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

#include <chrono>
#include <fstream>
#include <limits>
#include <map>
#include <set>
#include <sstream>
#include <string>
#include <itkImage.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>
#include <fpa/Filters/Image/Mori.h>
#include <fpa/Filters/Image/RandomWalker.h>
#include <fpa/Functors/Dijkstra/Image/Gaussian.h>
#include <CTBronchi/MoriLabelling.h>

// -------------------------------------------------------------------------
const unsigned int Dim = 3;
typedef short         TInputPixel;
typedef unsigned char TLabelPixel;
typedef float         TScalar;
typedef itk::Image< TInputPixel, Dim > TInputImage;
typedef itk::Image< TLabelPixel, Dim > TLabelImage;
typedef itk::Image< TScalar, Dim >     TScalarImage;
typedef std::map< std::string, std::string > TMap;

// -------------------------------------------------------------------------
TMap Args;
TInputImage::PointType global_seed;
TLabelPixel inside_value = TLabelPixel( 1 );
TLabelPixel outside_value = TLabelPixel( 0 );
TLabelPixel inside_label = TLabelPixel( 1 );
TLabelPixel outside_label = TLabelPixel( 2 );

// -------------------------------------------------------------------------
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 _TImagePtr >
void WriteImage( const _TImagePtr& image, const std::string& fname )
{
  typedef typename _TImagePtr::ObjectType _TImage;
  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;
}

