// Test of kinematic consistency: check if finite differences of velocities, accelerations
// match positions

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <iostream>

#include <gtest/gtest.h>

#include "../Extras/InverseDynamics/CoilCreator.hpp"
#include "../Extras/InverseDynamics/DillCreator.hpp"
#include "../Extras/InverseDynamics/SimpleTreeCreator.hpp"
#include "BulletInverseDynamics/MultiBodyTree.hpp"

using namespace btInverseDynamics;

const int kLevel = 5;
const int kNumBodies = BT_ID_POW(2, kLevel);

// template function for calculating the norm
template <typename T>
idScalar calculateNorm(T&);
// only implemented for vec3
template <>
idScalar calculateNorm(vec3& v) {
    return std::sqrt(BT_ID_POW(v(0), 2) + BT_ID_POW(v(1), 2) + BT_ID_POW(v(2), 2));
}

// template function to convert a DiffType (finite differences)
// to a ValueType. This is for angular velocity calculations
// via finite differences.
template <typename ValueType, typename DiffType>
DiffType toDiffType(ValueType& fd, ValueType& val);

// vector case: just return finite difference approximation
template <>
vec3 toDiffType(vec3& fd, vec3& val) {
    return fd;
}

// orientation case: calculate spin tensor and extract angular velocity
template <>
vec3 toDiffType(mat33& fd, mat33& val) {
    // spin tensor
    mat33 omega_tilde = fd * val.transpose();
    // extract vector from spin tensor
    vec3 omega;
    omega(0) = 0.5 * (omega_tilde(2, 1) - omega_tilde(1, 2));
    omega(1) = 0.5 * (omega_tilde(0, 2) - omega_tilde(2, 0));
    omega(2) = 0.5 * (omega_tilde(1, 0) - omega_tilde(0, 1));
    return omega;
}

/// Class for calculating finite difference approximation
/// of time derivatives and comparing it to an analytical solution
/// DiffType and ValueType can be different, to allow comparison
/// of angular velocity vectors and orientations given as transform matrices.
template <typename ValueType, typename DiffType>
class DiffFD {
public:
    DiffFD() : m_dt(0.0), m_num_updates(0), m_max_error(0.0), m_max_value(0.0), m_valid_fd(false) {}

    void init(std::string name, idScalar dt) {
        m_name = name;
        m_dt = dt;
        m_num_updates = 0;
        m_max_error = 0.0;
        m_max_value = 0.0;
        m_valid_fd = false;
    }

    void update(const ValueType& val, const DiffType& true_diff) {
        m_val = val;
        if (m_num_updates > 2) {
            // 2nd order finite difference approximation for d(value)/dt
            ValueType diff_value_fd = (val - m_older_val) / (2.0 * m_dt);
            // convert to analytical diff type. This is for angular velocities
            m_diff_fd = toDiffType<ValueType, DiffType>(diff_value_fd, m_old_val);
            // now, calculate the error
            DiffType error_value_type = m_diff_fd - m_old_true_diff;
            idScalar error = calculateNorm<DiffType>(error_value_type);
            if (error > m_max_error) {
                m_max_error = error;
            }

            idScalar value = calculateNorm<DiffType>(m_old_true_diff);
            if (value > m_max_value) {
                m_max_value = value;
            }

            m_valid_fd = true;
        }
        m_older_val = m_old_val;
        m_old_val = m_val;
        m_old_true_diff = true_diff;
        m_num_updates++;
        m_time += m_dt;
    }

    void printMaxError() {
        printf("max_error: %e dt= %e max_value= %e fraction= %e\n", m_max_error, m_dt, m_max_value,
               m_max_value > 0.0 ? m_max_error / m_max_value : 0.0);
    }
    void printCurrent() {
        if (m_valid_fd) {
            // note: m_old_true_diff already equals m_true_diff here, so values are not aligned.
            //      (but error calculation takes this into account)
            printf("%s time: %e fd: %e %e %e true: %e %e %e\n", m_name.c_str(), m_time,
                   m_diff_fd(0), m_diff_fd(1), m_diff_fd(2), m_old_true_diff(0), m_old_true_diff(1),
                   m_old_true_diff(2));
        }
    }

    idScalar getMaxError() const { return m_max_error; }
    idScalar getMaxValue() const { return m_max_value; }

private:
    idScalar m_dt;
    ValueType m_val;
    ValueType m_old_val;
    ValueType m_older_val;
    DiffType m_old_true_diff;
    DiffType m_diff_fd;
    int m_num_updates;
    idScalar m_max_error;
    idScalar m_max_value;
    idScalar m_time;
    std::string m_name;
    bool m_valid_fd;
};

