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

// BSD 3-Clause License

//

// Copyright (C) 2021-2024, Heriot-Watt University, University of Edinburgh

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

// All rights reserved.

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

#ifndef CROCODDYL_MULTIBODY_ACTIONS_CONTACT_INVDYN_HPP_

#define CROCODDYL_MULTIBODY_ACTIONS_CONTACT_INVDYN_HPP_

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

#include "crocoddyl/core/constraints/constraint-manager.hpp"

#include "crocoddyl/core/costs/cost-sum.hpp"

#include "crocoddyl/core/diff-action-base.hpp"

#include "crocoddyl/multibody/contacts/multiple-contacts.hpp"

#include "crocoddyl/multibody/data/contacts.hpp"

#include "crocoddyl/multibody/fwd.hpp"

#include "crocoddyl/multibody/states/multibody.hpp"

namespace crocoddyl {

/**

 * @brief Differential action model for contact inverse dynamics in multibody

 * systems.

 *

 * This class implements forward kinematic with holonomic contact constraints

 * (defined at the acceleration level) and inverse-dynamics computation using

 * the Recursive Newton Euler Algorithm (RNEA). The stack of cost and constraint

 * functions are implemented in `CostModelSumTpl` and

 * `ConstraintModelManagerTpl`, respectively. The acceleration and contact

 * forces are decision variables defined as the control inputs, and the

 * under-actuation and contact constraint are under the name `tau` and its frame

 * name, thus the user is not allowed to use it.

 *

 * Additionally, it is important to note that `calcDiff()` computes the

 * derivatives using the latest stored values by `calc()`. Thus, we need to

 * first run `calc()`.

 *

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

 * `createData()`

 */

template <typename _Scalar>

class DifferentialActionModelContactInvDynamicsTpl

