///////////////////////////////////////////////////////////////////////////////

// BSD 3-Clause License

//

// Copyright (C) 2019-2022, LAAS-CNRS, University of Edinburgh,

//                          University of Oxford, Heriot-Watt University

// Copyright note valid unless otherwise stated in individual files.

// All rights reserved.

///////////////////////////////////////////////////////////////////////////////

#ifndef CROCODDYL_CORE_SOLVERS_DDP_HPP_

#define CROCODDYL_CORE_SOLVERS_DDP_HPP_

#include <Eigen/Cholesky>

#include <vector>

#include "crocoddyl/core/mathbase.hpp"

#include "crocoddyl/core/solver-base.hpp"

#include "crocoddyl/core/utils/deprecate.hpp"

namespace crocoddyl {

/**

 * @brief Differential Dynamic Programming (DDP) solver

 *

 * The DDP solver computes an optimal trajectory and control commands by

 * iterates running `backwardPass()` and `forwardPass()`. The backward-pass

 * updates locally the quadratic approximation of the problem and computes

 * descent direction. If the warm-start is feasible, then it computes the gaps

 * \f$\mathbf{f}_s\f$ and run a modified Riccati sweep: \f{eqnarray*}

 *   \mathbf{Q}_{\mathbf{x}_k} &=& \mathbf{l}_{\mathbf{x}_k} +

 * \mathbf{f}^\top_{\mathbf{x}_k} (V_{\mathbf{x}_{k+1}} +

 * V_{\mathbf{xx}_{k+1}}\mathbf{\bar{f}}_{k+1}),\\

 *   \mathbf{Q}_{\mathbf{u}_k} &=& \mathbf{l}_{\mathbf{u}_k} +

 * \mathbf{f}^\top_{\mathbf{u}_k} (V_{\mathbf{x}_{k+1}} +

 * V_{\mathbf{xx}_{k+1}}\mathbf{\bar{f}}_{k+1}),\\

 *   \mathbf{Q}_{\mathbf{xx}_k} &=& \mathbf{l}_{\mathbf{xx}_k} +

 * \mathbf{f}^\top_{\mathbf{x}_k} V_{\mathbf{xx}_{k+1}}

 * \mathbf{f}_{\mathbf{x}_k},\\

 *   \mathbf{Q}_{\mathbf{xu}_k} &=& \mathbf{l}_{\mathbf{xu}_k} +

 * \mathbf{f}^\top_{\mathbf{x}_k} V_{\mathbf{xx}_{k+1}}

 * \mathbf{f}_{\mathbf{u}_k},\\

 *   \mathbf{Q}_{\mathbf{uu}_k} &=& \mathbf{l}_{\mathbf{uu}_k} +

 * \mathbf{f}^\top_{\mathbf{u}_k} V_{\mathbf{xx}_{k+1}}

 * \mathbf{f}_{\mathbf{u}_k}.

 * \f}

 * Then, the forward-pass rollouts this new policy by integrating the system

 * dynamics along a tuple of optimized control commands \f$\mathbf{u}^*_s\f$,

 * i.e. \f{eqnarray}

 *   \mathbf{\hat{x}}_0 &=& \mathbf{\tilde{x}}_0,\\

 *   \mathbf{\hat{u}}_k &=& \mathbf{u}_k + \alpha\mathbf{k}_k +

 * \mathbf{K}_k(\mathbf{\hat{x}}_k-\mathbf{x}_k),\\ \mathbf{\hat{x}}_{k+1} &=&

 * \mathbf{f}_k(\mathbf{\hat{x}}_k,\mathbf{\hat{u}}_k). \f}

 *

 * \sa SolverAbstract(), `backwardPass()` and `forwardPass()`

 */

class SolverDDP : public SolverAbstract {

 public:

  typedef typename MathBaseTpl<double>::MatrixXsRowMajor MatrixXdRowMajor;

  /**

   * @brief Initialize the DDP solver

   *

   * @param[in] problem  shooting problem

   */

  explicit SolverDDP(boost::shared_ptr<ShootingProblem> problem);

  virtual ~SolverDDP();

  virtual bool solve(

      const std::vector<Eigen::VectorXd>& init_xs = DEFAULT_VECTOR,

      const std::vector<Eigen::VectorXd>& init_us = DEFAULT_VECTOR,

      const std::size_t maxiter = 100, const bool is_feasible = false,

      const double init_reg = NAN);

  virtual void computeDirection(const bool recalc = true);

  virtual double tryStep(const double steplength = 1);

  virtual double stoppingCriteria();

  virtual const Eigen::Vector2d& expectedImprovement();

  virtual void resizeData();

