/*=========================================================================
 *
 *  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 itkRegularStepGradientDescentOptimizerv4_hxx
#define itkRegularStepGradientDescentOptimizerv4_hxx


namespace itk
{

template <typename TInternalComputationValueType>
RegularStepGradientDescentOptimizerv4<TInternalComputationValueType>::RegularStepGradientDescentOptimizerv4()
  : m_RelaxationFactor(0.5)
  , m_MinimumStepLength(1e-4)
  , // Initialize parameter for the convergence checker
  m_GradientMagnitudeTolerance(1e-4)
  , m_CurrentLearningRateRelaxation(0)
{
  this->m_DoEstimateLearningRateAtEachIteration = false;
  this->m_DoEstimateLearningRateOnce = false;
}

template <typename TInternalComputationValueType>
void
RegularStepGradientDescentOptimizerv4<TInternalComputationValueType>::StartOptimization(bool doOnlyInitialization)
{
  this->m_UseConvergenceMonitoring = false;

  // Call the grandparent version for basic validation and setup
  GradientDescentOptimizerBasev4Template<TInternalComputationValueType>::StartOptimization(doOnlyInitialization);

  if (this->m_ReturnBestParametersAndValue)
  {
    this->m_BestParameters = this->GetCurrentPosition();
    this->m_CurrentBestValue = NumericTraits<MeasureType>::max();
  }

  const SizeValueType spaceDimension = this->m_Metric->GetNumberOfParameters();

  this->m_Gradient = DerivativeType(spaceDimension);
  this->m_PreviousGradient = DerivativeType(spaceDimension);
  this->m_Gradient.Fill(0.0f);
  this->m_PreviousGradient.Fill(0.0f);

  // Reset the iterations and learning rate scale when the optimizer is started again.
  this->m_CurrentLearningRateRelaxation = 1.0;
  this->m_CurrentIteration = 0;

  // validity check for the value of GradientMagnitudeTolerance
  if (m_GradientMagnitudeTolerance < 0.0)
  {
    itkExceptionMacro("Gradient magnitude tolerance must be "
                      "greater or equal 0.0. Current value is "
                      << m_GradientMagnitudeTolerance);
  }

  if (!doOnlyInitialization)
  {
    this->ResumeOptimization();
  }
}

template <typename TInternalComputationValueType>
void
RegularStepGradientDescentOptimizerv4<TInternalComputationValueType>::AdvanceOneStep()
{
  itkDebugMacro("AdvanceOneStep");

  // Make sure the scales have been set properly
  if (this->m_Scales.size() != this->m_Gradient.Size())
  {
    itkExceptionMacro("The size of Scales is " << this->m_Scales.size()
                                               << ", but the NumberOfParameters for the CostFunction is "
                                               << this->m_Gradient.Size() << '.');
  }

  if (this->m_RelaxationFactor < 0.0)
  {
    itkExceptionMacro("Relaxation factor must be positive. Current value is " << this->m_RelaxationFactor);
  }

  if (this->m_RelaxationFactor >= 1.0)
  {
    itkExceptionMacro("Relaxation factor must be less than 1.0. Current value is " << this->m_RelaxationFactor);
  }

  //
  // Begin threaded gradient modification.
  // Scale gradient and previous gradient by scales.
  // The m_Gradient and m_PreviousGradient variables are modified in-place.
  //
  this->ModifyGradientByScales();

  CompensatedSummationType compensatedSummation;
  for (SizeValueType dim = 0; dim < this->m_Gradient.Size(); ++dim)
  {
    const double weighted = this->m_Gradient[dim];
    compensatedSummation += weighted * weighted;
  }
  const double gradientMagnitude = std::sqrt(compensatedSummation.GetSum());

  if (gradientMagnitude < this->m_GradientMagnitudeTolerance)
  {
    this->m_StopCondition = StopConditionObjectToObjectOptimizerEnum::GRADIENT_MAGNITUDE_TOLEARANCE;
    this->m_StopConditionDescription << "Gradient magnitude tolerance met after " << this->m_CurrentIteration
                                     << " iterations. Gradient magnitude (" << gradientMagnitude
                                     << ") is less than gradient magnitude tolerance ("
                                     << this->m_GradientMagnitudeTolerance << ").";
    this->StopOptimization();
    return;
  }

  compensatedSummation.ResetToZero();
  for (SizeValueType i = 0; i < this->m_Gradient.Size(); ++i)
  {
    const double weight1 = this->m_Gradient[i];
    const double weight2 = this->m_PreviousGradient[i];
    compensatedSummation += weight1 * weight2;
  }
  const double scalarProduct = compensatedSummation.GetSum();

  // Check for direction changes
  if (scalarProduct < 0)
  {
    this->m_CurrentLearningRateRelaxation *= this->m_RelaxationFactor;
  }

  const double stepLength = this->m_CurrentLearningRateRelaxation * this->m_LearningRate;

  if (stepLength < this->m_MinimumStepLength)
  {
    this->m_StopCondition = StopConditionObjectToObjectOptimizerEnum::STEP_TOO_SMALL;
    this->m_StopConditionDescription << "Step too small after " << this->m_CurrentIteration
                                     << " iterations. Current step (" << stepLength << ") is less than minimum step ("
                                     << this->m_MinimumStepLength << ").";
    this->StopOptimization();
    return;
  }

  this->EstimateLearningRate();
  this->ModifyGradientByLearningRate();
  const double factor = NumericTraits<typename ScalesType::ValueType>::OneValue() / gradientMagnitude;

  try
  {
    // Pass gradient to transform and let it do its own updating
    this->m_Metric->UpdateTransformParameters(this->m_Gradient, factor);
  }
  catch (const ExceptionObject &)
  {
    this->m_StopCondition = StopConditionObjectToObjectOptimizerEnum::UPDATE_PARAMETERS_ERROR;
    this->m_StopConditionDescription << "UpdateTransformParameters error";
    this->StopOptimization();

    // Pass exception to caller
    throw;
  }

  this->InvokeEvent(IterationEvent());
}

template <typename TInternalComputationValueType>
double
RegularStepGradientDescentOptimizerv4<TInternalComputationValueType>::GetCurrentStepLength() const
{
  return (this->m_CurrentLearningRateRelaxation * this->m_LearningRate);
}

template <typename TInternalComputationValueType>
void
RegularStepGradientDescentOptimizerv4<TInternalComputationValueType>::ModifyGradientByScalesOverSubRange(
  const IndexRangeType & subrange)
{
  const ScalesType & scales = this->GetScales();
  const ScalesType & weights = this->GetWeights();

  ScalesType factor(scales.Size());

  if (this->GetWeightsAreIdentity())
  {
    for (SizeValueType i = 0; i < factor.Size(); ++i)
    {
      factor[i] = NumericTraits<typename ScalesType::ValueType>::OneValue() / scales[i];
    }
  }
  else
  {
    for (SizeValueType i = 0; i < factor.Size(); ++i)
    {
      factor[i] = weights[i] / scales[i];
    }
  }

  // Loop over the range. It is inclusive.
  for (IndexValueType j = subrange[0]; j <= subrange[1]; ++j)
  {
    // Scale is checked during StartOptimization for values <=
    // machine epsilon.
    // Take the modulo of the index to handle gradients from transforms
    // with local support. The gradient array stores the gradient of local
    // parameters at each local index with linear packing.
    IndexValueType index = j % scales.Size();
    this->m_Gradient[j] = this->m_Gradient[j] * factor[index];
    this->m_PreviousGradient[j] = this->m_PreviousGradient[j] * factor[index];
  }
}

template <typename TInternalComputationValueType>
void
RegularStepGradientDescentOptimizerv4<TInternalComputationValueType>::ModifyGradientByLearningRateOverSubRange(
  const IndexRangeType & subrange)
{
  // Loop over the range. It is inclusive.
  for (IndexValueType j = subrange[0]; j <= subrange[1]; ++j)
  {
    this->m_Gradient[j] = this->m_Gradient[j] * this->m_CurrentLearningRateRelaxation * this->m_LearningRate;
  }
}

template <typename TInternalComputationValueType>
void
RegularStepGradientDescentOptimizerv4<TInternalComputationValueType>::EstimateLearningRate()
{
  if (this->m_ScalesEstimator.IsNull())
  {
    return;
  }
  if (this->m_DoEstimateLearningRateAtEachIteration ||
      (this->m_DoEstimateLearningRateOnce && this->m_CurrentIteration == 0))
  {
    TInternalComputationValueType stepScale = this->m_ScalesEstimator->EstimateStepScale(this->m_Gradient);

    if (stepScale <= NumericTraits<TInternalComputationValueType>::epsilon())
    {
      this->m_LearningRate = NumericTraits<TInternalComputationValueType>::OneValue();
    }
    else
    {
      this->m_LearningRate = this->m_MaximumStepSizeInPhysicalUnits / stepScale;
    }
    CompensatedSummationType compensatedSummation;
    for (SizeValueType dim = 0; dim < this->m_Gradient.Size(); ++dim)
    {
      const double weighted = this->m_Gradient[dim];
      compensatedSummation += weighted * weighted;
    }
    const double gradientMagnitude = std::sqrt(compensatedSummation.GetSum());

    //
    // Specialized to keep the step size regularized this additional
    // scale is needed to make the learning rate independent on the
    // gradient magnitude.
    //
    this->m_LearningRate *= gradientMagnitude;
  }
}
template <typename TInternalComputationValueType>
void
RegularStepGradientDescentOptimizerv4<TInternalComputationValueType>::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);

  os << indent << "Relaxation factor: " << this->m_RelaxationFactor << std::endl;
  os << indent << "Minimum step length: " << this->m_MinimumStepLength << std::endl;
  os << indent << "Gradient magnitude tolerance: " << this->m_GradientMagnitudeTolerance << std::endl;
  os << indent << "Current learning rate relaxation: " << this->m_CurrentLearningRateRelaxation << std::endl;
}

} // namespace itk

#endif