    : public DifferentialActionModelAbstractTpl<_Scalar> {

 public:

  typedef _Scalar Scalar;

  typedef DifferentialActionModelAbstractTpl<Scalar> Base;

  typedef DifferentialActionDataContactInvDynamicsTpl<Scalar> Data;

  typedef StateMultibodyTpl<Scalar> StateMultibody;

  typedef ActuationModelAbstractTpl<Scalar> ActuationModelAbstract;

  typedef CostModelSumTpl<Scalar> CostModelSum;

  typedef ConstraintModelManagerTpl<Scalar> ConstraintModelManager;

  typedef ContactModelMultipleTpl<Scalar> ContactModelMultiple;

  typedef DataCollectorAbstractTpl<Scalar> DataCollectorAbstract;

  typedef DifferentialActionDataAbstractTpl<Scalar>

      DifferentialActionDataAbstract;

  typedef ContactItemTpl<Scalar> ContactItem;

  typedef MathBaseTpl<Scalar> MathBase;

  typedef typename MathBase::VectorXs VectorXs;

  typedef typename MathBase::MatrixXs MatrixXs;

  /**

   * @brief Initialize the contact inverse-dynamics action model

   *

   * It describes the kinematic evolution of the multibody system with contacts,

   * and computes the needed torques using inverse-dynamics.

   *

   * @param[in] state      State of the multibody system

   * @param[in] actuation  Actuation model

   * @param[in] contacts   Multiple contacts

   * @param[in] costs      Cost model

   */

  DifferentialActionModelContactInvDynamicsTpl(

      boost::shared_ptr<StateMultibody> state,

      boost::shared_ptr<ActuationModelAbstract> actuation,

      boost::shared_ptr<ContactModelMultiple> contacts,

      boost::shared_ptr<CostModelSum> costs);

  /**

   * @brief Initialize the contact inverse-dynamics action model

   *

   * @param[in] state        State of the multibody system

   * @param[in] actuation    Actuation model

   * @param[in] contacts     Multiple contacts

   * @param[in] costs        Cost model

   * @param[in] constraints  Constraints model

   */

  DifferentialActionModelContactInvDynamicsTpl(

      boost::shared_ptr<StateMultibody> state,

      boost::shared_ptr<ActuationModelAbstract> actuation,

      boost::shared_ptr<ContactModelMultiple> contacts,

      boost::shared_ptr<CostModelSum> costs,

      boost::shared_ptr<ConstraintModelManager> constraints);

  virtual ~DifferentialActionModelContactInvDynamicsTpl();

  /**

   * @brief Compute the system acceleration, cost value and constraint residuals

   *

   * It extracts the acceleration value from control vector and also computes

   * the cost and constraints.

   *

   * @param[in] data  Contact inverse-dynamics data

   * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$

   * @param[in] u     Control input \f$\mathbf{u}\in\mathbb{R}^{nu}\f$

   */

  virtual void calc(

      const boost::shared_ptr<DifferentialActionDataAbstract>& data,

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

  /**

   * @brief @copydoc Base::calc(const

   * boost::shared_ptr<DifferentialActionDataAbstract>& data, const

   * Eigen::Ref<const VectorXs>& x)

   */

  virtual void calc(

      const boost::shared_ptr<DifferentialActionDataAbstract>& data,

      const Eigen::Ref<const VectorXs>& x);

  /**

   * @brief Compute the derivatives of the dynamics, cost and constraint

   * functions

   *

   * It computes the partial derivatives of the dynamical system and the cost

   * and contraint functions. It assumes that `calc()` has been run first. This

   * function builds a quadratic approximation of the time-continuous action

   * model (i.e., dynamical system, cost and constraint functions).

   *

   * @param[in] data  Contact inverse-dynamics data

   * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$

   * @param[in] u     Control input \f$\mathbf{u}\in\mathbb{R}^{nu}\f$

   */

  virtual void calcDiff(

      const boost::shared_ptr<DifferentialActionDataAbstract>& data,

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

  /**

   * @brief @copydoc Base::calcDiff(const

   * boost::shared_ptr<DifferentialActionDataAbstract>& data, const

   * Eigen::Ref<const VectorXs>& x)

   */

  virtual void calcDiff(

      const boost::shared_ptr<DifferentialActionDataAbstract>& data,

      const Eigen::Ref<const VectorXs>& x);

  /**

   * @brief Create the contact inverse-dynamics data

   *

   * @return contact inverse-dynamics data

   */

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

  /**

   * @brief Checks that a specific data belongs to the contact inverse-dynamics

   * model

   */

  virtual bool checkData(

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

  /**

   * @brief Computes the quasic static commands

   *

   * The quasic static commands are the ones produced for a reference posture as

   * an equilibrium point with zero acceleration, i.e., for

   * \f$\mathbf{f^q_x}\delta\mathbf{q}+\mathbf{f_u}\delta\mathbf{u}=\mathbf{0}\f$

   *

   * @param[in] data     Contact inverse-dynamics data

   * @param[out] u       Quasic-static commands

   * @param[in] x        State point (velocity has to be zero)

   * @param[in] maxiter  Maximum allowed number of iterations (default 100)

   * @param[in] tol      Tolerance (default 1e-9)

   */

  virtual void quasiStatic(

      const boost::shared_ptr<DifferentialActionDataAbstract>& data,

      Eigen::Ref<VectorXs> u, const Eigen::Ref<const VectorXs>& x,

      const std::size_t maxiter = 100, const Scalar tol = Scalar(1e-9));

  /**

   * @brief Return the number of inequality constraints

   */

  virtual std::size_t get_ng() const;

  /**

   * @brief Return the number of equality constraints

   */

  virtual std::size_t get_nh() const;

  /**

   * @brief Return the lower bound of the inequality constraints

   */

  virtual const VectorXs& get_g_lb() const;

  /**

   * @brief Return the upper bound of the inequality constraints

   */

  virtual const VectorXs& get_g_ub() const;

  /**

   * @brief Return the actuation model

   */

  const boost::shared_ptr<ActuationModelAbstract>& get_actuation() const;

  /**

   * @brief Return the contact model

   */

  const boost::shared_ptr<ContactModelMultiple>& get_contacts() const;

  /**

   * @brief Return the cost model

   */

  const boost::shared_ptr<CostModelSum>& get_costs() const;

  /**

   * @brief Return the constraint model manager

   */

  const boost::shared_ptr<ConstraintModelManager>& get_constraints() const;

  /**

   * @brief Return the Pinocchio model

   */

  pinocchio::ModelTpl<Scalar>& get_pinocchio() const;

  /**

   * @brief Print relevant information of the contact inverse-dynamics model

   * @param[out] os  Output stream object

   */

  virtual void print(std::ostream& os) const;

 protected:

  using Base::g_lb_;   //!< Lower bound of the inequality constraints

  using Base::g_ub_;   //!< Upper bound of the inequality constraints

  using Base::ng_;     //!< Number of inequality constraints

  using Base::nh_;     //!< Number of equality constraints

  using Base::nu_;     //!< Control dimension

  using Base::state_;  //!< Model of the state

 private:

  void init(const boost::shared_ptr<StateMultibody>& state);

  boost::shared_ptr<ActuationModelAbstract> actuation_;    //!< Actuation model

  boost::shared_ptr<ContactModelMultiple> contacts_;       //!< Contact model

  boost::shared_ptr<CostModelSum> costs_;                  //!< Cost model

  boost::shared_ptr<ConstraintModelManager> constraints_;  //!< Constraint model

  pinocchio::ModelTpl<Scalar>& pinocchio_;                 //!< Pinocchio model

 public:

  /**

   * @brief Actuation residual

   *

   * This residual function enforces the torques of under-actuated joints (e.g.,

   * floating-base joints) to be zero. We compute these torques and their

   * derivatives using RNEA inside

   * `DifferentialActionModelContactInvDynamicsTpl`.

   *

   * As described in `ResidualModelAbstractTpl`, the residual value and its

   * Jacobians are calculated by `calc` and `calcDiff`, respectively.

   *

   * \sa `ResidualModelAbstractTpl`, `calc()`, `calcDiff()`, `createData()`

   */

  class ResidualModelActuation : public ResidualModelAbstractTpl<_Scalar> {

   public:

    typedef _Scalar Scalar;

    typedef MathBaseTpl<Scalar> MathBase;

    typedef ResidualModelAbstractTpl<Scalar> Base;

    typedef StateMultibodyTpl<Scalar> StateMultibody;

    typedef ResidualDataAbstractTpl<Scalar> ResidualDataAbstract;

    typedef DataCollectorAbstractTpl<Scalar> DataCollectorAbstract;

    typedef typename MathBase::VectorXs VectorXs;

    /**

     * @brief Initialize the actuation residual model

     *

     * @param[in] state  State of the multibody system

     * @param[in] nu     Dimension of the joint torques

     * @param[in] nc     Dimension of all the contacts

     */

                           const std::size_t nu, const std::size_t nc)

               true),

          na_(nu),

    /**

     * @brief Compute the actuation residual

     *

     * @param[in] data  Actuation residual data

     * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$

     * @param[in] u     Control input \f$\mathbf{u}\in\mathbb{R}^{nv+nu}\f$

     */

                      const Eigen::Ref<const VectorXs>&,

                      const Eigen::Ref<const VectorXs>&) {

      typename Data::ResidualDataActuation* d =

      // Update the under-actuation set and residual

        }

      }

    /**

     * @brief @copydoc Base::calc(const boost::shared_ptr<ResidualDataAbstract>&

     * data, const Eigen::Ref<const VectorXs>& x)

     */

                      const Eigen::Ref<const VectorXs>&) {

    /**

     * @brief Compute the derivatives of the actuation residual

     *

     * @param[in] data  Actuation residual data

     * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$

     * @param[in] u     Control input \f$\mathbf{u}\in\mathbb{R}^{nu}\f$

     */

                          const Eigen::Ref<const VectorXs>&,

                          const Eigen::Ref<const VectorXs>&) {

      typename Data::ResidualDataActuation* d =

        }

      }

    /**

     * @brief @copydoc Base::calcDiff(const

     * boost::shared_ptr<ResidualDataAbstract>& data, const Eigen::Ref<const

     * VectorXs>& x)

     */

                          const Eigen::Ref<const VectorXs>&) {

    /**

     * @brief Create the actuation residual data

     *

     * @return Actuation residual data

     */

        DataCollectorAbstract* const data) {

      return boost::allocate_shared<typename Data::ResidualDataActuation>(

          Eigen::aligned_allocator<typename Data::ResidualDataActuation>(),

    }

    /**

     * @brief Print relevant information of the actuation residual model

     *

     * @param[out] os  Output stream object

     */

    }

   protected:

    using Base::nu_;

    using Base::state_;

   private:

    std::size_t na_;  //!< Number of actuated joints

    std::size_t nc_;  //!< Number of contacts

  };

 public:

  /**

   * @brief Contact-acceleration residual

   *

   * This residual function is defined as \f$\mathbf{r} = \mathbf{a_0}\f$, where

   * \f$\mathbf{a_0}\f$ defines the desired contact acceleration, which might

   * also include the Baumgarte stabilization gains. Furthermore, the Jacobians

   * of the residual function are computed analytically. This is used by

   * `ConstraintModelManagerTpl` inside parent

   * `DifferentialActionModelContactInvDynamicsTpl` class.

   *

   * As described in `ResidualModelAbstractTpl`, the residual value and its

   * Jacobians are calculated by `calc` and `calcDiff`, respectively.

   *

   * \sa `ResidualModelAbstractTpl`, `calc()`, `calcDiff()`, `createData()`

   */

  class ResidualModelContact : public ResidualModelAbstractTpl<_Scalar> {

   public:

    typedef _Scalar Scalar;

    typedef MathBaseTpl<Scalar> MathBase;

    typedef ResidualModelAbstractTpl<Scalar> Base;

    typedef StateMultibodyTpl<Scalar> StateMultibody;

    typedef ResidualDataAbstractTpl<Scalar> ResidualDataAbstract;

    typedef DataCollectorAbstractTpl<Scalar> DataCollectorAbstract;

    typedef typename MathBase::VectorXs VectorXs;

    typedef typename MathBase::MatrixXs MatrixXs;

    /**

     * @brief Initialize the contact-acceleration residual model

     *

     * @param[in] state  State of the multibody system

     * @param[in] id     Contact frame id

     * @param[in] nr     Dimension of the contact-acceleration residual

     * @param[in] nc     Dimension of all contacts

     */

                         const pinocchio::FrameIndex id, const std::size_t nr,

                         const std::size_t nc)

          id_(id),

    /**

     * @brief Compute the contact-acceleration residual

     *

     * @param[in] data  Contact-acceleration residual data

     * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$

     * @param[in] u     Control input \f$\mathbf{u}\in\mathbb{R}^{nv+nu}\f$

     */

              const Eigen::Ref<const VectorXs>&,

              const Eigen::Ref<const VectorXs>&) {

      typename Data::ResidualDataContact* d =

    /**

     * @brief @copydoc Base::calc(const boost::shared_ptr<ResidualDataAbstract>&

     * data, const Eigen::Ref<const VectorXs>& x)

     */

                      const Eigen::Ref<const VectorXs>&) {

    /**

     * @brief Compute the derivatives of the contact-acceleration residual

     *

     * @param[in] data  Contact-acceleration residual data

     * @param[in] x     State point \f$\mathbf{x}\in\mathbb{R}^{ndx}\f$

     * @param[in] u     Control input \f$\mathbf{u}\in\mathbb{R}^{nu}\f$

     */

                  const Eigen::Ref<const VectorXs>&,

                  const Eigen::Ref<const VectorXs>&) {

      typename Data::ResidualDataContact* d =

    /**

     * @brief @copydoc Base::calcDiff(const

     * boost::shared_ptr<ResidualDataAbstract>& data, const Eigen::Ref<const

     * VectorXs>& x)

     */

                          const Eigen::Ref<const VectorXs>&) {

    /**

     * @brief Create the contact-acceleration residual data

     *

     * @return contact-acceleration residual data

     */

        DataCollectorAbstract* const data) {

      return boost::allocate_shared<typename Data::ResidualDataContact>(

          Eigen::aligned_allocator<typename Data::ResidualDataContact>(), this,

    }

    /**

     * @brief Print relevant information of the contact-acceleration residual

     * model

     *

     * @param[out] os  Output stream object

     */

    }

   protected:

    using Base::nr_;

    using Base::nu_;

    using Base::state_;

   private:

    pinocchio::FrameIndex id_;  //!< Reference frame id

    std::string frame_name_;    //!< Reference frame name

  };

};

