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

#include <iDynTree/EigenHelpers.h>
#include <iDynTree/SpatialInertia.h>
#include <iDynTree/SpatialMomentum.h>
#include <iDynTree/Wrench.h>
#include <iDynTree/Transform.h>
#include <iDynTree/Utils.h>
#include <iDynTree/TestUtils.h>
#include <iDynTree/Twist.h>
#include <iDynTree/MatrixFixSize.h>
#include <iDynTree/SpatialAcc.h>
#include <iDynTree/InertiaNonLinearParametrization.h>

#include <Eigen/Dense>

#include <cstdio>
#include <cstdlib>

using namespace iDynTree;

SpatialInertia getNonPhysicalConsistentInertia()
{
    Vector16 paramsVec;
    paramsVec.zero();
    paramsVec(0) = -1.0;

    paramsVec(1) = 1.0;
    paramsVec(2) = 1.0;
    paramsVec(3) = 0.0;
    /*
    paramsVec(4) = 0.66474;
    paramsVec(5) = 0.534487;
    paramsVec(6) = -0.521962;
    paramsVec(7) = 0.124066;
    paramsVec(8) = 0.609993;
    paramsVec(9) = 0.782634;
    paramsVec(10) = 0.736701;
    paramsVec(11) = -0.585006;
    paramsVec(12) = 0.339175;
    paramsVec(13) = 0.434594;
    paramsVec(14) = -0.716795;
    paramsVec(15) = 0.213938;*/
    RigidBodyInertiaNonLinearParametrization params;
    params.fromVectorWithRotationAsVec(paramsVec);
    SpatialInertia I = params.toRigidBodyInertia();
    return I;
}

void checkInertiaTransformation(const Transform & trans, const SpatialInertia & inertia)
{
    SpatialInertia inertiaTranslated = trans*inertia;
    Matrix6x6      inertiaTranslatedCheck;

    Matrix6x6 adjWre = trans.asAdjointTransformWrench();
    Matrix6x6 I      = inertia.asMatrix();
    Matrix6x6 adj    = trans.inverse().asAdjointTransform();
    toEigen(inertiaTranslatedCheck) = toEigen(adjWre)*
                                      toEigen(I)*
                                      toEigen(adj);


    std::cerr << "inertiaTranslated :\n" << inertiaTranslated.asMatrix().toString() << std::endl;
    std::cerr << "inertiaTranslatedCheck :\n" << inertiaTranslatedCheck.toString() << std::endl;

    Matrix6x6 inertiaTranslatedRaw = inertiaTranslated.asMatrix();
    ASSERT_EQUAL_MATRIX(inertiaTranslatedCheck,inertiaTranslatedRaw);
}

void checkInertiaTwistProduct(const SpatialInertia & inertia, const Twist & twist)
{
    SpatialMomentum momentum = inertia*twist;
    Vector6         momentumCheck;

    Matrix6x6 I = inertia.asMatrix();

    Vector6 twistPlain = twist.asVector();
    toEigen(momentumCheck) = toEigen(I)*toEigen(twistPlain);

    ASSERT_EQUAL_VECTOR(momentum.asVector(),momentumCheck);
}

