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

#include <itkImage.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/IO/MinimumSpanningTreeWriter.h>
#include <fpa/IO/UniqueValuesContainerWriter.h>
#include <fpa/IO/MatrixValuesContainerWriter.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;

// -------------------------------------------------------------------------
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 < 12 )
  {
    std::cerr
      << "Usage: " << argv[ 0 ] << std::endl
      << " input_image" << std::endl
      << " seed_x seed_y seed_z" << std::endl
      << " output_distancemap" << 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 ];
  double seed_x = std::atof( argv[ 2 ] );
  double seed_y = std::atof( argv[ 3 ] );
  double seed_z = std::atof( argv[ 4 ] );
  std::string distancemap_fn = argv[ 5 ];
  std::string output_costmap_fn = argv[ 6 ];
  std::string output_labels_fn = argv[ 7 ];
  std::string mst_output_fn = argv[ 8 ];
  std::string endpoints_output_fn = argv[ 9 ];
  std::string bifurcations_output_fn = argv[ 10 ];
  std::string branches_output_fn = argv[ 11 ];

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

  // Get seed
  double p[ 3 ];
  TImage::PointType pnt;
  TImage::IndexType seed;
  pnt[ 0 ] = seed_x;
  pnt[ 1 ] = seed_y;
  pnt[ 2 ] = seed_z;
  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 ) );

  // 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( dmap.GetPointer( ), distancemap_fn );
  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( );

  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$