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

#include "itkQuasiNewtonOptimizerv4EstimateNewtonStepThreader.h"

namespace itk
{

template <typename TInternalComputationValueType>
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::QuasiNewtonOptimizerv4Template()
  : m_PreviousValue(0.0)
  , m_BestValue(0.0)
  , m_BestPosition(0)

{
  this->m_LearningRate = NumericTraits<TInternalComputationValueType>::OneValue();

  // m_MaximumNewtonStepSizeInPhysicalUnits is used for automatic learning
  // rate estimation. it may be initialized either by calling
  // SetMaximumNewtonStepSizeInPhysicalUnits manually or by using m_ScalesEstimator
  // automatically. and the former has higher priority than the latter.
  this->m_MaximumNewtonStepSizeInPhysicalUnits = TInternalComputationValueType{};

  /** Threader for Quasi-Newton method */
  using OptimizerType = QuasiNewtonOptimizerv4EstimateNewtonStepThreaderTemplate<TInternalComputationValueType>;
  using OptimizerPointer = typename OptimizerType::Pointer;
  OptimizerPointer estimateNewtonStepThreader = OptimizerType::New();

  this->m_EstimateNewtonStepThreader = estimateNewtonStepThreader;
}

template <typename TInternalComputationValueType>
void
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
}

template <typename TInternalComputationValueType>
void
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::StartOptimization(bool doOnlyInitialization)
{
  itkDebugMacro("StartOptimization");

  const SizeValueType numPara = this->m_Metric->GetNumberOfParameters();
  const SizeValueType numLocalPara = this->m_Metric->GetNumberOfLocalParameters();
  const SizeValueType numLocals = this->m_Metric->GetNumberOfParameters() / numLocalPara;

  /* Set up the data for Quasi-Newton method */
  this->m_NewtonStep.SetSize(numPara);

  this->m_NewtonStepValidFlags.resize(numLocals);

  this->m_HessianArray.resize(numLocals);
  for (SizeValueType loc = 0; loc < numLocals; ++loc)
  {
    this->m_HessianArray[loc].SetSize(numLocalPara, numLocalPara);
  }

  if (this->m_ScalesEstimator.IsNotNull())
  {
    if (this->m_MaximumNewtonStepSizeInPhysicalUnits <= NumericTraits<TInternalComputationValueType>::epsilon())
    {
      // Newton step size might be bigger than one voxel spacing.
      // empirically, we set it to 1~5 voxel spacings.
      this->m_MaximumNewtonStepSizeInPhysicalUnits = 3.0 * this->m_ScalesEstimator->EstimateMaximumStepSize();
    }
  }

  /* Must call the superclass version for basic validation, setup,
   * and to start the optimization loop. */
  Superclass::StartOptimization(doOnlyInitialization);
}

