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

#include <itkBinaryThresholdImageFilter.h>
#include <itkSmoothingRecursiveGaussianImageFilter.h>
#include <ivq/ITK/RegionOfInterestWithPaddingImageFilter.h>
#include <fpa/Filters/Image/ExtractAxis.h>
#include <fpa/Filters/Image/RandomWalker.h>
#include "DijkstraWithMeanAndVariance.h"
#include "RandomWalkLabelling.h"

/* TODO
   #include <itkImageRegionConstIterator.h>
   #include <itkImageRegionIterator.h>


   #include <fpa/Functors/Dijkstra/Image/Gaussian.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 unsigned char _TLabel;
  typedef typename TOutputImage::PixelType _TScalar;
  typedef itk::Image< _TLabel, TInputImage::ImageDimension > _TLabels;

  // Prepare initial seeds
  const TInputImage* input = this->GetInput( );
  this->_SynchSeed( input, this->m_StartSeed );
  this->_SynchSeed( input, this->m_EndSeed );
  for( TSeed& seed: this->m_Seeds )
    this->_SynchSeed( input, seed );

  // Smooth input
  typename TOutputImage::Pointer smooth_in;
  this->_Smooth( this->GetInput( ), smooth_in );

  // Initial segmentation
  typename TOutputImage::Pointer init_seg;
  typename TPath::Pointer init_axis;
  _TScalar init_mean, init_std;
  typename TInputImage::RegionType roi =
    this->_RawSegmentation(
      smooth_in.GetPointer( ), this->m_Seeds,
      this->m_StartSeed.Index, this->m_EndSeed.Index,
      this->m_Beta,
      init_seg, init_axis, init_mean, init_std
      );
  init_std *= _TScalar( this->m_Sigma );

  // Extract input ROIs
  typename TOutputImage::Pointer smooth_in_roi, init_seg_roi;
  roi = this->_ROI( smooth_in.GetPointer( ), roi, 10, smooth_in_roi );
  this->_ROI( init_seg.GetPointer( ), roi, 0, init_seg_roi );
  typename TPath::Pointer init_axis_roi = TPath::New( );
  init_axis_roi->SetReferenceImage( smooth_in_roi.GetPointer( ) );
  this->_AxisROI( init_axis.GetPointer( ), roi, init_axis_roi );

  // Labelling
  typename _TLabels::Pointer init_labels;
  _TScalar radius = _TScalar( this->m_Radius );
  this->_Label(
    smooth_in_roi.GetPointer( ),
    init_seg_roi.GetPointer( ),
    init_axis_roi.GetPointer( ),
    init_mean, init_std, radius,
    init_labels
    );

  // Random walker
  typename _TLabels::Pointer rw_seg;
  this->_RandomWalker(
    smooth_in_roi.GetPointer( ),
    init_labels.GetPointer( ),
    init_std / _TScalar( 2 ),
    rw_seg
    );

  // ROI outputs
  typename TOutputImage::Pointer out_dist;
  typename TPath::Pointer out_axis;
  this->_DistanceAndAxis(
    rw_seg.GetPointer( ),
    this->m_Seeds, this->m_StartSeed, this->m_EndSeed,
    out_dist, out_axis
    );

  // Put everything back to requested region
  /* TODO
     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 _TIn, class _TSeed >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_SynchSeed( const _TIn* in, _TSeed& seed )
{
  if( seed.IsPoint )
    in->TransformPhysicalPointToIndex( seed.Point, seed.Index );
  else
    in->TransformIndexToPhysicalPoint( seed.Index, seed.Point );
  seed.IsPoint = true;
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
template< class _TIn, class _TOutPtr >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_Smooth( const _TIn* in, _TOutPtr& out )
{
  typedef typename _TOutPtr::ObjectType _TOut;
  typedef itk::SmoothingRecursiveGaussianImageFilter< _TIn, _TOut > _TSmooth;

  typename _TSmooth::Pointer smooth = _TSmooth::New( );
  smooth->SetInput( in );
  smooth->SetNormalizeAcrossScale( true );
  smooth->SetSigmaArray( in->GetSpacing( ) * double( 2 ) );
  smooth->Update( );
  out = smooth->GetOutput( );
  out->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
template< class _TIn, class _TOutPtr, class _TAxisPtr, class _TSeeds >
typename _TIn::RegionType
RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_RawSegmentation(
  const _TIn* in, const _TSeeds& seeds,
  const typename _TIn::IndexType& s0,
  const typename _TIn::IndexType& s1,
  const typename _TOutPtr::ObjectType::PixelType& beta,
  _TOutPtr& out, _TAxisPtr& out_axis,
  typename _TOutPtr::ObjectType::PixelType& oMean,
  typename _TOutPtr::ObjectType::PixelType& oSTD
  )
{
  typedef typename _TOutPtr::ObjectType _TOut;
  typedef typename _TOut::PixelType  _TScalar;
  typedef DijkstraWithMeanAndVariance< _TIn, _TOut > _TInit;
  typedef fpa::Functors::Dijkstra::Image::Gaussian< _TIn, _TScalar > _TFun;

  typename _TFun::Pointer fun = _TFun::New( );
  fun->SetBeta( beta );

  typename _TInit::Pointer init = _TInit::New( );
  init->SetInput( in );
  init->SetWeightFunction( fun );
  init->StopAtOneFrontOn( );
  for( typename _TSeeds::value_type seed: seeds )
    init->AddSeed( seed.Point );
  init->Update( );

  // Get initial values
  oMean = _TScalar( init->GetMean( ) );
  oSTD = _TScalar( init->GetDeviation( ) );

  // Get initial objects
  init->GetMinimumSpanningTree( )->GetPath( out_axis, s0, s1 );
  out = init->GetOutput( );
  out->DisconnectPipeline( );

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

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
template< class _TIn, class _TOutPtr >
typename _TIn::RegionType
RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_ROI(
  const _TIn* in, const typename _TIn::RegionType& roi, unsigned int pad,
  _TOutPtr& out
  )
{
  typedef typename _TOutPtr::ObjectType _TOut;
  typedef ivq::ITK::RegionOfInterestWithPaddingImageFilter< _TIn, _TOut > _TROI;

  typename _TROI::Pointer filter = _TROI::New( );
  filter->SetInput( in );
  filter->SetRegionOfInterest( roi );
  filter->SetPadding( pad );
  filter->Update( );
  out = filter->GetOutput( );
  out->DisconnectPipeline( );
  return( filter->GetRegionOfInterest( ) );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
template< class _TAxisPtr, class _TRegion >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_AxisROI(
  const typename _TAxisPtr::ObjectType* in, const _TRegion& roi,
  _TAxisPtr& out
  )
{
  typedef typename _TAxisPtr::ObjectType _TAxis;

  for( unsigned long i = 0; i < in->GetSize( ); ++i )
  {
    typename _TAxis::TIndex v = in->GetVertex( i );
    for( unsigned int d = 0; d < _TAxis::Dimension; ++d )
      v[ d ] -= roi.GetIndex( )[ d ];
    out->AddVertex( v );

  } // rof
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
template< class _TInRaw, class _TInCosts, class _TAxis, class _TScalar, class _TOutPtr >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_Label(
  const _TInRaw* raw, const _TInCosts* costs, const _TAxis* axis,
  const _TScalar& mean, const _TScalar& dev, const _TScalar& radius,
  _TOutPtr& out
  )
{
  typedef typename _TOutPtr::ObjectType _TOut;
  typedef RandomWalkLabelling< _TInRaw, _TInCosts, _TOut > _TLabel;

  typename _TLabel::Pointer label = _TLabel::New( );
  label->SetInputImage( raw );
  label->SetInputCosts( costs );
  label->SetInputPath( axis );
  label->SetInsideLabel( 1 );
  label->SetOutsideLabel( 2 );
  label->SetLowerLabel( 3 );
  label->SetUpperLabel( 4 );
  label->SetRadius( radius );
  label->SetLowerThreshold( mean - dev );
  label->SetUpperThreshold( mean + dev );
  label->Update( );
  out = label->GetOutputLabels( );
  out->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
template< class _TIn, class _TOutPtr >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_RandomWalker(
  const _TIn* in, const typename _TOutPtr::ObjectType* labels,
  const typename _TIn::PixelType& beta,
  _TOutPtr& out
  )
{
  typedef typename _TIn::PixelType _TScalar;
  typedef typename _TOutPtr::ObjectType _TOut;
  typedef fpa::Functors::Dijkstra::Image::Gaussian< _TIn, _TScalar > _TFun;
  typedef fpa::Filters::Image::RandomWalker< _TIn, _TOut, _TScalar > _TRandomWalker;
  typedef itk::BinaryThresholdImageFilter< _TOut, _TOut > _TExtract;

  typename _TFun::Pointer fun = _TFun::New( );
  fun->SetBeta( beta );

  typename _TRandomWalker::Pointer rw = _TRandomWalker::New( );
  rw->SetInputImage( in );
  rw->SetInputLabels( labels );
  rw->SetWeightFunction( fun );

  typename _TExtract::Pointer extract = _TExtract::New( );
  extract->SetInput( rw->GetOutputLabels( ) );
  extract->SetInsideValue( 1 );
  extract->SetOutsideValue( 0 );
  extract->SetLowerThreshold( 1 );
  extract->SetUpperThreshold( 1 );
  extract->Update( );
  out = extract->GetOutput( );
  out->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
template< class _TInputImage, class _TOutputImage >
template< class _TIn, class _TSeeds, class _TSeed, class _TOutPtr, class _TAxisPtr >
void RandomWalkSegmentation< _TInputImage, _TOutputImage >::
_DistanceAndAxis(
  const _TIn* in, const _TSeeds& seeds,
  const _TSeed& p0, const _TSeed& p1,
  _TOutPtr& out_dist, _TAxisPtr& out_axis
  )
{
  typedef typename _TOutPtr::ObjectType _TOut;
  typedef typename _TOut::PixelType _TScalar;
  typedef fpa::Filters::Image::ExtractAxis< _TIn, _TScalar > _TExtract;

  // Prepare output values
  typename _TIn::IndexType s0, s1;
  in->TransformPhysicalPointToIndex( p0.Point, s0 );
  in->TransformPhysicalPointToIndex( p1.Point, s1 );

  typename _TExtract::Pointer extract = _TExtract::New( );
  extract->SetInput( const_cast< _TIn* >( in ) );
  extract->AddSeed( s0 );
  extract->AddSeed( s1 );
  for( typename _TSeeds::value_type seed: seeds )
    extract->AddSeed( seed.Point );
  extract->SetStartIndex( s0 );
  extract->SetEndIndex( s1 );
  extract->Update( );
  out_axis = TPath::New( );
  out_axis->Graft( extract->GetOutput( ) );
  this->_Smooth( extract->GetCenterness( )->GetOutput( ), out_dist );
}

#endif // __RandomWalkSegmentation__hxx__

// eof - $RCSfile$