// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Thomas Capricelli <orzel@freehackers.org>
// Copyright (C) 2012 Desire Nuentsa <desire.nuentsa_wakam@inria.fr>
//
// The algorithm of this class initially comes from MINPACK whose original authors are:
// Copyright Jorge More - Argonne National Laboratory
// Copyright Burt Garbow - Argonne National Laboratory
// Copyright Ken Hillstrom - Argonne National Laboratory
//
// This Source Code Form is subject to the terms of the Minpack license
// (a BSD-like license) described in the campaigned CopyrightMINPACK.txt file.
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_LEVENBERGMARQUARDT_H
#define EIGEN_LEVENBERGMARQUARDT_H

// IWYU pragma: private
#include "./InternalHeaderCheck.h"

namespace Eigen {
namespace LevenbergMarquardtSpace {
enum Status {
  NotStarted = -2,
  Running = -1,
  ImproperInputParameters = 0,
  RelativeReductionTooSmall = 1,
  RelativeErrorTooSmall = 2,
  RelativeErrorAndReductionTooSmall = 3,
  CosinusTooSmall = 4,
  TooManyFunctionEvaluation = 5,
  FtolTooSmall = 6,
  XtolTooSmall = 7,
  GtolTooSmall = 8,
  UserAsked = 9
};
}

template <typename Scalar_, int NX = Dynamic, int NY = Dynamic>
struct DenseFunctor {
  typedef Scalar_ Scalar;
  enum { InputsAtCompileTime = NX, ValuesAtCompileTime = NY };
  typedef Matrix<Scalar, InputsAtCompileTime, 1> InputType;
  typedef Matrix<Scalar, ValuesAtCompileTime, 1> ValueType;
  typedef Matrix<Scalar, ValuesAtCompileTime, InputsAtCompileTime> JacobianType;
  typedef ColPivHouseholderQR<JacobianType> QRSolver;
  const int m_inputs, m_values;

  DenseFunctor() : m_inputs(InputsAtCompileTime), m_values(ValuesAtCompileTime) {}
  DenseFunctor(int inputs, int values) : m_inputs(inputs), m_values(values) {}

  int inputs() const { return m_inputs; }
  int values() const { return m_values; }

  // int operator()(const InputType &x, ValueType& fvec) { }
  //  should be defined in derived classes

  // int df(const InputType &x, JacobianType& fjac) { }
  //  should be defined in derived classes
};

template <typename Scalar_, typename Index_>
struct SparseFunctor {
  typedef Scalar_ Scalar;
  typedef Index_ Index;
  typedef Matrix<Scalar, Dynamic, 1> InputType;
  typedef Matrix<Scalar, Dynamic, 1> ValueType;
  typedef SparseMatrix<Scalar, ColMajor, Index> JacobianType;
  typedef SparseQR<JacobianType, COLAMDOrdering<int> > QRSolver;
  enum { InputsAtCompileTime = Dynamic, ValuesAtCompileTime = Dynamic };

  SparseFunctor(int inputs, int values) : m_inputs(inputs), m_values(values) {}

  int inputs() const { return m_inputs; }
  int values() const { return m_values; }

  const int m_inputs, m_values;
  // int operator()(const InputType &x, ValueType& fvec) { }
  //  to be defined in the functor

  // int df(const InputType &x, JacobianType& fjac) { }
  //  to be defined in the functor if no automatic differentiation
};
namespace internal {
template <typename QRSolver, typename VectorType>
void lmpar2(const QRSolver &qr, const VectorType &diag, const VectorType &qtb, typename VectorType::Scalar m_delta,
            typename VectorType::Scalar &par, VectorType &x);
}
/**
 * \ingroup NonLinearOptimization_Module
 * \brief Performs non linear optimization over a non-linear function,
 * using a variant of the Levenberg Marquardt algorithm.
 *
 * Check wikipedia for more information.
 * http://en.wikipedia.org/wiki/Levenberg%E2%80%93Marquardt_algorithm
 */
template <typename FunctorType_>
class LevenbergMarquardt : internal::no_assignment_operator {
 public:
  typedef FunctorType_ FunctorType;
  typedef typename FunctorType::QRSolver QRSolver;
  typedef typename FunctorType::JacobianType JacobianType;
  typedef typename JacobianType::Scalar Scalar;
  typedef typename JacobianType::RealScalar RealScalar;
  typedef typename QRSolver::StorageIndex PermIndex;
  typedef Matrix<Scalar, Dynamic, 1> FVectorType;
  typedef PermutationMatrix<Dynamic, Dynamic, int> PermutationType;

