// SPDX-FileCopyrightText: Fondazione Istituto Italiano di Tecnologia (IIT)
// SPDX-License-Identifier: BSD-3-Clause

#ifndef IDYNTREE_ESTIMATION_EXTERNALWRENCHESTIMATION_H
#define IDYNTREE_ESTIMATION_EXTERNALWRENCHESTIMATION_H

#include <iDynTree/Direction.h>
#include <iDynTree/MatrixDynSize.h>
#include <iDynTree/Position.h>
#include <iDynTree/VectorDynSize.h>
#include <iDynTree/VectorFixSize.h>

#include <iDynTree/Indices.h>
#include <iDynTree/LinkState.h>

#include <vector>

namespace iDynTree
{
class Model;
class Traversal;
class SubModelDecomposition;
class LinkContactWrenches;
class SensorsMeasurements;
class SensorsList;
class LinkVelArray;
class LinkAccArray;
class JointPosDoubleArray;
class JointDOFsDoubleArray;

/**
 * Type of a UnknownWrenchContact.
 */
enum UnknownWrenchContactType
{
    /**
     * Contact for which the complete wrench is unknown.
     */
    FULL_WRENCH,

    /**
     * Contact assumed to be a pure force excerted on the contact point.
     */
    PURE_FORCE,

    /**
     * Contact assumed to be a pure force with a known direction excerted on the contact point
     */
    PURE_FORCE_WITH_KNOWN_DIRECTION,

    /**
     * The contact forces is assumed to be known.
     */
    NO_UNKNOWNS
};

/**
 * \brief A contact whose wrench is unknown.
 *
 * Class representing a contact for which
 * the approximate center of the contact surface is known,
 * but the wrench is unknown.
 *
 */
struct UnknownWrenchContact
{
    /**
     * Constructor
     */
    UnknownWrenchContact()
    {
    }

    UnknownWrenchContact(const UnknownWrenchContactType _unknownType,
                         const Position& _contactPoint,
                         const Direction& _forceDirection = iDynTree::Direction::Default(),
                         const Wrench& _knownWrench = iDynTree::Wrench(),
                         const unsigned long& _contactId = 0)
        : unknownType(_unknownType)
        , contactPoint(_contactPoint)
        , forceDirection(_forceDirection)
        , knownWrench(_knownWrench)
        , contactId(_contactId)
    {
    }

    /**
     * Type of the unknown contact.
     */
    UnknownWrenchContactType unknownType;

    /**
     * Position of the center of the contact, in the link frame.
     */
    Position contactPoint;

    /**
     * If unknownType is PURE_FORCE_WITH_KNOWN_DIRECTION,
     * contains the known direction (in link frame) of the force.
     */
    Direction forceDirection;

    /**
     * If unknownType is NO_UNKNOWNS,
     * contains the value of the contact force, with the orientation of the link frame,
     * and w.r.t. to the origin of the link frame, i.e. it ignores the contactPoint attribute.
     */
    Wrench knownWrench;

    /**
     * Unique id identifing the contact.
     * This id is propagated to the contact wrench data structure.
     * It is implemented mainly for compatibility with the skinDynLib library.
     */
    unsigned long contactId;
};

/**
 * A set of UnknownWrenchContact for each link, representing all the contacts
 * between the model and the external environment whose wrench is unkwnon.
 *
 */
class LinkUnknownWrenchContacts
{
private:
    std::vector<std::vector<UnknownWrenchContact>> m_linkUnknownWrenchContacts;

public:
    /**
     * Create a LinkWrenches vector, with the size given
     * by nrOfLinks .
     *
     * @param[in] nrOfLinks the size of the vector.
     */
    LinkUnknownWrenchContacts(unsigned int nrOfLinks = 0);
    LinkUnknownWrenchContacts(const Model& model);

    /**
     * Preserving the number of links, remove all the previously added unknowns.
     *
     */
    void clear();

