// =========================================================================
// @author Leonardo Florez Valencia
// @email florez-l@javeriana.edu.co
// =========================================================================
#ifndef __fpa__Common__SliceBySliceRandomWalker__hxx__
#define __fpa__Common__SliceBySliceRandomWalker__hxx__

#include <vector>
#include <itkBinaryThresholdImageFilter.h>
#include <itkExtractImageFilter.h>
#include <itkImageRegionConstIterator.h>
#include <itkImageRegionIterator.h>
#include <itkJoinSeriesImageFilter.h>
#include <itkMaximumImageFilter.h>
#include <itkMinimumMaximumImageCalculator.h>
#include <fpa/Common/RandomWalker.h>
#include <fpa/Functors/Dijkstra/Image/Gaussian.h>

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalarImage >
fpa::Common::SliceBySliceRandomWalker< _TImage, _TLabels, _TScalarImage >::
SliceBySliceRandomWalker( )
  : Superclass( ),
    m_Epsilon( TScalar( 1e-5 ) ),
    m_Beta( TScalar( 1 ) ),
    m_VesselnessThreshold( TScalar( 5 ) )
{
  ivqITKInputConfigureMacro( InputLabels, TLabels );
  ivqITKInputConfigureMacro( InputVesselness, TScalarImage );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalarImage >
fpa::Common::SliceBySliceRandomWalker< _TImage, _TLabels, _TScalarImage >::
~SliceBySliceRandomWalker( )
{
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalarImage >
void
fpa::Common::SliceBySliceRandomWalker< _TImage, _TLabels, _TScalarImage >::
GenerateData( )
{
  typedef itk::Image< TPixel, TImage::ImageDimension - 1 > _TSliceImage;
  typedef itk::Image< TLabel, TImage::ImageDimension - 1 > _TSliceLabels;
  typedef itk::Image< TScalar, TImage::ImageDimension - 1 > _TSliceScalarImage;

  // Preare I/O objects
  const TImage* input = this->GetInput( );
  const TLabels* labels = this->GetInputLabels( );
  const TScalarImage* vesselness = this->GetInputVesselness( );
  this->AllocateOutputs( );

  // Get vesselness range
  typedef itk::MinimumMaximumImageCalculator< TScalarImage > _TMinMax;
  typename _TMinMax::Pointer vMinMax = _TMinMax::New( );
  vMinMax->SetImage( vesselness );
  vMinMax->Compute( );
  TScalar vMax = vMinMax->GetMaximum( );

  // Some values
  typename TImage::RegionType region = input->GetRequestedRegion( );
  typename TImage::IndexType regionIndex = region.GetIndex( );
  typename TImage::SizeType sliceSize = region.GetSize( );
  int numSlices = sliceSize[ 2 ];
  sliceSize[ 2 ] = 0;
  this->m_VesselnessValue  = TScalar( 0.01 );
  this->m_VesselnessValue *= this->m_VesselnessThreshold * vMax;

  // Composite image
  typename TScalarImage::Pointer composite;
  this->_Composite( composite, labels, vesselness, vMax );

  // Extract slices
  std::vector< typename _TSliceImage::Pointer > data3D( numSlices );
  std::vector< typename _TSliceLabels::Pointer > binaryTree( numSlices );
  std::vector< typename _TSliceScalarImage::Pointer > fusion( numSlices );
  for( int i = 0; i < numSlices; ++i )
  {
    // Prepare extraction region
    typename TImage::IndexType sliceIndex = regionIndex;
    sliceIndex[ 2 ] += i;
    typename TImage::RegionType sliceRegion( sliceIndex, sliceSize );

    // Extract regions
    typename _TSliceImage::Pointer input_slice;
    typename _TSliceLabels::Pointer labels_slice;
    typename _TSliceScalarImage::Pointer composite_slice;
    this->_Slice( input_slice, input, sliceRegion );
    this->_Slice( labels_slice, labels, sliceRegion );
    this->_Slice( composite_slice, composite.GetPointer( ), sliceRegion );

    // Update vectors with each image
    data3D[ i ] = input_slice;
    binaryTree[ i ] = labels_slice;
    fusion[ i ] = composite_slice;

  } // rof

  // Random walker slice-by-slice
  this->_RandomWalker( binaryTree, data3D, fusion, true );
  this->_RandomWalker( binaryTree, data3D, fusion, false );

  // Join results
  typedef itk::JoinSeriesImageFilter< _TSliceLabels, TLabels > _TJoin;
  typename _TJoin::Pointer join = _TJoin::New( );
  for( int i = 0; i < numSlices; ++i )
    join->SetInput( i, binaryTree[ i ] );
  join->Update( );
  this->GetOutput( )->Graft( join->GetOutput( ) );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalarImage >
void
fpa::Common::SliceBySliceRandomWalker< _TImage, _TLabels, _TScalarImage >::
_Composite(
  typename TScalarImage::Pointer& composite,
  const TLabels* labels,
  const TScalarImage* vesselness,
  const TScalar& maxVess
  )
{
  // Some values
  typename TScalarImage::RegionType region = labels->GetRequestedRegion( );

  // Composite image
  composite = TScalarImage::New( );
  composite->CopyInformation( vesselness );
  composite->SetBufferedRegion( vesselness->GetBufferedRegion( ) );
  composite->Allocate( );
  itk::ImageRegionConstIterator< TScalarImage > v( vesselness, region );
  itk::ImageRegionConstIterator< TLabels > l( labels, region );
  itk::ImageRegionIterator< TScalarImage > c( composite, region );
  v.GoToBegin( );
  l.GoToBegin( );
  c.GoToBegin( );
  for( ; !v.IsAtEnd( ); ++v, ++l, ++c )
    c.Set( ( l.Get( ) != 0 )? maxVess: v.Get( ) );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalarImage >
template< class _TSlicePtr, class _TInput >
void
fpa::Common::SliceBySliceRandomWalker< _TImage, _TLabels, _TScalarImage >::
_Slice(
  _TSlicePtr& slice,
  const _TInput* input,
  typename _TInput::RegionType region
  )
{
  typedef typename _TSlicePtr::ObjectType _TSlice;
  typedef itk::ExtractImageFilter< _TInput, _TSlice > _TExtract;

  typename _TExtract::Pointer extract = _TExtract::New( );
  extract->SetInput( input );
  extract->SetExtractionRegion( region );
  extract->SetDirectionCollapseToIdentity( );
  extract->Update( );
  slice = extract->GetOutput( );
  slice->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalarImage >
template< class _TBinaryTree, class _TData3D, class _TFusion >
void
fpa::Common::SliceBySliceRandomWalker< _TImage, _TLabels, _TScalarImage >::
_RandomWalker(
  _TBinaryTree& binaryTree,
  const _TData3D& data3D,
  const _TFusion& fusion,
  bool down
  )
{
  typedef typename _TBinaryTree::value_type     _TSliceBinaryPtr;
  typedef typename _TData3D::value_type         _TSliceData3DPtr;
  typedef typename _TFusion::value_type         _TSliceFusionPtr;
  typedef typename _TSliceBinaryPtr::ObjectType _TSliceBinary;
  typedef typename _TSliceData3DPtr::ObjectType _TSliceData3D;
  typedef typename _TSliceFusionPtr::ObjectType _TSliceFusion;
  typedef typename _TSliceFusion::PixelType     _TSliceScalar;

  int numSlices = binaryTree.size( );
  int z = ( down )? -1: numSlices;
  int o = ( down )? 1: -1;
  while( ( down && z < numSlices - 2 ) || ( !down && z > 1 ) )
  {
    z += o;
    _TSliceBinaryPtr tmp = binaryTree[ z ];
    _TSliceBinaryPtr next = binaryTree[ z + o ];
    _TSliceData3DPtr data = data3D[ z + o ];
    _TSliceFusionPtr vess = fusion[ z + o ];
    typename _TSliceBinary::RegionType region = tmp->GetRequestedRegion( );

    // Create seeds and background
    _TSliceBinaryPtr seeds = _TSliceBinary::New( );
    seeds->CopyInformation( tmp );
    seeds->SetBufferedRegion( region );
    seeds->Allocate( );
    seeds->FillBuffer( 0 );

    itk::ImageRegionConstIterator< _TSliceBinary > tIt( tmp, region );
    itk::ImageRegionConstIterator< _TSliceFusion > vIt( vess, region );
    itk::ImageRegionIterator< _TSliceBinary > sIt( seeds, region );
    tIt.GoToBegin( );
    vIt.GoToBegin( );
    sIt.GoToBegin( );
    int len = 0;
    int numO = 0;
    int numB = 0;
    for( ; !tIt.IsAtEnd( ); ++tIt, ++vIt, ++sIt )
    {
      if( tIt.Get( ) != 0 )
      {
        len++;
        if( vIt.Get( ) > this->m_VesselnessValue )
        {
          sIt.Set( 1 );
          numO++;

        } // fi

      } // fi

      if( vIt.Get( ) < this->m_VesselnessValue )
      {
        sIt.Set( 2 );
        numB++;

      } // fi

    } // rof

    if( len == 0 )
      continue;

    // Random walker function
    typedef fpa::Functors::Dijkstra::Image::Gaussian< _TSliceData3D, _TSliceScalar > _TFunction;
    typename _TFunction::Pointer edge = _TFunction::New( );
    edge->SetBeta( this->m_Beta );
    edge->SetEpsilon( this->m_Epsilon );
    edge->TreatAsWeightOff( );

    // Real random walker
    typedef fpa::Common::RandomWalker< _TSliceData3D, _TSliceBinary, _TSliceScalar > _TRandomWalker;
    typename _TRandomWalker::Pointer rw = _TRandomWalker::New( );
    rw->SetInput( data );
    rw->SetInputLabels( seeds );
    rw->SetEdgeFunction( edge );

    // Extract label
    typedef itk::BinaryThresholdImageFilter< _TSliceBinary, _TSliceBinary > _TExtract;
    typename _TExtract::Pointer extract = _TExtract::New( );
    extract->SetInput( rw->GetOutput( ) );
    extract->SetInsideValue( 1 );
    extract->SetOutsideValue( 0 );
    extract->SetLowerThreshold( 1 );
    extract->SetUpperThreshold( 1 );

    // Keep maximum
    typedef itk::MaximumImageFilter< _TSliceBinary > _TMax;
    typename _TMax::Pointer max = _TMax::New();
    max->SetInput( 0, extract->GetOutput( ) );
    max->SetInput( 1, next );
    max->Update( );
    binaryTree[ z + o ] = max->GetOutput( );
    binaryTree[ z + o ]->DisconnectPipeline( );

  } // elihw
}

#endif // __fpa__Common__SliceBySliceRandomWalker__hxx__
// eof - $RCSfile$