/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         https://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef itkMutualInformationImageToImageMetric_hxx
#define itkMutualInformationImageToImageMetric_hxx

#include "itkImageRandomConstIteratorWithIndex.h"
#include "itkMath.h"
#include "itkGaussianKernelFunction.h"
#include "itkCompensatedSummation.h"

namespace itk
{

template <typename TFixedImage, typename TMovingImage>
MutualInformationImageToImageMetric<TFixedImage, TMovingImage>::MutualInformationImageToImageMetric()
{
  m_NumberOfSpatialSamples = 0;
  this->SetNumberOfSpatialSamples(50);

  m_KernelFunction = dynamic_cast<KernelFunctionType *>(GaussianKernelFunction<double>::New().GetPointer());

  m_FixedImageStandardDeviation = 0.4;
  m_MovingImageStandardDeviation = 0.4;

  m_MinProbability = 0.0001;

  //
  // Following initialization is related to
  // calculating image derivatives
  this->SetComputeGradient(false); // don't use the default gradient for now
  m_DerivativeCalculator = DerivativeFunctionType::New();
  m_DerivativeCalculator->UseImageDirectionOn();
}

template <typename TFixedImage, typename TMovingImage>
void
MutualInformationImageToImageMetric<TFixedImage, TMovingImage>::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
  os << indent << "NumberOfSpatialSamples: ";
  os << m_NumberOfSpatialSamples << std::endl;
  os << indent << "FixedImageStandardDeviation: ";
  os << m_FixedImageStandardDeviation << std::endl;
  os << indent << "MovingImageStandardDeviation: ";
  os << m_MovingImageStandardDeviation << std::endl;
  os << indent << "KernelFunction: ";
  os << m_KernelFunction.GetPointer() << std::endl;
}

template <typename TFixedImage, typename TMovingImage>
void
MutualInformationImageToImageMetric<TFixedImage, TMovingImage>::SetNumberOfSpatialSamples(unsigned int num)
{
  if (num == m_NumberOfSpatialSamples)
  {
    return;
  }

  this->Modified();

  // clamp to minimum of 1
  m_NumberOfSpatialSamples = ((num > 1) ? num : 1);

  // resize the storage vectors
  m_SampleA.resize(m_NumberOfSpatialSamples);
  m_SampleB.resize(m_NumberOfSpatialSamples);
}

template <typename TFixedImage, typename TMovingImage>
void
MutualInformationImageToImageMetric<TFixedImage, TMovingImage>::SampleFixedImageDomain(
  SpatialSampleContainer & samples) const
{
  using RandomIterator = ImageRandomConstIteratorWithIndex<FixedImageType>;
  RandomIterator randIter(this->m_FixedImage, this->GetFixedImageRegion());
  if (this->m_ReseedIterator)
  {
    randIter.ReinitializeSeed();
  }
  else
  {
    randIter.ReinitializeSeed(this->m_RandomSeed++);
  }
  randIter.SetNumberOfSamples(m_NumberOfSpatialSamples);
  randIter.GoToBegin();

  typename SpatialSampleContainer::iterator       iter;
  typename SpatialSampleContainer::const_iterator end = samples.end();

  bool allOutside = true;

  this->m_NumberOfPixelsCounted = 0; // Number of pixels that map into the
                                     // fixed and moving image mask, if
                                     // specified
                                     // and the resampled fixed grid after
                                     // transformation.

  // Number of random picks made from the portion of fixed image within the
  // fixed mask
  SizeValueType numberOfFixedImagePixelsVisited = 0;
  SizeValueType dryRunTolerance = this->GetFixedImageRegion().GetNumberOfPixels();
  for (iter = samples.begin(); iter != end; ++iter)
  {
    // Get sampled index
    FixedImageIndexType index = randIter.GetIndex();
    // Get sampled fixed image value
    iter->FixedImageValue = randIter.Get();
    // Translate index to point
    this->m_FixedImage->TransformIndexToPhysicalPoint(index, iter->FixedImagePointValue);

    // If not inside the fixed mask, ignore the point
    if (this->m_FixedImageMask && !this->m_FixedImageMask->IsInsideInWorldSpace(iter->FixedImagePointValue))
    {
      ++randIter; // jump to another random position
      continue;
    }

    if (allOutside)
    {
      ++numberOfFixedImagePixelsVisited;
      if (numberOfFixedImagePixelsVisited > dryRunTolerance)
      {
        // We randomly visited as many points as is the size of the fixed image
        // region.. Too may samples mapped outside.. go change your transform
        itkExceptionMacro("Too many samples mapped outside the moving buffer");
      }
    }

    MovingImagePointType mappedPoint = this->m_Transform->TransformPoint(iter->FixedImagePointValue);

    // If the transformed point after transformation does not lie within the
    // MovingImageMask, skip it.
    if (this->m_MovingImageMask && !this->m_MovingImageMask->IsInsideInWorldSpace(mappedPoint))
    {
      ++randIter;
      continue;
    }

    // The interpolator does not need to do bounds checking if we have masks,
    // since we know that the point is within the fixed and moving masks. But
    // a crazy user can specify masks that are bigger than the image. Then we
    // will need bounds checking.. So keep this anyway.
    if (this->m_Interpolator->IsInsideBuffer(mappedPoint))
    {
      iter->MovingImageValue = this->m_Interpolator->Evaluate(mappedPoint);
      this->m_NumberOfPixelsCounted++;
      allOutside = false;
    }
    else
    {
      iter->MovingImageValue = 0;
    }

    // Jump to random position
    ++randIter;
  }

  if (allOutside)
  {
    // if all the samples mapped to the outside throw an exception
    itkExceptionMacro("All the sampled point mapped to outside of the moving image");
  }
}

