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

  Program:   Insight Segmentation & Registration Toolkit
  Module:    $RCSfile: itkMeanSquaresImageToImageMetric2.txx,v $
  Language:  C++
  Date:      $Date: 2003/10/09 19:44:40 $
  Version:   $Revision: 1.45 $

  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 _itkMeanSquaresImageToImageMetric2_txx
#define _itkMeanSquaresImageToImageMetric2_txx

#include "itkMeanSquaresImageToImageMetric2.h"
#include "itkImageRegionConstIteratorWithIndex.h"
#include "itkBSplineDeformableTransform.h"

namespace itk
{

/*
 * Constructor
 */
template <class TFixedImage, class TMovingImage> 
MeanSquaresImageToImageMetric2<TFixedImage,TMovingImage>
::MeanSquaresImageToImageMetric2()
{
  itkDebugMacro("Constructor");

  m_ComputeGradient = false; // don't use the default gradient for now

  m_TransformIsBSpline    = false;
  typename BSplineTransformType::Pointer transformer = BSplineTransformType::New();
  this->SetTransform (transformer);
  m_BSplineTransform = NULL;

}

/*
 * Initialize
 */
template <class TFixedImage, class TMovingImage> 
void
MeanSquaresImageToImageMetric2<TFixedImage,TMovingImage>
::Initialize(void) throw ( ExceptionObject )
{
  this->Superclass::Initialize();

  // plug-in derivative calculator
  m_DerivativeCalculator = DerivativeFunctionType::New();
  m_DerivativeCalculator->SetInputImage( m_MovingImage );


  // Check if transform is BSpline
  m_TransformIsBSpline = true;

  BSplineTransformType * testPtr2 = dynamic_cast<BSplineTransformType *>(
    m_Transform.GetPointer() );
  if( !testPtr2 )
    {
    m_TransformIsBSpline = false;
    m_BSplineTransform = NULL;
    itkDebugMacro( "Transform is not BSplineDeformable" );
    }
  else
    {
    m_BSplineTransform = testPtr2;
    itkDebugMacro( "Transform is BSplineDeformable" );
    }

}


/*
 * Get the match Measure
 */
template <class TFixedImage, class TMovingImage> 
typename MeanSquaresImageToImageMetric2<TFixedImage,TMovingImage>::MeasureType
MeanSquaresImageToImageMetric2<TFixedImage,TMovingImage>
::GetValue( const TransformParametersType & parameters ) const
{

  itkDebugMacro("GetValue( " << parameters << " ) ");

  FixedImageConstPointer fixedImage = this->GetFixedImage();

  if( !fixedImage ) 
    {
    itkExceptionMacro( << "Fixed image has not been assigned" );
    }

  typedef  itk::ImageRegionConstIteratorWithIndex<FixedImageType> FixedIteratorType;


  FixedIteratorType ti( fixedImage, this->GetFixedImageRegion() );

  typename FixedImageType::IndexType index;

  MeasureType measure = NumericTraits< MeasureType >::Zero;

  m_NumberOfPixelsCounted = 0;

  this->SetTransformParameters( parameters );

  while(!ti.IsAtEnd())
    {

    index = ti.GetIndex();
    
    typename Superclass::InputPointType inputPoint;
    fixedImage->TransformIndexToPhysicalPoint( index, inputPoint );

    if( m_FixedImageMask && !m_FixedImageMask->IsInside( inputPoint ) )
      {
      ++ti;
      continue;
      }

    typename Superclass::OutputPointType transformedPoint = m_Transform->TransformPoint( inputPoint );

    if( m_MovingImageMask && !m_MovingImageMask->IsInside( transformedPoint ) )
      {
      ++ti;
      continue;
      }

    if( m_Interpolator->IsInsideBuffer( transformedPoint ) )
      {
      const RealType movingValue  = m_Interpolator->Evaluate( transformedPoint );
      const RealType fixedValue   = ti.Get();
      m_NumberOfPixelsCounted++;
      const RealType diff = movingValue - fixedValue; 
      measure += diff * diff; 
      }

    ++ti;
    }

  if( !m_NumberOfPixelsCounted )
    {
    itkExceptionMacro(<<"All the points mapped to outside of the moving image");
    }
  else
    {
    measure /= m_NumberOfPixelsCounted;
    }

  return measure;

}





/*
 * Get the Derivative Measure
 */
template < class TFixedImage, class TMovingImage> 
void
MeanSquaresImageToImageMetric2<TFixedImage,TMovingImage>
::GetDerivative( const TransformParametersType & parameters,
                 DerivativeType & derivative  ) const
{
  MeasureType value;
  // call the combined version
  this->GetValueAndDerivative( parameters, value, derivative );
}


/*
 * Get both the match Measure and theDerivative Measure 
 */
template <class TFixedImage, class TMovingImage> 
void
MeanSquaresImageToImageMetric2<TFixedImage,TMovingImage>
::GetValueAndDerivative(const TransformParametersType & parameters, 
                        MeasureType & value, DerivativeType  & derivative) const
{

  itkDebugMacro("GetValueAndDerivative( " << parameters << " ) ");

  FixedImageConstPointer fixedImage = this->GetFixedImage();

  if( !fixedImage ) 
    {
    itkExceptionMacro( << "Fixed image has not been assigned" );
    }

  const unsigned int ImageDimension = FixedImageType::ImageDimension;

  typedef  itk::ImageRegionConstIteratorWithIndex<
    FixedImageType> FixedIteratorType;

  typedef  itk::ImageRegionConstIteratorWithIndex<
    ITK_TYPENAME Superclass::GradientImageType> GradientIteratorType;


  FixedIteratorType ti( fixedImage, this->GetFixedImageRegion() );

  typename FixedImageType::IndexType index;

  MeasureType measure = NumericTraits< MeasureType >::Zero;

  m_NumberOfPixelsCounted = 0;

  this->SetTransformParameters( parameters );

  const unsigned int ParametersDimension = this->GetNumberOfParameters();
  derivative = DerivativeType( ParametersDimension );
  derivative.Fill( NumericTraits<ITK_TYPENAME DerivativeType::ValueType>::Zero );

  ti.GoToBegin();

  while(!ti.IsAtEnd())
    {

    index = ti.GetIndex();
    
    typename Superclass::InputPointType inputPoint;
    fixedImage->TransformIndexToPhysicalPoint( index, inputPoint );

    if( m_FixedImageMask && !m_FixedImageMask->IsInside( inputPoint ) )
      {
      ++ti;
      continue;
      }

    typename Superclass::OutputPointType transformedPoint;
    unsigned int numBSplineWeights;
    BSplineTransformWeightsType weights;
    BSplineTransformIndexArrayType indices;
    unsigned long numParametersPerDimension;
    bool isInside; 

    if ( !m_TransformIsBSpline )
      {
      // generic version: works for all transforms
      transformedPoint = m_Transform->TransformPoint( inputPoint );
      isInside = m_Interpolator->IsInsideBuffer( transformedPoint );
      }
    else
      {
      // optimized for BSplineDeformableTransform
      numBSplineWeights = m_BSplineTransform->GetNumberOfWeights();
      numParametersPerDimension = m_BSplineTransform->GetNumberOfParametersPerDimension();
      weights.SetSize( numBSplineWeights );
      indices.SetSize( numBSplineWeights );
      bool insideBSRegion;

      m_BSplineTransform->TransformPoint( inputPoint, transformedPoint,
                                          weights, indices, insideBSRegion );

      isInside = m_Interpolator->IsInsideBuffer( transformedPoint );
      isInside = isInside && insideBSRegion;

      }


    if( m_MovingImageMask && !m_MovingImageMask->IsInside( transformedPoint ) )
      {
      ++ti;
      continue;
      }

    if( isInside )
      {
      const RealType movingValue  = m_Interpolator->Evaluate( transformedPoint );
      const RealType fixedValue   = ti.Value();
      m_NumberOfPixelsCounted++;

      const RealType diff = movingValue - fixedValue; 
      measure += diff * diff;

      const GradientPixelType gradient = 
            m_DerivativeCalculator->Evaluate( transformedPoint );

      if ( !m_TransformIsBSpline )
        {
        // generic version: works for all transforms
        const TransformJacobianType & jacobian =
          m_Transform->GetJacobian( inputPoint ); 

        for(unsigned int par=0; par<ParametersDimension; par++)
          {
          RealType sum = NumericTraits< RealType >::Zero;
          for(unsigned int dim=0; dim<ImageDimension; dim++)
            {
            sum += 2.0 * diff * jacobian( dim, par ) * gradient[dim];
            } // for dim
          derivative[par] += sum;
          } // for par
        }
      else
        {
        // Optimized for BSplineDeformableTransform.
        // Only update the derivative of the parameters which is within
        // the BSpline support of this sample point.
        RealType sum = NumericTraits< RealType >::Zero;
        for ( unsigned int dim = 0; dim < FixedImageDimension; dim++ )
          {
          unsigned long offset = dim * numParametersPerDimension;
          for( unsigned int mu = 0; mu < numBSplineWeights; mu++ )
            {
            double value = 2.0 * diff * gradient[dim] * weights[mu];
            derivative[ indices[mu] + offset ] += value;
            } // for mu 
          } // for dim;
        }

      } // end if InsideBuffer

    ++ti;
    } // end while

  if( !m_NumberOfPixelsCounted )
    {
    itkExceptionMacro(<<"All the points mapped to outside of the moving image");
    }
  else
    {
    for(unsigned int i=0; i<ParametersDimension; i++)
      {
      derivative[i] /= m_NumberOfPixelsCounted;
      }
    measure /= m_NumberOfPixelsCounted;
    }

  value = measure;

}


} // end namespace itk


#endif