    /**
     *
     * @param[in] nrOfLinks the number of links used to resize
     */
    void resize(unsigned int nrOfLinks);
    void resize(const Model& model);

    /**
     * Get the number of external contacts for a given link.
     */
    size_t getNrOfContactsForLink(const LinkIndex linkIndex) const;

    /**
     * Set the number of external contacts for a given link.
     */
    void setNrOfContactsForLink(const LinkIndex linkIndex, const size_t nrOfContacts);

    /**
     * Add a new contact for a link.
     */
    void addNewContactForLink(const LinkIndex linkIndex, const UnknownWrenchContact& newContact);

    /**
     * Add a new contact for a frame.
     * If the specified frame is not the a link frame, the method automatically convert the unknown
     * wrench to the relative link frame.
     *
     * @param[in] model the model class for getting frame information.
     * @param[in] frameIndex the index of the frame in which you are expressing the new unknown
     * wrench.
     * @param[in] newContact the new unknown wrench to add.
     * @return true if all went well, false otherwise
     */
    bool addNewContactInFrame(const Model& model,
                              const FrameIndex frameIndex,
                              const UnknownWrenchContact& newContact);

    /**
     * Add a full wrench unknown at the origin of the specified frame.
     * Simplified version of a addNewContactInFrame, in which the contact
     * point is the origin of the frame and the unknown type is FULL_WRENCH.
     *
     * Equivalent to
     *   addNewContactInFrame(model,frame,UnknownWrenchContact(FULL_WRENCH,Position::Zero()))
     *
     * @param[in] model the model class for getting frame information.
     * @param[in] frameIndex the index of the frame in which you are expressing the new unknown
     * wrench.
     * @return true if all went well, false otherwise
     */
    bool addNewUnknownFullWrenchInFrameOrigin(const Model& model, const FrameIndex frameIndex);

    /**
     * Get a specific ContactWrench
     *
     * @param[in] linkIndex the index of the link for which the contact is retrieved
     * @param[in] contactIndex a index (between 0 and getNrOfContactsForLink(link)-1 ) identifing
     * the specific contact.
     */
    UnknownWrenchContact& contactWrench(const LinkIndex linkIndex, const size_t contactIndex);

    const UnknownWrenchContact&
    contactWrench(const LinkIndex linkIndex, const size_t contactIndex) const;

    /**
     * Get a human readable description of the LinkUnknownWrenchContacts (for debug)
     */
    std::string toString(const Model& model) const;
};

struct estimateExternalWrenchesBuffers
{
    estimateExternalWrenchesBuffers();
    estimateExternalWrenchesBuffers(const SubModelDecomposition& subModels);
    estimateExternalWrenchesBuffers(const size_t nrOfSubModels, const size_t nrOfLinks);

    /**
     * Resize the struct for the number of submodel
     */
    void resize(const SubModelDecomposition& subModels);
    void resize(const size_t nrOfSubModels, const size_t nrOfLinks);

    size_t getNrOfSubModels() const;
    size_t getNrOfLinks() const;

    /**
     * Check if the buffer size are consistent with the submodel
     * decomposition.
     */
    bool isConsistent(const SubModelDecomposition& subModels) const;

    /**
     * The problem of external wrenches estimation boils down to
     * solve a LS problem in the form argmin_x (Ax-b)^2 .
     */
    std::vector<MatrixDynSize> A;
    std::vector<VectorDynSize> x;
    std::vector<Vector6> b;

    /**
     * We compute the b term for each subtree
     * in a iterative way, so we need a buffer
     * to store it for each link
     */
    LinkWrenches b_contacts_subtree;