template <typename TFixedImage, typename TMovingImage>
auto
MutualInformationImageToImageMetric<TFixedImage, TMovingImage>::GetValue(const ParametersType & parameters) const
  -> MeasureType
{
  // make sure the transform has the current parameters
  this->m_Transform->SetParameters(parameters);

  // collect sample set A
  this->SampleFixedImageDomain(m_SampleA);

  // collect sample set B
  this->SampleFixedImageDomain(m_SampleB);

  // calculate the mutual information

  using SumType = CompensatedSummation<double>;
  SumType dLogSumFixed;
  SumType dLogSumMoving;
  SumType dLogSumJoint;
  SumType dSumFixed;
  SumType dSumMoving;
  SumType dSumJoint;

  typename SpatialSampleContainer::const_iterator aiter;
  typename SpatialSampleContainer::const_iterator aend = m_SampleA.end();
  typename SpatialSampleContainer::const_iterator biter;
  typename SpatialSampleContainer::const_iterator bend = m_SampleB.end();
  for (biter = m_SampleB.begin(); biter != bend; ++biter)
  {
    dSumFixed.ResetToZero();
    dSumMoving.ResetToZero();
    dSumJoint.ResetToZero();
    dSumFixed += m_MinProbability;
    dSumMoving += m_MinProbability;
    dSumJoint += m_MinProbability;
    for (aiter = m_SampleA.begin(); aiter != aend; ++aiter)
    {
      double valueFixed;
      double valueMoving;

      valueFixed = (biter->FixedImageValue - aiter->FixedImageValue) / m_FixedImageStandardDeviation;
      valueFixed = m_KernelFunction->Evaluate(valueFixed);

      valueMoving = (biter->MovingImageValue - aiter->MovingImageValue) / m_MovingImageStandardDeviation;
      valueMoving = m_KernelFunction->Evaluate(valueMoving);

      dSumFixed += valueFixed;
      dSumMoving += valueMoving;
      dSumJoint += valueFixed * valueMoving;
    } // end of sample A loop

    if (dSumFixed.GetSum() > 0.0)
    {
      dLogSumFixed -= std::log(dSumFixed.GetSum());
    }
    if (dSumMoving.GetSum() > 0.0)
    {
      dLogSumMoving -= std::log(dSumMoving.GetSum());
    }
    if (dSumJoint.GetSum() > 0.0)
    {
      dLogSumJoint -= std::log(dSumJoint.GetSum());
    }
  } // end of sample B loop

  auto nsamp = static_cast<double>(m_NumberOfSpatialSamples);

  double threshold = -0.5 * nsamp * std::log(m_MinProbability);
  if (dLogSumMoving.GetSum() > threshold || dLogSumFixed.GetSum() > threshold || dLogSumJoint.GetSum() > threshold)
  {
    // at least half the samples in B did not occur within
    // the Parzen window width of samples in A
    itkExceptionMacro("Standard deviation is too small");
  }

  MeasureType measure = dLogSumFixed.GetSum() + dLogSumMoving.GetSum() - dLogSumJoint.GetSum();
  measure /= nsamp;
  measure += std::log(nsamp);

  return measure;
}

