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

#include <fpa/Image/RandomWalker.h>
#include <fpa/Image/Functors/Dijkstra/Gaussian.h>
#include <ivq/ITK/ExtractLabelFunction.h>
#include <ivq/ITK/ImageUnaryFunctionFilter.h>

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

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

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

typedef ivq::ITK::ExtractLabelFunction< TLabel, TBinary > TLabelFunction;
typedef ivq::ITK::ImageUnaryFunctionFilter< TLabelImage, TBinaryImage > TLabelExtractor;

// -------------------------------------------------------------------------
int main( int argc, char* argv[] )
{
  // Get arguments
  if( argc < 7 )
  {
    std::cerr
      << "Usage: " << argv[ 0 ] << std::endl
      << "   input_image label_image output_image" << std::endl
      << "   label 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 ];
  TLabel label = TLabel( std::atoi( argv[ 4 ] ) );
  double alpha = std::atof( argv[ 5 ] );
  double beta = std::atof( argv[ 6 ] );

  // 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 );

  // Label extraction
  TLabelFunction::Pointer label_function = TLabelFunction::New( );
  label_function->SetLabel( label );
  label_function->SetInsideValue( 255 );
  label_function->SetOutsideValue( 0 );

  TLabelExtractor::Pointer label_extractor = TLabelExtractor::New( );
  label_extractor->SetInput( filter->GetMarks( ) );
  label_extractor->SetFunction( label_function );


  // 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;

    ts = std::chrono::high_resolution_clock::now( );
    label_extractor->Update( );
    te = std::chrono::high_resolution_clock::now( );
    tr = te - ts;
    std::cout
      << "Extract label 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< TBinaryImage >::Pointer output_image_writer =
    itk::ImageFileWriter< TBinaryImage >::New( );
  output_image_writer->SetInput( label_extractor->GetOutput( ) );
  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$