template <typename ValueType, typename DiffType>
class VecDiffFD {
public:
    VecDiffFD(std::string name, int dim, idScalar dt) : m_name(name), m_fd(dim), m_dt(dt) {
        for (int i = 0; i < m_fd.size(); i++) {
            char buf[256];
            BT_ID_SNPRINTF(buf, 256, "%s-%.2d", name.c_str(), i);
            m_fd[i].init(buf, dt);
        }
    }
    void update(int i, ValueType& val, DiffType& true_diff) { m_fd[i].update(val, true_diff); }
    idScalar getMaxError() const {
        idScalar max_error = 0;
        for (int i = 0; i < m_fd.size(); i++) {
            const idScalar error = m_fd[i].getMaxError();
            if (error > max_error) {
                max_error = error;
            }
        }
        return max_error;
    }
    idScalar getMaxValue() const {
        idScalar max_value = 0;
        for (int i = 0; i < m_fd.size(); i++) {
            const idScalar value = m_fd[i].getMaxValue();
            if (value > max_value) {
                max_value= value;
            }
        }
        return max_value;
    }
    void printMaxError() {
        printf("%s: total  dt= %e max_error= %e\n", m_name.c_str(), m_dt, getMaxError());
    }

    void printCurrent() {
        for (int i = 0; i < m_fd.size(); i++) {
            m_fd[i].printCurrent();
        }
    }

private:
    std::string m_name;
    std::vector<DiffFD<ValueType, DiffType> > m_fd;
    const idScalar m_dt;
    idScalar m_max_error;
};

// calculate maximum difference between finite difference and analytical differentiation
int calculateDifferentiationError(const MultiBodyTreeCreator& creator, idScalar deltaT,
                                  idScalar endTime, idScalar* max_linear_velocity_error,
                                  idScalar* max_angular_velocity_error,
                                  idScalar* max_linear_acceleration_error,
                                  idScalar* max_angular_acceleration_error) {
    // setup system
    MultiBodyTree* tree = CreateMultiBodyTree(creator);
    if (0x0 == tree) {
        return -1;
    }
    // set gravity to zero, so nothing is added to accelerations in forward kinematics
    vec3 gravity_zero;
    gravity_zero(0) = 0;
    gravity_zero(1) = 0;
    gravity_zero(2) = 0;
    tree->setGravityInWorldFrame(gravity_zero);
    //
    const idScalar kAmplitude = 1.0;
    const idScalar kFrequency = 1.0;

    vecx q(tree->numDoFs());
    vecx dot_q(tree->numDoFs());
    vecx ddot_q(tree->numDoFs());
    vecx joint_forces(tree->numDoFs());

    VecDiffFD<vec3, vec3> fd_vel("linear-velocity", tree->numBodies(), deltaT);
    VecDiffFD<vec3, vec3> fd_acc("linear-acceleration", tree->numBodies(), deltaT);
    VecDiffFD<mat33, vec3> fd_omg("angular-velocity", tree->numBodies(), deltaT);
    VecDiffFD<vec3, vec3> fd_omgd("angular-acceleration", tree->numBodies(), deltaT);

    for (idScalar t = 0.0; t < endTime; t += deltaT) {
        for (int body = 0; body < tree->numBodies(); body++) {
            q(body) = kAmplitude * sin(t * 2.0 * BT_ID_PI * kFrequency);
            dot_q(body) = kAmplitude * 2.0 * BT_ID_PI * kFrequency * cos(t * 2.0 * BT_ID_PI * kFrequency);
            ddot_q(body) =
                -kAmplitude * pow(2.0 * BT_ID_PI * kFrequency, 2) * sin(t * 2.0 * BT_ID_PI * kFrequency);
        }

        if (-1 == tree->calculateInverseDynamics(q, dot_q, ddot_q, &joint_forces)) {
            delete tree;
            return -1;
        }

        // position/velocity
        for (int body = 0; body < tree->numBodies(); body++) {
            vec3 pos;
            vec3 vel;
            mat33 world_T_body;
            vec3 omega;
            vec3 dot_omega;
            vec3 acc;

            tree->getBodyOrigin(body, &pos);
            tree->getBodyTransform(body, &world_T_body);
            tree->getBodyLinearVelocity(body, &vel);
            tree->getBodyAngularVelocity(body, &omega);
            tree->getBodyLinearAcceleration(body, &acc);
            tree->getBodyAngularAcceleration(body, &dot_omega);

            fd_vel.update(body, pos, vel);
            fd_omg.update(body, world_T_body, omega);
            fd_acc.update(body, vel, acc);
            fd_omgd.update(body, omega, dot_omega);


//        fd_vel.printCurrent();
//fd_acc.printCurrent();
//fd_omg.printCurrent();
//fd_omgd.printCurrent();
        }
    }

    *max_linear_velocity_error = fd_vel.getMaxError()/fd_vel.getMaxValue();
    *max_angular_velocity_error = fd_omg.getMaxError()/fd_omg.getMaxValue();
    *max_linear_acceleration_error = fd_acc.getMaxError()/fd_acc.getMaxValue();
    *max_angular_acceleration_error = fd_omgd.getMaxError()/fd_omgd.getMaxValue();

    delete tree;
    return 0;
}

