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

#include <itkImageToImageFilter.h>
#include <fpa/Config.h>
#include <fpa/DataStructures/Image/Path.h>

/**
 */
template< class _TInputImage, class _TOutputImage >
class RandomWalkSegmentation
  : public itk::ImageToImageFilter< _TInputImage, _TOutputImage >
{
public:
  typedef _TInputImage  TInputImage;
  typedef _TOutputImage TOutputImage;

  typedef itk::ImageToImageFilter< TInputImage, TOutputImage > Superclass;
  typedef RandomWalkSegmentation        Self;
  typedef itk::SmartPointer< Self >       Pointer;
  typedef itk::SmartPointer< const Self > ConstPointer;

  typedef typename TInputImage::IndexType TIndex;
  typedef typename TInputImage::PointType TPoint;

  struct TSeed
  {
    TIndex Index;
    TPoint Point;
    bool IsPoint;
  };
  typedef fpa::DataStructures::Image::Path< TInputImage::ImageDimension > TPath;

public:
  itkNewMacro( Self );
  itkTypeMacro( RandomWalkSegmentation, itk::ImageToImageFilter );

  itkGetConstMacro( StartSeed, TSeed );
  itkGetConstMacro( EndSeed, TSeed );

  itkGetConstMacro( Beta, double );
  itkSetMacro( Beta, double );

  itkGetConstMacro( Sigma, double );
  itkSetMacro( Sigma, double );

  itkGetConstMacro( Radius, double );
  itkSetMacro( Radius, double );

  fpaFilterOutputMacro( OutputAxis, TPath );

public:
  void AddSeed( const TIndex& s );
  void AddSeed( const TPoint& s );
  void ClearSeeds( );
  unsigned long GetNumberOfSeeds( ) const;

  void SetStartSeed( const TIndex& s );
  void SetStartSeed( const TPoint& s );

  void SetEndSeed( const TIndex& s );
  void SetEndSeed( const TPoint& s );

protected:
  RandomWalkSegmentation( );
  virtual ~RandomWalkSegmentation( );

  virtual void GenerateData( ) override;

private:
  template< class _TIn, class _TSeed >
  void _SynchSeed( const _TIn* in, _TSeed& seed );

  template< class _TIn, class _TOutPtr >
  void _Smooth( const _TIn* in, _TOutPtr& out );

  template< class _TIn, class _TOutPtr, class _TAxisPtr, class _TSeeds >
  typename _TIn::RegionType _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
    );

  template< class _TIn, class _TOutPtr >
  typename _TIn::RegionType _ROI(
    const _TIn* in, const typename _TIn::RegionType& roi, unsigned int pad,
    _TOutPtr& out
    );

  template< class _TAxisPtr, class _TRegion >
  void _AxisROI(
    const typename _TAxisPtr::ObjectType* in, const _TRegion& roi,
    _TAxisPtr& out
    );

  template< class _TInRaw, class _TInCosts, class _TAxis, class _TScalar, class _TOutPtr >
  void _Label(
    const _TInRaw* raw, const _TInCosts* costs, const _TAxis* axis,
    const _TScalar& mean, const _TScalar& dev, const _TScalar& radius,
    _TOutPtr& out
    );

  template< class _TIn, class _TOutPtr >
  void _RandomWalker(
    const _TIn* in, const typename _TOutPtr::ObjectType* labels,
    const typename _TIn::PixelType& beta,
    _TOutPtr& out
    );

  template< class _TIn, class _TSeeds, class _TSeed, class _TOutPtr, class _TAxisPtr >
  void _DistanceAndAxis(
    const _TIn* in, const _TSeeds& seeds,
    const _TSeed& p0, const _TSeed& p1,
    _TOutPtr& out_dist, _TAxisPtr& out_axis
    );

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

protected:
  TSeed m_StartSeed;
  TSeed m_EndSeed;
  std::vector< TSeed > m_Seeds;

  double m_Beta;
  double m_Sigma;
  double m_Radius;
};

#ifndef ITK_MANUAL_INSTANTIATION
#  include "RandomWalkSegmentation.hxx"
#endif // ITK_MANUAL_INSTANTIATION

#endif // __RandomWalkSegmentation__h__

// eof - $RCSfile$