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

#include <chrono>
#include <map>

#include <itkImage.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>
#include <itkRegionOfInterestImageFilter.h>
#include <itkMinimumMaximumImageCalculator.h>

#include <fpa/Filters/BaseMarksInterface.h>
#include <fpa/Filters/Image/SeedsFromLabelsInterface.h>
#include <fpa/Filters/Image/DefaultTraits.h>
#include <fpa/Filters/Image/RegionGrow.h>
#include <fpa/Functors/RegionGrow/BinaryThreshold.h>

// -------------------------------------------------------------------------
const unsigned int Dim = 3;
typedef short                                  TPixel;
typedef unsigned char                          TLabel;
typedef itk::NumericTraits< TPixel >::RealType TScalar;
typedef itk::Image< TPixel, Dim >              TImage;
typedef itk::Image< TLabel, Dim >              TLabels;
typedef itk::Image< TScalar, Dim >             TScalarImage;

/**
 */
class MoriLabellingTraits
  : public fpa::Filters::Image::DefaultTraits< TImage, TLabels, TLabel >
{
public:
  typedef fpa::Filters::Image::DefaultTraits< TImage, TLabels, TLabel > Superclass;
  typedef typename Superclass::TInternalTraits  TInternalTraits;
  typedef typename Superclass::TMarksImage      TMarksImage;
  typedef typename Superclass::TFilterInterface TFilterInterface;

  typedef fpa::Filters::BaseMarksInterface< TInternalTraits >  TMarksInterface;
  typedef fpa::Filters::Image::SeedsFromLabelsInterface< TInternalTraits > TSeedsInterface;
};

/**
 */
class MoriLabelling
  : public fpa::Filters::Image::RegionGrow< TImage, TLabels, TLabel, MoriLabellingTraits >
{
public:
  typedef MoriLabellingTraits TTraits;
  fpaTraitsMacro( TTraits );

  typedef fpa::Filters::Image::RegionGrow< TImage, TLabels, TMark, TTraits > Superclass;
  typedef MoriLabelling                   Self;
  typedef itk::SmartPointer< Self >       Pointer;
  typedef itk::SmartPointer< const Self > ConstPointer;
  typedef fpa::Functors::RegionGrow::BinaryThreshold< TInputValue > TFunctor;

public:
  itkNewMacro( Self );
  itkTypeMacro( MoriLabelling, fpa::Filters::Image::RegionGrow );

  itkGetConstMacro( LastThreshold, TInputValue );
  itkSetMacro( LastThreshold, TInputValue );

  itkGetConstMacro( MinVertex, TVertex );
  itkGetConstMacro( MaxVertex, TVertex );

  ivqITKInputMacro( InputLabels, TLabels );
  ivqITKInputMacro( InputVesselness, TScalarImage );

public:
  TInputValue GetUpperThreshold( ) const
    {
      return( this->m_Functor->GetUpperThreshold( ) );
    }
  void SetUpperThreshold( TInputValue t )
    {
      this->m_Functor->SetUpperThreshold( t );
    }

protected:
  MoriLabelling( )
  : Superclass( ),
    m_LastThreshold( TInputValue( 0 ) ),
    m_VesselnessThr( TScalar( 0.05 ) )
    {
      ivqITKInputConfigureMacro( InputLabels, TLabels );
      ivqITKInputConfigureMacro( InputVesselness, TScalarImage );
      this->m_Functor = TFunctor::New( );
      this->SetPredicate( this->m_Functor );
    }

  virtual ~MoriLabelling( )
    {
    }

  virtual const itk::DataObject* _GetReferenceInput( ) const override
    {
      return( this->GetInputLabels( ) );
    }
  virtual void _BeforeGenerateData( ) override
    {
      this->Superclass::_BeforeGenerateData( );
      this->m_FirstVertex = true;

      typedef itk::MinimumMaximumImageCalculator< TScalarImage > _TMinMax;
      _TMinMax::Pointer minMax = _TMinMax::New( );
      minMax->SetImage( this->GetInputVesselness( ) );
      minMax->Compute( );
      this->m_MaxVesselness = ( 1.0  - this->m_VesselnessThr ) * minMax->GetMaximum( );
    }

  virtual void _PostComputeOutputValue( TNode& n ) override
    {
      this->Superclass::_PostComputeOutputValue( n );
      if( n.Value == this->GetInsideValue( ) )
      {
        const TImage* input = this->GetInput( );
        const TLabels* labels = this->GetInputLabels( );
        const TScalarImage* vesselness = this->GetInputVesselness( );
        double x = input->GetPixel( n.Vertex );
        /* TODO
           double a = std::fabs( x - double( this->m_Functor->GetUpperThreshold( ) ) );
           double b = std::fabs( x - double( this->m_LastThreshold ) );
        */
        if( labels->GetPixel( n.Vertex ) == 0 )
        {
          if( vesselness->GetPixel( n.Vertex ) > this->m_MaxVesselness )
            n.Value = this->GetInsideValue( );
          else
            n.Value = 0;

        } // fi

        if( !( this->m_FirstVertex ) )
        {
          for( unsigned int d = 0; d < TImage::ImageDimension; ++d )
          {
            if( n.Vertex[ d ] < this->m_MinVertex[ d ] )
              this->m_MinVertex[ d ] = n.Vertex[ d ];
            if( this->m_MaxVertex[ d ] < n.Vertex[ d ] )
              this->m_MaxVertex[ d ] = n.Vertex[ d ];

          } // rof
        }
        else
        {
          this->m_MinVertex = n.Vertex;
          this->m_MaxVertex = n.Vertex;
          this->m_FirstVertex = false;

        } // fi

      } // fi
    }

private:
  // Purposely not implemented.
  MoriLabelling( const Self& other );
  Self& operator=( const Self& other );

protected:
  TFunctor::Pointer m_Functor;
  TInputValue m_LastThreshold;
  bool m_FirstVertex;
  TVertex m_MinVertex;
  TVertex m_MaxVertex;

  TScalar m_MaxVesselness;
  TScalar m_VesselnessThr;
};

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

