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

  Program:   Insight Segmentation & Registration Toolkit
  Module:    $RCSfile: itkFEMImageMetricLoad.h,v $
  Language:  C++
  Date:      $Date: 2009-01-29 21:28:16 $
  Version:   $Revision: 1.25 $

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm 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 __itkFEMImageMetricLoad_h
#define __itkFEMImageMetricLoad_h

#include "itkFEMLoadElementBase.h"

#include "itkImage.h"
#include "itkTranslationTransform.h"

#include "itkImageRegionIteratorWithIndex.h"
#include "itkNeighborhoodIterator.h"
#include "itkNeighborhoodIterator.h"
#include "itkNeighborhoodInnerProduct.h"
#include "itkDerivativeOperator.h"
#include "itkForwardDifferenceOperator.h"
#include "itkLinearInterpolateImageFunction.h"
#include "vnl/vnl_math.h"

#include <itkMutualInformationImageToImageMetric.h>
#include <itkMattesMutualInformationImageToImageMetric.h>
#include <itkMeanSquaresImageToImageMetric.h>
#include <itkNormalizedCorrelationImageToImageMetric.h>
//#include <itkMeanReciprocalSquareDifferenceImageToImageMetric.h>


namespace itk 
{
namespace fem
{

/**
 * \class ImageMetricLoad
 * \brief General image pair load that uses the itkImageToImageMetrics.
 *
 * LoadImageMetric computes FEM gravity loads by using derivatives provided 
 * by itkImageToImageMetrics (e.g. mean squares intensity difference.)
 * The function responsible for this is called Fg, as required by the FEMLoad
 * standards.  It takes a vnl_vector as input.
 * We assume the vector input is of size 2*ImageDimension.
 * The 0 to ImageDimension-1 elements contain the position, p,
 * in the reference (moving) image.  The next ImageDimension to 2*ImageDimension-1
 * elements contain the value of the vector field at that point, v(p).
 *
 * Then, we evaluate the derivative at the point p+v(p) with respect to
 * some region of the target (fixed) image by calling the metric with 
 * the translation parameters as provided by the vector field at p.
 * The metrics return both a scalar similarity value and vector-valued derivative.  
 * The derivative is what gives us the force to drive the FEM registration.
 * These values are computed with respect to some region in the Fixed image.
 * This region size may be set by the user by calling SetMetricRadius.
 * As the metric derivative computation evolves, performance should improve
 * and more functionality will be available (such as scale selection).
 */ 
template<class TMoving,class TFixed> 
class ImageMetricLoad : public LoadElement
{
  FEM_CLASS(ImageMetricLoad,LoadElement)
public:

  // Necessary typedefs for dealing with images BEGIN
  typedef typename LoadElement::Float Float;

  typedef TMoving                            MovingType;
  typedef typename MovingType::ConstPointer  MovingConstPointer;
  typedef MovingType*                        MovingPointer;
  typedef TFixed                             FixedType;
  typedef FixedType*                         FixedPointer;
  typedef typename FixedType::ConstPointer   FixedConstPointer;

  /** Dimensionality of input and output data is assumed to be the same. */
  itkStaticConstMacro(ImageDimension, unsigned int,
                      MovingType::ImageDimension);

  typedef ImageRegionIteratorWithIndex<MovingType> RefRegionIteratorType; 
  typedef ImageRegionIteratorWithIndex<FixedType>    TarRegionIteratorType; 
  