template <typename TInternalComputationValueType>
void
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::AdvanceOneStep()
{
  itkDebugMacro("AdvanceOneStep");
  const SizeValueType numPara = this->m_Metric->GetNumberOfParameters();
  this->m_CurrentPosition = this->m_Metric->GetParameters();

  if (this->GetCurrentIteration() == 0)
  {
    // initialize some information
    this->m_PreviousValue = this->GetCurrentMetricValue();
    this->m_PreviousPosition = this->GetCurrentPosition();
    this->m_PreviousGradient = this->GetGradient();

    this->m_BestValue = this->m_CurrentMetricValue;
    this->m_BestPosition = this->m_CurrentPosition;
    this->m_BestIteration = this->GetCurrentIteration();
  }
  else if (m_BestValue > this->m_CurrentMetricValue)
  {
    // store the best value and related information
    this->m_BestValue = this->m_CurrentMetricValue;
    this->m_BestPosition = this->m_CurrentPosition;
    this->m_BestIteration = this->GetCurrentIteration();
  }

  if (this->GetCurrentIteration() - m_BestIteration > m_MaximumIterationsWithoutProgress)
  {
    ParametersType backStep;
    backStep = m_BestPosition - this->m_Metric->GetParameters();
    this->m_Metric->UpdateTransformParameters(backStep);

    this->m_CurrentPosition = this->m_BestPosition;
    this->m_CurrentMetricValue = this->m_BestValue;

    this->m_StopCondition = StopConditionObjectToObjectOptimizerEnum::STEP_TOO_SMALL;
    this->m_StopConditionDescription << "Optimization stops after " << this->GetCurrentIteration()
                                     << " iterations since"
                                     << " there is no progress in the last " << m_MaximumIterationsWithoutProgress
                                     << " steps." << std::endl
                                     << " The best value is from Iteration " << m_BestIteration << '.';
    this->StopOptimization();
    return;
  }

  if (this->GetCurrentIteration() > 0)
  {
    ParametersType lastStep(numPara);
    lastStep = this->m_CurrentPosition - this->m_PreviousPosition;
    if (lastStep.squared_magnitude() < NumericTraits<TInternalComputationValueType>::epsilon())
    {
      this->m_StopCondition = StopConditionObjectToObjectOptimizerEnum::STEP_TOO_SMALL;
      this->m_StopConditionDescription << "Optimization stops after " << this->GetCurrentIteration()
                                       << " iterations since"
                                       << " the last step is almost zero.";
      this->StopOptimization();
      return;
    }
  }

  this->EstimateNewtonStep();

  /** Save for the next iteration */
  this->m_PreviousValue = this->GetCurrentMetricValue();
  this->m_PreviousPosition = this->GetCurrentPosition();
  this->m_PreviousGradient = this->GetGradient();

  this->CombineGradientNewtonStep();
  this->ModifyCombinedNewtonStep();

  try
  {
    /* Pass gradient to transform and let it do its own updating */
    this->m_Metric->UpdateTransformParameters(this->m_NewtonStep);
  }
  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>
void
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::CombineGradientNewtonStep()
{
  const SizeValueType numLocalPara = this->m_Metric->GetNumberOfLocalParameters();
  const SizeValueType numLocals = this->m_Metric->GetNumberOfParameters() / numLocalPara;

  bool validNewtonStepExists = false;
  for (SizeValueType loc = 0; loc < numLocals; ++loc)
  {
    if (this->m_NewtonStepValidFlags[loc])
    {
      validNewtonStepExists = true;
      break;
    }
  }

  TInternalComputationValueType ratio = NumericTraits<TInternalComputationValueType>::OneValue();
  if (validNewtonStepExists)
  {
    TInternalComputationValueType gradStepScale = this->m_ScalesEstimator->EstimateStepScale(this->m_Gradient);
    TInternalComputationValueType newtonStepScale = this->m_ScalesEstimator->EstimateStepScale(this->m_NewtonStep);
    if (gradStepScale > NumericTraits<TInternalComputationValueType>::epsilon())
    {
      ratio = newtonStepScale / gradStepScale;
    }
  }

  for (SizeValueType loc = 0; loc < numLocals; ++loc)
  {
    if (!this->m_NewtonStepValidFlags[loc])
    {
      // Using the Gradient step
      const SizeValueType offset = loc * numLocalPara;
      for (SizeValueType p = offset; p < (offset + numLocalPara); ++p)
      {
        this->m_NewtonStep[p] = this->m_Gradient[p] * ratio;
      }
    }
  }
}

template <typename TInternalComputationValueType>
void
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::ModifyCombinedNewtonStep()
{
  // If m_ScalesEstimator is not set, we will not change the Newton step.
  if (this->m_ScalesEstimator.IsNull())
  {
    return;
  }

  TInternalComputationValueType stepScale = this->m_ScalesEstimator->EstimateStepScale(this->m_NewtonStep);

  if (stepScale <= NumericTraits<TInternalComputationValueType>::epsilon())
  {
    this->m_LearningRate = NumericTraits<TInternalComputationValueType>::OneValue();
  }
  else
  {
    this->m_LearningRate = this->m_MaximumNewtonStepSizeInPhysicalUnits / stepScale;
    if (this->m_LearningRate > NumericTraits<TInternalComputationValueType>::OneValue())
    {
      // learning rate is at most 1 for a newton step
      this->m_LearningRate = NumericTraits<TInternalComputationValueType>::OneValue();
    }
  }

  if (itk::Math::abs(this->m_LearningRate - NumericTraits<TInternalComputationValueType>::OneValue()) > 0.01)
  {
    this->m_NewtonStep *= this->m_LearningRate;
  }
}

template <typename TInternalComputationValueType>
void
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::ResetNewtonStep(IndexValueType loc)
{
  const SizeValueType numLocalPara = this->m_Metric->GetNumberOfLocalParameters();

  // Initialize Hessian to identity matrix
  m_HessianArray[loc].Fill(TInternalComputationValueType{});

  for (unsigned int i = 0; i < numLocalPara; ++i)
  {
    m_HessianArray[loc][i][i] = NumericTraits<TInternalComputationValueType>::OneValue(); // identity matrix
  }

  IndexValueType offset = loc * numLocalPara;
  for (SizeValueType p = 0; p < numLocalPara; ++p)
  {
    // Set to zero for invalid Newton steps.
    // They must be defined since they will be used during step scale estimation.
    this->m_NewtonStep[offset + p] = TInternalComputationValueType{};
  }
}

template <typename TInternalComputationValueType>
void
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::EstimateNewtonStep()
{
  if (this->m_Gradient.GetSize() == 0)
  {
    return;
  }

  IndexRangeType fullrange;
  fullrange[0] = 0;
  fullrange[1] = this->m_Gradient.GetSize() - 1; // range is inclusive
  /* Perform the modification either with or without threading */

  if (this->m_Gradient.GetSize() > 10000)
  {
    /* This ends up calling EstimateNewtonStepOverSubRange from each thread */
    this->m_EstimateNewtonStepThreader->Execute(this, fullrange);
  }
  else
  {
    this->EstimateNewtonStepOverSubRange(fullrange);
  }
}

template <typename TInternalComputationValueType>
void
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::EstimateNewtonStepOverSubRange(
  const IndexRangeType & subrange)
{
  const SizeValueType numLocalPara = this->m_Metric->GetNumberOfLocalParameters();

  IndexValueType low = subrange[0] / numLocalPara;
  IndexValueType high = subrange[1] / numLocalPara;

  // let us denote the i-th thread's sub range by subrange_i
  // we assume subrange_i[1] + 1 = subrange_(i+1)[0] .
  // if the subrange_i doesn't start with the multiple of numLocalPara,
  // we assign this starting block of local parameters to thread_i
  // if the subrange_i doesn't end with the multiple of numLocalPara,
  // we assign this ending block of local parameters to thread_(i+1) .
  if ((subrange[1] + 1) % numLocalPara != 0)
  {
    --high;
  }

  for (IndexValueType loc = low; loc <= high; ++loc)
  {
    if (this->GetCurrentIteration() == 0)
    {
      this->m_NewtonStepValidFlags[loc] = false;
    }
    else
    {
      this->m_NewtonStepValidFlags[loc] = this->ComputeHessianAndStepWithBFGS(loc);
    }

    if (!this->m_NewtonStepValidFlags[loc])
    {
      this->ResetNewtonStep(loc);
    }

  } // end for loc
}

template <typename TInternalComputationValueType>
bool
QuasiNewtonOptimizerv4Template<TInternalComputationValueType>::ComputeHessianAndStepWithBFGS(IndexValueType loc)
{
  if (this->GetCurrentIteration() == 0)
  {
    return false;
  }

  const SizeValueType numLocalPara = this->m_Metric->GetNumberOfLocalParameters();
  IndexValueType      offset = loc * numLocalPara;

  ParametersType dx(numLocalPara);  // delta of position x: x_k+1 - x_k
  DerivativeType dg(numLocalPara);  // delta of gradient: g_k+1 - g_k
  DerivativeType edg(numLocalPara); // estimated delta of gradient: hessian_k * dx

  for (SizeValueType p = 0; p < numLocalPara; ++p)
  {
    dx[p] = this->m_CurrentPosition[offset + p] - this->m_PreviousPosition[offset + p];
    // gradient is already negated
    dg[p] = this->m_PreviousGradient[offset + p] - this->m_Gradient[offset + p];
  }

  edg = this->m_HessianArray[loc] * dx;

  TInternalComputationValueType dot_dg_dx = inner_product(dg, dx);
  TInternalComputationValueType dot_edg_dx = inner_product(edg, dx);

  if (itk::Math::abs(dot_dg_dx) <= NumericTraits<TInternalComputationValueType>::epsilon() ||
      itk::Math::abs(dot_edg_dx) <= NumericTraits<TInternalComputationValueType>::epsilon())
  {
    return false;
  }

  vnl_matrix<TInternalComputationValueType> plus = outer_product(dg, dg) / dot_dg_dx;
  vnl_matrix<TInternalComputationValueType> minus = outer_product(edg, edg) / dot_edg_dx;
  vnl_matrix<TInternalComputationValueType> newHessian = this->m_HessianArray[loc] + plus - minus;

  this->m_HessianArray[loc] = newHessian;

  for (SizeValueType p = 0; p < numLocalPara; ++p)
  {
    if (newHessian[p][p] < 0)
    {
      return false;
    }
  }

  TInternalComputationValueType threshold = NumericTraits<TInternalComputationValueType>::epsilon();

  if (itk::Math::abs(vnl_determinant(newHessian)) <= threshold)
  {
    return false;
  }
  else
  {
    DerivativeType gradient(numLocalPara);
    for (SizeValueType p = 0; p < numLocalPara; ++p)
    {
      gradient[p] = this->m_Gradient[offset + p];
    }

    const vnl_matrix<TInternalComputationValueType> hessianInverse{
      vnl_matrix_inverse<TInternalComputationValueType>(newHessian).as_matrix()
    };
    // gradient is already negated
    const DerivativeType newtonStep{ hessianInverse * gradient };
    for (SizeValueType p = 0; p < numLocalPara; ++p)
    {
      this->m_NewtonStep[offset + p] = newtonStep[p];
    }
  }

  return true;
}

} // end namespace itk

#endif