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

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

#include <itkImageFileReader.h>
#include <itkMinimumMaximumImageCalculator.h>
#include <itkShiftScaleImageFilter.h>

#include <itkSignedMaurerDistanceMapImageFilter.h>

#include <itkImageFileWriter.h>

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

#include <vtkCellArray.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 I, class O >
void DistanceMap(
  const typename I::Pointer& input, typename O::Pointer& output
  );

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  if( argc < 5 )
  {
    std::cerr
      << "Usage: " << argv[ 0 ]
      << " input_image distancemap output_costmap output_labels"
      << std::endl;
    return( 1 );

  } // fi
  std::string input_image_fn = argv[ 1 ];
  std::string distancemap_fn = argv[ 2 ];
  std::string output_costmap_fn = argv[ 3 ];
  std::string output_labels_fn = 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( 800, 800 );
  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( );
  while( view.GetNumberOfSeeds( ) == 0 )
    view.Start( );

  // Get seed
  double p[ 3 ];
  view.GetSeed( 0, p );
  TImage::PointType pnt;
  TImage::IndexType seed;
  pnt[ 0 ] = TImage::PointType::ValueType( p[ 0 ] );
  pnt[ 1 ] = TImage::PointType::ValueType( p[ 1 ] );
  pnt[ 2 ] = TImage::PointType::ValueType( p[ 2 ] );
  input_image->TransformPhysicalPointToIndex( pnt, seed );

  // Compute squared distance map
  TScalarImage::Pointer dmap;
  DistanceMap< TImage, TScalarImage >( input_image, dmap );

  // Prepare cost conversion function
  typedef fpa::Base::Functors::InvertCostFunction< TScalar > TFunction;
  TFunction::Pointer function = TFunction::New( );

  // Prepare Dijkstra filter
  typedef fpa::Image::DijkstraWithEndPointDetection< TScalarImage, TScalarImage > TFilter;
  TFilter::Pointer filter = TFilter::New( );
  filter->SetInput( dmap );
  filter->SetConversionFunction( function );
  filter->SetNeighborhoodOrder( 2 );
  filter->StopAtOneFrontOff( );
  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.0001 );
  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;

  // Minimum spanning tree
  const TFilter::TMinimumSpanningTree* mst =
    filter->GetMinimumSpanningTree( );

  // Build branches from endpoints to seed
  vtkSmartPointer< vtkPoints > endpoints_points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > endpoints_cells =
    vtkSmartPointer< vtkCellArray >::New( );
  vtkSmartPointer< vtkCellArray > endpoints_vertices =
    vtkSmartPointer< vtkCellArray >::New( );

  const TFilter::TUniqueVertices& endpoints = filter->GetEndPoints( );
  TFilter::TUniqueVertices::const_iterator eIt = endpoints.begin( );
  for( ; eIt != endpoints.end( ); ++eIt )
  {
    TFilter::TVertices path;
    mst->GetPath( path, *eIt, filter->GetSeed( 0 ) );

    std::cout << *eIt << std::endl;
    TFilter::TVertices::const_iterator pIt = path.begin( );
    for( ; pIt != path.end( ); ++pIt )
    {
      TImage::PointType pnt;
      input_image->TransformIndexToPhysicalPoint( *pIt, pnt );
      endpoints_points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );

      if( pIt != path.begin( ) )
      {
        endpoints_cells->InsertNextCell( 2 );
        endpoints_cells->InsertCellPoint( endpoints_points->GetNumberOfPoints( ) - 2 );
        endpoints_cells->InsertCellPoint( endpoints_points->GetNumberOfPoints( ) - 1 );
      }
      else
      {
        std::cout << "ok" << std::endl;
        endpoints_vertices->InsertNextCell( 1 );
        endpoints_vertices->InsertCellPoint( endpoints_points->GetNumberOfPoints( ) - 1 );

      } // fi




    } // rof

  } // rof
  std::cout << endpoints.size( ) << std::endl;
      
  vtkSmartPointer< vtkPolyData > endpoints_polydata =
    vtkSmartPointer< vtkPolyData >::New( );
  endpoints_polydata->SetPoints( endpoints_points );
  endpoints_polydata->SetLines( endpoints_cells );
  endpoints_polydata->SetVerts( endpoints_vertices );











  // Bifurcations
  vtkSmartPointer< vtkPoints > bifurcations_points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > bifurcations_vertices =
    vtkSmartPointer< vtkCellArray >::New( );

  const TFilter::TUniqueVertices& bifurcations = filter->GetBifurcationPoints( );
  TFilter::TUniqueVertices::const_iterator bIt = bifurcations.begin( );
  for( ; bIt != bifurcations.end( ); ++bIt )
  {
    std::cout << *bIt << std::endl;

    TImage::PointType pnt;
    input_image->TransformIndexToPhysicalPoint( *bIt, pnt );
    bifurcations_points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );
    bifurcations_vertices->InsertNextCell( 1 );
    bifurcations_vertices->InsertCellPoint( bifurcations_points->GetNumberOfPoints( ) - 1 );

  } // rof

  vtkSmartPointer< vtkPolyData > bifurcations_polydata =
    vtkSmartPointer< vtkPolyData >::New( );
  bifurcations_polydata->SetPoints( bifurcations_points );
  bifurcations_polydata->SetVerts( bifurcations_vertices );





  // Build branches from endpoints to seed
  vtkSmartPointer< vtkPoints > branches_points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > branches_cells =
    vtkSmartPointer< vtkCellArray >::New( );

  const TFilter::TBranches& branches = filter->GetBranches( );
  TFilter::TBranches::const_iterator brIt = branches.begin( );
  for( ; brIt != branches.end( ); ++brIt )
  {
    TFilter::TBranch::const_iterator brsIt = brIt->second.begin( );
    for( ; brsIt != brIt->second.end( ); ++brsIt )
    {
      TImage::PointType pnt0, pnt1;
      input_image->TransformIndexToPhysicalPoint( brIt->first, pnt0 );
      input_image->TransformIndexToPhysicalPoint( brsIt->first, pnt1 );
      branches_points->InsertNextPoint( pnt0[ 0 ], pnt0[ 1 ], pnt0[ 2 ] );
      branches_points->InsertNextPoint( pnt1[ 0 ], pnt1[ 1 ], pnt1[ 2 ] );

      branches_cells->InsertNextCell( 2 );
      branches_cells->InsertCellPoint( branches_points->GetNumberOfPoints( ) - 2 );
      branches_cells->InsertCellPoint( branches_points->GetNumberOfPoints( ) - 1 );

    } // rof

  } // rof

  vtkSmartPointer< vtkPolyData > branches_polydata =
    vtkSmartPointer< vtkPolyData >::New( );
  branches_polydata->SetPoints( branches_points );
  branches_polydata->SetLines( branches_cells );

  view.AddPolyData( endpoints_polydata, 1, 0, 0, 1 );
  view.AddPolyData( bifurcations_polydata, 0, 0, 1, 1 );
  view.AddPolyData( branches_polydata, 0, 1, 0, 1 );
  view.Render( );
  view.Start( );

  // Save final total cost map
  SaveImage( filter->GetOutput( ), output_costmap_fn );
  SaveImage( dmap.GetPointer( ), distancemap_fn );
  SaveImage( filter->GetLabelImage( ), output_labels_fn );
  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( );

  typename itk::MinimumMaximumImageCalculator< I >::Pointer minmax =
    itk::MinimumMaximumImageCalculator< I >::New( );
  minmax->SetImage( reader->GetOutput( ) );
  minmax->Compute( );
  double vmin = double( minmax->GetMinimum( ) );
  double vmax = double( minmax->GetMaximum( ) );

  typename itk::ShiftScaleImageFilter< I, I >::Pointer shift =
    itk::ShiftScaleImageFilter< I, I >::New( );
  shift->SetInput( reader->GetOutput( ) );
  shift->SetScale( vmax - vmin );
  shift->SetShift( vmin / ( vmax - vmin ) );
  shift->Update( );

  image = shift->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
}

// -------------------------------------------------------------------------
template< class I, class O >
void DistanceMap(
  const typename I::Pointer& input, typename O::Pointer& output
  )
{
  typename itk::SignedMaurerDistanceMapImageFilter< I, O >::Pointer dmap =
    itk::SignedMaurerDistanceMapImageFilter< I, O >::New( );
  dmap->SetInput( input );
  dmap->SetBackgroundValue( ( typename I::PixelType )( 0 ) );
  dmap->InsideIsPositiveOn( );
  dmap->SquaredDistanceOn( );
  dmap->UseImageSpacingOn( );

  std::time_t start, end;
  std::time( &start );
  dmap->Update( );
  std::time( &end );
  std::cout
    << "Distance map time = "
    << std::difftime( end, start )
    << " s." << std::endl;

  output = dmap->GetOutput( );
  output->DisconnectPipeline( );
}

// eof - $RCSfile$