
#ifndef _ContrastEnhancementFilter_h
#define _ContrastEnhancementFilter_h

#include "itkImage.h"
#include "itkPoint.h"
#include "itkVector.h"

#include "itkMultiThreader.h"
#include "itkSimpleFastMutexLock.h"

#include "DynArray.h"

class ContrastEnhancementFilter: public itk::Object
{
public:

  // Standard ITK class typedefs
  typedef ContrastEnhancementFilter Self;
  typedef itk::SmartPointer<Self> Pointer;
  typedef itk::SmartPointer<const Self> ConstPointer;

  // Type of tumor enhancement
  // Note: vessel regions always appear enhanced
  typedef enum{RingEnhancement, UniformEnhancement, NoEnhancement}
    EnhancementMode;

  // Method for creation through the object factory
  itkNewMacro(Self);

  typedef itk::Image<float, 3> ProbabilityImageType;
  typedef ProbabilityImageType::Pointer ProbabilityImagePointer;
  typedef ProbabilityImageType::IndexType ProbabilityImageIndexType;
  typedef ProbabilityImageType::OffsetType ProbabilityImageOffsetType;
  typedef ProbabilityImageType::PixelType ProbabilityImagePixelType;
  typedef ProbabilityImageType::PointType ProbabilityImagePointType;
  typedef ProbabilityImageType::RegionType ProbabilityImageRegionType;
  typedef ProbabilityImageType::SizeType ProbabilityImageSizeType;
  typedef ProbabilityImageType::SpacingType ProbabilityImageSpacingType;

  typedef itk::Image<unsigned char, 3> ByteImageType;
  typedef ByteImageType::Pointer ByteImagePointer;
  typedef ByteImageType::IndexType ByteImageIndexType;
  typedef ByteImageType::OffsetType ByteImageOffsetType;
  typedef ByteImageType::PixelType ByteImagePixelType;
  typedef ByteImageType::RegionType ByteImageRegionType;
  typedef ByteImageType::SizeType ByteImageSizeType;
  typedef ByteImageType::SpacingType ByteImageSpacingType;

  typedef itk::Image<unsigned int, 3> UIntImageType;
  typedef UIntImageType::Pointer UIntImagePointer;
  typedef UIntImageType::IndexType UIntImageIndexType;
  typedef UIntImageType::OffsetType UIntImageOffsetType;
  typedef UIntImageType::PixelType UIntImagePixelType;
  typedef UIntImageType::RegionType UIntImageRegionType;
  typedef UIntImageType::SizeType UIntImageSizeType;
  typedef UIntImageType::SpacingType UIntImageSpacingType;

  typedef itk::Point<float, 3> PointType;
  typedef itk::Vector<float, 3> VectorType;

  DynArray<ProbabilityImagePointer> GetOutput();

  void SetInputProbabilities(DynArray<ProbabilityImagePointer> plist);

  itkSetMacro(VesselProbability, ProbabilityImagePointer);

  itkGetConstMacro(MaximumIterations, unsigned int);
  itkSetMacro(MaximumIterations, unsigned int);

  itkGetConstMacro(TimeStep, float);
  itkSetMacro(TimeStep, float);

  itkGetConstMacro(GrowCoefficient, float);
  itkSetMacro(GrowCoefficient, float);

  itkGetConstMacro(DeathCoefficient, float);
  itkSetMacro(DeathCoefficient, float);

  itkGetConstMacro(Mode, EnhancementMode);
  itkSetMacro(Mode, EnhancementMode);

  itkGetConstMacro(NumberOfThreads, unsigned int);
  itkSetMacro(NumberOfThreads, unsigned int);

  void Update();

protected:

  ContrastEnhancementFilter();
  ~ContrastEnhancementFilter();

  void CheckInput();

  void Initialize();

  float Step();

  float ThreadedStep();
  static ITK_THREAD_RETURN_TYPE _stepThread(void* arg);

  void Finish();

  void ComputeSourceSinkMask();

private:

  bool m_Modified;

  DynArray<ProbabilityImagePointer> m_InputProbabilities;
  DynArray<ProbabilityImagePointer> m_OutputProbabilities;

  ProbabilityImagePointer m_VesselProbability;

  unsigned int m_MaximumIterations;
  float m_TimeStep;

  float m_BlurVariance;

  float m_SeedCountFraction;

  unsigned int m_DilationRadius;

  float m_InputVolume;
  float m_VolumeFraction;

  ByteImagePointer m_SourceMask;
  ByteImagePointer m_SinkMask;

  ProbabilityImagePointer m_DiffusionCoeffImage;
  ProbabilityImagePointer m_Phi;
  ProbabilityImagePointer m_NewPhi;

  EnhancementMode m_Mode;

  float m_GrowCoefficient;
  float m_DeathCoefficient;

  DynArray<ProbabilityImageRegionType> m_SplitRegions;

  itk::SimpleFastMutexLock m_Mutex;

  unsigned int m_NumberOfThreads;

};

#endif
