#include <ctime>
#include <iostream>
#include <limits>
#include <map>
#include <string>
#include <deque>

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

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

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

// -------------------------------------------------------------------------
const unsigned int Dim = 3;
typedef double                               TPixel;
typedef double                               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::Dijkstra< TImage, TScalar > TBaseDijkstra;

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

// -------------------------------------------------------------------------
class TDijkstra
  : public TBaseDijkstra
{
public:
  typedef TDijkstra                       Self;
  typedef TBaseDijkstra                   Superclass;
  typedef itk::SmartPointer< Self >       Pointer;
  typedef itk::SmartPointer< const Self > ConstPointer;

  typedef Superclass::TCost          TCost;
  typedef Superclass::TVertex        TVertex;
  typedef Superclass::InputImageType TImage;
  typedef std::deque< TVertex >      TVertices;

protected:
  typedef std::pair< TCost, TVertex >     _TCandidate;
  typedef std::multimap< TCost, TVertex > _TCandidates;
  typedef Superclass::_TNode _TNode;

public:
  itkNewMacro( Self );
  itkTypeMacro( TDijkstra, TBaseDijkstra );

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

  virtual void _BeforeMainLoop( )
    {
      const TImage* img = this->GetInput( );

      // Correct seeds
      TImage::SizeType radius;
      radius.Fill( 3 );
      itk::ConstNeighborhoodIterator< TImage > it(
        radius, img, img->GetRequestedRegion( )
        );
      for( unsigned int s = 0; s < this->m_Seeds.size( ); ++s )
      {
        _TNode seed = this->m_Seeds[ s ];
        it.SetLocation( seed.Vertex );

        TImage::SizeType size = it.GetSize( );
        unsigned int nN = 1;
        for( unsigned int d = 0; d < TImage::ImageDimension; ++d )
          nN *= size[ d ];
        TImage::PixelType max_value = img->GetPixel( seed.Vertex );
        for( unsigned int i = 0; i < nN; ++i )
        {
          if( it.GetPixel( i ) > max_value )
          {
            seed.Vertex = it.GetIndex( i );
            seed.Parent = seed.Vertex;
            max_value = it.GetPixel( i );

          } // fi

        } // rof
        this->m_Seeds[ s ] = seed;

      } // rof

      // End initialization
      this->Superclass::_BeforeMainLoop( );
      this->m_Candidates.clear( );
    }

  virtual void _AfterMainLoop( )
    {
      this->Superclass::_AfterMainLoop( );
      if( this->m_Candidates.size( ) == 0 )
        return;
      
      this->m_Axes = vtkSmartPointer< vtkPolyData >::New( );
      vtkSmartPointer< vtkPoints > points =
        vtkSmartPointer< vtkPoints >::New( );
      vtkSmartPointer< vtkCellArray > lines =
        vtkSmartPointer< vtkCellArray >::New( );

      _TCandidates::const_reverse_iterator cIt =
        this->m_Candidates.rbegin( );

      TVertices pS;
      TVertex vIt = cIt->second;
      TVertex parent = this->_Parent( vIt );
      while( parent != vIt )
      {
        pS.push_front( vIt );
        vIt = parent;
        parent = this->_Parent( vIt );

        TImage::PointType pnt;
        this->GetInput( )->TransformIndexToPhysicalPoint( vIt, pnt );
        points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );
        if( points->GetNumberOfPoints( ) > 1 )
        {
          lines->InsertNextCell( 2 );
          lines->InsertCellPoint( points->GetNumberOfPoints( ) - 2 );
          lines->InsertCellPoint( points->GetNumberOfPoints( ) - 1 );

        } // fi

      } // elihw
      pS.push_front( vIt );
      TImage::PointType pnt;
      this->GetInput( )->TransformIndexToPhysicalPoint( vIt, pnt );
      points->InsertNextPoint( pnt[ 0 ], pnt[ 1 ], pnt[ 2 ] );
        if( points->GetNumberOfPoints( ) > 1 )
        {
          lines->InsertNextCell( 2 );
          lines->InsertCellPoint( points->GetNumberOfPoints( ) - 2 );
          lines->InsertCellPoint( points->GetNumberOfPoints( ) - 1 );

        } // fi

        this->m_Axes->SetPoints( points );
        this->m_Axes->SetLines( lines );
    }

  virtual bool _UpdateNeigh( _TNode& nn, const _TNode& n )
    {
      TCost nc = this->_Cost( nn.Vertex, n.Vertex );
      if( TCost( 0 ) < nc )
      {
        nn.Cost = n.Cost + ( TCost( 1 ) / nc );
        nn.Result = nn.Cost;
        return( true );
      }
      else
        return( false );
    }

  virtual bool _UpdateResult( _TNode& n )
    {
      bool ret = this->Superclass::_UpdateResult( n );
      if( ret )
      {
        TCost nc = this->_Cost( n.Vertex, n.Parent );
        if( TCost( 0 ) < nc )
        {
          n.Result = n.Cost / ( nc * nc );
          this->GetOutput( )->SetPixel( n.Vertex, n.Result );
          this->m_Candidates.insert( _TCandidate( n.Result, n.Vertex ) );

        } // fi

      } // fi
      return( ret );
    }

private:
  TDijkstra( const Self& other );
  Self& operator=( const Self& other );

protected:
  _TCandidates m_Candidates;
public:
  vtkSmartPointer< vtkPolyData > m_Axes;
};

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

  } // fi
  std::string input_image_fn = argv[ 1 ];
  TImage::IndexType seed_idx;
  seed_idx[ 0 ] = std::atoi( argv[ 2 ] );
  seed_idx[ 1 ] = std::atoi( argv[ 3 ] );
  seed_idx[ 2 ] = std::atoi( argv[ 4 ] );
  std::string output_image_fn = argv[ 5 ];
  bool visual_debug = false;
  if( argc > 6 )
    visual_debug = ( std::atoi( argv[ 6 ] ) == 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( ) );

  // Allow some interaction
  view.Start( );

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

  if( visual_debug )
  {
    // Configure observer
    TDijkstraObs::Pointer obs = TDijkstraObs::New( );
    obs->SetImage( input_image, 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;

  view.AddPolyData( paths->m_Axes, 1, 0, 0 );
  view.Start( );

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

  // Show result
  /*
    TVTKImage::Pointer output_image_vtk = TVTKImage::New( );
    output_image_vtk->SetInput( segmentor->GetOutput( ) );
    output_image_vtk->Update( );

    vtkSmartPointer< vtkImageMarchingCubes > mc =
    vtkSmartPointer< vtkImageMarchingCubes >::New( );
    mc->SetInputData( output_image_vtk->GetOutput( ) );
    mc->SetValue(
    0,
    double( segmentor->GetInsideValue( ) ) * double( 0.95 )
    );
    mc->Update( );

    // Let some interaction and close program
    view.AddPolyData( mc->GetOutput( ), 0.1, 0.6, 0.8, 0.5 );
    view.Start( );

    // Write resulting image
    TImageWriter::Pointer output_image_writer = TImageWriter::New( );
    output_image_writer->SetInput( segmentor->GetOutput( ) );
    output_image_writer->SetFileName( output_image_fn );
    try
    {
    output_image_writer->Update( );
    }
    catch( itk::ExceptionObject& err )
    {
    std::cerr << "Error caught: " << err << std::endl;
    return( 1 );

    } // yrt
  */
  return( 0 );
}

// eof - $RCSfile$