template <typename TFixedImage, typename TMovingImage>
void
MutualInformationImageToImageMetric<TFixedImage, TMovingImage>::GetValueAndDerivative(const ParametersType & parameters,
                                                                                      MeasureType &          value,
                                                                                      DerivativeType & derivative) const
{
  value = MeasureType{};
  unsigned int   numberOfParameters = this->m_Transform->GetNumberOfParameters();
  DerivativeType temp(numberOfParameters);
  temp.Fill(0);
  derivative = temp;

  // make sure the transform has the current parameters
  this->m_Transform->SetParameters(parameters);

  // set the DerivativeCalculator
  m_DerivativeCalculator->SetInputImage(this->m_MovingImage);

  // collect sample set A
  this->SampleFixedImageDomain(m_SampleA);

  // collect sample set B
  this->SampleFixedImageDomain(m_SampleB);

  // calculate the mutual information
  using SumType = CompensatedSummation<double>;
  SumType dLogSumFixed;
  SumType dLogSumMoving;
  SumType dLogSumJoint;
  SumType dSumFixed;
  SumType dDenominatorMoving;
  SumType dDenominatorJoint;

  typename SpatialSampleContainer::iterator       aiter;
  typename SpatialSampleContainer::const_iterator aend = m_SampleA.end();
  typename SpatialSampleContainer::iterator       biter;
  typename SpatialSampleContainer::const_iterator bend = m_SampleB.end();

  // precalculate all the image derivatives for sample A
  using DerivativeContainer = std::vector<DerivativeType>;
  DerivativeContainer sampleADerivatives;
  sampleADerivatives.resize(m_NumberOfSpatialSamples);

  typename DerivativeContainer::iterator aditer;
  DerivativeType                         tempDeriv(numberOfParameters);
  TransformJacobianType                  jacobian(numberOfParameters, numberOfParameters);
  for (aiter = m_SampleA.begin(), aditer = sampleADerivatives.begin(); aiter != aend; ++aiter, ++aditer)
  {
    /** FIXME: is there a way to avoid the extra copying step? */
    this->CalculateDerivatives(aiter->FixedImagePointValue, tempDeriv, jacobian);
    (*aditer) = tempDeriv;
  }

  DerivativeType derivB(numberOfParameters);
  for (biter = m_SampleB.begin(); biter != bend; ++biter)
  {
    dDenominatorMoving.ResetToZero();
    dDenominatorMoving += m_MinProbability;
    dDenominatorJoint.ResetToZero();
    dDenominatorJoint += m_MinProbability;
    dSumFixed.ResetToZero();
    dSumFixed += m_MinProbability;
    for (aiter = m_SampleA.begin(); aiter != aend; ++aiter)
    {
      double valueFixed;
      double valueMoving;

      valueFixed = (biter->FixedImageValue - aiter->FixedImageValue) / m_FixedImageStandardDeviation;
      valueFixed = m_KernelFunction->Evaluate(valueFixed);

      valueMoving = (biter->MovingImageValue - aiter->MovingImageValue) / m_MovingImageStandardDeviation;
      valueMoving = m_KernelFunction->Evaluate(valueMoving);

      dDenominatorMoving += valueMoving;
      dDenominatorJoint += valueMoving * valueFixed;

      dSumFixed += valueFixed;
    } // end of sample A loop

    if (dSumFixed.GetSum() > 0.0)
    {
      dLogSumFixed -= std::log(dSumFixed.GetSum());
    }
    if (dDenominatorMoving.GetSum() > 0.0)
    {
      dLogSumMoving -= std::log(dDenominatorMoving.GetSum());
    }
    if (dDenominatorJoint.GetSum() > 0.0)
    {
      dLogSumJoint -= std::log(dDenominatorJoint.GetSum());
    }

    /** get the image derivative for this B sample */
    this->CalculateDerivatives(biter->FixedImagePointValue, derivB, jacobian);

    SumType totalWeight;
    for (aiter = m_SampleA.begin(), aditer = sampleADerivatives.begin(); aiter != aend; ++aiter, ++aditer)
    {
      double valueFixed;
      double valueMoving;
      double weightMoving;
      double weightJoint;
      double weight;

      valueFixed = (biter->FixedImageValue - aiter->FixedImageValue) / m_FixedImageStandardDeviation;
      valueFixed = m_KernelFunction->Evaluate(valueFixed);

      valueMoving = (biter->MovingImageValue - aiter->MovingImageValue) / m_MovingImageStandardDeviation;
      valueMoving = m_KernelFunction->Evaluate(valueMoving);

      weightMoving = valueMoving / dDenominatorMoving.GetSum();
      weightJoint = valueMoving * valueFixed / dDenominatorJoint.GetSum();

      weight = (weightMoving - weightJoint);
      weight *= biter->MovingImageValue - aiter->MovingImageValue;

      totalWeight += weight;
      derivative -= (*aditer) * weight;
    } // end of sample A loop

    derivative += derivB * totalWeight.GetSum();
  } // end of sample B loop

  auto nsamp = static_cast<double>(m_NumberOfSpatialSamples);

  double threshold = -0.5 * nsamp * std::log(m_MinProbability);
  if (dLogSumMoving.GetSum() > threshold || dLogSumFixed.GetSum() > threshold || dLogSumJoint.GetSum() > threshold)
  {
    // at least half the samples in B did not occur within
    // the Parzen window width of samples in A
    itkExceptionMacro("Standard deviation is too small");
  }

  value = dLogSumFixed.GetSum() + dLogSumMoving.GetSum() - dLogSumJoint.GetSum();
  value /= nsamp;
  value += std::log(nsamp);

  derivative /= nsamp;
  derivative /= itk::Math::sqr(m_MovingImageStandardDeviation);
}

