#ifndef __FPA__IMAGE__DIJKSTRAWITHENDPOINTDETECTION__HXX__
#define __FPA__IMAGE__DIJKSTRAWITHENDPOINTDETECTION__HXX__

#include <vector>
#include <set>
#include <itkImageRegionConstIteratorWithIndex.h>
#include <itkImageRegionIteratorWithIndex.h>

// -------------------------------------------------------------------------
template< class I, class O >
typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TLabelImage* fpa::Image::DijkstraWithEndPointDetection< I, O >::
GetLabelImage( )
{
  return(
    dynamic_cast< TLabelImage* >(
      this->itk::ProcessObject::GetOutput( this->m_LabelImageIndex )
      )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
const typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TLabelImage* fpa::Image::DijkstraWithEndPointDetection< I, O >::
GetLabelImage( ) const
{
  return(
    dynamic_cast< const TLabelImage* >(
      this->itk::ProcessObject::GetOutput( this->m_LabelImageIndex )
      )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
GraftLabelImage( itk::DataObject* obj )
{
  TLabelImage* lbl =
    dynamic_cast< TLabelImage* >(
      this->itk::ProcessObject::GetOutput( this->m_LabelImageIndex )
      );
  if( lbl != NULL )
    this->GraftNthOutput( this->m_LabelImageIndex, lbl );
}

// -------------------------------------------------------------------------
template< class I, class O >
typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TUniqueVertices* fpa::Image::DijkstraWithEndPointDetection< I, O >::
GetEndPoints( )
{
  return(
    dynamic_cast< TUniqueVertices* >(
      this->itk::ProcessObject::GetOutput( this->m_EndPointsIndex )
      )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
const typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TUniqueVertices* fpa::Image::DijkstraWithEndPointDetection< I, O >::
GetEndPoints( ) const
{
  return(
    dynamic_cast< const TUniqueVertices* >(
      this->itk::ProcessObject::GetOutput( this->m_EndPointsIndex )
      )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
GraftEndPoints( itk::DataObject* obj )
{
  TUniqueVertices* lbl =
    dynamic_cast< TUniqueVertices* >(
      this->itk::ProcessObject::GetOutput( this->m_EndPointsIndex )
      );
  if( lbl != NULL )
    this->GraftNthOutput( this->m_EndPointsIndex, lbl );
}

// -------------------------------------------------------------------------
template< class I, class O >
typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TUniqueVertices* fpa::Image::DijkstraWithEndPointDetection< I, O >::
GetBifurcations( )
{
  return(
    dynamic_cast< TUniqueVertices* >(
      this->itk::ProcessObject::GetOutput( this->m_BifurcationsIndex )
      )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
const typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TUniqueVertices* fpa::Image::DijkstraWithEndPointDetection< I, O >::
GetBifurcations( ) const
{
  return(
    dynamic_cast< const TUniqueVertices* >(
      this->itk::ProcessObject::GetOutput( this->m_BifurcationsIndex )
      )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
GraftBifurcations( itk::DataObject* obj )
{
  TUniqueVertices* lbl =
    dynamic_cast< TUniqueVertices* >(
      this->itk::ProcessObject::GetOutput( this->m_BifurcationsIndex )
      );
  if( lbl != NULL )
    this->GraftNthOutput( this->m_BifurcationsIndex, lbl );
}


// -------------------------------------------------------------------------
template< class I, class O >
typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TBranches* fpa::Image::DijkstraWithEndPointDetection< I, O >::
GetBranches( )
{
  return(
    dynamic_cast< TBranches* >(
      this->itk::ProcessObject::GetOutput( this->m_BranchesIndex )
      )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
const typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TBranches* fpa::Image::DijkstraWithEndPointDetection< I, O >::
GetBranches( ) const
{
  return(
    dynamic_cast< const TBranches* >(
      this->itk::ProcessObject::GetOutput( this->m_BranchesIndex )
      )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
GraftBranches( itk::DataObject* obj )
{
  TBranches* lbl =
    dynamic_cast< TBranches* >(
      this->itk::ProcessObject::GetOutput( this->m_BranchesIndex )
      );
  if( lbl != NULL )
    this->GraftNthOutput( this->m_BranchesIndex, lbl );
}

// -------------------------------------------------------------------------
template< class I, class O >
fpa::Image::DijkstraWithEndPointDetection< I, O >::
DijkstraWithEndPointDetection( )
  : Superclass( ),
    m_CorrectSeeds( true ),
    m_CorrectEndPoints( true ),
    m_SafetyNeighborhoodSize( 3 )
{
  this->m_LabelImageIndex = this->GetNumberOfRequiredOutputs( );
  this->m_EndPointsIndex = this->m_LabelImageIndex + 1;
  this->m_BifurcationsIndex = this->m_LabelImageIndex + 2;
  this->m_BranchesIndex = this->m_LabelImageIndex + 3;
  this->SetNumberOfRequiredOutputs( this->m_LabelImageIndex + 4 );
  this->itk::ProcessObject::SetNthOutput(
    this->m_LabelImageIndex, TLabelImage::New( )
    );
  this->itk::ProcessObject::SetNthOutput(
    this->m_EndPointsIndex, TUniqueVertices::New( )
    );
  this->itk::ProcessObject::SetNthOutput(
    this->m_BifurcationsIndex, TUniqueVertices::New( )
    );
  this->itk::ProcessObject::SetNthOutput(
    this->m_BranchesIndex, TBranches::New( )
    );
}

// -------------------------------------------------------------------------
template< class I, class O >
fpa::Image::DijkstraWithEndPointDetection< I, O >::
~DijkstraWithEndPointDetection( )
{
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
_BeforeGenerateData( )
{
  if( this->m_CorrectSeeds )
  {
    const I* input = this->GetInput( );
    typename I::SpacingType spac = input->GetSpacing( );
    double max_spac = double( spac[ 0 ] );
    for( unsigned int d = 1; d < I::ImageDimension; ++d )
      max_spac =
        ( double( spac[ d ] ) > max_spac )? double( spac[ d ] ): max_spac;
    max_spac *= double( 3 );

    // Correct seeds
    for( unsigned int s = 0; s < this->m_SeedVertices.size( ); ++s )
    {
      TVertex seed = this->m_SeedVertices[ s ];
      _TNode n = this->m_Seeds[ seed ];
      _TRegion region = this->_Region( seed, max_spac );
      itk::ImageRegionConstIteratorWithIndex< I > iIt( input, region );

      iIt.GoToBegin( );
      _TPixel max_value = iIt.Get( );
      for( ++iIt; !iIt.IsAtEnd( ); ++iIt )
      {
        if( iIt.Get( ) > max_value )
        {
          this->m_SeedVertices[ s ] = iIt.GetIndex( );
          max_value = iIt.Get( );

        } // fi

      } // rof
      this->m_Seeds.erase( seed );
      n.Parent = this->m_SeedVertices[ s ];
      this->m_Seeds[ this->m_SeedVertices[ s ] ] = n;

    } // rof

  } // fi

  // End initialization
  this->Superclass::_BeforeGenerateData( );
  this->m_Candidates.clear( );
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
_AfterGenerateData( )
{
  // Finish base algorithm
  this->Superclass::_AfterGenerateData( );

  // Check if backtracking has any sense
  if( this->m_Candidates.size( ) == 0 )
    return;
  this->InvokeEvent( TEndEvent( ) );
  this->InvokeEvent( TStartBacktrackingEvent( ) );

  // Detect endpoints and bifurcations
  this->_EndPointsAndBifurcations( );

  // Find branches
  this->_FindBranches( );

  // Label pixels and branches
  this->_LabelAll( );
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
_SetResult( const TVertex& v, const _TNode& n )
{
  this->Superclass::_SetResult( v, n );
  this->m_Candidates.insert( _TCandidate( n.Result, v ) );
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
_EndPointsAndBifurcations( )
{
  // Prepare backtracking objects
  TUniqueVertices* endpoints = this->GetEndPoints( );
  TUniqueVertices* bifurcations = this->GetBifurcations( );
  endpoints->Clear( );
  bifurcations->Clear( );

  // Get some input values
  TVertex seed = this->GetSeed( 0 );
  const I* input = this->GetInput( );
  const TMinimumSpanningTree* mst = this->GetMinimumSpanningTree( );

  // Prepare a pixel marks structure
  typedef std::map< TVertex, bool, TVertexCompare > _TPixelMarks;
  _TPixelMarks pixel_marks;

  // Object to hold tree-nodes marks
  typedef std::set< TVertex, TVertexCompare > _TTreeMarks;
  _TTreeMarks tree_marks;

  // Iterate over the candidates, starting from the candidates that
  // are near thin branches (high accumulated cost)
  typename _TCandidates::const_reverse_iterator cIt =
    this->m_Candidates.rbegin( );
  for( ; cIt != this->m_Candidates.rend( ); ++cIt )
  {
    // If pixel has already been visited, pass
    TVertex v = cIt->second;
    if( pixel_marks[ v ] )
      continue;

    // Correct it to nearest start candidate (high distance value)
    double vr = std::sqrt( double( input->GetPixel( v ) ) );
    if( this->m_CorrectEndPoints )
      v = this->_MaxInRegion( input, v, vr * double( 1.5 ) );

    // Now, check for real marking conditions
    //   1. Has it been visited by dijkstra?
    if( this->_Node( v ).Label == Self::FarLabel )
      continue;
    //   2. Is it already marked?
    if( pixel_marks[ v ] )
      continue;
    //   3. Is it already an endpoint?
    if( endpoints->Find( v ) )
      continue;
    //   4. Ok, it is completely new!
    endpoints->Insert( v );

    // Get the path all the way to global seed
    TVertices path = mst->GetPath( v, seed );

    // Backtracking to find endpoints and bifurcations
    bool adding_new_points = true;
    typename TVertices::const_iterator pIt = path.begin( );
    for( ; pIt != path.end( ) && adding_new_points; ++pIt )
    {
      this->InvokeEvent( TBacktrackingEvent( *pIt, 1 ) );
      if( tree_marks.find( *pIt ) == tree_marks.end( ) )
      {
        // Mark current point as a tree point
        tree_marks.insert( *pIt );

        // Mark a region around current point as visited
        vr = std::sqrt( double( input->GetPixel( *pIt ) ) );
        _TRegion region = this->_Region( *pIt, vr );
        if( region.GetNumberOfPixels( ) > 0 )
        {
          itk::ImageRegionConstIteratorWithIndex< I > iIt( input, region );
          for( iIt.GoToBegin( ); !iIt.IsAtEnd( ); ++iIt )
            pixel_marks[ iIt.GetIndex( ) ] = true;
        }
        else
          pixel_marks[ *pIt ] = true;
      }
      else
      {
        // A bifurcation point has been reached!
        if( *pIt != seed )
        {
          bifurcations->Insert( *pIt );
          adding_new_points = false;

        } // fi

      } // fi

    } // rof
    this->InvokeEvent( TEndBacktrackingEvent( ) );

  } // rof
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
_FindBranches( )
{
  // Configure and obtain inputs
  TVertex seed = this->GetSeed( 0 );
  const TMinimumSpanningTree* mst = this->GetMinimumSpanningTree( );
  const TUniqueVertices* endpoints = this->GetEndPoints( );
  const TUniqueVertices* bifurcations = this->GetBifurcations( );
  TBranches* branches = this->GetBranches( );
  branches->Clear( );

  // Reconstruct pixels
  typename TUniqueVertices::ConstIterator eIt = endpoints->Begin( );
  for( ; eIt != endpoints->End( ); ++eIt )
  {
    // Get the path all the way to global seed
    TVertices path = mst->GetPath( *eIt, seed );

    TVertex start_vertex = *eIt;
    typename TVertices::const_iterator pIt = path.begin( );
    for( ; pIt != path.end( ); ++pIt )
    {
      if( bifurcations->Find( *pIt ) )
      {
        branches->SetValue( start_vertex, *pIt, 1 );
        start_vertex = *pIt;

      } // fi

    } // rof

    // Finish with branches to global seed
    branches->SetValue( start_vertex, seed, 1 );

  } // rof
}

// -------------------------------------------------------------------------
template< class I, class O >
void fpa::Image::DijkstraWithEndPointDetection< I, O >::
_LabelAll( )
{
  // Configure and obtain inputs
  const I* input = this->GetInput( );
  const TMinimumSpanningTree* mst = this->GetMinimumSpanningTree( );
  const TUniqueVertices* bifurcations = this->GetBifurcations( );
  TBranches* branches = this->GetBranches( );
  TLabelImage* label = this->GetLabelImage( );
  label->FillBuffer( 0 );

  // Label branches
  typename TBranches::Iterator bIt = branches->Begin( );
  TLabel actual_label = 0;
  for( ; bIt != branches->End( ); ++bIt )
  {
    typename TBranches::RowIterator brIt = branches->Begin( bIt );
    for( ; brIt != branches->End( bIt ); ++brIt )
    {
      actual_label++;
      brIt->second = actual_label;

      TVertices path = mst->GetPath( bIt->first, brIt->first );
      typename TVertices::const_iterator pIt = path.begin( );
      for( ; pIt != path.end( ); ++pIt )
      {
        double d = std::sqrt( double( input->GetPixel( *pIt ) ) );
        _TRegion region = this->_Region( *pIt, d );
        itk::ImageRegionConstIteratorWithIndex< I > iIt( input, region );
        itk::ImageRegionIteratorWithIndex< TLabelImage > lIt( label, region );
        iIt.GoToBegin( );
        lIt.GoToBegin( );
        for( ; !iIt.IsAtEnd( ); ++iIt, ++lIt )
        {
          // Mask real input image
          if(
            !( iIt.Get( ) < ( typename I::PixelType )( 0 ) ) &&
            lIt.Get( ) == TLabel( 0 )
            )
            lIt.Set( actual_label );
          
        } // rof

      } // rof

    } // rof

  } // rof
  this->m_NumberOfBranches = actual_label;

  // Label bifurcations
  typename TUniqueVertices::ConstIterator bifIt =
    bifurcations->Begin( );
  for( ; bifIt != bifurcations->End( ); ++bifIt )
  {
    actual_label++;
    double d = std::sqrt( double( input->GetPixel( *bifIt ) ) );
    _TRegion region = this->_Region( *bifIt, d * 1.5 );
    itk::ImageRegionConstIteratorWithIndex< I > iIt( input, region );
    itk::ImageRegionIteratorWithIndex< TLabelImage > lIt( label, region );
    iIt.GoToBegin( );
    lIt.GoToBegin( );
    for( ; !iIt.IsAtEnd( ); ++iIt, ++lIt )
    {
      // Mask real input image
      if( !( iIt.Get( ) < ( typename I::PixelType )( 0 ) ) )
        lIt.Set( actual_label );

    } // rof

  } // rof
}

// -------------------------------------------------------------------------
template< class I, class O >
typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
_TRegion fpa::Image::DijkstraWithEndPointDetection< I, O >::
_Region( const TVertex& c, const double& r )
{
  typename I::ConstPointer input = this->GetInput( );
  typename I::SpacingType spac = input->GetSpacing( );
  _TRegion region = input->GetLargestPossibleRegion( );
  typename I::IndexType idx0 = region.GetIndex( );
  typename I::IndexType idx1 = idx0 + region.GetSize( );

  // Compute region size and index
  typename I::IndexType i0, i1;
  _TSize size;
  for( unsigned int d = 0; d < I::ImageDimension; ++d )
  {
    // NOTE: 3 is a minimum neighborhood size
    long s =
      long( std::ceil( r / double( spac[ d ] ) ) ) +
      long( this->m_SafetyNeighborhoodSize );
    i0[ d ] = c[ d ] - s;
    i1[ d ] = c[ d ] + s;

    if( i0[ d ] < idx0[ d ] ) i0[ d ] = idx0[ d ];
    if( i1[ d ] < idx0[ d ] ) i1[ d ] = idx0[ d ];
    if( i0[ d ] > idx1[ d ] ) i0[ d ] = idx1[ d ];
    if( i1[ d ] > idx1[ d ] ) i1[ d ] = idx1[ d ];
    size[ d ] = i1[ d ] - i0[ d ];

  }  // rof

  // Prepare region and return it
  region.SetIndex( i0 );
  region.SetSize( size );

  return( region );
}

// -------------------------------------------------------------------------
template< class I, class O >
template< class _T > 
typename fpa::Image::DijkstraWithEndPointDetection< I, O >::
TVertex fpa::Image::DijkstraWithEndPointDetection< I, O >::
_MaxInRegion( const _T* image, const TVertex& v, const double& r )
{
  typedef itk::ImageRegionConstIteratorWithIndex< _T > _TIt;

  _TIt iIt( image, this->_Region( v, r ) );
  iIt.GoToBegin( );
  TVertex max_vertex = iIt.GetIndex( );
  typename _T::PixelType max_value = iIt.Get( );
  for( ++iIt; !iIt.IsAtEnd( ); ++iIt )
  {
    typename _T::PixelType value = iIt.Get( );
    if( value > max_value )
    {
      max_value = value;
      max_vertex = iIt.GetIndex( );

    } // fi

  } // rof
  return( max_vertex );
}

#endif // __FPA__IMAGE__DIJKSTRAWITHENDPOINTDETECTION__HXX__

// eof - $RCSfile$