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

#include "iDynTree/GravityCompensationHelpers.h"

namespace iDynTree
{

GravityCompensationHelper::GravityCompensationHelper()
    : m_isModelValid(false)
    , m_isKinematicsUpdated(false)
    , m_model()
    , m_dynamicTraversal()
    , m_kinematicTraversals()
    , m_jointPos()
    , m_linkVels()
    , m_linkProperAccs()
    , m_linkNetExternalWrenchesZero()
    , m_linkIntWrenches()
    , m_generalizedTorques()
{
}

GravityCompensationHelper::~GravityCompensationHelper()
{
    freeKinematicTraversals();
}

void GravityCompensationHelper::allocKinematicTraversals(const size_t nrOfLinks)
{
    m_kinematicTraversals.resize(nrOfLinks);
    for (size_t link = 0; link < m_kinematicTraversals.size(); link++)
    {
        m_kinematicTraversals[link] = new iDynTree::Traversal;
    }
}

void GravityCompensationHelper::freeKinematicTraversals()
{
    for (size_t link = 0; link < m_kinematicTraversals.size(); link++)
    {
        delete m_kinematicTraversals[link];
    }
    m_kinematicTraversals.resize(0);
}

bool GravityCompensationHelper::loadModel(const Model& _model, const std::string dynamicBase)
{
    m_model = _model;

    // resize the data structures
    iDynTree::LinkIndex dynamicBaseIndex = m_model.getLinkIndex(dynamicBase);

    bool ok = m_model.computeFullTreeTraversal(m_dynamicTraversal, dynamicBaseIndex);
    if (!ok)
    {
        m_isModelValid = false;
        return false;
    }

    freeKinematicTraversals();
    allocKinematicTraversals(m_model.getNrOfLinks());

    m_jointPos.resize(m_model);
    m_jointDofsZero.resize(m_model);
    m_jointDofsZero.zero();
    m_linkVels.resize(m_model);
    m_linkProperAccs.resize(m_model);
    m_linkIntWrenches.resize(m_model);
    m_linkNetExternalWrenchesZero.resize(m_model);
    m_linkNetExternalWrenchesZero.zero();
    m_generalizedTorques.resize(m_model);

    // set the model valid
    m_isModelValid = true;
    return true;
}

bool GravityCompensationHelper::updateKinematicsFromGravity(const JointPosDoubleArray& jointPos,
                                                            const FrameIndex& floatingFrame,
                                                            const Vector3& gravity)
{
    if (!m_isModelValid)
    {
        iDynTree::reportError("GravityCompensationHelper",
                              "updateKinematicsFromGravity",
                              "Model and sensors information is not set");
        return false;
    }

    iDynTree::Vector3 properClassicalAcceleration;
    properClassicalAcceleration(0) = -gravity(0);
    properClassicalAcceleration(1) = -gravity(1);
    properClassicalAcceleration(2) = -gravity(2);

    return updateKinematicsFromProperAcceleration(jointPos,
                                                  floatingFrame,
                                                  properClassicalAcceleration);
}

bool GravityCompensationHelper::updateKinematicsFromProperAcceleration(
    const JointPosDoubleArray& jointPos,
    const FrameIndex& floatingFrame,
    const Vector3& properClassicalLinearAcceleration)
{
    if (!m_isModelValid)
    {
        iDynTree::reportError("GravityCompensationHelper",
                              "updateKinematicsFromProperAcceleration",
                              "Model and sensors information is not set");
        return false;
    }

    if (floatingFrame == FRAME_INVALID_INDEX || floatingFrame < 0
        || floatingFrame >= (int)m_model.getNrOfFrames())
    {
        iDynTree::reportError("GravityCompensationHelper",
                              "updateKinematicsFromProperAcceleration",
                              "Unknown frame index specified");
        return false;
    }

    // Get link index of the specified frame
    iDynTree::LinkIndex floatingLinkIndex = m_model.getFrameLink(floatingFrame);

    // Build the traversal if it is not present
    if (m_kinematicTraversals[floatingLinkIndex]->getNrOfVisitedLinks() == 0)
    {
        m_model.computeFullTreeTraversal(*m_kinematicTraversals[floatingLinkIndex],
                                         floatingLinkIndex);
    }

    // To initialize the kinmatic propagation, we should first convert the kinematics
    // information from the frame in which they are specified to the main frame of the link
    iDynTree::Transform link_H_frame = m_model.getFrameTransform(floatingFrame);

    iDynTree::Twist base_vel_frame, base_vel_link;
    iDynTree::SpatialAcc base_acc_frame, base_acc_link;
    iDynTree::ClassicalAcc base_classical_acc_link;

    iDynTree::Vector3 zero3;
    zero3.zero();

    base_vel_frame.setLinearVec3(zero3);
    base_vel_frame.setAngularVec3(zero3);
    base_vel_link = link_H_frame * base_vel_frame;

    base_acc_frame.setLinearVec3(properClassicalLinearAcceleration);
    base_acc_frame.setAngularVec3(zero3);
    base_acc_link = link_H_frame * base_acc_frame;

    base_classical_acc_link.fromSpatial(base_acc_link, base_vel_link);

    // Propagate the kinematics information
    bool ok = iDynTree::
        dynamicsEstimationForwardVelAccKinematics(m_model,
                                                  *(m_kinematicTraversals[floatingLinkIndex]),
                                                  base_classical_acc_link.getLinearVec3(),
                                                  base_vel_link.getAngularVec3(),
                                                  base_classical_acc_link.getAngularVec3(),
                                                  jointPos,
                                                  m_jointDofsZero,
                                                  m_jointDofsZero,
                                                  m_linkVels,
                                                  m_linkProperAccs);
    // store joint positions
    m_jointPos = jointPos;

    if (!ok)
    {
        return false;
    } else
    {
        m_isKinematicsUpdated = true;
        return true;
    }
}

bool GravityCompensationHelper::getGravityCompensationTorques(JointDOFsDoubleArray& jointTrqs)
{
    if (!m_isModelValid)
    {
        iDynTree::reportError("GravityCompensationHelper",
                              "getGravityCompensationTorques",
                              "Model not set");
        return false;
    }

    if (!m_isKinematicsUpdated)
    {
        iDynTree::reportError("GravityCompensationHelper",
                              "getGravityCompensationTorques",
                              "Kinematic information not set");
        return false;
    }

    // Compute joint torques
    bool ok = iDynTree::RNEADynamicPhase(m_model,
                                         m_dynamicTraversal,
                                         m_jointPos,
                                         m_linkVels,
                                         m_linkProperAccs,
                                         m_linkNetExternalWrenchesZero,
                                         m_linkIntWrenches,
                                         m_generalizedTorques);
    if (!ok)
    {
        iDynTree::reportError("GravityCompensationHelper",
                              "getGravityCompensationTorques",
                              "Error in computing RNEADynamicPhase");
        return false;
    }

    // store computed torques
    jointTrqs = m_generalizedTorques.jointTorques();

    return true;
}

} // namespace iDynTree