 public:
  LevenbergMarquardt(FunctorType &functor)
      : m_functor(functor),
        m_nfev(0),
        m_njev(0),
        m_fnorm(0.0),
        m_gnorm(0),
        m_isInitialized(false),
        m_info(InvalidInput) {
    resetParameters();
    m_useExternalScaling = false;
  }

  LevenbergMarquardtSpace::Status minimize(FVectorType &x);
  LevenbergMarquardtSpace::Status minimizeInit(FVectorType &x);
  LevenbergMarquardtSpace::Status minimizeOneStep(FVectorType &x);
  LevenbergMarquardtSpace::Status lmder1(FVectorType &x, const Scalar tol = std::sqrt(NumTraits<Scalar>::epsilon()));
  static LevenbergMarquardtSpace::Status lmdif1(FunctorType &functor, FVectorType &x, Index *nfev,
                                                const Scalar tol = std::sqrt(NumTraits<Scalar>::epsilon()));

  /** Sets the default parameters */
  void resetParameters() {
    using std::sqrt;

    m_factor = 100.;
    m_maxfev = 400;
    m_ftol = sqrt(NumTraits<RealScalar>::epsilon());
    m_xtol = sqrt(NumTraits<RealScalar>::epsilon());
    m_gtol = 0.;
    m_epsfcn = 0.;
  }

  /** Sets the tolerance for the norm of the solution vector*/
  void setXtol(RealScalar xtol) { m_xtol = xtol; }

  /** Sets the tolerance for the norm of the vector function*/
  void setFtol(RealScalar ftol) { m_ftol = ftol; }

  /** Sets the tolerance for the norm of the gradient of the error vector*/
  void setGtol(RealScalar gtol) { m_gtol = gtol; }

  /** Sets the step bound for the diagonal shift */
  void setFactor(RealScalar factor) { m_factor = factor; }

  /** Sets the error precision  */
  void setEpsilon(RealScalar epsfcn) { m_epsfcn = epsfcn; }

  /** Sets the maximum number of function evaluation */
  void setMaxfev(Index maxfev) { m_maxfev = maxfev; }

  /** Use an external Scaling. If set to true, pass a nonzero diagonal to diag() */
  void setExternalScaling(bool value) { m_useExternalScaling = value; }

  /** \returns the tolerance for the norm of the solution vector */
  RealScalar xtol() const { return m_xtol; }

  /** \returns the tolerance for the norm of the vector function */
  RealScalar ftol() const { return m_ftol; }

  /** \returns the tolerance for the norm of the gradient of the error vector */
  RealScalar gtol() const { return m_gtol; }

  /** \returns the step bound for the diagonal shift */
  RealScalar factor() const { return m_factor; }

  /** \returns the error precision */
  RealScalar epsilon() const { return m_epsfcn; }

  /** \returns the maximum number of function evaluation */
  Index maxfev() const { return m_maxfev; }

  /** \returns a reference to the diagonal of the jacobian */
  FVectorType &diag() { return m_diag; }

  /** \returns the number of iterations performed */
  Index iterations() const { return m_iter; }

  /** \returns the number of functions evaluation */
  Index nfev() const { return m_nfev; }

  /** \returns the number of jacobian evaluation */
  Index njev() const { return m_njev; }

  /** \returns the norm of current vector function */
  RealScalar fnorm() const { return m_fnorm; }

  /** \returns the norm of the gradient of the error */
  RealScalar gnorm() const { return m_gnorm; }

  /** \returns the LevenbergMarquardt parameter */
  RealScalar lm_param(void) const { return m_par; }

  /** \returns a reference to the  current vector function
   */
  FVectorType &fvec() { return m_fvec; }

  /** \returns a reference to the matrix where the current Jacobian matrix is stored
   */
  JacobianType &jacobian() { return m_fjac; }

  /** \returns a reference to the triangular matrix R from the QR of the jacobian matrix.
   * \sa jacobian()
   */
  JacobianType &matrixR() { return m_rfactor; }

  /** the permutation used in the QR factorization
   */
  PermutationType permutation() const { return m_permutation; }

  /**
   * \brief Reports whether the minimization was successful
   * \returns \c Success if the minimization was successful,
   *         \c NumericalIssue if a numerical problem arises during the
   *          minimization process, for example during the QR factorization
   *         \c NoConvergence if the minimization did not converge after
   *          the maximum number of function evaluation allowed
   *          \c InvalidInput if the input matrix is invalid
   */
  ComputationInfo info() const { return m_info; }