void checkInvariance(const Transform & trans, const SpatialInertia & inertia, const Twist & twist)
{
    Transform invTrans = trans.inverse();
    SpatialMomentum momentumTranslated = trans*(inertia*twist);
    SpatialMomentum momentumTranslatedCheck = (trans*inertia)*(trans*twist);

    Twist           twistTranslated         = trans*twist;
    SpatialInertia  inertiaTranslated       = trans*inertia;
    Vector6 momentumTranslatedCheck2;
    Vector6 momentumTranslatedCheck3;
    Vector6 twistTranslatedCheck;
    Matrix6x6 transAdjWrench = trans.asAdjointTransformWrench();
    Matrix6x6 inertiaRaw     = inertia.asMatrix();
    Matrix6x6 transInvAdj    = trans.inverse().asAdjointTransform();
    Matrix6x6 transAdj       = trans.asAdjointTransform();
    Matrix6x6 inertiaTranslatedCheck;
    Vector6 twistPlain = twist.asVector();

    toEigen(momentumTranslatedCheck2) = toEigen(transAdjWrench)*
                                        toEigen(inertiaRaw)*
                                        toEigen(twistPlain);

    toEigen(momentumTranslatedCheck3) = toEigen(transAdjWrench)*
                                        toEigen(inertiaRaw)*
                                        toEigen(transInvAdj)*
                                        toEigen(transAdj)*
                                        toEigen(twistPlain);

    toEigen(twistTranslatedCheck)     = toEigen(transAdj)*
                                        toEigen(twistPlain);

    toEigen(inertiaTranslatedCheck)   = toEigen(transAdjWrench)*
                                        toEigen(inertiaRaw)*
                                        toEigen(transInvAdj);

    ASSERT_EQUAL_MATRIX(inertiaTranslatedCheck,inertiaTranslated.asMatrix());
    ASSERT_EQUAL_VECTOR(twistTranslated.asVector(),twistTranslatedCheck);
    ASSERT_EQUAL_VECTOR(momentumTranslated.asVector(),momentumTranslatedCheck3);
    ASSERT_EQUAL_VECTOR(momentumTranslated.asVector(),momentumTranslatedCheck2);
    ASSERT_EQUAL_VECTOR(momentumTranslated.asVector(),momentumTranslatedCheck.asVector());

    SpatialMomentum momentum = invTrans*momentumTranslated;
    SpatialMomentum momentumCheck = (invTrans*(trans*inertia))*(invTrans*(trans*twist));

    ASSERT_EQUAL_VECTOR(momentum.asVector(),momentumCheck.asVector());
}

void checkBiasWrench(const SpatialInertia & inertia, const Twist & twist)
{
    Wrench biasWrench = inertia.biasWrench(twist);
    Wrench biasWrenchCheck = twist*(inertia*twist);

    ASSERT_EQUAL_VECTOR(biasWrench.asVector(),biasWrenchCheck.asVector());
}

void checkRegressors(const SpatialInertia & I,
                     const Twist & v,
                     const SpatialAcc & a,
                     const Twist & vRef)
{
    //////////
    /// Check the momentum regressor
    /////////
    // Compute momentum with classical operation
    SpatialMomentum h = I*v;

    // Compute the inertial parameters for the given spatial inertia
    Vector10 inertialParams = I.asVector();

    // Compute momentum with regressors
    Vector6 hRegr;

    toEigen(hRegr) = toEigen(SpatialInertia::momentumRegressor(v))*toEigen(inertialParams);

    // Check that they are equal
    ASSERT_EQUAL_VECTOR(h.asVector(),hRegr);

    //////////
    /// Check the momentum derivative (net wrench) regressor
    /////////

    // Compute momentum derivative with classical operation
    Wrench dotH = I*a+v*(I*v);

    // Compute momentum derivative with regressors
    Vector6 dotHregr;

    toEigen(dotHregr) = toEigen(SpatialInertia::momentumDerivativeRegressor(v,a))*toEigen(inertialParams);

    // Check that they are equal
    ASSERT_EQUAL_VECTOR(dotH.asVector(),dotHregr);

    //////////
    /// Check the momentum derivative (net wrench) regressor, Slotine Li version
    /// Check also that the operator is simmetric with respect to v, vDes
    /////////

    Wrench dotH_SL = I*a+v*(I*vRef)-I*(v*vRef);

    Vector6 dotH_SL_regr;

    toEigen(dotH_SL_regr) =
        toEigen(SpatialInertia::momentumDerivativeSlotineLiRegressor(v,vRef,a))*toEigen(inertialParams);

    ASSERT_EQUAL_VECTOR(dotH_SL.asVector(),dotH_SL_regr);

    // Check that the regressor matches the momentumDerivativeRegressor one if v = vRef
    Wrench dotH_SL_check = I*a+v*(I*v)-I*(v*v);

    Vector6 dotH_SL_regr_check;

    toEigen(dotH_SL_regr_check) =
        toEigen(SpatialInertia::momentumDerivativeSlotineLiRegressor(v,v,a))*toEigen(inertialParams);

    ASSERT_EQUAL_VECTOR(dotH_SL_check.asVector(),dotH_SL_regr_check);
    ASSERT_EQUAL_VECTOR(dotH.asVector(),dotH_SL_check.asVector());

}

