/*=========================================================================

  Program:   Advanced Normalization Tools

  Copyright (c) ConsortiumOfANTS. All rights reserved.
  See accompanying COPYING.txt or
 https://github.com/stnava/ANTs/blob/master/ANTSCopyright.txt for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#ifndef _SurfaceImageCurvature_h
#define _SurfaceImageCurvature_h

#include "itkNeighborhoodIterator.h"

#include "itkSurfaceCurvatureBase.h"
#include "itkGradientRecursiveGaussianImageFilter.h"
#include "itkGradientImageFilter.h"
#include "itkVectorLinearInterpolateImageFunction.h"

namespace itk
{
/** \class SurfaceImageCurvature
 *
 * This class takes a surface as input and creates a local
 * geometric frame for each surface point.
 *
 *
 */
template <typename TSurface>
class SurfaceImageCurvature :
  public SurfaceCurvatureBase<TSurface,
                              3>
{
public:

  /** Standard class typedefs. */
  typedef SurfaceImageCurvature          Self;
  typedef SurfaceCurvatureBase<TSurface> Superclass;
  typedef SmartPointer<Self>             Pointer;
  typedef SmartPointer<const Self>       ConstPointer;

  /** Run-time type information (and related methods). */
  itkTypeMacro(SurfaceImageCurvature, SurfaceCurvatureBase);

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

  /** Image related types. */
  typedef typename TSurface::PixelType PixelType;
  enum { ImageDimension = TSurface::ImageDimension };
  typedef Image<PixelType, itkGetStaticConstMacro(ImageDimension)>
    ImageType;
  typedef typename ImageType::IndexType           IndexType;
  typedef typename ImageType::SizeType            SizeType;
  typedef ImageRegionIteratorWithIndex<ImageType> ImageIteratorType;
  /** Image dimension. */
  itkStaticConstMacro(SurfaceDimension, unsigned int, TSurface::ImageDimension);

  typedef typename Superclass::RealType   RealType;
  typedef typename Superclass::PointType  VectorType;
  typedef typename Superclass::PointType  FixedVectorType;
  typedef typename Superclass::PointType  PointType;
  typedef typename Superclass::MatrixType MatrixType;
  typedef typename ImageType::PointType ImagePointType;

  typedef  Image<PixelType, itkGetStaticConstMacro(ImageDimension)>
    OutputImageType;
  typedef ImageRegionIteratorWithIndex<OutputImageType> OutputImageIteratorType;

  typedef typename OutputImageType::Pointer OutputImagePointer;

  typedef Image<MatrixType, itkGetStaticConstMacro(ImageDimension)>
    FrameImageType;

  /** Gradient filtering */
  typedef CovariantVector<RealType,
                          itkGetStaticConstMacro(ImageDimension)> GradientPixelType;
  typedef Image<GradientPixelType,
                itkGetStaticConstMacro(ImageDimension)> GradientImageType;
  typedef itk::VectorLinearInterpolateImageFunction<GradientImageType, RealType> VectorInterpolatorType;
  typedef SmartPointer<GradientImageType> GradientImagePointer;
  typedef GradientRecursiveGaussianImageFilter<OutputImageType, GradientImageType>
    GradientImageFilterType;
  typedef GradientImageFilter<OutputImageType>
    GradientImageFilterType2;
  typedef typename GradientImageFilterType::Pointer GradientImageFilterPointer;

  typedef NeighborhoodIterator<ImageType> NeighborhoodIteratorType;

  /** Find all points within some distance of the origin.
    * The argument gives the number of times to apply the
    * mean shift algorithm to find the best neighborhood.
    */
  void FindNeighborhood(unsigned int numMeanShifts = 0) ITK_OVERRIDE;

  void  FindEuclideanNeighborhood(PointType p);

  void  FindGeodesicNeighborhood();

  /** This applies one of the algorithms for finding the local curvature
      and frame.  The default is joshi. */
  void ComputeFrameOverDomain(unsigned int which = 0) ITK_OVERRIDE;

  ImageType * GetInput();

  virtual void SetInputImage(typename ImageType::Pointer & input);
  OutputImageType * GetOutput();

  /** Apply the level set curvature equation over the whole image */
  void LevelSetMeanCurvature();

  /** Use the gradient of the image to estimate the normals everywhere.
   *  Also compute area, if boolean is set.
  */
  void EstimateNormalsFromGradient();

  /** Use the Weingarten map to estimate the curvature.*/
  void WeingartenMap();

  /** Use the Weingarten map to estimate the curvature.*/
  void WeingartenMapGradients();

  /** Computes a neighborhood surface area function everywhere*/
  void ComputeSurfaceArea();

  /** Use the gradient estimated normal to get the local frame.
      Requires call to SetNormal to find tangents. */
  void EstimateFrameFromGradient(IndexType);
  void EstimateFrameFromGradient(ImagePointType);

  /** Implemented version of virtual function from parent class.
      Here, we just sum the computed function, held in CurvatureImage,
      over a neighborhood */
  RealType IntegrateFunctionOverNeighborhood(bool norm = false) ITK_OVERRIDE;

  /** Get the neighborhood integral for every surface point.*/
  RealType IntegrateFunctionOverSurface(bool norm = false);

  /** Postprocess the curvature function by, e.g., gaussian
      smoothing of the curvature (and perhaps frame)
      in the local neighbhorhood. */
  void PostProcessGeometry();

  itkSetMacro(NeighborhoodRadius, RealType);
  itkGetMacro(NeighborhoodRadius, RealType);

  itkSetMacro(UseLabel, bool);
  itkGetMacro(UseLabel, bool);

  itkSetMacro(kSign, float);
  itkSetMacro(Sigma, float);

  itkSetMacro(Threshold, float);

  itkGetMacro(FunctionImage, OutputImagePointer);
  itkSetMacro(FunctionImage, OutputImagePointer);

  void ProcessLabelImage();

  float CurvatureAtIndex(IndexType index)
  {
    PointType p;

    for( unsigned int k = 0; k < ImageDimension; k++ )
      {
      p[k] = (RealType) index[k];
      }
    this->SetOrigin(p);
    this->EstimateFrameFromGradient(index);
    this->FindNeighborhood();
    this->WeingartenMap();
    float fval = this->m_GaussianKappa;
    float kpix;
    fval = this->m_MeanKappa;
    if( fabs(fval) > 1 )
      {
      fval /= fval;
      }
    kpix = kpix + m_kSign * fval;

    return kpix;
  }

  inline RealType innerProduct( PointType v1, PointType v2 )
    {
    return ( v1[0]*v2[0] + v1[1]*v2[1] + v1[2]*v2[2] );
    }

  inline bool IsValidIndex( IndexType ind )
  {
    for( unsigned int i = 0; i < ImageDimension; i++ )
      {
      float shifted = ind[i];
      if( shifted < 1 || shifted > ( this->m_ImageSize[i] - 1  ) )
        {
        return false;
        }
      }
    return true;
  }

  inline bool IsValidSurface(PixelType pix, IndexType ind )
  {
    for( unsigned int i = 0; i < ImageDimension; i++ )
      {
      float shifted = ind[i];
      if( shifted < 1 || shifted > ( this->m_ImageSize[i] - 1  ) )
        {
        return false;
        }
      }

    if( this->m_UseLabel )
      {
      if( pix == this->m_SurfaceLabel )
        {
        return true;
        }
      else
        {
        return false;
        }
      }
    else
      {
      if( pix > this->m_Threshold )
        {
        return true;
        }
      else
        {
        return false;
        }
      }
  }

protected:

  SurfaceImageCurvature();
  ~SurfaceImageCurvature()
  {
  };

  void CopyImageToFunctionImage( OutputImagePointer, OutputImagePointer);

  /** This function changes the values of the label image for use with
      the fast marching image filter. */
private:

  PixelType                m_SurfaceLabel;
  OutputImagePointer       m_FunctionImage;
  RealType                 m_NeighborhoodRadius;
  SizeType                 m_ImageSize;
  GradientImagePointer     m_GradientImage;
  NeighborhoodIteratorType m_ti;
  NeighborhoodIteratorType m_ti2;
  bool                     m_UseLabel;
  float                    m_kSign;
  float                    m_Sigma;
  float                    m_Threshold;
  float                    m_Area;
  RealType                 m_MinSpacing;
  typename VectorInterpolatorType::Pointer m_Vinterp;
};
} // namespace itk

#ifndef ITK_MANUAL_INSTANTIATION
#include "itkSurfaceImageCurvature.hxx"
#endif

#endif
