#include <cstdlib>
#include <ctime>
#include <iostream>
#include <limits>
#include <string>

#include <itkImage.h>
#include <itkImageToVTKImageFilter.h>

#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>

#include <fpa/Image/Functors/ImageCostFunction.h>
#include <fpa/Image/DijkstraWithEndPointDetection.h>
#include <fpa/VTK/ImageMPR.h>
#include <fpa/VTK/Image3DObserver.h>

#include <fpa/IO/MinimumSpanningTreeWriter.h>
#include <fpa/IO/UniqueValuesContainerWriter.h>
#include <fpa/IO/MatrixValuesContainerWriter.h>

#include <vtkCellArray.h>
#include <vtkFloatArray.h>
#include <vtkImageMarchingCubes.h>
#include <vtkPoints.h>
#include <vtkPolyData.h>
#include <vtkSmartPointer.h>

// -------------------------------------------------------------------------
const unsigned int Dim = 3;
typedef unsigned char TPixel;
typedef float         TScalar;

typedef itk::Image< TPixel, Dim >            TImage;
typedef itk::Image< TScalar, Dim >           TScalarImage;
typedef itk::ImageToVTKImageFilter< TImage > TVTKInputImage;

// -------------------------------------------------------------------------
template< class I >
void ReadImage( typename I::Pointer& image, const std::string& filename );

template< class I >
void SaveImage( const I* image, const std::string& filename );