// -------------------------------------------------------------------------
template< class _TImagePtr, class _TRegion >
void ROI(
  _TImagePtr& output,
  const typename _TImagePtr::ObjectType* input,
  const _TRegion& roi
  )
{
  typedef typename _TImagePtr::ObjectType _TImage;
  typedef itk::RegionOfInterestImageFilter< _TImage, _TImage > _TFilter;

  typename _TFilter::Pointer filter = _TFilter::New( );
  filter->SetInput( input );
  filter->SetRegionOfInterest( roi );
  double t = MeasureTime( filter );
  std::cout << "ROI computed in " << t << " s" << std::endl;
  output = filter->GetOutput( );
  output->DisconnectPipeline( );
}

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

  args[ "upper_threshold" ] = "-650";

  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-labels filename" << std::endl
      << "\t-vesselness filename" << std::endl
      << "\t[-roi] filename" << std::endl
      << "\t[-upper_threshold value]" << 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" ] );

      // Read labels image
      TLabels::Pointer input_labels;
      ReadImage( input_labels, args[ "labels" ] );

      // Read vesselness image
      TScalarImage::Pointer input_vesselness;
      ReadImage( input_vesselness, args[ "vesselness" ] );

      // Mori labelling
      MoriLabelling::Pointer labelling = MoriLabelling::New( );
      labelling->SetInput( input_image );
      labelling->SetInputLabels( input_labels );
      labelling->SetInputVesselness( input_vesselness );
      labelling->SetOutsideValue( 2 );
      labelling->SetFillValue( 2 );
      labelling->SetInsideValue( 1 );
      labelling->SetUpperThreshold(
        TPixel( std::atof( args[ "upper_threshold" ].c_str( ) ) )
        );
      /* TODO
         labelling->SetLastThreshold( last_thr );
      */
      double t = MeasureTime( labelling );
      std::cout << "Labelling executed in " << t << " s" << std::endl;

      std::map< std::string, std::string >::const_iterator i =
        args.find( "roi" );
      if( i != args.end( ) )
      {
        TImage::IndexType minV = labelling->GetMinVertex( );
        TImage::IndexType maxV = labelling->GetMaxVertex( );
        TImage::SizeType roiS;
        for( unsigned d = 0; d < TImage::ImageDimension; ++d )
          roiS[ d ] = maxV[ d ] - minV[ d ] + 1;
        TImage::RegionType roi;
        roi.SetIndex( minV );
        roi.SetSize( roiS );

        // ROI input image
        TImage::Pointer input_image_roi;
        ROI( input_image_roi, input_image.GetPointer( ), roi );
        WriteImage( input_image_roi.GetPointer( ), i->second );

        // ROI output image
        TLabels::Pointer output_roi;
        ROI( output_roi, labelling->GetOutput( ), roi );
        WriteImage( output_roi.GetPointer( ), args[ "out" ] );
      }
      else
        WriteImage( labelling->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$