  /**

   * @brief Update the Jacobian, Hessian and feasibility of the optimal control

   * problem

   *

   * These derivatives are computed around the guess state and control

   * trajectory. These trajectory can be set by using `setCandidate()`.

   *

   * @return the total cost around the guess trajectory

   */

  virtual double calcDiff();

  /**

   * @brief Run the backward pass (Riccati sweep)

   *

   * It assumes that the Jacobian and Hessians of the optimal control problem

   * have been compute (i.e., `calcDiff()`). The backward pass handles

   * infeasible guess through a modified Riccati sweep: \f{eqnarray*}

   *   \mathbf{Q}_{\mathbf{x}_k} &=& \mathbf{l}_{\mathbf{x}_k} +

   * \mathbf{f}^\top_{\mathbf{x}_k} (V_{\mathbf{x}_{k+1}}

   * +

   * V_{\mathbf{xx}_{k+1}}\mathbf{\bar{f}}_{k+1}),\\

   *   \mathbf{Q}_{\mathbf{u}_k} &=& \mathbf{l}_{\mathbf{u}_k} +

   * \mathbf{f}^\top_{\mathbf{u}_k} (V_{\mathbf{x}_{k+1}}

   * +

   * V_{\mathbf{xx}_{k+1}}\mathbf{\bar{f}}_{k+1}),\\

   *   \mathbf{Q}_{\mathbf{xx}_k} &=& \mathbf{l}_{\mathbf{xx}_k} +

   * \mathbf{f}^\top_{\mathbf{x}_k} V_{\mathbf{xx}_{k+1}}

   * \mathbf{f}_{\mathbf{x}_k},\\

   *   \mathbf{Q}_{\mathbf{xu}_k} &=& \mathbf{l}_{\mathbf{xu}_k} +

   * \mathbf{f}^\top_{\mathbf{x}_k} V_{\mathbf{xx}_{k+1}}

   * \mathbf{f}_{\mathbf{u}_k},\\

   *   \mathbf{Q}_{\mathbf{uu}_k} &=& \mathbf{l}_{\mathbf{uu}_k} +

   * \mathbf{f}^\top_{\mathbf{u}_k} V_{\mathbf{xx}_{k+1}}

   * \mathbf{f}_{\mathbf{u}_k}, \f} where

   * \f$\mathbf{l}_{\mathbf{x}_k}\f$,\f$\mathbf{l}_{\mathbf{u}_k}\f$,\f$\mathbf{f}_{\mathbf{x}_k}\f$

   * and \f$\mathbf{f}_{\mathbf{u}_k}\f$ are the Jacobians of the cost function

   * and dynamics,

   * \f$\mathbf{l}_{\mathbf{xx}_k}\f$,\f$\mathbf{l}_{\mathbf{xu}_k}\f$ and

   * \f$\mathbf{l}_{\mathbf{uu}_k}\f$ are the Hessians of the cost function,

   * \f$V_{\mathbf{x}_{k+1}}\f$ and \f$V_{\mathbf{xx}_{k+1}}\f$ defines the

   * linear-quadratic approximation of the Value function, and

   * \f$\mathbf{\bar{f}}_{k+1}\f$ describes the gaps of the dynamics.

   */

  virtual void backwardPass();

  /**

   * @brief Run the forward pass or rollout

   *

   * It rollouts the action model given the computed policy (feedforward terns

   * and feedback gains) by the `backwardPass()`: \f{eqnarray}

   *   \mathbf{\hat{x}}_0 &=& \mathbf{\tilde{x}}_0,\\

   *   \mathbf{\hat{u}}_k &=& \mathbf{u}_k + \alpha\mathbf{k}_k +

   * \mathbf{K}_k(\mathbf{\hat{x}}_k-\mathbf{x}_k),\\ \mathbf{\hat{x}}_{k+1} &=&

   * \mathbf{f}_k(\mathbf{\hat{x}}_k,\mathbf{\hat{u}}_k). \f} We can define

   * different step lengths \f$\alpha\f$.

   *

   * @param stepLength  applied step length (\f$0\leq\alpha\leq1\f$)

   */

  virtual void forwardPass(const double stepLength);

  /**

   * @brief Compute the linear-quadratic approximation of the control

   * Hamiltonian function

   *

   * @param[in] t      Time instance

   * @param[in] model  Action model in the given time instance

   * @param[in] data   Action data in the given time instance

   */

  virtual void computeActionValueFunction(

      const std::size_t t, const boost::shared_ptr<ActionModelAbstract>& model,

      const boost::shared_ptr<ActionDataAbstract>& data);

  /**

   * @brief Compute the linear-quadratic approximation of the Value function

   *

   * This function is called in the backward pass after updating the local

   * Hamiltonian.

   *

   * @param[in] t      Time instance

   * @param[in] model  Action model in the given time instance

   */