    /**
     * We compute the transform between each link
     * and the submodel base, for computing the
     * A matrices
     */
    LinkPositions subModelBase_H_link;
};

/**
 * \brief Estimate the external contact wrenches using the MultiBody Newton-Euler equations.
 *
 * This function is used to estimate the external contacts forces **without** using any measurement
 * of the internal FT sensors. It is tipically used to get data for calibrating the offset of
 * the internal FT sensors.
 */
bool estimateExternalWrenchesWithoutInternalFT(const Model& model,
                                               const Traversal& traversal,
                                               const LinkUnknownWrenchContacts& unknownWrenches,
                                               const JointPosDoubleArray& jointPos,
                                               const LinkVelArray& linkVel,
                                               const LinkAccArray& linkProperAcc,
                                               estimateExternalWrenchesBuffers& bufs,
                                               LinkContactWrenches& outputContactWrenches);

/**
 * \brief Estimate the external wrenches trasmitted by the contacts between the model and the
 * external environment.
 *
 * This function exploits the measurements of internal FT sensors (whose structure is contained
 * in the sensors parameters and which measurements are contained in the ftSensorsMeasurements
 * parameters) to compute an estimation of the values of the unknown wrenches specified in the
 * unknownWrenches parameter.
 *
 * @param[in] model the considered model.
 * @param[in] subModels a decomposition of the model along the joint of the six axis F/T sensors.
 * @param[in] sensors a description of the sensors available in the model.
 * @param[in] unknownWrenches a description of the contacts for which the contact wrench is unknown.
 * @param[in] linkVel a vector of link twists, expressed w.r.t to the link orientation and the link
 * origin
 * @param[in] linkProperAcc a vector of link spatial (in the Featherstone sense) and proper
 * accelerations, expressed w.r.t to the link orientation and the link origin
 * @param[in] ftSensorsMeasurements the measurements of the internal six axis F/T sensors.
 * @param[out] outputContactWrenches the estimated contact wrenches.
 * @return true if all went well (the dimension of the inputs are consistent), false otherwise
 *
 */
bool estimateExternalWrenches(const Model& model,
                              const SubModelDecomposition& subModels,
                              const SensorsList& sensors,
                              const LinkUnknownWrenchContacts& unknownWrenches,
                              const JointPosDoubleArray& jointPos,
                              const LinkVelArray& linkVel,
                              const LinkAccArray& linkProperAcc,
                              const SensorsMeasurements& ftSensorsMeasurements,
                              estimateExternalWrenchesBuffers& bufs,
                              LinkContactWrenches& outputContactWrenches);

/**
 * \brief Modified forward kinematics for torque/force estimation.
 *
 * This is a version of forward kinematics modified to fit the needs
 * of joint torques/external wrenches estimation.
 *
 * There are several difference with respect to the classical
 * forward kinematics.
 * The first one is that the only inputs necessary related to the base link are
 * the base link classical proper acceleration, the base link angular velocity
 * and the base link angular acceleration. This is because the dynamics
 * of an articulated system does not depend on an offset in linear velocity, and hence
 * the estimation of joint torques/external wrenches is not affected by the base link
 * linear velocity. This will mean that the link velocitity computed by this
 * algorithm are not the velocity of the links with respect to an inertial frame.
 * Nevertherless they can still be used for estimation.
 *
 * There are two main ways in which the base information is computed: one is
 * exploiting the knoledge that a link is not moving with respect to an inertial frame:
 * in this case the classical proper acceleration boils down to the inverted gravitational
 * acceleration, while the angular velocity and angular accelerations are equal to zero.
 * The other way is to exploit the measure of an accelerometer and of a gyroscope
 * mounted on the base link of the traversal: the accelerometer will then measure
 * directly the classical proper acceleration, while the gyroscope will measure the angular
 * velocity. The angular acceleration can be computed by numerical derivation, or simply neglected
 * if its effect on the estimation is minimal.
 *
 *
 * \param[in] model the input model
 * \param[in] traversal the traversal used to propagate the velocity and the proper acceleration
 * \param[in] base_classicalProperAcc classical proper acceleration of the base origin
 * \param[in] base_angularVel angular velocity of the base link frame
 * \param[in] base_angularAcc angular acceleration of the base link frame
 * \param[in] jointPos joint positions
 * \param[in] jointVel joint velocities
 * \param[in] jointAcc joint accelerations
 * \param[out] linkVel vector of link twists, expressed in the link frame for both orientation and
 * origin
 * \param[out] linkProperAcc vector of link proper spatial acceleration, expressed in the link frame
 * for both orientation and origin
 * @return true if all went well, false otherwise
 *
 */
bool dynamicsEstimationForwardVelAccKinematics(const Model& model,
                                               const Traversal& traversal,
                                               const Vector3& base_classicalProperAcc,
                                               const Vector3& base_angularVel,
                                               const Vector3& base_angularAcc,
                                               const JointPosDoubleArray& jointPos,
                                               const JointDOFsDoubleArray& jointVel,
                                               const JointDOFsDoubleArray& jointAcc,
                                               LinkVelArray& linkVel,
                                               LinkAccArray& linkProperAcc);

/**
 * \brief Modified forward kinematics for floating basedynamics estimation.
 *
 * This is a version of velocity forward kinematics modified to fit the needs
 * of free floating dynamics estimation.
 *
 * There are several difference with respect to the classical
 * forward kinematics.
 * The first one is that the only inputs necessary related to the base link is
 *   the base link angular velocity. This is because the dynamics
 * of an articulated system does not depend on an offset in linear velocity.
 * This will mean that the link velocities computed by this
 * algorithm are not the velocity of the links with respect to an inertial frame.
 * Nevertherless they can still be used for estimation.
 *
 * There are two main ways in which the base information is computed: one is
 * exploiting the knoledge that a link is not moving with respect to an inertial frame:
 * in this case the angular velocity is equal to zero.
 * The other way is to exploit the measure of a gyroscope
 * mounted on the base link of the traversal: the gyroscope will measure the link angular velocity.
 *
 *
 * \param[in] model the input model
 * \param[in] traversal the traversal used to propagate the velocity and the proper acceleration
 * \param[in] base_angularVel angular velocity of the base link frame
 * \param[in] jointPos joint positions
 * \param[in] jointVel joint velocities
 * \param[out] linkVel vector of link twists, expressed in the link frame for both orientation and
 * origin
 * @return true if all went well, false otherwise
 */
bool dynamicsEstimationForwardVelKinematics(const Model& model,
                                            const Traversal& traversal,
                                            const Vector3& base_angularVel,
                                            const JointPosDoubleArray& jointPos,
                                            const JointDOFsDoubleArray& jointVel,
                                            LinkVelArray& linkVel);

/**
 * \brief Compute the net internal and external wrenches (excluding gravity forces) acting on the
 * links.
 * @param[in] model the input model
 * @param[in] linkVel a vector of link twists, expressed w.r.t to the link orientation and the link
 * origin
 * @param[in] linkProperAcc a vector of link spatial (in the Featherstone sense) and proper
 * accelerations, expressed w.r.t to the link orientation and the link origin
 * @param[in] linkNetWrenchesWithoutGravity the vector of the sum of all the wrenches (both internal
 * and external, excluding gravity) acting on link i, expressed (both orientation and point) with
 * respect to the reference frame of link i
 */
bool computeLinkNetWrenchesWithoutGravity(
    const Model& model,
    const LinkVelArray& linkVel,
    const LinkAccArray& linkProperAcc,
    LinkNetTotalWrenchesWithoutGravity& linkNetWrenchesWithoutGravity);

/**
 * Compute the link contact wrenches from the net external wrenches
 *
 * If there are more than 6 unknows for link, the problem becomes ill-defined
 * and the function just assign all the external wrench to the first contact.
 * \todo(traversaro): support arbitrary LinkUnknownWrenchContacts by performing
 *                    a least square fitting, similar to what implemented in
 *                    the estimateExternalWrenches .
 */
bool estimateLinkContactWrenchesFromLinkNetExternalWrenches(
    const Model& model,
    const LinkUnknownWrenchContacts& unknownWrenches,
    const LinkNetExternalWrenches& netExtWrenches,
    LinkContactWrenches& outputContactWrenches);

} // namespace iDynTree

#endif