  typedef NeighborhoodIterator<MovingType> 
                                     MovingNeighborhoodIteratorType; 
  typedef typename MovingNeighborhoodIteratorType::IndexType  
                                     MovingNeighborhoodIndexType;
  typedef typename MovingNeighborhoodIteratorType::RadiusType 
                                     MovingRadiusType;
  typedef NeighborhoodIterator<FixedType> 
                                     FixedNeighborhoodIteratorType; 
  typedef typename FixedNeighborhoodIteratorType::IndexType  
                                     FixedNeighborhoodIndexType;
  typedef typename FixedNeighborhoodIteratorType::RadiusType 
                                     FixedRadiusType;


// IMAGE DATA
  typedef   typename  MovingType::PixelType   RefPixelType;
  typedef   typename  FixedType::PixelType    TarPixelType;
  typedef   Float                             PixelType;
  typedef   Float                             ComputationType;
  typedef   Image< RefPixelType, itkGetStaticConstMacro(ImageDimension) >       RefImageType;
  typedef   Image< TarPixelType, itkGetStaticConstMacro(ImageDimension) >       TarImageType;
  typedef   Image< PixelType, itkGetStaticConstMacro(ImageDimension) >            ImageType;
  typedef   vnl_vector<Float>                             VectorType;

// Necessary typedefs for dealing with images END
 
//------------------------------------------------------------
// Set up the metrics
//------------------------------------------------------------
  typedef double                   CoordinateRepresentationType;
  typedef Transform< CoordinateRepresentationType,itkGetStaticConstMacro(ImageDimension), itkGetStaticConstMacro(ImageDimension) > TransformBaseType;
  typedef TranslationTransform<CoordinateRepresentationType,  itkGetStaticConstMacro(ImageDimension) >  DefaultTransformType;

 /**  Type of supported metrics. */
  typedef   ImageToImageMetric<FixedType,MovingType > MetricBaseType;
  typedef typename MetricBaseType::Pointer            MetricBaseTypePointer;

  typedef   MutualInformationImageToImageMetric<  MovingType, FixedType   > MutualInformationMetricType;

  typedef   MeanSquaresImageToImageMetric< MovingType, FixedType   > MeanSquaresMetricType;

  typedef   NormalizedCorrelationImageToImageMetric< MovingType, FixedType  > NormalizedCorrelationMetricType;

  typedef  MeanSquaresMetricType                                DefaultMetricType;
  typedef typename DefaultTransformType::ParametersType         ParametersType;
  typedef typename DefaultTransformType::JacobianType           JacobianType;


//------------------------------------------------------------
// Set up an Interpolator
//------------------------------------------------------------
  typedef LinearInterpolateImageFunction< MovingType, double > InterpolatorType;

  /** Gradient filtering */
  typedef float RealType;
  typedef CovariantVector<RealType,
          itkGetStaticConstMacro(ImageDimension)> GradientPixelType;
  typedef Image<GradientPixelType,
               itkGetStaticConstMacro(ImageDimension)> GradientImageType;
  typedef SmartPointer<GradientImageType>     GradientImagePointer;
  typedef GradientRecursiveGaussianImageFilter< ImageType,
                                                GradientImageType >
          GradientImageFilterType;  
  //  typedef typename GradientImageFilterType::Pointer GradientImageFilterPointer;


// FUNCTIONS

  /** Set/Get the Metric.  */
  void SetMetric(MetricBaseTypePointer MP)
    { m_Metric=MP; }
  
 /** Define the reference (moving) image. */
  void SetMovingImage(MovingType* R)
    { 
    m_RefImage = R; 
    m_RefSize=m_RefImage->GetLargestPossibleRegion().GetSize();
    }

  void SetMetricMovingImage(MovingType* R)  
    { 
    m_Metric->SetMovingImage( R ); 
    m_RefSize=R->GetLargestPossibleRegion().GetSize(); 
    }

  /** Define the target (fixed) image. */ 
  void SetFixedImage(FixedType* T)
    { 
    m_TarImage=T; 
    m_TarSize=T->GetLargestPossibleRegion().GetSize(); 
    }
  void SetMetricFixedImage(FixedType* T)  
    { 
    m_Metric->SetFixedImage( T );
    m_TarSize=T->GetLargestPossibleRegion().GetSize(); 
    }

  MovingPointer GetMovingImage()
    { return m_RefImage; }
  FixedPointer GetFixedImage() { return m_TarImage; }

  /** Define the metric region size. */ 
  void SetMetricRadius(MovingRadiusType T) {m_MetricRadius  = T; }   
  /** Get the metric region size. */ 
  MovingRadiusType GetMetricRadius() { return m_MetricRadius; }
  