// first test: absolute difference between numerical and numerial
// differentiation should be small
TEST(InvDynKinematicsDifferentiation, errorAbsolute) {
    //CAVEAT:these values are hand-tuned to work for the specific trajectory defined above.
#ifdef BT_ID_USE_DOUBLE_PRECISION
    const idScalar kDeltaT = 1e-7;
	const idScalar kAcceptableError = 1e-4;
#else
    const idScalar kDeltaT = 1e-4;
	const idScalar kAcceptableError = 5e-3;
#endif
    const idScalar kDuration = 0.01;
    

    CoilCreator coil_creator(kNumBodies);
    DillCreator dill_creator(kLevel);
    SimpleTreeCreator simple_creator(kNumBodies);

    idScalar max_linear_velocity_error;
    idScalar max_angular_velocity_error;
    idScalar max_linear_acceleration_error;
    idScalar max_angular_acceleration_error;

    // test serial chain
    calculateDifferentiationError(coil_creator, kDeltaT, kDuration, &max_linear_velocity_error,
                                  &max_angular_velocity_error, &max_linear_acceleration_error,
                                  &max_angular_acceleration_error);

    EXPECT_LT(max_linear_velocity_error, kAcceptableError);
    EXPECT_LT(max_angular_velocity_error, kAcceptableError);
    EXPECT_LT(max_linear_acceleration_error, kAcceptableError);
    EXPECT_LT(max_angular_acceleration_error, kAcceptableError);

    // test branched tree
    calculateDifferentiationError(dill_creator, kDeltaT, kDuration, &max_linear_velocity_error,
                                  &max_angular_velocity_error, &max_linear_acceleration_error,
                                  &max_angular_acceleration_error);

    EXPECT_LT(max_linear_velocity_error, kAcceptableError);
    EXPECT_LT(max_angular_velocity_error, kAcceptableError);
    EXPECT_LT(max_linear_acceleration_error, kAcceptableError);
    EXPECT_LT(max_angular_acceleration_error, kAcceptableError);

    // test system with different joint types
    calculateDifferentiationError(simple_creator, kDeltaT, kDuration, &max_linear_velocity_error,
                                  &max_angular_velocity_error, &max_linear_acceleration_error,
                                  &max_angular_acceleration_error);

    EXPECT_LT(max_linear_velocity_error, kAcceptableError);
    EXPECT_LT(max_angular_velocity_error, kAcceptableError);
    EXPECT_LT(max_linear_acceleration_error, kAcceptableError);
    EXPECT_LT(max_angular_acceleration_error, kAcceptableError);
}

// second test: check if the change in the differentiation error
// is consitent with the second order approximation, ie, error ~ O(dt^2)
TEST(InvDynKinematicsDifferentiation, errorOrder) {
    const idScalar kDeltaTs[2] = {1e-4, 1e-5};
    const idScalar kDuration = 1e-2;

    CoilCreator coil_creator(kNumBodies);
    //    DillCreator dill_creator(kLevel);
    //    SimpleTreeCreator simple_creator(kNumBodies);

    idScalar max_linear_velocity_error[2];
    idScalar max_angular_velocity_error[2];
    idScalar max_linear_acceleration_error[2];
    idScalar max_angular_acceleration_error[2];

    // test serial chain
    calculateDifferentiationError(coil_creator, kDeltaTs[0], kDuration,
                                  &max_linear_velocity_error[0], &max_angular_velocity_error[0],
                                  &max_linear_acceleration_error[0],
                                  &max_angular_acceleration_error[0]);

    calculateDifferentiationError(coil_creator, kDeltaTs[1], kDuration,
                                  &max_linear_velocity_error[1], &max_angular_velocity_error[1],
                                  &max_linear_acceleration_error[1],
                                  &max_angular_acceleration_error[1]);

    const idScalar expected_linear_velocity_error_1 =
        max_linear_velocity_error[0] * pow(kDeltaTs[1] / kDeltaTs[0], 2);
    const idScalar expected_angular_velocity_error_1 =
        max_angular_velocity_error[0] * pow(kDeltaTs[1] / kDeltaTs[0], 2);
    const idScalar expected_linear_acceleration_error_1 =
        max_linear_acceleration_error[0] * pow(kDeltaTs[1] / kDeltaTs[0], 2);
    const idScalar expected_angular_acceleration_error_1 =
        max_angular_acceleration_error[0] * pow(kDeltaTs[1] / kDeltaTs[0], 2);

/*    printf("linear vel error: %e %e  %e\n", max_linear_velocity_error[1],
           expected_linear_velocity_error_1,
           max_linear_velocity_error[1] - expected_linear_velocity_error_1);
    printf("angular vel error: %e %e  %e\n", max_angular_velocity_error[1],
           expected_angular_velocity_error_1,
           max_angular_velocity_error[1] - expected_angular_velocity_error_1);
    printf("linear acc error: %e %e  %e\n", max_linear_acceleration_error[1],
           expected_linear_acceleration_error_1,
           max_linear_acceleration_error[1] - expected_linear_acceleration_error_1);
    printf("angular acc error: %e %e  %e\n", max_angular_acceleration_error[1],
           expected_angular_acceleration_error_1,
           max_angular_acceleration_error[1] - expected_angular_acceleration_error_1);
*/
}

int main(int argc, char** argv) {

    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();

    return EXIT_SUCCESS;
}