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

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

#include <itkImageFileReader.h>
#include <fpa/VTK/ImageMPR.h>

#include <fpa/Image/DijkstraWithEndPointDetection.h>
#include <fpa/IO/MinimumSpanningTreeReader.h>
#include <fpa/IO/UniqueValuesContainerReader.h>
#include <fpa/IO/MatrixValuesContainerReader.h>

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

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

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

typedef fpa::Image::DijkstraWithEndPointDetection< TImage, TImage > TFilter;
typedef TFilter::TMinimumSpanningTree TMinimumSpanningTree;
typedef TFilter::TUniqueVertices TUniqueVertices;
typedef TFilter::TBranches TBranches;
typedef TFilter::TVertices TVertices;

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  if( argc < 6 )
  {
    std::cerr
      << "Usage: " << argv[ 0 ] << std::endl
      << " input_image" << std::endl
      << " input_minimum_spanning_tree" << std::endl
      << " input_endpoints" << std::endl
      << " input_bifurcations" << std::endl
      << " input_branches" << std::endl
      << std::endl;
    return( 1 );

  } // fi
  std::string input_image_fn = argv[ 1 ];
  std::string mst_input_fn = argv[ 2 ];
  std::string endpoints_input_fn = argv[ 3 ];
  std::string bifurcations_input_fn = argv[ 4 ];
  std::string branches_input_fn = argv[ 5 ];

  // Read image
  itk::ImageFileReader< TImage >::Pointer input_image_reader =
    itk::ImageFileReader< TImage >::New( );
  input_image_reader->SetFileName( input_image_fn );
  input_image_reader->Update( );
  TImage::ConstPointer input_image = input_image_reader->GetOutput( );

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

  // Read minimum spanning tree
  fpa::IO::MinimumSpanningTreeReader< TMinimumSpanningTree >::Pointer
    mst_input_reader =
    fpa::IO::MinimumSpanningTreeReader< TMinimumSpanningTree >::New( );
  mst_input_reader->SetFileName( mst_input_fn );
  mst_input_reader->Update( );
  TMinimumSpanningTree::ConstPointer mst_input = mst_input_reader->GetOutput( );

  fpa::IO::UniqueValuesContainerReader< TUniqueVertices >::Pointer
    endpoints_input_reader =
    fpa::IO::UniqueValuesContainerReader< TUniqueVertices >::New( );
  endpoints_input_reader->SetFileName( endpoints_input_fn );
  endpoints_input_reader->Update( );
  TUniqueVertices::ConstPointer endpoints_input = endpoints_input_reader->GetOutput( );

  fpa::IO::UniqueValuesContainerReader< TUniqueVertices >::Pointer
    bifurcations_input_reader =
    fpa::IO::UniqueValuesContainerReader< TUniqueVertices >::New( );
  bifurcations_input_reader->SetFileName( bifurcations_input_fn );
  bifurcations_input_reader->Update( );
  TUniqueVertices::ConstPointer bifurcations_input =
    bifurcations_input_reader->GetOutput( );

  fpa::IO::MatrixValuesContainerReader< TBranches >::Pointer
    branches_input_reader =
    fpa::IO::MatrixValuesContainerReader< TBranches >::New( );
  branches_input_reader->SetFileName( branches_input_fn );
  branches_input_reader->Update( );
  TBranches::ConstPointer branches_input = branches_input_reader->GetOutput( );
  unsigned int nBranches = branches_input_reader->GetNumberOfLabels( );

  // 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 );

  // Show endpoints
  vtkSmartPointer< vtkPoints > endpoints_points =
    vtkSmartPointer< vtkPoints >::New( );
  vtkSmartPointer< vtkCellArray > endpoints_cells =
    vtkSmartPointer< vtkCellArray >::New( );
  for(
    TUniqueVertices::ConstIterator epIt = endpoints_input->Begin( );
    epIt != endpoints_input->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(
    TUniqueVertices::ConstIterator bfIt = bifurcations_input->Begin( );
    bfIt != bifurcations_input->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( );

  TBranches::ConstIterator brIt = branches_input->Begin( );
  for( ; brIt != branches_input->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 ]
      );

    TBranches::ConstRowIterator brRowIt = branches_input->Begin( brIt );
    for( ; brRowIt != branches_input->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 );
      TVertices path;
      mst_input->GetPath( path, brIt->first, brRowIt->first );
      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 );

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

  return( 0 );
}

// eof - $RCSfile$