// Check the nonlinear parametrization of inertia
void checkInertiaNonLinearParametrization()
{
    RigidBodyInertiaNonLinearParametrization inertiaParams, inertiaParamsNotConsistent, inertiaParamsCheck;
    // populate the inertia parametrization with random (but physical consistent data)
    inertiaParams.mass = getRandomDouble(0.3,10.0);
    inertiaParams.com = getRandomPosition();
    inertiaParams.link_R_centroidal = getRandomRotation();
    inertiaParams.centralSecondMomentOfMass(2) = getRandomDouble(0.0,10.0);
    inertiaParams.centralSecondMomentOfMass(1) = getRandomDouble(0.0,inertiaParams.centralSecondMomentOfMass(2));
    inertiaParams.centralSecondMomentOfMass(0) = getRandomDouble(0.0,inertiaParams.centralSecondMomentOfMass(1));

    ASSERT_IS_TRUE(inertiaParams.isPhysicallyConsistent());

    inertiaParamsNotConsistent.mass = getRandomDouble(-10.0,-0.3);
    inertiaParamsNotConsistent.com.zero();
    inertiaParamsNotConsistent.link_R_centroidal =  getRandomRotation(); //iDynTree::Rotation::Identity();
    inertiaParamsNotConsistent.centralSecondMomentOfMass(0) = 1.0; //getRandomDouble(0.0,10.0);
    inertiaParamsNotConsistent.centralSecondMomentOfMass(1) = 2.0; //getRandomDouble(0.0,inertiaParams.centralSecondMomentOfMass(0));
    inertiaParamsNotConsistent.centralSecondMomentOfMass(2) = 3.0; //getRandomDouble(0.0,inertiaParams.centralSecondMomentOfMass(1));

    ASSERT_IS_FALSE(inertiaParamsNotConsistent.isPhysicallyConsistent());

    // Check if the two nonlinear-->inertia--> nonlinear works fine
    inertiaParamsCheck.fromRigidBodyInertia(inertiaParams.toRigidBodyInertia());

    ASSERT_EQUAL_DOUBLE(inertiaParams.mass,inertiaParamsCheck.mass);
    ASSERT_EQUAL_VECTOR(inertiaParams.com,inertiaParamsCheck.com);

    //ASSERT_EQUAL_MATRIX(inertiaParams.link_R_centroidal,inertiaParamsCheck.link_R_centroidal);

    ASSERT_EQUAL_VECTOR(inertiaParams.centralSecondMomentOfMass,inertiaParamsCheck.centralSecondMomentOfMass);
    ASSERT_EQUAL_MATRIX(inertiaParams.toRigidBodyInertia().asMatrix(),inertiaParamsCheck.toRigidBodyInertia().asMatrix());

    // Check mapping to and from usual inertial parameters
    RigidBodyInertiaNonLinearParametrization inertiaParamsCheck2;

    inertiaParamsCheck2.fromInertialParameters(inertiaParams.toRigidBodyInertia().asVector());

    ASSERT_EQUAL_MATRIX(inertiaParams.toRigidBodyInertia().asMatrix(),inertiaParamsCheck2.toRigidBodyInertia().asMatrix());
}

void checkInertiaNonLinearParametrizationGradient()
{
    RigidBodyInertiaNonLinearParametrization inertiaParams;

    inertiaParams.mass = getRandomDouble(0.3,10.0);
    inertiaParams.com = getRandomPosition();
    inertiaParams.link_R_centroidal =  getRandomRotation();
    inertiaParams.centralSecondMomentOfMass(2) = getRandomDouble(0.0,10.0);
    inertiaParams.centralSecondMomentOfMass(1) = getRandomDouble(0.0,inertiaParams.centralSecondMomentOfMass(2));
    inertiaParams.centralSecondMomentOfMass(0) = getRandomDouble(0.0,inertiaParams.centralSecondMomentOfMass(1));

    // Analytical gradient
    Matrix10x16 gradient = inertiaParams.getGradientWithRotationAsVec();

    // Numerical gradient
    Matrix10x16 numericalGradient;

    // Let's compute the numericalGradient
    double delta = 1e-4;
    Vector16 currentValue = inertiaParams.asVectorWithRotationAsVec();
    Vector10 currentOutput = inertiaParams.toRigidBodyInertia().asVector();

    Vector16 upperPerturbation;
    Vector16 lowerPerturbation;
    Vector10 upperOutput;
    Vector10 lowerOutput;
    for(int i=0; i < 16; i++ )
    {
        upperPerturbation = currentValue;
        lowerPerturbation = currentValue;

        upperPerturbation(i) = currentValue(i) + delta;
        lowerPerturbation(i) = currentValue(i) - delta;

        RigidBodyInertiaNonLinearParametrization upperPertubed,lowerPertubed;
        upperPertubed.fromVectorWithRotationAsVec(upperPerturbation);
        lowerPertubed.fromVectorWithRotationAsVec(lowerPerturbation);

        upperOutput = upperPertubed.toRigidBodyInertia().asVector();
        lowerOutput = lowerPertubed.toRigidBodyInertia().asVector();


        for(int j = 0; j < 10; j++)
        {
            numericalGradient(j,i) = (upperOutput(j)-lowerOutput(j))/(2*delta);
        }
    }

    // Check that the two are consistent
    ASSERT_EQUAL_MATRIX_TOL(gradient,numericalGradient,1e-8);

}

