// =========================================================================
// @author Leonardo Florez-Valencia (florez-l@javeriana.edu.co)
// =========================================================================
#ifndef __RandomWalkSegmentation__hxx__
#define __RandomWalkSegmentation__hxx__

#include <itkBinaryThresholdImageFilter.h>
#include <itkImageRegionConstIterator.h>
#include <itkImageRegionIterator.h>
#include <itkSmoothingRecursiveGaussianImageFilter.h>

#include <ivq/ITK/RegionOfInterestWithPaddingImageFilter.h>

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

#include "DijkstraWithMeanAndVariance.h"
#include "RandomWalkLabelling.h"

#include <itkImageFileWriter.h>

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
AddSeed( const TIndex& s )
{
  TSeed seed;
  seed.Index = s;
  seed.IsPoint = false;
  this->m_Seeds.push_back( seed );
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
AddSeed( const TPoint& s )
{
  TSeed seed;
  seed.Point = s;
  seed.IsPoint = true;
  this->m_Seeds.push_back( seed );
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
ClearSeeds( )
{
  this->m_Seeds.clear( );
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
unsigned long RandomWalkSegmentation< _TInputImage, _TOutputImage >::
GetNumberOfSeeds( ) const
{
  return( this->m_Seeds.size( ) );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
SetStartSeed( const TIndex& s )
{
  this->m_StartSeed.Index = s;
  this->m_StartSeed.IsPoint = false;
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
SetStartSeed( const TPoint& s )
{
  this->m_StartSeed.Point = s;
  this->m_StartSeed.IsPoint = true;
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
SetEndSeed( const TIndex& s )
{
  this->m_EndSeed.Index = s;
  this->m_EndSeed.IsPoint = false;
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
SetEndSeed( const TPoint& s )
{
  this->m_EndSeed.Point = s;
  this->m_EndSeed.IsPoint = true;
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
RandomWalkSegmentation< _TInputImage, _TOutputImage >::
RandomWalkSegmentation( )
  : Superclass( ),
    m_Beta( double( 20 ) ),
    m_Sigma( double( 10 ) ),
    m_Radius( double( 5 ) )
{
  fpaFilterOutputConfigureMacro( OutputAxis, TPath );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
RandomWalkSegmentation< _TInputImage, _TOutputImage >::
~RandomWalkSegmentation( )
{
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
GenerateData( )
{
  typedef typename TOutputImage::PixelType _TScalar;
  typedef DijkstraWithMeanAndVariance< TOutputImage, TOutputImage > _TInit;
  typedef typename _TInit::TMarksImage _TMarksImage;
  typedef fpa::Functors::Dijkstra::Image::Gaussian< TOutputImage, _TScalar > _TGaussFun;
  typedef RandomWalkLabelling< TOutputImage, _TMarksImage > _TLabelling;
  typedef typename _TLabelling::TLabelsImage _TLabelsImage;
  typedef fpa::Filters::Image::RandomWalker< TOutputImage, _TLabelsImage, _TScalar > _TRandomWalker;
  typedef itk::BinaryThresholdImageFilter< _TLabelsImage, _TLabelsImage > _TLabelExtract;
  typedef fpa::Filters::Image::ExtractAxis< _TLabelsImage, _TScalar > _TAxisExtract;
  typedef ivq::ITK::RegionOfInterestWithPaddingImageFilter< TOutputImage, TOutputImage > _TInputROI;
  typedef ivq::ITK::RegionOfInterestWithPaddingImageFilter< _TMarksImage, _TMarksImage > _TMarksROI;
  typedef itk::SmoothingRecursiveGaussianImageFilter< TInputImage, TOutputImage > _TSmooth;

  // Prepare initial seeds
  const TInputImage* input = this->GetInput( );
  if( this->m_StartSeed.IsPoint )
    input->TransformPhysicalPointToIndex(
      this->m_StartSeed.Point,
      this->m_StartSeed.Index
      );
  else
    input->TransformIndexToPhysicalPoint(
      this->m_StartSeed.Index,
      this->m_StartSeed.Point
      );
  this->m_StartSeed.IsPoint = true;
  if( this->m_EndSeed.IsPoint )
    input->TransformPhysicalPointToIndex(
      this->m_EndSeed.Point,
      this->m_EndSeed.Index
      );
  else
    input->TransformIndexToPhysicalPoint(
      this->m_EndSeed.Index,
      this->m_EndSeed.Point
      );
  this->m_EndSeed.IsPoint = true;
  for( TSeed& seed: this->m_Seeds )
  {
    if( seed.IsPoint )
      input->TransformPhysicalPointToIndex( seed.Point, seed.Index );
    else
      input->TransformIndexToPhysicalPoint( seed.Index, seed.Point );
    seed.IsPoint = true;

  } // rof

  // Intermediary objects
  typename TOutputImage::Pointer smooth_input;
  typename TOutputImage::Pointer input_roi;
  typename TPath::Pointer init_axis;
  typename TPath::Pointer init_axis_roi;
  typename TPath::Pointer output_axis_roi;
  typename _TMarksImage::Pointer init_marks;
  typename _TMarksImage::Pointer init_marks_roi;
  typename _TLabelsImage::Pointer init_labels;
  typename _TLabelsImage::Pointer final_labels;
  typename _TLabelsImage::Pointer inside_labels;
  typename TInputImage::RegionType roi;
  typename TOutputImage::Pointer output_roi;
  double init_mean, init_std;

  // Smooth input
  std::cout << "0" << std::endl;
  { // begin
    typename _TSmooth::Pointer smooth = _TSmooth::New( );
    smooth->SetInput( input );
    smooth->SetNormalizeAcrossScale( true );
    smooth->SetSigmaArray( input->GetSpacing( ) * double( 2 ) );
    smooth->Update( );
    smooth_input = smooth->GetOutput( );
    smooth_input->DisconnectPipeline( );

  } // end

  // Initial segmentation
  std::cout << "1" << std::endl;
  { // begin
    typename _TGaussFun::Pointer init_fun = _TGaussFun::New( );
    init_fun->SetBeta( this->m_Beta );

    typename _TInit::Pointer init = _TInit::New( );
    init->SetInput( smooth_input );
    init->SetWeightFunction( init_fun );
    init->StopAtOneFrontOn( );
    for( TSeed seed: this->m_Seeds )
      init->AddSeed( seed.Point );
    init->Update( );

    // Get initial values
    init_mean = init->GetMean( );
    init_std = init->GetDeviation( ) * this->m_Sigma;

    // Get initial objects
    init->GetMinimumSpanningTree( )->
      GetPath( init_axis, this->m_StartSeed.Index, this->m_EndSeed.Index );
    init_marks = init->GetMarks( );
    init_marks->DisconnectPipeline( );

    typename TInputImage::IndexType min_idx = init->GetMinVertex( );
    typename TInputImage::IndexType max_idx = init->GetMaxVertex( );
    typename TInputImage::SizeType roi_size;
    for( unsigned int i = 0; i < TInputImage::ImageDimension; ++i )
      roi_size[ i ] = max_idx[ i ] - min_idx[ i ] + 1;
    roi.SetIndex( min_idx );
    roi.SetSize( roi_size );

  } // end

  // Extract input ROIs
  unsigned int pad = 10;
  { // begin
    typename _TInputROI::Pointer input_roi_filter = _TInputROI::New( );
    input_roi_filter->SetInput( smooth_input );
    input_roi_filter->SetRegionOfInterest( roi );
    input_roi_filter->SetPadding( pad );
    input_roi_filter->Update( );
    input_roi = input_roi_filter->GetOutput( );
    input_roi->DisconnectPipeline( );
    roi = input_roi_filter->GetRegionOfInterest( );

  } // end

  { // begin
    typename _TMarksROI::Pointer init_marks_roi_filter = _TMarksROI::New( );
    init_marks_roi_filter->SetInput( init_marks );
    init_marks_roi_filter->SetRegionOfInterest( roi );
    init_marks_roi_filter->SetPadding( 0 );
    init_marks_roi_filter->Update( );
    init_marks_roi = init_marks_roi_filter->GetOutput( );
    init_marks_roi->DisconnectPipeline( );

  } // end
  
  // Convert initial axis
  { // begin
    init_axis_roi = TPath::New( );
    init_axis_roi->SetReferenceImage( input_roi.GetPointer( ) );
    for( unsigned long i = 0; i < init_axis->GetSize( ); ++i )
    {
      TIndex v = init_axis->GetVertex( i );
      for( unsigned int d = 0; d < TInputImage::ImageDimension; ++d )
        v[ d ] -= roi.GetIndex( )[ d ];
      init_axis_roi->AddVertex( v );

    } // rof

  } // end

  // Labelling
  std::cout << "2" << std::endl;
  { // begin
    typename _TLabelling::Pointer labelling = _TLabelling::New( );
    labelling->SetInputImage( input_roi );
    labelling->SetInputMarks( init_marks_roi );
    labelling->SetInputPath( init_axis_roi );
    labelling->SetInsideLabel( 1 );
    labelling->SetOutsideLabel( 2 );
    labelling->SetLowerLabel( 3 );
    labelling->SetUpperLabel( 4 );
    labelling->SetRadius( this->m_Radius );
    labelling->SetLowerThreshold( init_mean - init_std );
    labelling->SetUpperThreshold( init_mean + init_std );
    labelling->Update( );
    init_labels = labelling->GetOutputLabels( );
    init_labels->DisconnectPipeline( );

  } // end

  // Random walker
  std::cout << "3" << std::endl;
  { // begin
    typename _TGaussFun::Pointer rw_fun = _TGaussFun::New( );
    rw_fun->SetBeta( init_std / double( 2 ) );

    typename _TRandomWalker::Pointer rw = _TRandomWalker::New( );
    rw->SetInputImage( input_roi );
    rw->SetInputLabels( init_labels );
    rw->SetWeightFunction( rw_fun );
    rw->Update( );
    final_labels = rw->GetOutputLabels( );
    final_labels->DisconnectPipeline( );

  } // end

  // Extract inside label
  std::cout << "4" << std::endl;
  { // begin
    typename _TLabelExtract::Pointer label_extract = _TLabelExtract::New( );
    label_extract->SetInput( final_labels );
    label_extract->SetInsideValue( 1 );
    label_extract->SetOutsideValue( 0 );
    label_extract->SetLowerThreshold( 1 );
    label_extract->SetUpperThreshold( 1 );
    label_extract->Update( );
    inside_labels = label_extract->GetOutput( );
    inside_labels->DisconnectPipeline( );

  } // end

  // Prepare output values
  std::cout << "5" << std::endl;
  { // begin
    TIndex start_seed, end_seed;
    inside_labels->TransformPhysicalPointToIndex( this->m_StartSeed.Point, start_seed );
    inside_labels->TransformPhysicalPointToIndex( this->m_EndSeed.Point, end_seed );

    typename _TAxisExtract::Pointer axis_extract = _TAxisExtract::New( );
    axis_extract->SetInput( inside_labels );
    axis_extract->AddSeed( start_seed );
    axis_extract->AddSeed( end_seed );
    for( TSeed seed: this->m_Seeds )
      axis_extract->AddSeed( seed.Point );
    axis_extract->SetStartIndex( start_seed );
    axis_extract->SetEndIndex( end_seed );
    axis_extract->Update( );
    output_axis_roi = TPath::New( );
    output_axis_roi->Graft( axis_extract->GetOutput( ) );
    output_roi = axis_extract->GetCenterness( )->GetOutput( );
    output_roi->DisconnectPipeline( );

  } // end

  // Put everything back to requested region
  std::cout << "6" << std::endl;
  { // begin
    TOutputImage* output = this->GetOutput( );
    output->SetBufferedRegion( output->GetRequestedRegion( ) );
    output->Allocate( );
    output->FillBuffer( -std::numeric_limits< _TScalar >::max( ) );

    itk::ImageRegionConstIterator< TOutputImage > rIt(
      output_roi, output_roi->GetRequestedRegion( )
      );
    itk::ImageRegionIterator< TOutputImage > oIt( output, roi );
    rIt.GoToBegin( );
    oIt.GoToBegin( );
    for( ; !rIt.IsAtEnd( ); ++rIt, ++oIt )
      oIt.Set( rIt.Get( ) );

    TPath* output_axis = this->GetOutputAxis( );
    output_axis->SetReferenceImage( output );
    for( unsigned long i = 0; i < output_axis_roi->GetSize( ); ++i )
    {
      TIndex v = output_axis_roi->GetVertex( i );
      for( unsigned int d = 0; d < TInputImage::ImageDimension; ++d )
        v[ d ] += roi.GetIndex( )[ d ];
      output_axis->AddVertex( v );

    } // rof

  } // end
  std::cout << "7" << std::endl;
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
template< class _TImage >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_save( _TImage* image, const std::string& fname )
{
  typename itk::ImageFileWriter< _TImage >::Pointer w =
    itk::ImageFileWriter< _TImage >::New( );
  w->SetInput( image );
  w->SetFileName( fname );
  w->Update( );
}


#endif // __RandomWalkSegmentation__hxx__

// eof - $RCSfile$