#include <cmath>
#include <ctime>
#include <deque>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <limits>
#include <map>
#include <string>

#include <itkConstNeighborhoodIterator.h>
#include <itkNeighborhoodIterator.h>
#include <itkDanielssonDistanceMapImageFilter.h>
#include <itkImage.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>
#include <itkImageToVTKImageFilter.h>

#include <itkMinimumMaximumImageCalculator.h>
#include <itkInvertIntensityImageFilter.h>

#include <vtkPoints.h>
#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>
#include <fpa/VTK/Image3DObserver.h>

// -------------------------------------------------------------------------
const unsigned int Dim = 3;
typedef float                                TPixel;
typedef float                                TScalar;
typedef itk::Image< TPixel, Dim >            TImage;
typedef itk::ImageToVTKImageFilter< TImage > TVTKImage;

typedef itk::ImageFileReader< TImage >          TImageReader;
typedef itk::ImageFileWriter< TImage >    TImageWriter;
typedef fpa::Image::DijkstraWithSphereBacktracking< TImage, TScalar > TDijkstra;

typedef fpa::VTK::ImageMPR TMPR;
typedef fpa::VTK::Image3DObserver< TDijkstra, vtkRenderWindow > TDijkstraObs;

struct TBranch
{
  double Length;
  TImage::PointType::VectorType V1;
  TImage::PointType::VectorType V2;

