/**
 * @file aug_lagrangian_impl.hpp
 * @author Ryan Curtin
 *
 * Implementation of AugLagrangian class (Augmented Lagrangian optimization
 * method).
 *
 * ensmallen is free software; you may redistribute it and/or modify it under
 * the terms of the 3-clause BSD license.  You should have received a copy of
 * the 3-clause BSD license along with ensmallen.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#ifndef ENSMALLEN_AUG_LAGRANGIAN_AUG_LAGRANGIAN_IMPL_HPP
#define ENSMALLEN_AUG_LAGRANGIAN_AUG_LAGRANGIAN_IMPL_HPP

#include <ensmallen_bits/lbfgs/lbfgs.hpp>
#include <ensmallen_bits/function.hpp>
#include "aug_lagrangian_function.hpp"

namespace ens {

template<typename VecType>
inline AugLagrangianType<VecType>::AugLagrangianType(
    const size_t maxIterations,
    const double penaltyThresholdFactor,
    const double sigmaUpdateFactor,
    const L_BFGS& lbfgs) :
    maxIterations(maxIterations),
    penaltyThresholdFactor(penaltyThresholdFactor),
    sigmaUpdateFactor(sigmaUpdateFactor),
    lbfgs(lbfgs),
    terminate(false),
    deprecatedSigma(0.0)
{
}

template<typename VecType>
template<typename LagrangianFunctionType,
         typename MatType,
         typename InVecType,
         typename GradType,
         typename... CallbackTypes>
typename std::enable_if<IsMatrixType<GradType>::value, bool>::type
AugLagrangianType<VecType>::Optimize(
    LagrangianFunctionType& function,
    MatType& coordinates,
    InVecType& lambda,
    double& sigma,
    CallbackTypes&&... callbacks)
{
  AugLagrangianFunction<LagrangianFunctionType, InVecType> augfunc(
      function, lambda, sigma);

  return Optimize(augfunc, coordinates, callbacks...);
}

template<typename VecType>
template<typename LagrangianFunctionType,
         typename MatType,
         typename GradType,
         typename... CallbackTypes>
typename std::enable_if<IsMatrixType<GradType>::value &&
                        IsAllNonMatrix<CallbackTypes...>::value, bool>::type
AugLagrangianType<VecType>::Optimize(LagrangianFunctionType& function,
                                     MatType& coordinates,
                                     CallbackTypes&&... callbacks)
{
  typedef typename ForwardType<MatType>::bvec InVecType;

  // If the user did not specify the right size for sigma and lambda, we will
  // use defaults.
  // TODO: remove this when ensmallen 4.x is released!
  if (!deprecatedLambda.is_empty())
  {
    InVecType lambda(conv_to<InVecType>::from(deprecatedLambda));

    AugLagrangianFunction<LagrangianFunctionType, InVecType> augfunc(function,
        lambda, deprecatedSigma);
    const bool result = Optimize(augfunc, coordinates, callbacks...);
    deprecatedLambda = conv_to<VecType>::from(lambda);

    return result;
  }
  else
  {
    // Use default values.
    InVecType lambda(function.NumConstraints());
    lambda.zeros();
    double sigma = 10;

    AugLagrangianFunction<LagrangianFunctionType, InVecType> augfunc(
        function, lambda, sigma);
    return Optimize(augfunc, coordinates, callbacks...);
  }
}

template<typename VecType>
template<typename LagrangianFunctionType,
         typename MatType,
         typename InVecType,
         typename GradType,
         typename... CallbackTypes>
typename std::enable_if<IsMatrixType<GradType>::value, bool>::type
AugLagrangianType<VecType>::Optimize(
    AugLagrangianFunction<LagrangianFunctionType, InVecType>& augfunc,
    MatType& coordinatesIn,
    CallbackTypes&&... callbacks)
{
  // Convenience typedefs.
  typedef typename MatType::elem_type ElemType;
  typedef typename MatTypeTraits<MatType>::BaseMatType BaseMatType;
  typedef typename MatTypeTraits<GradType>::BaseMatType BaseGradType;

  BaseMatType& coordinates = (BaseMatType&) coordinatesIn;

  // Check that the types satisfy our needs.
  traits::CheckConstrainedFunctionTypeAPI<LagrangianFunctionType, BaseMatType,
      BaseGradType>();
  RequireFloatingPointType<BaseMatType>();
  RequireFloatingPointType<BaseGradType>();
  RequireSameInternalTypes<BaseMatType, BaseGradType>();

  LagrangianFunctionType& function = augfunc.Function();

  // Ensure that we update lambda immediately.
  ElemType penaltyThreshold = std::numeric_limits<ElemType>::max();

  // Track the last objective to compare for convergence.
  ElemType lastObjective = function.Evaluate(coordinates);

  // Convergence tolerance---depends on the epsilon of the type we are using for
  // optimization.
  ElemType tolerance = 1000 * std::numeric_limits<ElemType>::epsilon();

  // Then, calculate the current penalty.
  ElemType penalty = 0;
  for (size_t i = 0; i < function.NumConstraints(); i++)
  {
    const ElemType p = std::pow(function.EvaluateConstraint(i, coordinates),
        ElemType(2));
    terminate |= Callback::EvaluateConstraint(*this, function, coordinates, i,
        p, callbacks...);

    penalty += p;
  }

  Info << "Penalty is " << penalty << " (threshold " << penaltyThreshold
      << ")." << std::endl;

  // The odd comparison allows user to pass maxIterations = 0 (i.e. no limit on
  // number of iterations).
  size_t it;
  Callback::BeginOptimization(*this, function, coordinates, callbacks...);
  for (it = 0; it != (maxIterations - 1) && !terminate; it++)
  {
    Info << "AugLagrangian on iteration " << it
        << ", starting with objective "  << lastObjective << "." << std::endl;

    if (!lbfgs.Optimize(augfunc, coordinates, callbacks...))
      Info << "L-BFGS reported an error during optimization." << std::endl;
    Info << "Done with L-BFGS." << std::endl;

    const ElemType objective = function.Evaluate(coordinates);

    terminate |= Callback::Evaluate(*this, function, coordinates, objective,
        callbacks...);

    // Check if we are done with the entire optimization (the threshold we are
    // comparing with is arbitrary).
    if (std::abs(lastObjective - objective) < tolerance &&
        augfunc.Sigma() > 500000)
    {
      Callback::EndOptimization(*this, function, coordinates, callbacks...);
      return true;
    }

    lastObjective = objective;

    // Assuming that the optimization has converged to a new set of coordinates,
    // we now update either lambda or sigma.  We update sigma if the penalty
    // term is too high, and we update lambda otherwise.

    // First, calculate the current penalty.
    ElemType penalty = 0;
    for (size_t i = 0; i < function.NumConstraints(); i++)
    {
      const ElemType p = std::pow(function.EvaluateConstraint(i, coordinates),
          ElemType(2));
      terminate |= Callback::EvaluateConstraint(*this, function, coordinates, i,
          p, callbacks...);

      penalty += p;
    }

    Info << "Penalty is " << penalty << " (threshold " << penaltyThreshold
        << ")." << std::endl;

    if (terminate)
      break;

    if (penalty < penaltyThreshold) // We update lambda.
    {
      // We use the update: lambda_{k + 1} = lambda_k - sigma * c(coordinates),
      // but we have to write a loop to do this for each constraint.
      for (size_t i = 0; i < function.NumConstraints(); i++)
      {
        const ElemType p = function.EvaluateConstraint(i, coordinates);
        terminate |= Callback::EvaluateConstraint(*this, function, coordinates,
            i, p, callbacks...);

        augfunc.Lambda()[i] -= ElemType(augfunc.Sigma()) * p;
      }

      // We also update the penalty threshold to be a factor of the current
      // penalty.
      penaltyThreshold = ElemType(penaltyThresholdFactor) * penalty;
      Info << "Lagrange multiplier estimates updated." << std::endl;
    }
    else
    {
      // We multiply sigma by a constant value.
      augfunc.Sigma() *= sigmaUpdateFactor;
      Info << "Updated sigma to " << augfunc.Sigma() << "." << std::endl;
      if (augfunc.Sigma() >= std::numeric_limits<ElemType>::max() / 2.0)
      {
        Warn << "AugLagrangian::Optimize(): sigma too large for element type; "
            << "terminating." << std::endl;
        Callback::EndOptimization(*this, function, coordinates, callbacks...);
        return true;
      }
    }

    terminate |= Callback::StepTaken(*this, function, coordinates,
        callbacks...);
  }

  Callback::EndOptimization(*this, function, coordinates, callbacks...);
  return false;
}

} // namespace ens

#endif // ENSMALLEN_AUG_LAGRANGIAN_AUG_LAGRANGIAN_IMPL_HPP