template <typename TFixedImage, typename TMovingImage>
void
MutualInformationImageToImageMetric<TFixedImage, TMovingImage>::GetDerivative(const ParametersType & parameters,
                                                                              DerivativeType &       derivative) const
{
  MeasureType value;

  // call the combined version
  this->GetValueAndDerivative(parameters, value, derivative);
}

template <typename TFixedImage, typename TMovingImage>
void
MutualInformationImageToImageMetric<TFixedImage, TMovingImage>::CalculateDerivatives(
  const FixedImagePointType & point,
  DerivativeType &            derivatives,
  TransformJacobianType &     jacobian) const
{
  MovingImagePointType mappedPoint = this->m_Transform->TransformPoint(point);

  CovariantVector<double, MovingImageDimension> imageDerivatives;

  if (m_DerivativeCalculator->IsInsideBuffer(mappedPoint))
  {
    imageDerivatives = m_DerivativeCalculator->Evaluate(mappedPoint);
  }
  else
  {
    derivatives.Fill(0.0);
    return;
  }

  this->m_Transform->ComputeJacobianWithRespectToParameters(point, jacobian);

  unsigned int numberOfParameters = this->m_Transform->GetNumberOfParameters();
  for (unsigned int k = 0; k < numberOfParameters; ++k)
  {
    derivatives[k] = 0.0;
    for (unsigned int j = 0; j < MovingImageDimension; ++j)
    {
      derivatives[k] += jacobian[j][k] * imageDerivatives[j];
    }
  }
}

} // end namespace itk

#endif