  bool operator<( const TBranch& other ) const
    {
      return( other.Length < this->Length );
    }
};
typedef std::multiset< TBranch > TBranches;

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  if( argc < 8 )
  {
    std::cerr
      << "Usage: " << argv[ 0 ]
      << " input_image x y z output_image output_imagej output_vtk"
      << " visual_debug"
      << std::endl;
    return( 1 );

  } // fi
  std::string input_image_fn = argv[ 1 ];
  TImage::PointType seed_pnt;
  seed_pnt[ 0 ] = std::atof( argv[ 2 ] );
  seed_pnt[ 1 ] = std::atof( argv[ 3 ] );
  seed_pnt[ 2 ] = std::atof( argv[ 4 ] );
  std::string output_image_fn = argv[ 5 ];
  std::string output_imagej_fn = argv[ 6 ];
  std::string output_vtk_fn = argv[ 7 ];
  bool visual_debug = false;
  if( argc > 8 )
    visual_debug = ( std::atoi( argv[ 8 ] ) == 1 );

  // Read image
  TImageReader::Pointer input_image_reader = TImageReader::New( );
  input_image_reader->SetFileName( input_image_fn );
  try
  {
    input_image_reader->Update( );
  }
  catch( itk::ExceptionObject& err )
  {
    std::cerr << "Error caught: " << err << std::endl;
    return( 1 );

  } // yrt
  TImage::ConstPointer input_image = input_image_reader->GetOutput( );

  // Show image
  TVTKImage::Pointer vtk_image = TVTKImage::New( );
  vtk_image->SetInput( input_image );
  vtk_image->Update( );

  TMPR view;
  view.SetBackground( 0.3, 0.2, 0.8 );
  view.SetSize( 800, 800 );
  view.SetImage( vtk_image->GetOutput( ) );

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

  // Allow some interaction
  view.Render( );
  if( visual_debug )
    view.Start( );

  TImage::IndexType seed_idx;
  input_image->TransformPhysicalPointToIndex( seed_pnt, seed_idx );
  /* TODO
     std::cout << seed_idx << " " << seed_pnt << std::endl;
     seed_idx[ 0 ] = 256;
     seed_idx[ 1 ] = 313;
     seed_idx[ 2 ] = 381;
  */

  // Extract paths
  TDijkstra::Pointer paths = TDijkstra::New( );
  paths->AddSeed( seed_idx, TScalar( 0 ) );
  paths->SetInput( input_image );
  paths->SetNeighborhoodOrder( 2 );

  if( visual_debug )
  {
    // Configure observer
    TDijkstraObs::Pointer obs = TDijkstraObs::New( );
    obs->SetRenderWindow( view.GetWindow( ) );
    paths->AddObserver( itk::AnyEvent( ), obs );
    paths->ThrowEventsOn( );
  }
  else
    paths->ThrowEventsOff( );
  std::clock_t start = std::clock( );
  paths->Update( );
  std::clock_t end = std::clock( );
  double seconds = double( end - start ) / double( CLOCKS_PER_SEC );
  std::cout << "Paths extraction time = " << seconds << std::endl;

  // Create polydata
  vtkSmartPointer< vtkPoints > points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > cells =
    vtkSmartPointer< vtkCellArray >::New( );
  vtkSmartPointer< vtkFloatArray > scalars =
    vtkSmartPointer< vtkFloatArray >::New( );

  TDijkstra::TMarkImage::Pointer new_marks = TDijkstra::TMarkImage::New( );
  new_marks->SetLargestPossibleRegion( input_image->GetLargestPossibleRegion( ) );
  new_marks->SetRequestedRegion( input_image->GetRequestedRegion( ) );
  new_marks->SetBufferedRegion( input_image->GetBufferedRegion( ) );
  new_marks->SetOrigin( input_image->GetOrigin( ) );
  new_marks->SetSpacing( input_image->GetSpacing( ) );
  new_marks->SetDirection( input_image->GetDirection( ) );
  new_marks->Allocate( );
  new_marks->FillBuffer( 0 );

  const TDijkstra::TMarkImage* marks = paths->GetOutputMarkImage( );
  TDijkstra::TMark max_mark = paths->GetNumberOfBranches( );
  const TDijkstra::TVertices& endpoints = paths->GetEndPoints( );
  const TDijkstra::TVertices& bifurcations = paths->GetBifurcationPoints( );
  const TDijkstra::TTree& tree = paths->GetFullTree( );
  const TDijkstra::TTree& reduced_tree = paths->GetReducedTree( );
  TDijkstra::TVertices::const_iterator epIt = endpoints.begin( );

  TBranches branches;
  TImage::PointType ori = input_image->GetOrigin( );

  TDijkstra::TTree::const_iterator rIt = reduced_tree.begin( );
  for( ; rIt != reduced_tree.end( ); ++rIt )
  {
    TDijkstra::TTree::const_iterator tIt = tree.find( rIt->first );

    // Prepare branch "a la ImageJ"
    TBranch branch;
    TImage::PointType p1, p2;
    input_image->TransformIndexToPhysicalPoint( rIt->second.first, p1 );
    input_image->TransformIndexToPhysicalPoint( rIt->first, p2 );
    branch.V1 = p1 - ori;
    branch.V2 = p2 - ori;
    branch.Length = double( 0 );
      
    double pd =
      ( double( tIt->second.second ) - double( 1 ) ) /
      ( double( max_mark ) - double( 1 ) );
    bool start = true;
    TImage::PointType prev;
    do
    {
      TDijkstra::TVertex idx = tIt->first;
      new_marks->SetPixel( idx, 255 );
      TImage::PointType pnt;
      input_image->TransformIndexToPhysicalPoint( idx, pnt );
      points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );
      scalars->InsertNextTuple1( pd );
      if( !start )
      {
        cells->InsertNextCell( 2 );
        cells->InsertCellPoint( points->GetNumberOfPoints( ) - 2 );
        cells->InsertCellPoint( points->GetNumberOfPoints( ) - 1 );
        branch.Length += prev.EuclideanDistanceTo( pnt );

      } // fi
      start = false;
      prev = pnt;
      tIt = tree.find( tIt->second.first );

    } while( tIt->first != rIt->second.first );

    // Last point
    TDijkstra::TVertex idx = tIt->first;
    new_marks->SetPixel( idx, 255 );
    TImage::PointType pnt;
    input_image->TransformIndexToPhysicalPoint( idx, pnt );
    points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );
    scalars->InsertNextTuple1( pd );
    if( !start )
    {
      cells->InsertNextCell( 2 );
      cells->InsertCellPoint( points->GetNumberOfPoints( ) - 2 );
      cells->InsertCellPoint( points->GetNumberOfPoints( ) - 1 );
      branch.Length += prev.EuclideanDistanceTo( pnt );

    } // fi
    branches.insert( branch );

  } // rof

  std::ofstream bout( output_imagej_fn.c_str( ) );
  if( !bout )
  {
    std::cerr << "Error: could not open " << output_imagej_fn << std::endl;
    return( 1 );

  } // fi
  int i = 1;
  for( TBranches::const_iterator bIt = branches.begin( ); bIt != branches.end( ); ++bIt, ++i )
  {
    bout
      << std::setiosflags( std::ios::fixed) << std::setprecision( 3 )
      << i << "\t1.000\t"
      << bIt->Length << "\t"
      << bIt->V1[ 0 ] << "\t"
      << bIt->V1[ 1 ] << "\t"
      << bIt->V1[ 2 ] << "\t"
      << bIt->V2[ 0 ] << "\t"
      << bIt->V2[ 1 ] << "\t"
      << bIt->V2[ 2 ] << "\t"
      << ( bIt->V2 - bIt->V1 ).GetNorm( )
      << std::endl;

  } // rof
  bout.close( );

  vtkSmartPointer< vtkPolyData > vtk_tree =
    vtkSmartPointer< vtkPolyData >::New( );
  vtk_tree->SetPoints( points );
  vtk_tree->SetLines( cells );
  vtk_tree->GetPointData( )->SetScalars( scalars );

  view.AddPolyData( vtk_tree );
  view.Render( );
  if( visual_debug )
    view.Start( );

  vtkSmartPointer< vtkPolyDataWriter > writer =
    vtkSmartPointer< vtkPolyDataWriter >::New( );
  writer->SetInputData( vtk_tree );
  writer->SetFileName( output_vtk_fn.c_str( ) );
  writer->Update( );

  itk::ImageFileWriter< TDijkstra::TMarkImage >::Pointer marks_writer =
    itk::ImageFileWriter< TDijkstra::TMarkImage >::New( );
  marks_writer->SetInput( new_marks );
  marks_writer->SetFileName( output_image_fn );
  marks_writer->Update( );

  TImageWriter::Pointer dijkstra_writer =
    TImageWriter::New( );
  dijkstra_writer->SetInput( paths->GetOutput( ) );
  dijkstra_writer->SetFileName( "dijkstra.mhd" );
  dijkstra_writer->Update( );

  return( 0 );
}

// eof - $RCSfile$