  /** Set/Get methods for the number of integration points to use 
   * in each 1-dimensional line integral when evaluating the load.
   * This value is passed to the load implementation.
   */
  void SetNumberOfIntegrationPoints(unsigned int i)
    { m_NumberOfIntegrationPoints=i;}
  unsigned int GetNumberOfIntegrationPoints()
    { return m_NumberOfIntegrationPoints;}

  /** Set the direction of the gradient (uphill or downhill). 
    * E.g. the mean squares metric should be minimized while NCC and PR should be maximized.
    */ 
  void SetSign(Float s)
    {m_Sign=s;}
  
  /** Set the sigma in a gaussian measure. */
  void SetTemp(Float s)
    {m_Temp=s;}

  /** Scaling of the similarity energy term */
  void SetGamma(Float s)
    {m_Gamma=s;}

  void SetSolution(Solution::ConstPointer ptr)
    {  m_Solution=ptr; }
  Solution::ConstPointer GetSolution()
    { return m_Solution; }

  /**
   *  This method returns the total metric evaluated over the image with respect to the current solution.
   */
  Float GetMetric (VectorType  InVec);
  VectorType GetPolynomialFitToMetric(VectorType PositionInElement, VectorType SolutionAtPosition);

  VectorType MetricFiniteDiff(VectorType PositionInElement, VectorType SolutionAtPosition);

  // FIXME - WE ASSUME THE 2ND VECTOR (INDEX 1) HAS THE INFORMATION WE WANT
  Float GetSolution(unsigned int i,unsigned int which=0)
    {  
    return m_Solution->GetSolutionValue(i,which); 
    }
  
// define the copy constructor 
//  ImageMetricLoad(const ImageMetricLoad& LMS);

  void InitializeMetric(void);
  ImageMetricLoad(); // cannot be private until we always use smart pointers
  Float EvaluateMetricGivenSolution ( Element::ArrayType* el, Float step=1.0);
 
  /**
   * Compute the image based load - implemented with ITK metric derivatives.
   */
  VectorType Fe1(VectorType);
  VectorType Fe(VectorType,VectorType);
 
  static Baseclass* NewImageMetricLoad(void)
    { return new ImageMetricLoad; }

  /** Set/Get the metric gradient image */
  //void InitializeGradientImage();
  void SetMetricGradientImage(GradientImageType* g)
    { m_MetricGradientImage=g;}
  GradientImageType* GetMetricGradientImage()
    { return  m_MetricGradientImage;}


  void PrintCurrentEnergy()
    { std:: cout << " energy " << m_Energy << std::endl;}
  double GetCurrentEnergy()
    { return m_Energy; }
  void  SetCurrentEnergy( double e )
    { m_Energy=e; }

protected:


private:
  GradientImageType*            m_MetricGradientImage;
  MovingPointer                 m_RefImage;
  FixedPointer                  m_TarImage;
  MovingRadiusType              m_MetricRadius; /** used by the metric to set region size for fixed image*/ 
  typename MovingType::SizeType m_RefSize;
  typename FixedType::SizeType  m_TarSize;
  unsigned int                  m_NumberOfIntegrationPoints;
  unsigned int                  m_SolutionIndex;
  unsigned int                  m_SolutionIndex2;
  Float                         m_Sign;
  Float                         m_Temp;
  Float                         m_Gamma;

  typename Solution::ConstPointer     m_Solution;
  MetricBaseTypePointer               m_Metric;
  typename TransformBaseType::Pointer m_Transform;
  typename InterpolatorType::Pointer  m_Interpolator;

  mutable double                       m_Energy;
private:
  /** Dummy static int that enables automatic registration
      with FEMObjectFactory. */
  static const int m_DummyCLID;

};

}} // end namespace fem/itk

#ifndef ITK_MANUAL_INSTANTIATION
#include "itkFEMImageMetricLoad.txx"
#endif

#endif
