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

// BSD 3-Clause License

//

// Copyright (C) 2021, University of Edinburgh, University of Trento

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

// All rights reserved.

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

#ifndef CROCODDYL_CORE_CONTROLS_POLY_TWO_RK_HPP_

#define CROCODDYL_CORE_CONTROLS_POLY_TWO_RK_HPP_

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

#include "crocoddyl/core/fwd.hpp"

#include "crocoddyl/core/integrator/rk.hpp"

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

namespace crocoddyl {

/**

 * @brief A polynomial function of time of degree two, that is a quadratic

 * function

 *

 * The size of the parameters \f$\mathbf{u}\f$ is 3 times the size of the

 * control input \f$\mathbf{w}\f$. It defines a polynomial of degree two,

 * customized for the RK4 and RK4 integrators (even though it can be used with

 * whatever integration scheme). The first third of \f$\mathbf{u}\f$ represents

 * the value of \f$\mathbf{w}\f$ at time 0. The second third of \f$\mathbf{u}\f$

 * represents the value of \f$\mathbf{w}\f$ at time 0.5 or 1/3 for RK4 and RK3

 * parametrization, respectively. The last third of \f$\mathbf{u}\f$ represents

 * the value of \f$\mathbf{w}\f$ at time 1 or 2/3 for the RK4 and RK3

 * parametrization, respectively. This parametrization is suitable to be used

 * with the RK-4 or RK-3 integration schemes, because they require the value of

 * \f$\mathbf{w}\f$ exactly at 0, 0.5, 1 (for RK4) or 0, 1/3, 2/3 (for RK3).

 *

 * The main computations are carried out in `calc`, `multiplyByJacobian` and

 * `multiplyJacobianTransposeBy`, where the former computes control input

 * \f$\mathbf{w}\in\mathbb{R}^{nw}\f$ from a set of control parameters

 * \f$\mathbf{u}\in\mathbb{R}^{nu}\f$ where `nw` and `nu` represent the

 * dimension of the control inputs and parameters, respectively, and the latter

 * defines useful operations across the Jacobian of the control-parametrization

 * model. Finally, `params` allows us to obtain the control parameters from a

 * the control input, i.e., it is the inverse of `calc`. Note that

 * `multiplyByJacobian` and `multiplyJacobianTransposeBy` requires to run `calc`

 * first.

 *

 * \sa `ControlParametrizationAbstractTpl`, `calc()`, `calcDiff()`,

 * `createData()`, `params`, `multiplyByJacobian`, `multiplyJacobianTransposeBy`

 */

template <typename _Scalar>

class ControlParametrizationModelPolyTwoRKTpl