  virtual void computeValueFunction(

      const std::size_t t, const boost::shared_ptr<ActionModelAbstract>& model);

  /**

   * @brief Compute the feedforward and feedback terms using a Cholesky

   * decomposition

   *

   * To compute the feedforward \f$\mathbf{k}_k\f$ and feedback

   * \f$\mathbf{K}_k\f$ terms, we use a Cholesky decomposition to solve

   * \f$\mathbf{Q}_{\mathbf{uu}_k}^{-1}\f$ term: \f{eqnarray}

   * \mathbf{k}_k &=& \mathbf{Q}_{\mathbf{uu}_k}^{-1}\mathbf{Q}_{\mathbf{u}},\\

   * \mathbf{K}_k &=& \mathbf{Q}_{\mathbf{uu}_k}^{-1}\mathbf{Q}_{\mathbf{ux}}.

   * \f}

   *

   * Note that if the Cholesky decomposition fails, then we re-start the

   * backward pass and increase the state and control regularization values.

   */

  virtual void computeGains(const std::size_t t);

  /**

   * @brief Increase the state and control regularization values by a

   * `regfactor_` factor

   */

  void increaseRegularization();

  /**

   * @brief Decrease the state and control regularization values by a

   * `regfactor_` factor

   */

  void decreaseRegularization();

  /**

   * @brief Allocate all the internal data needed for the solver

   */

  virtual void allocateData();

  /**

   * @brief Return the regularization factor used to increase the damping value

   */

  double get_reg_incfactor() const;

  /**

   * @brief Return the regularization factor used to decrease the damping value

   */

  double get_reg_decfactor() const;

  /**

   * @brief Return the regularization factor used to decrease / increase it

   */

  DEPRECATED("Use get_reg_incfactor() or get_reg_decfactor()",

             double get_regfactor() const;)

  /**

   * @brief Return the minimum regularization value

   */

  double get_reg_min() const;

  DEPRECATED("Use get_reg_min()", double get_regmin() const);

  /**

   * @brief Return the maximum regularization value

   */

  double get_reg_max() const;

  DEPRECATED("Use get_reg_max()", double get_regmax() const);

  /**

   * @brief Return the set of step lengths using by the line-search procedure

   */

  const std::vector<double>& get_alphas() const;

  /**

   * @brief Return the step-length threshold used to decrease regularization

   */

  double get_th_stepdec() const;

  /**

   * @brief Return the step-length threshold used to increase regularization

   */

  double get_th_stepinc() const;

  /**

   * @brief Return the tolerance of the expected gradient used for testing the

   * step

   */

  double get_th_grad() const;

  /**

   * @brief Return the Hessian of the Value function \f$V_{\mathbf{xx}_s}\f$

   */

  const std::vector<Eigen::MatrixXd>& get_Vxx() const;

  /**

   * @brief Return the Hessian of the Value function \f$V_{\mathbf{x}_s}\f$

   */

  const std::vector<Eigen::VectorXd>& get_Vx() const;

  /**

   * @brief Return the Hessian of the Hamiltonian function

   * \f$\mathbf{Q}_{\mathbf{xx}_s}\f$

   */

  const std::vector<Eigen::MatrixXd>& get_Qxx() const;

  /**

   * @brief Return the Hessian of the Hamiltonian function

   * \f$\mathbf{Q}_{\mathbf{xu}_s}\f$

   */

  const std::vector<Eigen::MatrixXd>& get_Qxu() const;

  /**

   * @brief Return the Hessian of the Hamiltonian function

   * \f$\mathbf{Q}_{\mathbf{uu}_s}\f$

   */

  const std::vector<Eigen::MatrixXd>& get_Quu() const;

  /**

   * @brief Return the Jacobian of the Hamiltonian function

   * \f$\mathbf{Q}_{\mathbf{x}_s}\f$

   */

  const std::vector<Eigen::VectorXd>& get_Qx() const;

  /**

   * @brief Return the Jacobian of the Hamiltonian function

   * \f$\mathbf{Q}_{\mathbf{u}_s}\f$

   */

  const std::vector<Eigen::VectorXd>& get_Qu() const;

  /**

   * @brief Return the feedback gains \f$\mathbf{K}_{s}\f$

   */

  const std::vector<MatrixXdRowMajor>& get_K() const;

  /**

   * @brief Return the feedforward gains \f$\mathbf{k}_{s}\f$

   */

  const std::vector<Eigen::VectorXd>& get_k() const;

  /**

   * @brief Modify the regularization factor used to increase the damping value

   */

  void set_reg_incfactor(const double reg_factor);

  /**

   * @brief Modify the regularization factor used to decrease the damping value

   */

  void set_reg_decfactor(const double reg_factor);