void checkRegressorsAsGradients()
{
    SpatialAcc a = getRandomTwist();
    Twist v = getRandomTwist();
    SpatialInertia I = getRandomInertia();

    Matrix6x10 momentumGradient, momentumDerivativeGradient;
    Matrix6x10 momentumGradientNum, momentumDerivativeGradientNum;

    momentumGradient = SpatialInertia::momentumRegressor(v);
    momentumDerivativeGradient = SpatialInertia::momentumDerivativeRegressor(v,a);

    double delta = 1e-4;
    Vector10 currentValue = I.asVector();
    Vector6  currentOutputMom = (I*v).asVector();
    Vector6  currentOutputMomDer = (I*a+v.cross(I*v)).asVector();

    Vector10 upperPerturbation;
    Vector10 lowerPerturbation;
    Vector6 upperOutputMom, lowerOutputMom, upperOutputMomDer, lowerOutputMomDer;
    for(int i=0; i < 10; i++ )
    {
        upperPerturbation = currentValue;
        lowerPerturbation = currentValue;

        upperPerturbation(i) = currentValue(i) + delta;
        lowerPerturbation(i) = currentValue(i) - delta;

        SpatialInertia upperPertubed,lowerPertubed;
        upperPertubed.fromVector(upperPerturbation);
        lowerPertubed.fromVector(lowerPerturbation);

        upperOutputMom = (upperPertubed*v).asVector();
        lowerOutputMom = (lowerPertubed*v).asVector();

        upperOutputMomDer = (upperPertubed*a+v.cross(upperPertubed*v)).asVector();
        lowerOutputMomDer = (lowerPertubed*a+v.cross(lowerPertubed*v)).asVector();

        for(int j = 0; j < 6; j++)
        {
            momentumGradientNum(j,i) = (upperOutputMom(j)-lowerOutputMom(j))/(2*delta);
            momentumDerivativeGradientNum(j,i) =  (upperOutputMomDer(j)-lowerOutputMomDer(j))/(2*delta);
        }
    }

    // Check that the two are consistent
    ASSERT_EQUAL_MATRIX_TOL(momentumGradient,momentumGradientNum,1e-8);
    ASSERT_EQUAL_MATRIX_TOL(momentumDerivativeGradient,momentumDerivativeGradientNum,1e-8);

}