    : public ControlParametrizationModelAbstractTpl<_Scalar> {

 public:

  typedef _Scalar Scalar;

  typedef MathBaseTpl<Scalar> MathBase;

  typedef ControlParametrizationDataAbstractTpl<Scalar>

      ControlParametrizationDataAbstract;

  typedef ControlParametrizationModelAbstractTpl<Scalar> Base;

  typedef ControlParametrizationDataPolyTwoRKTpl<Scalar> Data;

  typedef typename MathBase::VectorXs VectorXs;

  typedef typename MathBase::MatrixXs MatrixXs;

  /**

   * @brief Initialize the poly-two RK control parametrization

   *

   * @param[in] nw      Dimension of control vector

   * @param[in] rktype  Type of RK parametrization

   */

  explicit ControlParametrizationModelPolyTwoRKTpl(const std::size_t nw,

                                                   const RKType rktype);

  virtual ~ControlParametrizationModelPolyTwoRKTpl();

  /**

   * @brief Get the value of the control at the specified time

   *

   * @param[in]  data   Poly-two-RK data

   * @param[in]  t      Time in [0,1]

   * @param[in]  u      Control parameters

   */

  virtual void calc(

      const boost::shared_ptr<ControlParametrizationDataAbstract>& data,

      const Scalar t, const Eigen::Ref<const VectorXs>& u) const;

  /**

   * @brief Get the value of the Jacobian of the control with respect to the

   * parameters

   *

   * It assumes that `calc()` has been run first

   *

   * @param[in]  data   Poly-two-RK data

   * @param[in]  t      Time in [0,1]

   * @param[in]  u      Control parameters

   */

  virtual void calcDiff(

      const boost::shared_ptr<ControlParametrizationDataAbstract>& data,

      const Scalar t, const Eigen::Ref<const VectorXs>& u) const;

  /**

   * @brief Create the control-parametrization data

   *

   * @return the control-parametrization data

   */

  virtual boost::shared_ptr<ControlParametrizationDataAbstract> createData();

  /**

   * @brief Get a value of the control parameters u such that the control at the

   * specified time t is equal to the specified value w

   *

   * @param[in]  data   Poly-two-RK data

   * @param[in]  t      Time in [0,1]

   * @param[in]  w      Control values

   */

  virtual void params(

      const boost::shared_ptr<ControlParametrizationDataAbstract>& data,

      const Scalar t, const Eigen::Ref<const VectorXs>& w) const;

  /**

   * @brief Map the specified bounds from the control space to the parameter

   * space

   *

   * @param[in]  w_lb   Control lower bound

   * @param[in]  w_ub   Control lower bound

   * @param[out] u_lb   Control parameters lower bound

   * @param[out] u_ub   Control parameters upper bound

   */

  virtual void convertBounds(const Eigen::Ref<const VectorXs>& w_lb,

                             const Eigen::Ref<const VectorXs>& w_ub,

                             Eigen::Ref<VectorXs> u_lb,

                             Eigen::Ref<VectorXs> u_ub) const;

  /**

   * @brief Compute the product between a specified matrix and the Jacobian of

   * the control (with respect to the parameters)

   *

   * It assumes that `calc()` has been run first

   *

   * @param[in]  data   Poly-two-RK data

   * @param[in]  A      A matrix to multiply times the Jacobian

   * @param[out] out    Product between the matrix A and the Jacobian of the

   * control with respect to the parameters

   * @param[in] op      Assignment operator which sets, adds, or removes the

   * given results

   */

  virtual void multiplyByJacobian(

      const boost::shared_ptr<ControlParametrizationDataAbstract>& data,

      const Eigen::Ref<const MatrixXs>& A, Eigen::Ref<MatrixXs> out,

      const AssignmentOp = setto) const;

  /**

   * @brief Compute the product between the transposed Jacobian of the control

   * (with respect to the parameters) and a specified matrix

   *

   * It assumes that `calc()` has been run first

   *

   * @param[in]  data   Poly-two-RK data

   * @param[in]  A      A matrix to multiply times the Jacobian

   * @param[out] out    Product between the transposed Jacobian of the control

   * with respect to the parameters and the matrix A

   * @param[in] op      Assignment operator which sets, adds, or removes the

   * given results

   */

  virtual void multiplyJacobianTransposeBy(

      const boost::shared_ptr<ControlParametrizationDataAbstract>& data,

      const Eigen::Ref<const MatrixXs>& A, Eigen::Ref<MatrixXs> out,

      const AssignmentOp = setto) const;

 protected:

  using Base::nu_;

  using Base::nw_;

 private:

  RKType rktype_;

};

template <typename _Scalar>

struct ControlParametrizationDataPolyTwoRKTpl

    : public ControlParametrizationDataAbstractTpl<_Scalar> {

  typedef _Scalar Scalar;

  typedef MathBaseTpl<Scalar> MathBase;

  typedef ControlParametrizationDataAbstractTpl<Scalar> Base;

  typedef typename MathBase::Vector3s Vector3s;

  template <template <typename Scalar> class Model>

  Vector3s c;     //!< Polynomial coefficients of the second-order control model

                  //!< that depends on time

  Scalar tmp_t2;  //!< Temporary variable to store the square of the time

};

}  // namespace crocoddyl

/* --- Details -------------------------------------------------------------- */

/* --- Details -------------------------------------------------------------- */

/* --- Details -------------------------------------------------------------- */

#include "crocoddyl/core/controls/poly-two-rk.hxx"

#endif  // CROCODDYL_CORE_CONTROLS_POLY_TWO_RK_HPP_