// =========================================================================
// @author Leonardo Florez Valencia
// @email florez-l@javeriana.edu.co
// =========================================================================
#ifndef __fpa__Filters__Image__ExtractSkeleton__hxx__
#define __fpa__Filters__Image__ExtractSkeleton__hxx__

#include <itkImage.h>
#include <itkImageRegionIteratorWithIndex.h>
#include <itkMinimumMaximumImageCalculator.h>
#include <fpa/Functors/Dijkstra/Invert.h>

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
_TDijkstra::_TDijkstra( )
  : Superclass( )
{
  // Prepare weight function
  typedef fpa::Functors::Dijkstra::Invert< TScalar > _TWeight;
  typename _TWeight::Pointer weight = _TWeight::New( );
  weight->SetAlpha( 1 );
  weight->SetBeta( 1 );
  this->SetWeightFunction( weight );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
_TDijkstra::~_TDijkstra( )
{
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
void fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
_TDijkstra::_BeforeGenerateData( )
{
  this->Superclass::_BeforeGenerateData( );
  this->m_SkeletonQueue.clear( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
void fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
_TDijkstra::_PostComputeOutputValue( TNode& n )
{
  typedef typename _TSkeletonQueue::value_type _TSkeletonQueueValue;

  this->Superclass::_PostComputeOutputValue( n );
  double d = double( this->GetInput( )->GetPixel( n.Vertex ) );
  if( d >= double( 0 ) )
  {
    // Update skeleton candidates
    d += double( 1e-5 );
    double v = double( n.Value ) / ( d * d );
    this->m_SkeletonQueue.insert( _TSkeletonQueueValue( v, n.Vertex ) );

  } // fi
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
itk::ModifiedTimeType
fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
GetMTime( ) const
{
  itk::ModifiedTimeType t = this->Superclass::GetMTime( );
  itk::ModifiedTimeType q = this->m_DistanceMap->GetMTime( );
  return( ( q < t )? q: t );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
void fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
SetInput( TInputImage* input )
{
  this->Superclass::SetNthInput( 0, input );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
typename fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
TInputImage* fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
GetInput( )
{
  return( dynamic_cast< TInputImage* >( this->Superclass::GetInput( 0 ) ) );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
const typename fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
TInputImage* fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
GetInput( ) const
{
  return(
    dynamic_cast< const TInputImage* >( this->Superclass::GetInput( 0 ) )
    );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
typename fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
TSkeleton* fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
GetOutput( )
{
  return( dynamic_cast< TSkeleton* >( this->Superclass::GetOutput( 0 ) ) );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
const typename fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
TSkeleton* fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
GetOutput( ) const
{
  return(
    dynamic_cast< const TSkeleton* >( this->Superclass::GetOutput( 0 ) )
    );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
ExtractSkeleton( )
  : Superclass( ),
    m_SeedFromMaximumDistance( false )
{
  this->SetNumberOfRequiredInputs( 1 );
  this->SetNumberOfRequiredOutputs( 1 );
  this->SetNthOutput( 0, TSkeleton::New( ) );

  this->m_DistanceMap = TDistanceMap::New( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
~ExtractSkeleton( )
{
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
void fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
GenerateData( )
{
  // Update distance map
  this->m_DistanceMap->SetInput( this->GetInput( ) );
  this->m_DistanceMap->Update( );

  // Correct seed
  if( this->m_SeedFromMaximumDistance )
  {
    typedef itk::MinimumMaximumImageCalculator< TOutputImage > _TMinMax;
    typename _TMinMax::Pointer minmax = _TMinMax::New( );
    minmax->SetImage( this->m_DistanceMap->GetOutput( ) );
    minmax->Compute( );
    this->m_Seed = minmax->GetIndexOfMaximum( );

  } // fi

  // Compute MST
  typename _TDijkstra::Pointer dijkstra = _TDijkstra::New( );
  dijkstra->SetInput( this->m_DistanceMap->GetOutput( ) );
  dijkstra->AddSeed( this->m_Seed );
  dijkstra->Update( );

  // Compute end-points
  const _TMST* mst = dijkstra->GetMinimumSpanningTree( );
  std::vector< TIndex > end_points;
  this->_EndPoints(
    end_points, this->m_DistanceMap->GetOutput( ),
    mst, dijkstra->GetSkeletonQueue( )
    );

  // Compute symbolic branches
  typedef std::map< TIndex, TIndex, typename TIndex::LexicographicCompare > _TTags;
  _TTags tags, branches;
  typename std::vector< TIndex >::const_iterator eIt = end_points.begin( );
  for( ; eIt != end_points.end( ); ++eIt )
  {
    // Tag path
    TIndex it = *eIt;
    TIndex p = mst->GetParent( it );
    typename _TTags::iterator bIt = tags.end( );
    while( it != p && bIt == tags.end( ) )
    {
      typename _TTags::iterator tIt = tags.find( it );
      if( tIt != tags.end( ) )
      {
        // Ok, a bifurcation point has been found
        // branch1: tIt->second <-> it (ok)
        branches[ tIt->second ] = it;

        // branch2: *eit <-> it (ok)
        branches[ *eIt ] = it;

        // branch3: it <-> until next bifurcation
        bIt = tIt;
      }
      else
        tags[ it ] = *eIt;
      it = p;
      p = mst->GetParent( it );

    } // elihw
    if( bIt != tags.end( ) )
    {
      TIndex pTag = bIt->second;
      TIndex nTag = bIt->first;
      it = bIt->first;
      p = it;
      while( tags[ it ] == pTag )
      {
        tags[ it ] = nTag;
        p = it;
        it = mst->GetParent( it );

      } // elihw
      tags[ it ] = nTag;
      branches[ bIt->first ] = p;
    }
    else
    {
      tags[ it ] = *eIt;
      branches[ *eIt ] = it;

    } // fi

  } // rof

  // Fill full branches
  typedef typename _TMST::TVertices _TVertices;
  typedef typename TSkeleton::TPath _TPath;

  TSkeleton* sk = this->GetOutput( );
  typename _TTags::const_iterator bIt = branches.begin( );
  for( ; bIt != branches.end( ); ++bIt )
  {
    _TVertices v = mst->GetAxis( bIt->first, bIt->second );
    typename _TPath::Pointer path = _TPath::New( );
    path->SetReferenceImage( this->GetInput( ) );
    typename _TVertices::const_reverse_iterator vIt = v.rbegin( );
    for( ; vIt != v.rend( ); ++vIt )
      path->AddVertex( *vIt );
    sk->AddBranch( path );

  } // rof
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
template< class _TMarksPointer >
void fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
_MarkSphere(
  _TMarksPointer& marks, const TOutputImage* dmap, const TIndex& center
  )
{
  typedef typename _TMarksPointer::ObjectType          _TMarks;
  typedef itk::ImageRegionIteratorWithIndex< _TMarks > _TMarksIt;

  static const double _eps = std::sqrt( double( Self::Dimension + 1 ) );
  typename _TMarks::SpacingType spac = dmap->GetSpacing( );
  typename _TMarks::RegionType region = dmap->GetRequestedRegion( );

  typename _TMarks::PointType cnt;
  dmap->TransformIndexToPhysicalPoint( center, cnt );
  double r = double( dmap->GetPixel( center ) ) * _eps;

  TIndex i0, i1;
  for( unsigned int d = 0; d < Self::Dimension; ++d )
  {
    long off = long( std::ceil( r / double( spac[ d ] ) ) );
    if( off < 3 )
      off = 3;
    i0[ d ] = center[ d ] - off;
    i1[ d ] = center[ d ] + off;

    if( i0[ d ] < region.GetIndex( )[ d ] )
      i0[ d ] = region.GetIndex( )[ d ];

    if( i1[ d ] >= region.GetIndex( )[ d ] + region.GetSize( )[ d ] )
      i1[ d ] = region.GetIndex( )[ d ] + region.GetSize( )[ d ] - 1;

  } // rof

  typename _TMarks::SizeType size;
  for( unsigned int d = 0; d < Self::Dimension; ++d )
    size[ d ] = i1[ d ] - i0[ d ] + 1;

  typename _TMarks::RegionType neighRegion;
  neighRegion.SetIndex( i0 );
  neighRegion.SetSize( size );

  _TMarksIt mIt( marks, neighRegion );
  for( mIt.GoToBegin( ); !mIt.IsAtEnd( ); ++mIt )
    mIt.Set( true );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TDistanceMap >
void fpa::Filters::Image::ExtractSkeleton< _TInputImage, _TDistanceMap >::
_EndPoints(
  std::vector< TIndex >& end_points,
  const TOutputImage* dmap,
  const _TMST* mst,
  const _TSkeletonQueue& queue
  )
{
  typedef typename _TSkeletonQueue::value_type _TSkeletonQueueValue;

  // Some values
  typedef itk::Image< bool, Self::Dimension > _TMarks;
  typename _TMarks::Pointer marks = _TMarks::New( );
  marks->SetLargestPossibleRegion( mst->GetLargestPossibleRegion( ) );
  marks->SetRequestedRegion( mst->GetRequestedRegion( ) );
  marks->SetBufferedRegion( mst->GetBufferedRegion( ) );
  marks->SetSpacing( mst->GetSpacing( ) );
  marks->SetOrigin( mst->GetOrigin( ) );
  marks->SetDirection( mst->GetDirection( ) );
  marks->Allocate( );
  marks->FillBuffer( false );

  // Get candidates in maximum to minimum iteration
  typename _TSkeletonQueue::const_reverse_iterator nIt = queue.rbegin( );
  for( ; nIt != queue.rend( ); ++nIt )
  {
    // Mark it and update end-points
    if( !( marks->GetPixel( nIt->second ) ) )
    {
      marks->SetPixel( nIt->second, true );
      end_points.push_back( nIt->second );

      // Mark path
      TIndex it = nIt->second;
      TIndex p = mst->GetParent( it );
      while( it != p )
      {
        this->_MarkSphere( marks, dmap, it );
        it = p;
        p = mst->GetParent( it );

      } // elihw
      this->_MarkSphere( marks, dmap, it );

    } // fi

  } // rof
}

#endif // __fpa__Filters__Image__ExtractSkeleton__hxx__

// eof - $RCSfile$