template <typename _Scalar>

struct DifferentialActionDataContactInvDynamicsTpl

    : public DifferentialActionDataAbstractTpl<_Scalar> {

  typedef _Scalar Scalar;

  typedef MathBaseTpl<Scalar> MathBase;

  typedef DifferentialActionDataAbstractTpl<Scalar> Base;

  typedef JointDataAbstractTpl<Scalar> JointDataAbstract;

  typedef DataCollectorJointActMultibodyInContactTpl<Scalar>

      DataCollectorJointActMultibodyInContact;

  typedef CostDataSumTpl<Scalar> CostDataSum;

  typedef ConstraintDataManagerTpl<Scalar> ConstraintDataManager;

  typedef ContactModelMultipleTpl<Scalar> ContactModelMultiple;

  typedef typename MathBase::VectorXs VectorXs;

  typedef typename MathBase::MatrixXs MatrixXs;

  template <template <typename Scalar> class Model>

      Model<Scalar>* const model)

      : Base(model),

        multibody(

                model->get_state(), model->get_actuation(), model->get_nu()),

    // Set constant values for Fu, df_dx, and df_du

    }

  pinocchio::DataTpl<Scalar> pinocchio;                  //!< Pinocchio data

  DataCollectorJointActMultibodyInContact multibody;     //!< Multibody data

  boost::shared_ptr<CostDataSum> costs;                  //!< Costs data

  boost::shared_ptr<ConstraintDataManager> constraints;  //!< Constraints data

  VectorXs

      tmp_xstatic;  //!< State point used for computing the quasi-static input

  VectorXs

      tmp_rstatic;  //!< Factorization used for computing the quasi-static input

  MatrixXs tmp_Jstatic;  //!< Jacobian used for computing the quasi-static input

  using Base::cost;

  using Base::Fu;

  using Base::Fx;

  using Base::Lu;

  using Base::Luu;

  using Base::Lx;

  using Base::Lxu;

  using Base::Lxx;

  using Base::r;

  using Base::xout;

  struct ResidualDataActuation : public ResidualDataAbstractTpl<_Scalar> {

    typedef _Scalar Scalar;

    typedef MathBaseTpl<Scalar> MathBase;

    typedef ResidualDataAbstractTpl<Scalar> Base;

    typedef DataCollectorAbstractTpl<Scalar> DataCollectorAbstract;

    typedef DataCollectorActMultibodyInContactTpl<Scalar>

        DataCollectorActMultibodyInContact;

    typedef ActuationDataAbstractTpl<Scalar> ActuationDataAbstract;

    typedef ContactDataMultipleTpl<Scalar> ContactDataMultiple;

    typedef typename MathBase::MatrixXs MatrixXs;

    template <template <typename Scalar> class Model>

                          DataCollectorAbstract* const data)

        : Base(model, data),

      // Check that proper shared data has been passed

            "Invalid argument: the shared data should be derived from "

            "DataCollectorActMultibodyInContact");

      }

      // Avoids data casting at runtime

    pinocchio::DataTpl<Scalar>* pinocchio;               //!< Pinocchio data

    boost::shared_ptr<ActuationDataAbstract> actuation;  //!< Actuation data

    boost::shared_ptr<ContactDataMultiple> contact;      //!< Contact data

    MatrixXs dtau_dx;

    MatrixXs dtau_du;

    using Base::r;

    using Base::Ru;

    using Base::Rx;

    using Base::shared;

  };

  struct ResidualDataContact : public ResidualDataAbstractTpl<_Scalar> {

    typedef _Scalar Scalar;

    typedef MathBaseTpl<Scalar> MathBase;

    typedef ResidualDataAbstractTpl<Scalar> Base;

    typedef DataCollectorAbstractTpl<Scalar> DataCollectorAbstract;

    typedef DataCollectorMultibodyInContactTpl<Scalar>

        DataCollectorMultibodyInContact;

    typedef ContactModelMultipleTpl<Scalar> ContactModelMultiple;

    template <template <typename Scalar> class Model>

                        DataCollectorAbstract* const data, const std::size_t id)

            "Invalid argument: the shared data should be derived from "

            "DataCollectorMultibodyInContact");

      }

                                        // and avoid to pass id in constructor

        }

      }

    ContactDataAbstractTpl<Scalar>* contact;  //!< Contact force data

    using Base::r;

    using Base::Ru;

    using Base::Rx;

    using Base::shared;

  };

};

}  // namespace crocoddyl

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

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

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

#include <crocoddyl/multibody/actions/contact-invdyn.hxx>

#endif  // CROCODDYL_MULTIBODY_ACTIONS_CONTACT_INVDYN_HPP_