 private:
  JacobianType m_fjac;
  JacobianType m_rfactor;  // The triangular matrix R from the QR of the jacobian matrix m_fjac
  FunctorType &m_functor;
  FVectorType m_fvec, m_qtf, m_diag;
  Index n;
  Index m;
  Index m_nfev;
  Index m_njev;
  RealScalar m_fnorm;   // Norm of the current vector function
  RealScalar m_gnorm;   // Norm of the gradient of the error
  RealScalar m_factor;  //
  Index m_maxfev;       // Maximum number of function evaluation
  RealScalar m_ftol;    // Tolerance in the norm of the vector function
  RealScalar m_xtol;    //
  RealScalar m_gtol;    // tolerance of the norm of the error gradient
  RealScalar m_epsfcn;  //
  Index m_iter;         // Number of iterations performed
  RealScalar m_delta;
  bool m_useExternalScaling;
  PermutationType m_permutation;
  FVectorType m_wa1, m_wa2, m_wa3, m_wa4;  // Temporary vectors
  RealScalar m_par;
  bool m_isInitialized;  // Check whether the minimization step has been called
  ComputationInfo m_info;
};

template <typename FunctorType>
LevenbergMarquardtSpace::Status LevenbergMarquardt<FunctorType>::minimize(FVectorType &x) {
  LevenbergMarquardtSpace::Status status = minimizeInit(x);
  if (status == LevenbergMarquardtSpace::ImproperInputParameters) {
    m_isInitialized = true;
    return status;
  }
  do {
    status = minimizeOneStep(x);
  } while (status == LevenbergMarquardtSpace::Running);
  m_isInitialized = true;
  return status;
}

template <typename FunctorType>
LevenbergMarquardtSpace::Status LevenbergMarquardt<FunctorType>::minimizeInit(FVectorType &x) {
  n = x.size();
  m = m_functor.values();

  m_wa1.resize(n);
  m_wa2.resize(n);
  m_wa3.resize(n);
  m_wa4.resize(m);
  m_fvec.resize(m);
  // FIXME: Sparse case: allocate space for the Jacobian.
  m_fjac.resize(m, n);
  if (!m_useExternalScaling) m_diag.resize(n);
  eigen_assert((!m_useExternalScaling || m_diag.size() == n) &&
               "When m_useExternalScaling is set, the caller must provide a valid 'm_diag'");
  m_qtf.resize(n);

  /* Function Body */
  m_nfev = 0;
  m_njev = 0;

  /*     check the input parameters for errors. */
  if (n <= 0 || m < n || m_ftol < 0. || m_xtol < 0. || m_gtol < 0. || m_maxfev <= 0 || m_factor <= 0.) {
    m_info = InvalidInput;
    return LevenbergMarquardtSpace::ImproperInputParameters;
  }

  if (m_useExternalScaling)
    for (Index j = 0; j < n; ++j)
      if (m_diag[j] <= 0.) {
        m_info = InvalidInput;
        return LevenbergMarquardtSpace::ImproperInputParameters;
      }

  /*     evaluate the function at the starting point */
  /*     and calculate its norm. */
  m_nfev = 1;
  if (m_functor(x, m_fvec) < 0) return LevenbergMarquardtSpace::UserAsked;
  m_fnorm = m_fvec.stableNorm();

  /*     initialize levenberg-marquardt parameter and iteration counter. */
  m_par = 0.;
  m_iter = 1;

  return LevenbergMarquardtSpace::NotStarted;
}

template <typename FunctorType>
LevenbergMarquardtSpace::Status LevenbergMarquardt<FunctorType>::lmder1(FVectorType &x, const Scalar tol) {
  n = x.size();
  m = m_functor.values();

  /* check the input parameters for errors. */
  if (n <= 0 || m < n || tol < 0.) return LevenbergMarquardtSpace::ImproperInputParameters;

  resetParameters();
  m_ftol = tol;
  m_xtol = tol;
  m_maxfev = 100 * (n + 1);

  return minimize(x);
}

template <typename FunctorType>
LevenbergMarquardtSpace::Status LevenbergMarquardt<FunctorType>::lmdif1(FunctorType &functor, FVectorType &x,
                                                                        Index *nfev, const Scalar tol) {
  Index n = x.size();
  Index m = functor.values();

  /* check the input parameters for errors. */
  if (n <= 0 || m < n || tol < 0.) return LevenbergMarquardtSpace::ImproperInputParameters;

  NumericalDiff<FunctorType> numDiff(functor);
  // embedded LevenbergMarquardt
  LevenbergMarquardt<NumericalDiff<FunctorType> > lm(numDiff);
  lm.setFtol(tol);
  lm.setXtol(tol);
  lm.setMaxfev(200 * (n + 1));

  LevenbergMarquardtSpace::Status info = LevenbergMarquardtSpace::Status(lm.minimize(x));
  if (nfev) *nfev = lm.nfev();
  return info;
}

}  // end namespace Eigen

#endif  // EIGEN_LEVENBERGMARQUARDT_H