  /**

   * @brief Modify the regularization factor used to decrease / increase it

   */

  DEPRECATED("Use set_reg_incfactor() or set_reg_decfactor()",

             void set_regfactor(const double reg_factor);)

  /**

   * @brief Modify the minimum regularization value

   */

  void set_reg_min(const double regmin);

  DEPRECATED("Use set_reg_min()", void set_regmin(const double regmin));

  /**

   * @brief Modify the maximum regularization value

   */

  void set_reg_max(const double regmax);

  DEPRECATED("Use set_reg_max()", void set_regmax(const double regmax));

  /**

   * @brief Modify the set of step lengths using by the line-search procedure

   */

  void set_alphas(const std::vector<double>& alphas);

  /**

   * @brief Modify the step-length threshold used to decrease regularization

   */

  void set_th_stepdec(const double th_step);

  /**

   * @brief Modify the step-length threshold used to increase regularization

   */

  void set_th_stepinc(const double th_step);

  /**

   * @brief Modify the tolerance of the expected gradient used for testing the

   * step

   */

  void set_th_grad(const double th_grad);

 protected:

  double reg_incfactor_;  //!< Regularization factor used to increase the

                          //!< damping value

  double reg_decfactor_;  //!< Regularization factor used to decrease the

                          //!< damping value

  double reg_min_;        //!< Minimum allowed regularization value

  double reg_max_;        //!< Maximum allowed regularization value

  double cost_try_;  //!< Total cost computed by line-search procedure

  std::vector<Eigen::VectorXd>

      xs_try_;  //!< State trajectory computed by line-search procedure

  std::vector<Eigen::VectorXd>

      us_try_;  //!< Control trajectory computed by line-search procedure

  std::vector<Eigen::VectorXd>

      dx_;  //!< State error during the roll-out/forward-pass (size T)

  // allocate data

  std::vector<Eigen::MatrixXd>

      Vxx_;  //!< Hessian of the Value function \f$\mathbf{V_{xx}}\f$

  Eigen::MatrixXd

      Vxx_tmp_;  //!< Temporary variable for ensuring symmetry of Vxx

  std::vector<Eigen::VectorXd>

      Vx_;  //!< Gradient of the Value function \f$\mathbf{V_x}\f$

  std::vector<Eigen::MatrixXd>

      Qxx_;  //!< Hessian of the Hamiltonian \f$\mathbf{Q_{xx}}\f$

  std::vector<Eigen::MatrixXd>

      Qxu_;  //!< Hessian of the Hamiltonian \f$\mathbf{Q_{xu}}\f$

  std::vector<Eigen::MatrixXd>

      Quu_;  //!< Hessian of the Hamiltonian \f$\mathbf{Q_{uu}}\f$

  std::vector<Eigen::VectorXd>

      Qx_;  //!< Gradient of the Hamiltonian \f$\mathbf{Q_x}\f$

  std::vector<Eigen::VectorXd>

      Qu_;  //!< Gradient of the Hamiltonian \f$\mathbf{Q_u}\f$

  std::vector<MatrixXdRowMajor> K_;  //!< Feedback gains \f$\mathbf{K}\f$

  std::vector<Eigen::VectorXd> k_;   //!< Feed-forward terms \f$\mathbf{l}\f$

  Eigen::VectorXd xnext_;      //!< Next state \f$\mathbf{x}^{'}\f$

  MatrixXdRowMajor FxTVxx_p_;  //!< Store the value of

                               //!< \f$\mathbf{f_x}^T\mathbf{V_{xx}}^{'}\f$

  std::vector<MatrixXdRowMajor>

      FuTVxx_p_;             //!< Store the values of

                             //!< \f$\mathbf{f_u}^T\mathbf{V_{xx}}^{'}\f$

                             //!< per each running node

  Eigen::VectorXd fTVxx_p_;  //!< Store the value of

                             //!< \f$\mathbf{\bar{f}}^T\mathbf{V_{xx}}^{'}\f$

  std::vector<Eigen::LLT<Eigen::MatrixXd> > Quu_llt_;  //!< Cholesky LLT solver

  std::vector<Eigen::VectorXd>

      Quuk_;  //!< Store the values of \f$\mathbf{Q_{uu}\mathbf{k}} per each

              //!< running node

  std::vector<double>

      alphas_;      //!< Set of step lengths using by the line-search procedure

  double th_grad_;  //!< Tolerance of the expected gradient used for testing the

                    //!< step

  double

      th_stepdec_;  //!< Step-length threshold used to decrease regularization

  double

      th_stepinc_;  //!< Step-length threshold used to increase regularization

};

}  // namespace crocoddyl

#endif  // CROCODDYL_CORE_SOLVERS_DDP_HPP_