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

#include <chrono>
#include <map>
#include <itkImage.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>
#include <itkMinimumMaximumImageCalculator.h>
#include <itkInvertIntensityImageFilter.h>
#include <itkHessianRecursiveGaussianImageFilter.h>
#include <itkHessian3DToVesselnessMeasureImageFilter.h>

// -------------------------------------------------------------------------
const unsigned int Dim = 3;
typedef short TPixel;
typedef itk::NumericTraits< TPixel >::RealType TScalar;
typedef itk::Image< TPixel, Dim >  TImage;
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 typename _TImagePtr::ObjectType _TImage;
  typedef itk::ImageFileReader< _TImage > _TReader;
  typename _TReader::Pointer reader = _TReader::New( );
  reader->SetFileName( fname );
  double t = MeasureTime( reader );
  std::cout << "Read " << fname << " in " << t << " s" << std::endl;
  image = reader->GetOutput( );
  image->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
template< class _TImage >
void WriteImage( const _TImage* image, const std::string& fname )
{
  typedef itk::ImageFileWriter< _TImage > _TWriter;
  typename _TWriter::Pointer writer = _TWriter::New( );
  writer->SetFileName( fname );
  writer->SetInput( image );
  double t = MeasureTime( writer );
  std::cout << "Wrote " << fname << " in " << t << " s" << std::endl;
}

// -------------------------------------------------------------------------
bool ParseArgs(
  std::map< std::string, std::string >& args, int argc, char* argv[]
  )
{
  std::set< std::string > mandatory;
  mandatory.insert( "in" );
  mandatory.insert( "out" );

  args[ "sigma" ] = "0.5";
  args[ "alpha1" ] = "0.5";
  args[ "alpha2" ] = "2";

  int i = 1;
  while( i < argc )
  {
    std::string cmd = argv[ i ] + 1;
    args[ cmd ] = argv[ i + 1 ];
    i += 2;

  } // elihw

  bool complete = true;
  for( std::string t: mandatory )
    complete &= ( args.find( t ) != args.end( ) );

  if( !complete )
  {
    std::cerr
      << "Usage: " << argv[ 0 ] << std::endl
      << "\t-in filename" << std::endl
      << "\t-out filename" << std::endl
      << "\t[-sigma val]" << std::endl
      << "\t[-alpha1 val]" << std::endl
      << "\t[-alpha2 val]" << std::endl;
    return( false );

  } // fi
  return( true );
}

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  std::map< std::string, std::string > args;
  try
  {
    if( ParseArgs( args, argc, argv ) )
    {
      // Read input image
      TImage::Pointer input_image;
      ReadImage( input_image, args[ "in" ] );

      // Vesselness
      typedef itk::MinimumMaximumImageCalculator< TImage >            TMinMax;
      typedef itk::InvertIntensityImageFilter< TImage >               TInverter;
      typedef itk::HessianRecursiveGaussianImageFilter< TImage >      THessian;
      typedef itk::Hessian3DToVesselnessMeasureImageFilter< TScalar > TVesselness;

      TMinMax::Pointer minMax = TMinMax::New( );
      minMax->SetImage( input_image );
      minMax->Compute( );

      TInverter::Pointer inverter = TInverter::New( );
      inverter->SetInput( input_image );
      inverter->SetMaximum( minMax->GetMaximum( ) );
      double t = MeasureTime( inverter );
      std::cout << "Inversion executed in " << t << " s" << std::endl;

      THessian::Pointer hessian = THessian::New( );
      hessian->SetInput( inverter->GetOutput( ) );
      hessian->SetSigma(
        std::atof( args[ "sigma" ].c_str( ) )
        );
      t = MeasureTime( hessian );
      std::cout << "Hessian executed in " << t << " s" << std::endl;

      TVesselness::Pointer vesselness = TVesselness::New( );
      vesselness->SetInput( hessian->GetOutput( ) );
      vesselness->SetAlpha1(
        std::atof( args[ "alpha1" ].c_str( ) )
        );
      vesselness->SetAlpha2(
        std::atof( args[ "alpha2" ].c_str( ) )
        );
      t = MeasureTime( vesselness );
      std::cout << "Vesselness executed in " << t << " s" << std::endl;

      WriteImage( vesselness->GetOutput( ), args[ "out" ] );
    }
    else
      return( 1 );
  }
  catch( std::exception& err )
  {
    std::cerr
      << "===============================" << std::endl
      << "Error caught: " << std::endl
      << err.what( )
      << "===============================" << std::endl
      << std::endl;
    return( 1 );

  } // yrt
  return( 0 );
}

// eof - $RCSfile$