#include <chrono>
#include <itkImage.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>

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

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

typedef itk::Image< TPixel, Dim > TInputImage;
typedef itk::Image< TLabel, Dim > TLabelImage;

typedef fpa::Image::RandomWalker< TInputImage, TLabelImage, TScalar > TFilter;
typedef fpa::Image::Functors::Dijkstra::Gaussian< TInputImage, TScalar > TWeight;

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  // Get arguments
  if( argc < 6 )
  {
    std::cerr
      << "Usage: " << argv[ 0 ] << std::endl
      << "   input_image label_image output_image" << std::endl
      << "   alpha(0) beta(100)"
      << std::endl;
    return( 1 );

  } // fi
  std::string input_image_filename = argv[ 1 ];
  std::string label_image_filename = argv[ 2 ];
  std::string output_image_filename = argv[ 3 ];
  double alpha = std::atof( argv[ 4 ] );
  double beta = std::atof( argv[ 5 ] );

  // Read images
  itk::ImageFileReader< TInputImage >::Pointer input_image_reader =
    itk::ImageFileReader< TInputImage >::New( );
  input_image_reader->SetFileName( input_image_filename );

  itk::ImageFileReader< TLabelImage >::Pointer label_image_reader =
    itk::ImageFileReader< TLabelImage >::New( );
  label_image_reader->SetFileName( label_image_filename );

  // Prepare weight
  TWeight::Pointer weight = TWeight::New( );
  weight->SetAlpha( alpha );
  weight->SetBeta( beta );

  // Prepare filter
  TFilter::Pointer filter = TFilter::New( );
  filter->SetInput( input_image_reader->GetOutput( ) );
  filter->SetLabels( label_image_reader->GetOutput( ) );
  filter->SetWeightFunction( weight );

  // Show some information
  std::cout << "----------------------------------------------" << std::endl;
  std::cout << "Image: " << input_image_filename << std::endl;

  // Execute pipeline
  std::chrono::time_point< std::chrono::high_resolution_clock > ts, te;
  std::chrono::duration< double > tr;
  try
  {
    ts = std::chrono::high_resolution_clock::now( );
    input_image_reader->Update( );
    te = std::chrono::high_resolution_clock::now( );
    tr = te - ts;
    std::cout << "Raw read time: " << tr.count( ) << " s" << std::endl;

    ts = std::chrono::high_resolution_clock::now( );
    label_image_reader->Update( );
    te = std::chrono::high_resolution_clock::now( );
    tr = te - ts;
    std::cout << "Label read time: " << tr.count( ) << " s" << std::endl;

    ts = std::chrono::high_resolution_clock::now( );
    filter->Update( );
    te = std::chrono::high_resolution_clock::now( );
    tr = te - ts;
    std::cout
      << "Labelling time: " << tr.count( ) << " s" << std::endl;
  }
  catch( std::exception& err )
  {
    std::cerr << "Error caught: " << err.what( ) << std::endl;
    return( 1 );

  } // yrt

  // Save output image
  itk::ImageFileWriter< TLabelImage >::Pointer output_image_writer =
    itk::ImageFileWriter< TLabelImage >::New( );
  output_image_writer->SetInput( filter->GetMarks( ) );
  output_image_writer->SetFileName( output_image_filename );
  try
  {
    ts = std::chrono::high_resolution_clock::now( );
    output_image_writer->Update( );
    te = std::chrono::high_resolution_clock::now( );
    tr = te - ts;
    std::cout << "Write time: " << tr.count( ) << " s" << std::endl;
  }
  catch( std::exception& err )
  {
    std::cerr << "Error caught: " << err.what( ) << std::endl;
    return( 1 );

  } // yrt
  std::cout << "----------------------------------------------" << std::endl;

  return( 0 );
}

// eof - $RCSfile$