/*=========================================================================
 *
 *  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
 *
 *         http://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 itkRegistrationParameterScalesFromJacobian_hxx
#define itkRegistrationParameterScalesFromJacobian_hxx

#include "itkRegistrationParameterScalesFromJacobian.h"

namespace itk
{

/** Compute parameter scales from average jacobian norms.
 *  For each parameter, compute the squared norm of its transform Jacobian,
 *  then average the squared norm over the sample points. This average is
 *  used as the scale of this parameter.
 */
template <typename TMetric>
void
RegistrationParameterScalesFromJacobian<TMetric>::EstimateScales(ScalesType & parameterScales)
{
  this->CheckAndSetInputs();
  this->SetScalesSamplingStrategy();
  this->SampleVirtualDomain();

  const SizeValueType numPara = this->GetNumberOfLocalParameters();

  parameterScales.SetSize(numPara);

  ParametersType norms(numPara);

  const auto numSamples = static_cast<const SizeValueType>(this->m_SamplePoints.size());

  norms.Fill(NumericTraits<typename ParametersType::ValueType>::ZeroValue());
  parameterScales.Fill(NumericTraits<typename ScalesType::ValueType>::OneValue());

  // checking each sample point
  for (SizeValueType c = 0; c < numSamples; c++)
  {
    const VirtualPointType point = this->m_SamplePoints[c];

    ParametersType squaredNorms(numPara);
    this->ComputeSquaredJacobianNorms(point, squaredNorms);

    norms = norms + squaredNorms;
  } // for numSamples

  if (numSamples > 0)
  {
    for (SizeValueType p = 0; p < numPara; p++)
    {
      parameterScales[p] = norms[p] / numSamples;
    }
  }
}

/**
 *  Compute the scale for a STEP, the impact of a STEP on the transform.
 */
template <typename TMetric>
typename RegistrationParameterScalesFromJacobian<TMetric>::FloatType
RegistrationParameterScalesFromJacobian<TMetric>::EstimateStepScale(const ParametersType & step)
{
  this->CheckAndSetInputs();
  this->SetStepScaleSamplingStrategy();
  this->SampleVirtualDomain();

  ScalesType sampleScales;
  this->ComputeSampleStepScales(step, sampleScales);

  const auto numSamples = static_cast<const SizeValueType>(this->m_SamplePoints.size());
  FloatType  scaleSum = NumericTraits<FloatType>::ZeroValue();

  // checking each sample point
  for (SizeValueType c = 0; c < numSamples; c++)
  {
    scaleSum += sampleScales[c];
  }

  return scaleSum / numSamples;
}

/**
 * Estimate the scales of local steps. For each voxel, we compute the impact
 * of a STEP on its location as in EstimateStepScale. Then we attribute this
 * impact to the corresponding local parameters.
 */
template <typename TMetric>
void
RegistrationParameterScalesFromJacobian<TMetric>::EstimateLocalStepScales(const ParametersType & step,
                                                                          ScalesType &           localStepScales)
{
  if (!this->IsDisplacementFieldTransform())
  {
    itkExceptionMacro(<< "EstimateLocalStepScales: the transform doesn't have local support.");
  }

  this->CheckAndSetInputs();
  this->SetStepScaleSamplingStrategy();
  this->SampleVirtualDomain();

  ScalesType sampleScales;
  this->ComputeSampleStepScales(step, sampleScales);

  const auto          numSamples = static_cast<const SizeValueType>(this->m_SamplePoints.size());
  const SizeValueType numPara = this->GetNumberOfLocalParameters();
  const SizeValueType numAllPara = this->GetTransform()->GetNumberOfParameters();
  const SizeValueType numLocals = numAllPara / numPara;

  localStepScales.SetSize(numLocals);
  localStepScales.Fill(NumericTraits<typename ScalesType::ValueType>::ZeroValue());

  // checking each sample point
  for (SizeValueType c = 0; c < numSamples; c++)
  {
    VirtualPointType & point = this->m_SamplePoints[c];
    IndexValueType     localId =
      this->m_Metric->ComputeParameterOffsetFromVirtualPoint(point, NumericTraits<SizeValueType>::OneValue());
    localStepScales[localId] = sampleScales[c];
  }
}

/**
 *  Compute the step scales for samples, i.e. the impacts on each sampled
 *  voxel from a change on the transform.
 */
template <typename TMetric>
void
RegistrationParameterScalesFromJacobian<TMetric>::ComputeSampleStepScales(const ParametersType & step,
                                                                          ScalesType &           sampleScales)
{
  const auto          numSamples = static_cast<const SizeValueType>(this->m_SamplePoints.size());
  const SizeValueType dim = this->GetDimension();
  const SizeValueType numPara = this->GetNumberOfLocalParameters();

  sampleScales.SetSize(numSamples);

  itk::Array<FloatType> dTdt(dim);

  JacobianType jacobianCache;
  JacobianType jacobian(dim,
                        (this->GetTransformForward() ? this->m_Metric->GetMovingTransform()->GetNumberOfParameters()
                                                     : this->m_Metric->GetFixedTransform()->GetNumberOfParameters()));


  // checking each sample point
  for (SizeValueType c = 0; c < numSamples; c++)
  {
    const VirtualPointType & point = this->m_SamplePoints[c];

    if (this->GetTransformForward())
    {
      this->m_Metric->GetMovingTransform()->ComputeJacobianWithRespectToParametersCachedTemporaries(
        point, jacobian, jacobianCache);
    }
    else
    {
      this->m_Metric->GetFixedTransform()->ComputeJacobianWithRespectToParametersCachedTemporaries(
        point, jacobian, jacobianCache);
    }

    if (!this->IsDisplacementFieldTransform())
    {
      dTdt = jacobian * step;
    }
    else
    {
      SizeValueType offset = this->m_Metric->ComputeParameterOffsetFromVirtualPoint(point, numPara);

      ParametersType localStep(numPara);
      for (SizeValueType p = 0; p < numPara; p++)
      {
        localStep[p] = step[offset + p];
      }
      dTdt = jacobian * localStep;
    }

    sampleScales[c] = dTdt.two_norm();
  }
}

/** Print the information about this class */
template <typename TMetric>
void
RegistrationParameterScalesFromJacobian<TMetric>::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
}

} // namespace itk

#endif /* itkRegistrationParameterScalesFromJacobian_hxx */