
#ifndef _BrainMeshGenerator_h
#define _BrainMeshGenerator_h

#include "itkBSplineInterpolateImageFunction.h"
#include "itkImage.h"

#include "vnl/vnl_vector.h"

#include "vtkPointLocator.h"
#include "vtkSmartPointer.h"
#include "vtkUnstructuredGrid.h"

#include "DynArray.h"

class BrainMeshGenerator
{

public:

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

  typedef itk::Image<float, 3> FloatImageType;
  typedef FloatImageType::Pointer FloatImagePointer;
  typedef FloatImageType::IndexType FloatImageIndexType;
  typedef FloatImageType::PointType FloatImagePointType;
  typedef FloatImageType::OffsetType FloatImageOffsetType;
  typedef FloatImageType::RegionType FloatImageRegionType;
  typedef FloatImageType::SizeType FloatImageSizeType;
  typedef FloatImageType::SpacingType FloatImageSpacingType;

  typedef itk::BSplineInterpolateImageFunction<FloatImageType, double>
    BSplineInterpolatorType;

  typedef itk::ContinuousIndex<double, 3> FloatImageContinuousIndexType;

  typedef vnl_vector<float> VectorType;
  typedef vnl_matrix<float> MatrixType;

  BrainMeshGenerator();
  ~BrainMeshGenerator() { }

  void SetInputProbabilities(const DynArray<FloatImagePointer>& probs);
  void SetLabelImage(const ByteImagePointer& labelImg);

  vtkSmartPointer<vtkUnstructuredGrid> GenerateMesh();

  float GetInitialSpacing() const { return m_InitialSpacing; }
  void SetInitialSpacing(float s) { m_InitialSpacing = s; }

  unsigned int GetMaximumIterations() const { return m_MaximumIterations; }
  void SetMaximumIterations(unsigned int n) { m_MaximumIterations = n; }

  float GetTimeStep() const { return m_TimeStep; }
  void SetTimeStep(float d) { m_TimeStep = d; }

protected:

  void ComputeInternalImages();

  float EvaluateHFunction(float x, float y, float z);
  float EvaluateInternalHFunction(float x, float y, float z, unsigned int i);

  // Support method for sequentially processing labels, process internal regions
  // individually
  void ProcessRegion(vtkPointLocator* pLoc, unsigned int which);

  vtkSmartPointer<vtkUnstructuredGrid> GenerateMeshAllAtOnce();
  vtkSmartPointer<vtkUnstructuredGrid> GenerateMeshSequentially();

  // Compute gradient of a float image (for distance image gradients)
  VectorType ComputeGradient(
    const FloatImageType* img, float x, float y, float z);

private:

  float m_InitialSpacing;

  float m_FScale;

  float m_TimeStep;

  unsigned int m_MaximumIterations;

  DynArray<FloatImagePointer> m_InputProbabilities;
  ByteImagePointer m_LabelImage;

  FloatImagePointer m_BrainDistanceImage;

  BSplineInterpolatorType::Pointer m_BrainDistanceInterpolator;
  DynArray<BSplineInterpolatorType::Pointer> m_InternalDistanceInterpolators;

  FloatImagePointer m_CurvatureImage;

  unsigned int m_MaxLabel;

  float m_PointTolerance;

  unsigned int m_SplineInterpolationOrder;

};

#endif