// -------------------------------------------------------------------------
template< class F >
class SkeletonCostFunction
  : public fpa::Image::Functors::ImageCostFunction< typename F::TInputImage, typename F::TResult >
{
public:
  typedef fpa::Image::Functors::ImageCostFunction< typename F::TInputImage, typename F::TResult > Superclass;
  typedef SkeletonCostFunction            Self;
  typedef itk::SmartPointer< Self >       Pointer;
  typedef itk::SmartPointer< const Self > ConstPointer;

  typedef typename Superclass::TInputImage TInputImage;
  typedef typename Superclass::TResult     TResult;
  typedef typename Superclass::TIndex      TIndex;

public:
  itkNewMacro( Self );
  itkTypeMacro( SkeletonCostFunction, fpa_Image_Functors_ImageCostFunction );

public:
  virtual void SetInputImage( const TInputImage* img )
    {
      this->Superclass::SetInputImage( img );

      typename TInputImage::SpacingType spac = img->GetSpacing( );
      typename TInputImage::SpacingType::ValueType min_spac = spac[ 0 ];
      for( unsigned int d = 1; d < TInputImage::ImageDimension; ++d )
        min_spac = ( spac[ d ] < min_spac )? spac[ d ]: min_spac;
      this->m_Distance = TResult( min_spac );
    }

  virtual TResult Evaluate( const TIndex& v, const TIndex& p ) const
    {
      TResult res = this->Superclass::Evaluate( v, p );
      if( res > TResult( 0 ) )
        return( this->m_Distance );
      else
        return( TResult( -1 ) );
    }

protected:
  SkeletonCostFunction( )
    : Superclass( )
    { }
  virtual ~SkeletonCostFunction( )
    { }

private:
  // Purposely not implemented
  SkeletonCostFunction( const Self& );
  void operator=( const Self& );

protected:
  TResult m_Distance;
};

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  if( argc < 11 )
  {
    std::cerr
      << "Usage: " << argv[ 0 ] << std::endl
      << " input_image" << std::endl
      << " seed_x seed_y seed_z" << std::endl
      << " output_costmap" << std::endl
      << " output_labels" << std::endl
      << " output_minimum_spanning_tree" << std::endl
      << " output_endpoints" << std::endl
      << " output_bifurcations" << std::endl
      << " output_branches" << std::endl
      << std::endl;
    return( 1 );

  } // fi
  std::string input_image_fn = argv[ 1 ];
  std::string output_costmap_fn = argv[ 5 ];
  std::string output_labels_fn = argv[ 6 ];
  std::string mst_output_fn = argv[ 7 ];
  std::string endpoints_output_fn = argv[ 8 ];
  std::string bifurcations_output_fn = argv[ 9 ];
  std::string branches_output_fn = argv[ 10 ];

  // Get seed
  TImage::PointType pnt;
  pnt[ 0 ] = TImage::PointType::ValueType( std::atof( argv[ 2 ] ) );
  pnt[ 1 ] = TImage::PointType::ValueType( std::atof( argv[ 3 ] ) );
  pnt[ 2 ] = TImage::PointType::ValueType( std::atof( argv[ 4 ] ) );

  // Read image
  TImage::Pointer input_image;
  try
  {
    ReadImage< TImage >( input_image, input_image_fn );
  }
  catch( itk::ExceptionObject& err )
  {
    std::cerr
      << "Error caught while reading \""
      << input_image_fn << "\": " << err
      << std::endl;
    return( 1 );

  } // yrt

  TVTKInputImage::Pointer vtk_input_image = TVTKInputImage::New( );
  vtk_input_image->SetInput( input_image );
  vtk_input_image->Update( );

  // Show input image and let some interaction
  fpa::VTK::ImageMPR view;
  view.SetBackground( 0.3, 0.2, 0.8 );
  view.SetSize( 500, 500 );
  view.SetImage( vtk_input_image->GetOutput( ) );

  vtkSmartPointer< vtkImageMarchingCubes > mc =
    vtkSmartPointer< vtkImageMarchingCubes >::New( );
  mc->SetInputData( vtk_input_image->GetOutput( ) );
  mc->SetValue( 0, 1e-1 );
  mc->Update( );
  view.AddPolyData( mc->GetOutput( ), 1, 1, 1, 0.4 );

  // Allow some interaction and wait for, at least, one seed
  view.Render( );
  view.Start( );

  // Transform seed
  TImage::IndexType seed;
  input_image->TransformPhysicalPointToIndex( pnt, seed );

  // Prepare Dijkstra filter
  typedef fpa::Image::DijkstraWithEndPointDetection< TImage, TScalarImage > TFilter;
  typedef SkeletonCostFunction< TFilter > TFunction;

  TFunction::Pointer function = TFunction::New( );

  TFilter::Pointer filter = TFilter::New( );
  filter->SetInput( input_image );
  filter->SetCostFunction( function );
  filter->SetNeighborhoodOrder( 2 );
  filter->SetSafetyNeighborhoodSize( 0 );
  filter->StopAtOneFrontOff( );
  filter->CorrectSeedsOff( );
  filter->CorrectEndPointsOff( );
  filter->AddSeed( seed, TScalar( 0 ) );

  // Prepare graphical debugger
  typedef fpa::VTK::Image3DObserver< TFilter, vtkRenderWindow > TDebugger;
  TDebugger::Pointer debugger = TDebugger::New( );
  debugger->SetRenderWindow( view.GetWindow( ) );
  debugger->SetRenderPercentage( 0.0000001 );
  filter->AddObserver( itk::AnyEvent( ), debugger );
  filter->ThrowEventsOn( );

  // Go!
  std::time_t start, end;
  std::time( &start );
  filter->Update( );
  std::time( &end );
  std::cout
    << "Extraction time = "
    << std::difftime( end, start )
    << " s." << std::endl;

  // Outputs
  const TFilter::TOutputImage* accumulated_costs = filter->GetOutput( );
  const TFilter::TLabelImage* labeled_image = filter->GetLabelImage( );
  const TFilter::TMinimumSpanningTree* mst = filter->GetMinimumSpanningTree( );
  const TFilter::TUniqueVertices* endpoints = filter->GetEndPoints( );
  const TFilter::TUniqueVertices* bifurcations = filter->GetBifurcations( );
  const TFilter::TBranches* branches = filter->GetBranches( );
  unsigned long nBranches = filter->GetNumberOfBranches( );

  // Save outputs
  SaveImage( accumulated_costs, output_costmap_fn );
  SaveImage( labeled_image, output_labels_fn );

  fpa::IO::MinimumSpanningTreeWriter< TFilter::TMinimumSpanningTree >::Pointer
    mst_writer =
    fpa::IO::MinimumSpanningTreeWriter< TFilter::TMinimumSpanningTree >::New( );
  mst_writer->SetInput( mst );
  mst_writer->SetFileName( mst_output_fn );
  mst_writer->Update( );

  fpa::IO::UniqueValuesContainerWriter< TFilter::TUniqueVertices >::Pointer
    endpoints_writer =
    fpa::IO::UniqueValuesContainerWriter< TFilter::TUniqueVertices >::New( );
  endpoints_writer->SetInput( endpoints );
  endpoints_writer->SetFileName( endpoints_output_fn );
  endpoints_writer->Update( );

  fpa::IO::UniqueValuesContainerWriter< TFilter::TUniqueVertices >::Pointer
    bifurcations_writer =
    fpa::IO::UniqueValuesContainerWriter< TFilter::TUniqueVertices >::New( );
  bifurcations_writer->SetInput( bifurcations );
  bifurcations_writer->SetFileName( bifurcations_output_fn );
  bifurcations_writer->Update( );

  fpa::IO::MatrixValuesContainerWriter< TFilter::TBranches >::Pointer
    branches_writer =
    fpa::IO::MatrixValuesContainerWriter< TFilter::TBranches >::New( );
  branches_writer->SetInput( branches );
  branches_writer->SetFileName( branches_output_fn );
  branches_writer->Update( );

  // Show endpoints
  vtkSmartPointer< vtkPoints > endpoints_points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > endpoints_cells =
    vtkSmartPointer< vtkCellArray >::New( );
  for(
    TFilter::TUniqueVertices::ConstIterator epIt = endpoints->Begin( );
    epIt != endpoints->End( );
    ++epIt
    )
  {
    TImage::PointType pnt;
    input_image->TransformIndexToPhysicalPoint( *epIt, pnt );
    endpoints_points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );
    endpoints_cells->InsertNextCell( 1 );
    endpoints_cells->
      InsertCellPoint( endpoints_points->GetNumberOfPoints( ) - 1 );

  } // rof
  vtkSmartPointer< vtkPolyData > endpoints_polydata =
    vtkSmartPointer< vtkPolyData >::New( );
  endpoints_polydata->SetPoints( endpoints_points );
  endpoints_polydata->SetVerts( endpoints_cells );
  view.AddPolyData( endpoints_polydata, 0, 0, 1, 1 );

  // Show bifurcations
  vtkSmartPointer< vtkPoints > bifurcations_points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > bifurcations_cells =
    vtkSmartPointer< vtkCellArray >::New( );
  for(
    TFilter::TUniqueVertices::ConstIterator bfIt = bifurcations->Begin( );
    bfIt != bifurcations->End( );
    ++bfIt
    )
  {
    TImage::PointType pnt;
    input_image->TransformIndexToPhysicalPoint( *bfIt, pnt );
    bifurcations_points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );
    bifurcations_cells->InsertNextCell( 1 );
    bifurcations_cells->
      InsertCellPoint( bifurcations_points->GetNumberOfPoints( ) - 1 );

  } // rof
  vtkSmartPointer< vtkPolyData > bifurcations_polydata =
    vtkSmartPointer< vtkPolyData >::New( );
  bifurcations_polydata->SetPoints( bifurcations_points );
  bifurcations_polydata->SetVerts( bifurcations_cells );
  view.AddPolyData( bifurcations_polydata, 0, 1, 0, 1 );

  // Show branches (simple and detailed)
  vtkSmartPointer< vtkPoints > simple_branches_points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > simple_branches_cells =
    vtkSmartPointer< vtkCellArray >::New( );

  vtkSmartPointer< vtkPoints > detailed_branches_points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > detailed_branches_cells =
    vtkSmartPointer< vtkCellArray >::New( );
  vtkSmartPointer< vtkFloatArray > detailed_branches_scalars =
    vtkSmartPointer< vtkFloatArray >::New( );

  TFilter::TBranches::ConstIterator brIt = branches->Begin( );
  for( ; brIt != branches->End( ); ++brIt )
  {
    // Branch's first point
    TImage::PointType first_point;
    input_image->TransformIndexToPhysicalPoint( brIt->first, first_point );
    unsigned long first_id = simple_branches_points->GetNumberOfPoints( );
    simple_branches_points->InsertNextPoint(
      first_point[ 0 ], first_point[ 1 ], first_point[ 2 ]
      );

    TFilter::TBranches::ConstRowIterator brRowIt = branches->Begin( brIt );
    for( ; brRowIt != branches->End( brIt ); ++brRowIt )
    {
      // Branch's second point
      TImage::PointType second_point;
      input_image->
        TransformIndexToPhysicalPoint( brRowIt->first, second_point );
      unsigned long second_id = simple_branches_points->GetNumberOfPoints( );
      simple_branches_points->InsertNextPoint(
        second_point[ 0 ], second_point[ 1 ], second_point[ 2 ]
        );
      simple_branches_cells->InsertNextCell( 2 );
      simple_branches_cells->InsertCellPoint( first_id );
      simple_branches_cells->InsertCellPoint( second_id );

      // Detailed path
      double pathId = double( brRowIt->second - 1 ) / double( nBranches - 1 );
      TFilter::TVertices path;
      mst->GetPath( path, brIt->first, brRowIt->first );
      TFilter::TVertices::const_iterator pIt = path.begin( );
      for( ; pIt != path.end( ); ++pIt )
      {
        TImage::PointType path_point;
        input_image->TransformIndexToPhysicalPoint( *pIt, path_point );
        detailed_branches_points->InsertNextPoint(
          path_point[ 0 ], path_point[ 1 ], path_point[ 2 ]
          );
        detailed_branches_scalars->InsertNextTuple1( pathId );
        if( pIt != path.begin( ) )
        {
          unsigned long nPoints =
            detailed_branches_points->GetNumberOfPoints( );
          detailed_branches_cells->InsertNextCell( 2 );
          detailed_branches_cells->InsertCellPoint( nPoints - 2 );
          detailed_branches_cells->InsertCellPoint( nPoints - 1 );

        } // fi

      } // rof

    } // rof

  } // rof
  vtkSmartPointer< vtkPolyData > simple_branches_polydata =
    vtkSmartPointer< vtkPolyData >::New( );
  simple_branches_polydata->SetPoints( simple_branches_points );
  simple_branches_polydata->SetLines( simple_branches_cells );
  view.AddPolyData( simple_branches_polydata, 1, 0, 1, 1 );

  vtkSmartPointer< vtkPolyData > detailed_branches_polydata =
    vtkSmartPointer< vtkPolyData >::New( );
  detailed_branches_polydata->SetPoints( detailed_branches_points );
  detailed_branches_polydata->SetLines( detailed_branches_cells );
  detailed_branches_polydata->
    GetPointData( )->SetScalars( detailed_branches_scalars );
  view.AddPolyData( detailed_branches_polydata, 1 );

  // Let some more interaction
  view.Render( );
  view.Start( );

  return( 0 );
}

// -------------------------------------------------------------------------
template< class I >
void ReadImage( typename I::Pointer& image, const std::string& filename )
{
  typename itk::ImageFileReader< I >::Pointer reader =
    itk::ImageFileReader< I >::New( );
  reader->SetFileName( filename );
  reader->Update( );
  image = reader->GetOutput( );
  image->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
template< class I >
void SaveImage( const I* image, const std::string& filename )
{
  typename itk::ImageFileWriter< I >::Pointer writer =
    itk::ImageFileWriter< I >::New( );
  writer->SetInput( image );
  writer->SetFileName( filename );
  try
  {
    writer->Update( );
  }
  catch( itk::ExceptionObject& err )
  {
    std::cerr
      << "Error saving \"" << filename << "\": " << err
      << std::endl;

  } // yrt
}

// eof - $RCSfile$