// =========================================================================
// @author Leonardo Florez Valencia
// @email florez-l@javeriana.edu.co
// =========================================================================

#ifndef __fpa__Image__SkeletonFilter__hxx__
#define __fpa__Image__SkeletonFilter__hxx__

#include <itkImageRegionIteratorWithIndex.h>
#include <itkMinimumMaximumImageCalculator.h>

#include <fpa/Image/Functors/VertexIdentity.h>
#include <fpa/Base/Functors/InvertValue.h>

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
void fpa::Image::SkeletonFilter< _TImage, _TScalar >::
SetInput( const TScalarImage* image )
{
  // Do nothing
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
void fpa::Image::SkeletonFilter< _TImage, _TScalar >::
SetInput( const TImage* image )
{
  this->m_DistanceMap->SetInput( image );
  this->Superclass::SetInput( this->m_DistanceMap->GetOutput( ) );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
_TImage* fpa::Image::SkeletonFilter< _TImage, _TScalar >::
GetInput( )
{
  return( this->m_DistanceMap->GetInput( ) );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
const _TImage* fpa::Image::SkeletonFilter< _TImage, _TScalar >::
GetInput( ) const
{
  return( this->m_DistanceMap->GetInput( ) );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
typename fpa::Image::SkeletonFilter< _TImage, _TScalar >::
TSkeleton* fpa::Image::SkeletonFilter< _TImage, _TScalar >::
GetSkeleton( )
{
  return(
    dynamic_cast< TSkeleton* >(
      this->itk::ProcessObject::GetOutput( this->m_SkeletonIdx )
      )
    );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
typename fpa::Image::SkeletonFilter< _TImage, _TScalar >::
TMarks* fpa::Image::SkeletonFilter< _TImage, _TScalar >::
GetMarks( )
{
  return(
    dynamic_cast< TMarks* >(
      this->itk::ProcessObject::GetOutput( this->m_MarksIdx )
      )
    );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
fpa::Image::SkeletonFilter< _TImage, _TScalar >::
SkeletonFilter( )
  : Superclass( ),
    m_SeedFromMaximumDistance( false )
{
  typedef fpa::Image::Functors::VertexIdentity< TScalarImage, _TScalar > _TVertexFunc;
  typedef fpa::Base::Functors::InvertValue< _TScalar, _TScalar > _TValueFunc;

  typename _TVertexFunc::Pointer vertex_func = _TVertexFunc::New( );
  typename _TValueFunc::Pointer value_func = _TValueFunc::New( );
  value_func->SetAlpha( 1 );
  value_func->SetBeta( 1 );

  this->SetFunctor( vertex_func );
  this->SetFunctor( value_func );

  this->m_DistanceMap = TDistanceMap::New( );
  this->m_DistanceMap->InsideIsPositiveOn( );
  this->m_DistanceMap->SquaredDistanceOff( );
  this->m_DistanceMap->UseImageSpacingOn( );

  unsigned int nOutputs = this->GetNumberOfRequiredOutputs( );
  this->SetNumberOfRequiredOutputs( nOutputs + 2 );
  this->m_SkeletonIdx = nOutputs;
  this->m_MarksIdx = nOutputs + 1;

  typename TSkeleton::Pointer skeleton = TSkeleton::New( );
  typename TMarks::Pointer marks = TMarks::New( );
  this->SetNthOutput( this->m_SkeletonIdx, skeleton.GetPointer( ) );
  this->SetNthOutput( this->m_MarksIdx, marks.GetPointer( ) );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
fpa::Image::SkeletonFilter< _TImage, _TScalar >::
~SkeletonFilter( )
{
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
void fpa::Image::SkeletonFilter< _TImage, _TScalar >::
GenerateData( )
{
  // Prepare distance map
  this->m_DistanceMap->Update( );

  // Update start seed
  TVertex seed;
  if( this->m_SeedFromMaximumDistance )
  {
    typedef itk::MinimumMaximumImageCalculator< TScalarImage > _TMinMax;
    typename _TMinMax::Pointer minmax = _TMinMax::New( );
    minmax->SetImage( this->m_DistanceMap->GetOutput( ) );
    minmax->Compute( );
    seed = minmax->GetIndexOfMaximum( );
  }
  else
    seed = *( this->BeginSeeds( ) );
  this->m_Seeds.clear( );
  this->m_Seeds.insert( seed );

  // Prepare skeleton candidates queue
  this->m_SkeletonQueue.clear( );

  // Go!
  this->Superclass::GenerateData( );

  // Backtracking
  _TAdjacencies A;
  std::vector< TVertex > end_points;
  this->_EndPoints( end_points, A );
  this->_Skeleton( end_points, A );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
void fpa::Image::SkeletonFilter< _TImage, _TScalar >::
_SetOutputValue( const TVertex& vertex, const TOutputValue& value )
{
  typedef typename _TSkeletonQueue::value_type _TSkeletonQueueValue;

  this->Superclass::_SetOutputValue( vertex, value );
  double d = double( this->m_DistanceMap->GetOutput( )->GetPixel( vertex ) );
  if( d >= double( 0 ) )
  {
    // Update skeleton candidates
    d += double( 1e-5 );
    double v = double( value ) / ( d * d );
    this->m_SkeletonQueue.insert( _TSkeletonQueueValue( v, vertex ) );

  } // fi
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
void fpa::Image::SkeletonFilter< _TImage, _TScalar >::
_EndPoints( std::vector< TVertex >& end_points, _TAdjacencies& A )
{
  typedef typename _TSkeletonQueue::value_type _TSkeletonQueueValue;

  TMarks* marks = this->GetMarks( );
  TMST* mst = this->GetMinimumSpanningTree( );
  typename TImage::SpacingType spac = marks->GetSpacing( );

  // Some values
  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( 0 );

  // BFS from maximum queue
  while( this->m_SkeletonQueue.size( ) > 0 )
  {
    // Get node
    typename _TSkeletonQueue::iterator nIt = this->m_SkeletonQueue.begin( );
    _TSkeletonQueueValue n = *nIt;
    this->m_SkeletonQueue.erase( nIt );

    // Mark it and update end-points
    unsigned char m = marks->GetPixel( n.second );
    if( m != 0 )
      continue;
    marks->SetPixel( n.second, 1 );
    end_points.push_back( n.second );

    // Mark path
    TVertex it = n.second;
    TVertex p = mst->GetParent( it );
    while( it != p )
    {
      this->_MarkSphere( it );
      it = p;
      p = mst->GetParent( it );

    } // elihw
    this->_MarkSphere( it );
    A[ n.second ] = it;

  } // elihw
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
void fpa::Image::SkeletonFilter< _TImage, _TScalar >::
_Skeleton( const std::vector< TVertex >& end_points, _TAdjacencies& A )
{
  typedef typename TSkeleton::TPath _TPath;
  typedef itk::Image< unsigned long, TImage::ImageDimension > _TTagsImage;

  TMST* mst = this->GetMinimumSpanningTree( );
  TSkeleton* skeleton = this->GetSkeleton( );

  // Tag branches
  typename _TTagsImage::Pointer tags = _TTagsImage::New( );
  tags->SetLargestPossibleRegion( mst->GetLargestPossibleRegion( ) );
  tags->SetRequestedRegion( mst->GetRequestedRegion( ) );
  tags->SetBufferedRegion( mst->GetBufferedRegion( ) );
  tags->Allocate( );
  tags->FillBuffer( 0 );
  typename std::vector< TVertex >::const_iterator eIt = end_points.begin( );
  for( ; eIt != end_points.end( ); ++eIt )
  {
    TVertex it = *eIt;
    TVertex p = mst->GetParent( it );
    while( it != p )
    {
      tags->SetPixel( it, tags->GetPixel( it ) + 1 );
      it = p;
      p = mst->GetParent( it );

    } // elihw
    tags->SetPixel( it, tags->GetPixel( it ) + 1 );

  } // rof

  // Build paths (branches)
  eIt = end_points.begin( );
  for( ; eIt != end_points.end( ); ++eIt )
  {
    TVertex it = *eIt;
    TVertex p = mst->GetParent( it );
    TVertex sIdx = it;
    typename _TPath::Pointer path = _TPath::New( );
    path->SetReferenceImage( mst );
    while( it != p )
    {
      if( tags->GetPixel( sIdx ) != tags->GetPixel( it ) )
      {
        // Ok, a new branch has been added
        path->AddVertex( it );
        skeleton->AddBranch( path );

        // Update a new starting index
        path = _TPath::New( );
        path->SetReferenceImage( mst );
        sIdx = it;
      }
      else
        path->AddVertex( it );
      it = p;
      p = mst->GetParent( it );

    } // elihw

    // Finally, add last branch
    path->AddVertex( it );
    skeleton->AddBranch( path );

  } // rof
}

// -------------------------------------------------------------------------
template< class _TImage, class _TScalar >
void fpa::Image::SkeletonFilter< _TImage, _TScalar >::
_MarkSphere( const TVertex& idx )
{
  typedef itk::ImageRegionIteratorWithIndex< TMarks > _TMarksIt;

  static const double _eps = std::sqrt( double( TImage::ImageDimension + 1 ) );
  const TScalarImage* dmap = this->m_DistanceMap->GetOutput( );
  TMarks* marks = this->GetMarks( );
  typename TImage::SpacingType spac = dmap->GetSpacing( );
  typename TImage::RegionType region = dmap->GetRequestedRegion( );

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

  TVertex i0, i1;
  for( unsigned int d = 0; d < TImage::ImageDimension; ++d )
  {
    long off = long( std::ceil( r / double( spac[ d ] ) ) );
    if( off < 3 )
      off = 3;
    i0[ d ] = idx[ d ] - off;
    i1[ d ] = idx[ 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 TImage::SizeType size;
  for( unsigned int d = 0; d < TImage::ImageDimension; ++d )
    size[ d ] = i1[ d ] - i0[ d ] + 1;

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

  _TMarksIt mIt( marks, neighRegion );
  for( mIt.GoToBegin( ); !mIt.IsAtEnd( ); ++mIt )
  {
    TVertex w = mIt.GetIndex( );
    typename TImage::PointType p;
    marks->TransformIndexToPhysicalPoint( w, p );
    mIt.Set( 1 );

  } // rof
}

#endif // __fpa__Image__SkeletonFilter__hxx__

// eof - $RCSfile$