// -------------------------------------------------------------------------
template< class _TInputPtr, class _TOutputPtr >
void Mori( _TOutputPtr& output, const _TInputPtr& input )
{
  typedef typename _TInputPtr::ObjectType  _TInput;
  typedef typename _TOutputPtr::ObjectType _TOutput;
  typedef fpa::Filters::Image::Mori< _TInput, _TOutput > _TMori;

  typename _TMori::Pointer mori = _TMori::New( );
  mori->SetInput( input );
  mori->SetSeed( global_seed );
  mori->SetInsideValue( inside_value );
  mori->SetOutsideValue( outside_value );
  mori->SetMinimumThreshold(
    TInputPixel( std::atof( Args[ "mori_minimum_threshold" ].c_str( ) ) )
    );
  mori->SetSignalKernelSize(
    std::atoi( Args[ "mori_signal_kernel_size" ].c_str( ) )
    );
  mori->SetSignalThreshold(
    std::atof( Args[ "mori_signal_threshold" ].c_str( ) )
    );
  mori->SetSignalInfluence(
    std::atof( Args[ "mori_signal_influence" ].c_str( ) )
    );
  mori->SetThresholds(
    TInputPixel( std::atof( Args[ "mori_lower_threshold" ].c_str( ) ) ),
    TInputPixel( std::atof( Args[ "mori_upper_threshold" ].c_str( ) ) ),
    TInputPixel( std::atof( Args[ "mori_delta_threshold" ].c_str( ) ) )
    );
  double t = MeasureTime( mori );
  std::cout << "Mori executed in " << t << " s" << std::endl;
  output = mori->GetOutput( );

  TMap::const_iterator i = Args.find( "out_mori" );
  if( i != Args.end( ) )
    WriteImage( output, i->second );

  i = Args.find( "out_signal" );
  if( i != Args.end( ) )
  {
    std::stringstream signal;
    unsigned long nthr = mori->GetNumberOfEvaluatedThresholds( );
    signal << "# nThr = " << nthr << std::endl;
    signal << "# Opt  = " << mori->GetOptimumThreshold( ) << std::endl;
    for( unsigned long j = 0; j < nthr; ++j )
    {
      typename _TMori::TPeak p;
      double x, y;
      mori->GetSignalValues( j, x, y, p );
      signal << x << " " << y << std::endl;

    } // rof

    std::ofstream signals_str( i->second.c_str( ) );
    signals_str << signal.str( );
    signals_str.close( );

  } // fi
  output->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
template< class _TInputPtr, class _TOutputPtr >
void Label( _TOutputPtr& output, const _TInputPtr& input, const _TOutputPtr& labels )
{
  typedef typename _TInputPtr::ObjectType  _TInput;
  typedef typename _TOutputPtr::ObjectType _TOutput;
  typedef CTBronchi::MoriLabelling< _TInput, _TOutput > _TLabelling;

  typename _TLabelling::Pointer labelling = _TLabelling::New( );
  labelling->SetInput( input );
  labelling->SetInputLabels( labels );
  labelling->SetOutsideValue( outside_label );
  labelling->SetInsideValue( inside_label );
  labelling->SetInputInsideValue( inside_value );
  labelling->SetUpperThreshold(
    TInputPixel( std::atof( Args[ "labelling_upper_threshold" ].c_str( ) ) )
    );
  double t = MeasureTime( labelling );
  std::cout << "Labelling executed in " << t << " s" << std::endl;
  output = labelling->GetOutput( );
  TMap::const_iterator i = Args.find( "out_labels" );
  if( i != Args.end( ) )
    WriteImage( output, i->second );
  output->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
template< class _TRawPtr, class _TLabelPtr >
void RandomWalker( _TLabelPtr& output, const _TRawPtr& raw, const _TLabelPtr& labels )
{
  typedef typename _TRawPtr::ObjectType  _TRaw;
  typedef typename _TLabelPtr::ObjectType _TLabel;
}

// -------------------------------------------------------------------------
bool ParseArgs( int argc, char* argv[] )
{
  std::set< std::string > mandatory;
  mandatory.insert( "in" );
  mandatory.insert( "out_segmentation" );
  mandatory.insert( "seed_x" );
  mandatory.insert( "seed_y" );
  mandatory.insert( "seed_z" );

  Args[ "mori_minimum_threshold" ] = "-850";
  Args[ "mori_signal_kernel_size" ] = "20";
  Args[ "mori_signal_threshold" ] = "100";
  Args[ "mori_signal_influence" ] = "0.5";
  Args[ "mori_lower_threshold" ] = "-1024";
  Args[ "mori_upper_threshold" ] = "0";
  Args[ "mori_delta_threshold" ] = "1";
  Args[ "labelling_upper_threshold" ] = "-400";

  for( int i = 1; i < argc; i += 2 )
    Args[ argv[ i ] + 1 ] = argv[ i + 1 ];

  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_segmentation filename" << std::endl
      << "\t-seed_x value -seed_y value -seed_z value" << std::endl
      << "\t[-out_mori filename]" << std::endl
      << "\t[-out_signal filename]" << std::endl
      << "\t[-out_labels filename]" << std::endl
      << "\t[-mori_minimum_threshold value]" << std::endl
      << "\t[-mori_signal_kernel_size value]" << std::endl
      << "\t[-mori_signal_threshold value]" << std::endl
      << "\t[-mori_signal_influence value]" << std::endl
      << "\t[-mori_lower_threshold value]" << std::endl
      << "\t[-mori_upper_threshold value]" << std::endl
      << "\t[-mori_delta_threshold value]" << std::endl
      << "\t[-labelling_upper_threshold value]" << std::endl;
    return( false );

  } // fi
  return( true );
}

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  try
  {
    if( ParseArgs( argc, argv ) )
    {
      // Parse seed
      global_seed[ 0 ] = std::atof( Args[ "seed_x" ].c_str( ) );
      global_seed[ 1 ] = std::atof( Args[ "seed_y" ].c_str( ) );
      global_seed[ 2 ] = std::atof( Args[ "seed_z" ].c_str( ) );

      // Read input image
      TInputImage::Pointer input_image;
      ReadImage( input_image, Args[ "in" ] );

      // Mori segmentation
      TLabelImage::Pointer mori;
      Mori( mori, input_image );

      // Create labels
      TLabelImage::Pointer labels;
      Label( labels, input_image, mori );

      return( 0 );

    } // fi
    return( 1 );
  }
  catch( std::exception& err )
  {
    std::cerr
      << "===============================" << std::endl
      << "Error caught: " << std::endl
      << err.what( )
      << "===============================" << std::endl
      << std::endl;
    return( 1 );

  } // yrt
  /* TODO
     if [ -z "$label_upper_threshold" ]; then label_upper_threshold="-600"; fi
     if [ -z "$label_inside" ]; then label_inside="1"; fi
     if [ -z "$label_outside" ]; then label_outside="2"; fi
     if [ -z "$random_walker_beta" ]; then random_walker_beta="20"; fi
  */
}

// eof - $RCSfile$