// =========================================================================
// @author Leonardo Florez Valencia
// @email florez-l@javeriana.edu.co
// =========================================================================

#include <chrono>
#include <iostream>
#include <string>
#include <sstream>
#include <itkImage.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>

#include <fpa/Common/OriginalRandomWalker.h>
#include <fpa/Filters/Image/RandomWalker.h>
#include <fpa/Functors/Dijkstra/Image/Gaussian.h>

// -------------------------------------------------------------------------
const unsigned int Dim = 3;
typedef short         TPixel;
typedef unsigned char TLabel;
typedef double        TScalar;

typedef itk::Image< TPixel, Dim > TImage;
typedef itk::Image< TLabel, Dim > TLabels;
typedef itk::Image< TScalar, Dim > TScalarImage;

// -------------------------------------------------------------------------
double MeasureTime( itk::ProcessObject* f )
{
  std::chrono::time_point< std::chrono::high_resolution_clock > s, e;
  std::chrono::duration< double > t;
  s = std::chrono::high_resolution_clock::now( );
  f->Update( );
  e = std::chrono::high_resolution_clock::now( );
  t = e - s;
  return( t.count( ) );
}

// -------------------------------------------------------------------------
template< class _TImagePtr >
void ReadImage( _TImagePtr& image, const std::string& fname )
{
  typedef itk::ImageFileReader< typename _TImagePtr::ObjectType > _TReader;
  typename _TReader::Pointer r = _TReader::New( );
  r->SetFileName( fname );
  double t = MeasureTime( r );
  image = r->GetOutput( );
  image->DisconnectPipeline( );
  std::cout << "Image \"" << fname << "\" read in " << t << "s" << std::endl;
}

// -------------------------------------------------------------------------
template< class _TImagePtr >
void WriteImage( const _TImagePtr& image, const std::string& fname )
{
  typedef itk::ImageFileWriter< typename _TImagePtr::ObjectType > _TWriter;
  typename _TWriter::Pointer w = _TWriter::New( );
  w->SetInput( image );
  w->SetFileName( fname );
  double t = MeasureTime( w );
  std::cout << "Image \"" << fname << "\" written in " << t << "s" << std::endl;
}

// -------------------------------------------------------------------------
template< class _TImagePtr, class _TLabelsPtr, class _TScalarImagePtr >
void Original(
  _TLabelsPtr& output_labels,
  _TScalarImagePtr& output_probabilities,
  const _TImagePtr& input,
  const _TLabelsPtr& labels,
  const typename _TScalarImagePtr::ObjectType::PixelType& beta,
  const typename _TScalarImagePtr::ObjectType::PixelType& epsilon
  )
{
  typedef typename _TImagePtr::ObjectType       _TImage;
  typedef typename _TLabelsPtr::ObjectType      _TLabels;
  typedef typename _TScalarImagePtr::ObjectType _TScalarImage;
  typedef typename _TScalarImage::PixelType     _TScalar;
  typedef
    fpa::Functors::Dijkstra::Image::Gaussian< _TImage, _TScalar >
    _TFunction;
  typedef
    fpa::Common::OriginalRandomWalker< _TImage, _TLabels, _TScalar >
    _TFilter;

  // Random walker function
  typename _TFunction::Pointer edge = _TFunction::New( );
  edge->SetBeta( beta );
  edge->SetEpsilon( epsilon );
  edge->TreatAsWeightOff( );

  // Random walker
  typename _TFilter::Pointer filter = _TFilter::New( );
  filter->SetInput( input );
  filter->SetInputLabels( labels );
  filter->SetEdgeFunction( edge );
  filter->DebugOn( );
  double t = MeasureTime( filter );
  output_labels = filter->GetOutput( );
  output_probabilities = filter->GetOutputProbabilities( );
  output_labels->DisconnectPipeline( );
  output_probabilities->DisconnectPipeline( );
  std::cout << "Original RW executed in " << t << "s" << std::endl;
}

// -------------------------------------------------------------------------
template< class _TImagePtr, class _TLabelsPtr, class _TScalarImagePtr >
void FrontPropagation(
  _TLabelsPtr& output_labels,
  _TScalarImagePtr& output_probabilities,
  const _TImagePtr& input,
  const _TLabelsPtr& labels,
  const typename _TScalarImagePtr::ObjectType::PixelType& beta,
  const typename _TScalarImagePtr::ObjectType::PixelType& epsilon
  )
{
  typedef typename _TImagePtr::ObjectType       _TImage;
  typedef typename _TLabelsPtr::ObjectType      _TLabels;
  typedef typename _TScalarImagePtr::ObjectType _TScalarImage;
  typedef typename _TScalarImage::PixelType     _TScalar;
  typedef
    fpa::Functors::Dijkstra::Image::Gaussian< _TImage, _TScalar >
    _TFunction;
  typedef
    fpa::Filters::Image::RandomWalker< _TImage, _TLabels, _TScalar >
    _TFilter;

  // Random walker function
  typename _TFunction::Pointer edge = _TFunction::New( );
  edge->SetBeta( beta );
  edge->SetEpsilon( epsilon );
  edge->TreatAsWeightOn( );

  // Random walker
  typename _TFilter::Pointer filter = _TFilter::New( );
  filter->SetInputImage( input );
  filter->SetInputLabels( labels );
  filter->SetWeightFunction( edge );
  double t = MeasureTime( filter );
  output_labels = filter->GetOutputLabels( );
  output_probabilities = filter->GetOutputCosts( );
  output_labels->DisconnectPipeline( );
  output_probabilities->DisconnectPipeline( );
  std::cout << "Front propagation RW executed in " << t << "s" << std::endl;
}

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  // Get input arguments
  if( argc < 5 )
  {
    std::cerr
      << "Usage: " << argv[ 0 ]
      << " input_image input_labels beta epsilon"
      << std::endl;
    return( 1 );

  } // fi
  std::string input_image_filename = argv[ 1 ];
  std::string input_labels_filename = argv[ 2 ];
  std::istringstream values( std::string( argv[ 3 ] ) + " " + argv[ 4 ] );
  TScalar beta, epsilon;
  values >> beta >> epsilon;

  // Read inputs
  TImage::Pointer input_image;
  TLabels::Pointer input_labels;
  ReadImage( input_image, input_image_filename );
  ReadImage( input_labels, input_labels_filename );

  // Execute algorithms
  TLabels::Pointer original_labels, fp_labels;
  TScalarImage::Pointer original_probabilities, fp_probabilities;
  Original(
    original_labels, original_probabilities,
    input_image, input_labels,
    beta, epsilon
    );
  FrontPropagation(
    fp_labels, fp_probabilities,
    input_image, input_labels,
    beta, epsilon
    );

  WriteImage( original_labels, "orig.mhd" );
  WriteImage( fp_labels, "fp.mhd" );

  return( 0 );
}

// eof - $RCSfile$