void checkInertiaIdentificationCostFunction()
{
    SpatialAcc a;
    a.zero();
    a(3) = 1.0;
    Twist v;
    v.zero();
    Vector16 paramsVec;
    paramsVec.zero();
    paramsVec(0) = -1.0;

    paramsVec(1) = 1.0;
    paramsVec(2) = 1.0;
    paramsVec(3) = 0.0;
    /*
    paramsVec(4) = 0.66474;
    paramsVec(5) = 0.534487;
    paramsVec(6) = -0.521962;
    paramsVec(7) = 0.124066;
    paramsVec(8) = 0.609993;
    paramsVec(9) = 0.782634;
    paramsVec(10) = 0.736701;
    paramsVec(11) = -0.585006;
    paramsVec(12) = 0.339175;
    paramsVec(13) = 0.434594;
    paramsVec(14) = -0.716795;
    paramsVec(15) = 0.213938;*/
    RigidBodyInertiaNonLinearParametrization params;
    params.fromVectorWithRotationAsVec(paramsVec);
    SpatialInertia I = params.toRigidBodyInertia();

    Wrench f;
    f.zero();
    f(3) = 1.0;

    Vector6 residual = (f-I*a).asVector();
    double cost = toEigen(residual).transpose()*toEigen(residual);
    Vector16 costGradient;
    toEigen(costGradient) = 2*toEigen(I*a-f).transpose()*toEigen(SpatialInertia::momentumDerivativeRegressor(v,a))*toEigen(params.getGradientWithRotationAsVec());

    double delta = 1e-4;
    Vector16 currentValue = paramsVec;
    double currentOutput = cost;

    Vector16 upperPerturbation;
    Vector16 lowerPerturbation;
    Vector16 costGradientNumerical;
    for(int i=0; i < 16; i++ )
    {
        upperPerturbation = currentValue;
        lowerPerturbation = currentValue;

        upperPerturbation(i) = currentValue(i) + delta;
        lowerPerturbation(i) = currentValue(i) - delta;

        SpatialInertia upperPertubed,lowerPertubed;
        RigidBodyInertiaNonLinearParametrization upperParams, lowerParams;
        upperParams.fromVectorWithRotationAsVec(upperPerturbation);
        lowerParams.fromVectorWithRotationAsVec(lowerPerturbation);
        upperPertubed = upperParams.toRigidBodyInertia();
        lowerPertubed = lowerParams.toRigidBodyInertia();

        Vector6 upperResidual = (f-upperPertubed*a).asVector();
        double upperOutput = toEigen(upperResidual).transpose()*toEigen(upperResidual);

        Vector6 lowerResidual = (f-lowerPertubed*a).asVector();
        double lowerOutput = toEigen(lowerResidual).transpose()*toEigen(lowerResidual);

        costGradientNumerical(i) = (upperOutput-lowerOutput)/(2*delta);
    }

    /*
    std::cerr << "Residual                " << residual.toString() << std::endl;
    std::cerr << "I                     \n" << I.asMatrix().toString() << std::endl;
    std::cerr << "Cost                    " << cost << std::endl;
    std::cerr << "Cost gradient           " << costGradient.toString() << std::endl;
    std::cerr << "Cost gradient numerical " << costGradientNumerical.toString() << std::endl;
    std::cerr << "toEigen(SpatialInertia::momentumDerivativeRegressor(v,a)):\n" << SpatialInertia::momentumDerivativeRegressor(v,a).toString() << std::endl;
    std::cerr << "toEigen(params.getGradientWithRotationAsVec()):\n" << (params.getGradientWithRotationAsVec()).toString() << std::endl;
    */

    // Check that the two are consistent
    ASSERT_EQUAL_VECTOR_TOL(costGradient,costGradientNumerical,1e-6);

}

int main()
{
    Transform trans(Rotation::RPY(5.0,6.0,-4.0),Position(10,6,-4));

    double twistData[6] = {-5.0,-6.0,-5.0,1.0,2.0,3.0};
    Twist twist(LinVelocity(twistData,3),AngVelocity(twistData+3,3));

    ASSERT_EQUAL_DOUBLE(twist.asVector()(0),twistData[0]);

    double rotInertiaData[3*3] = {10.0,0.04,0.04,
                                  0.04,20.0,0.04,
                                  0.04,0.04,24.0};
    SpatialInertia inertia(1.0,Position(100,-5,10),RotationalInertia(rotInertiaData,3,3));

    checkInertiaTwistProduct(inertia,twist);
    checkInertiaTransformation(trans,inertia);
    checkInvariance(trans,inertia,twist);
    checkBiasWrench(inertia,twist);

    inertia = getNonPhysicalConsistentInertia();

    checkInertiaTwistProduct(inertia,twist);
    checkInertiaTransformation(trans,inertia);
    checkInvariance(trans,inertia,twist);
    checkBiasWrench(inertia,twist);

    int nrOfChecks = 30;
    for(int i=0; i < nrOfChecks; i++ )
    {
        inertia = getRandomInertia();
        twist   = getRandomTwist();
        Twist twist2  = getRandomTwist();
        SpatialAcc a  = getRandomTwist();

        checkRegressors(inertia,twist,a,twist2);

        checkInertiaNonLinearParametrization();
        checkInertiaNonLinearParametrizationGradient();
        checkRegressorsAsGradients();
    }

    checkInertiaIdentificationCostFunction();

    return EXIT_SUCCESS;
}