#include <vtkCellArray.h>
 #include <vtkFloatArray.h>
 #include <vtkPolyData.h>
+#include <vtkPolyDataWriter.h>
 #include <vtkSmartPointer.h>
 #include <vtkImageMarchingCubes.h>
+#include <vtkLookupTable.h>
 
 #include <fpa/Image/DijkstraWithSphereBacktracking.h>
 #include <fpa/VTK/ImageMPR.h>
   vtkSmartPointer< vtkFloatArray > scalars =
     vtkSmartPointer< vtkFloatArray >::New( );
 
+  const TDijkstra::TMarkImage* marks = paths->GetOutputMarkImage( );
+  TDijkstra::TMark max_mark = paths->GetNumberOfBranches( );
   const TDijkstra::TVertices& endpoints = paths->GetEndPoints( );
   const TDijkstra::TTree& tree = paths->GetFinalTree( );
   TDijkstra::TVertices::const_iterator epIt = endpoints.begin( );
   for( unsigned int epId = 0; epIt != endpoints.end( ); ++epIt, ++epId )
   {
-    double pd = double( epId ) / double( endpoints.size( ) - 1 );
-
     TDijkstra::TVertex idx = *epIt;
     do
     {
       TImage::PointType pnt;
       input_image->TransformIndexToPhysicalPoint( idx, pnt );
 
+      TDijkstra::TMark mark = marks->GetPixel( idx );
+      double pd = double( mark );
+
       points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );
       scalars->InsertNextTuple1( pd );
       if( idx != *epIt )
   vtk_tree->SetLines( cells );
   vtk_tree->GetPointData( )->SetScalars( scalars );
 
-  view.AddPolyData( vtk_tree );
+  vtkSmartPointer<vtkLookupTable> lut =
+    vtkSmartPointer<vtkLookupTable>::New( );
+  lut->SetNumberOfTableValues( max_mark + 1 );
+  lut->SetTableRange( 0, max_mark );
+  lut->Build( );
+
+  view.AddPolyData( vtk_tree, lut );
   view.Render( );
   view.Start( );
 
+  vtkSmartPointer< vtkPolyDataWriter > writer =
+    vtkSmartPointer< vtkPolyDataWriter >::New( );
+  writer->SetInputData( vtk_tree );
+  writer->SetFileName( output_image_fn.c_str( ) );
+  writer->Update( );
+
+  itk::ImageFileWriter< TDijkstra::TMarkImage >::Pointer marks_writer =
+    itk::ImageFileWriter< TDijkstra::TMarkImage >::New( );
+  marks_writer->SetInput( marks );
+  marks_writer->SetFileName( "marks.mhd" );
+  marks_writer->Update( );
+
   /* TODO
      TImageWriter::Pointer dijkstra_writer =
      TImageWriter::New( );
 
 #include <itkImageRegionConstIteratorWithIndex.h>
 #include <itkImageRegionIteratorWithIndex.h>
 
+// -------------------------------------------------------------------------
+template< class I, class C >
+typename fpa::Image::DijkstraWithSphereBacktracking< I, C >::
+TMarkImage* fpa::Image::DijkstraWithSphereBacktracking< I, C >::
+GetOutputMarkImage( )
+{
+  return(
+    dynamic_cast< TMarkImage* >(
+      this->itk::ProcessObject::GetOutput( 1 )
+      )
+    );
+}
+
+// -------------------------------------------------------------------------
+template< class I, class C >
+const typename fpa::Image::DijkstraWithSphereBacktracking< I, C >::
+TMarkImage* fpa::Image::DijkstraWithSphereBacktracking< I, C >::
+GetOutputMarkImage( ) const
+{
+  return(
+    dynamic_cast< const TMarkImage* >(
+      this->itk::ProcessObject::GetOutput( 1 )
+      )
+    );
+}
+
 // -------------------------------------------------------------------------
 template< class I, class C >
 fpa::Image::DijkstraWithSphereBacktracking< I, C >::
 DijkstraWithSphereBacktracking( )
-  : Superclass( )
+  : Superclass( ),
+    m_NumberOfBranches( 0 )
 {
+  this->SetNumberOfRequiredOutputs( 2 );
+  this->SetNthOutput( 0, I::New( ) );
+  this->SetNthOutput( 1, TMarkImage::New( ) );
 }
 
 // -------------------------------------------------------------------------
 void fpa::Image::DijkstraWithSphereBacktracking< I, C >::
 _AfterMainLoop( )
 {
-  typedef itk::Image< bool, I::ImageDimension > _TMarkImage;
-
   // Finish base algorithm
   this->Superclass::_AfterMainLoop( );
   this->m_FinalTree.clear( );
   max_spac *= double( 3 );
 
   // Prepare an object to hold marks
-  typename _TMarkImage::Pointer marks = _TMarkImage::New( );
-  marks->SetLargestPossibleRegion( input->GetLargestPossibleRegion( ) );
-  marks->SetRequestedRegion( input->GetRequestedRegion( ) );
-  marks->SetBufferedRegion( input->GetBufferedRegion( ) );
-  marks->SetOrigin( input->GetOrigin( ) );
-  marks->SetSpacing( input->GetSpacing( ) );
-  marks->SetDirection( input->GetDirection( ) );
-  marks->Allocate( );
-  marks->FillBuffer( false );
+  typename TMarkImage::Pointer marks = this->GetOutputMarkImage( );
+  marks->FillBuffer( 0 );
 
   // Iterate over the candidates, starting from the candidates that
   // are near thin branches
   typename _TCandidates::const_reverse_iterator cIt =
     this->m_Candidates.rbegin( );
-  for( int backId = 0; cIt != this->m_Candidates.rend( ); ++cIt )
+  this->m_NumberOfBranches = 0;
+  for( ; cIt != this->m_Candidates.rend( ); ++cIt )
   {
     // If pixel hasn't been visited, continue
     TVertex v = cIt->second;
-    if( marks->GetPixel( v ) )
+    if( marks->GetPixel( v ) != 0 )
       continue;
 
     // Compute nearest start candidate
     } // rof
 
     // Keep endpoint
-    if( marks->GetPixel( max_vertex ) )
+    if( marks->GetPixel( max_vertex ) != 0 )
       continue;
-    backId++;
+    this->m_NumberOfBranches++;
     this->m_EndPoints.push_back( max_vertex );
 
     // Construct path (at least the part that hasn't been iterated)
-    bool terminate = false;
     do
     {
-      if( this->m_FinalTree.find( max_vertex ) == this->m_FinalTree.end( ) )
+      if(
+        std::find(
+          this->m_BifurcationPoints.begin( ),
+          this->m_BifurcationPoints.end( ),
+          max_vertex
+          ) == this->m_BifurcationPoints.end( )
+        )
       {
         // Mark a sphere around current point as visited
         double dist = std::sqrt( double( input->GetPixel( max_vertex ) ) );
-        region = this->_Region( max_vertex, dist * double( 1.25 ) );
-        itk::ImageRegionIteratorWithIndex< _TMarkImage >
+        region = this->_Region( max_vertex, dist * double( 1.1 ) );
+        itk::ImageRegionIteratorWithIndex< TMarkImage >
           mIt( marks, region );
         for( mIt.GoToBegin( ); !mIt.IsAtEnd( ); ++mIt )
-          mIt.Set( true );
+          mIt.Set( this->m_NumberOfBranches );
 
         // Next vertex in current path
-        this->InvokeEvent( TBacktrackingEvent( max_vertex, backId ) );
+        this->InvokeEvent( TBacktrackingEvent( max_vertex, this->m_NumberOfBranches ) );
         this->m_FinalTree[ max_vertex ] = this->_Parent( max_vertex );
       }
       else
       {
         // A bifurcation point has been reached!
         this->m_BifurcationPoints.push_back( max_vertex );
-        terminate = true;
+        this->m_NumberOfBranches++;
 
       } // fi
       max_vertex = this->_Parent( max_vertex );
 
-    } while( max_vertex != this->_Parent( max_vertex ) && !terminate );
-
-    if( !terminate )
-    {
-      this->m_FinalTree[ max_vertex ] = max_vertex;
-      this->InvokeEvent( TEndBacktrackingEvent( backId ) );
-
-    } // fi
+    } while( max_vertex != this->_Parent( max_vertex ) );
+
+    /* TODO
+       bool terminate = false;
+       do
+       {
+       if( this->m_FinalTree.find( max_vertex ) == this->m_FinalTree.end( ) )
+       {
+       // Mark a sphere around current point as visited
+       double dist = std::sqrt( double( input->GetPixel( max_vertex ) ) );
+       region = this->_Region( max_vertex, dist * double( 1.25 ) );
+       itk::ImageRegionIteratorWithIndex< TMarkImage >
+       mIt( marks, region );
+       for( mIt.GoToBegin( ); !mIt.IsAtEnd( ); ++mIt )
+       mIt.Set( true );
+
+       // Next vertex in current path
+       this->InvokeEvent( TBacktrackingEvent( max_vertex, this->m_NumberOfBranches ) );
+       this->m_FinalTree[ max_vertex ] = this->_Parent( max_vertex );
+       }
+       else
+       {
+       // A bifurcation point has been reached!
+       this->m_BifurcationPoints.push_back( max_vertex );
+       terminate = true;
+
+       } // fi
+       max_vertex = this->_Parent( max_vertex );
+
+       } while( max_vertex != this->_Parent( max_vertex ) && !terminate );
+
+       if( !terminate )
+       {
+       this->m_FinalTree[ max_vertex ] = max_vertex;
+       this->InvokeEvent( TEndBacktrackingEvent( this->m_NumberOfBranches ) );
+
+       } // fi
+    */
 
   } // rof
 }