/*
 * Copyright (C) 2016 Fondazione Istituto Italiano di Tecnologia
 *
 * Licensed under either the GNU Lesser General Public License v3.0 :
 * https://www.gnu.org/licenses/lgpl-3.0.html
 * or the GNU Lesser General Public License v2.1 :
 * https://www.gnu.org/licenses/old-licenses/lgpl-2.1.html
 * at your option.
 */
#include "iDynTree/Estimation/BerdySparseMAPSolver.h"
#include "iDynTree/Estimation/BerdyHelper.h"

#include <iDynTree/Core/VectorDynSize.h>
#include <iDynTree/Core/VectorFixSize.h>
#include <iDynTree/Core/SparseMatrix.h>
#include <iDynTree/Core/EigenHelpers.h>
#include <iDynTree/Core/EigenSparseHelpers.h>

#include <Eigen/SparseCore>
#include <Eigen/SparseCholesky>

#include <cassert>

namespace iDynTree {

    class BerdySparseMAPSolver::BerdySparseMAPSolverPimpl
    {
    public:
        iDynTree::BerdyHelper& berdy;

        bool valid;

        // Priors on regularization, measurements and dynamics constraints
        // These variables should be filled in the "initialization" phase
        iDynTree::SparseMatrix<iDynTree::ColumnMajor> priorDynamicsConstraintsCovarianceInverse; // Sigma_D^-1
        iDynTree::SparseMatrix<iDynTree::ColumnMajor> priorDynamicsRegularizationCovarianceInverse; // Sigma_d^-1
        iDynTree::VectorDynSize priorDynamicsRegularizationExpectedValue; // mu_d
        iDynTree::SparseMatrix<iDynTree::ColumnMajor> priorMeasurementsCovarianceInverse; // Sigma_y^-1

        // The following variables are mainly buffers to contain intermediate results

        // These matrices contain the elements representing the berdy system
        iDynTree::SparseMatrix<iDynTree::ColumnMajor> dynamicsConstraintsMatrix;
        iDynTree::SparseMatrix<iDynTree::ColumnMajor> measurementsMatrix;
        iDynTree::VectorDynSize dynamicsConstraintsBias;
        iDynTree::VectorDynSize measurementsBias;

        // Measurements vector and system state
        iDynTree::JointPosDoubleArray jointsConfiguration;
        iDynTree::JointDOFsDoubleArray jointsVelocity;
        iDynTree::VectorDynSize measurements;

        // Expected value and variance of the prior on the dynamics
        iDynTree::VectorDynSize expectedDynamicsPrior;
        Eigen::SparseMatrix<double, Eigen::ColMajor> covarianceDynamicsPriorInverse;

        iDynTree::VectorDynSize expectedDynamicsPriorRHS;

        // Expected value and variance of the a-posteriori on the dynamics
        iDynTree::VectorDynSize expectedDynamicsAPosteriori;
        Eigen::SparseMatrix<double, Eigen::ColMajor> covarianceDynamicsAPosterioriInverse;
        iDynTree::VectorDynSize expectedDynamicsAPosterioriRHS;

        // Decomposition buffers
        Eigen::SimplicialLDLT<Eigen::SparseMatrix<double, Eigen::ColMajor> > covarianceDynamicsPriorInverseDecomposition;
        Eigen::SimplicialLDLT<Eigen::SparseMatrix<double, Eigen::ColMajor> > covarianceDynamicsAPosterioriInverseDecomposition;

        BerdySparseMAPSolverPimpl(BerdyHelper& berdyHelper)
        : berdy(berdyHelper)
        , valid(false)
        {
            initialize();
        }

        bool initialize();
        void computeMAP(bool computePermutation);
        static bool invertSparseMatrix(const iDynTree::SparseMatrix<iDynTree::ColumnMajor>&in, iDynTree::SparseMatrix<iDynTree::ColumnMajor>& inverted);
    };

    BerdySparseMAPSolver::BerdySparseMAPSolver(BerdyHelper& berdyHelper)
    : m_pimpl(new BerdySparseMAPSolverPimpl(berdyHelper))
    {
        assert(m_pimpl);
    }

    BerdySparseMAPSolver::~BerdySparseMAPSolver()
    {
        assert(m_pimpl);
        delete m_pimpl;
    }

    void BerdySparseMAPSolver::setDynamicsConstraintsPriorCovariance(const iDynTree::SparseMatrix<iDynTree::ColumnMajor>& covariance)
    {
        assert(m_pimpl);
        assert(m_pimpl->berdy.getOptions().berdyVariant!=BERDY_FLOATING_BASE_NON_COLLOCATED_EXT_WRENCHES);
        assert(covariance.rows() == m_pimpl->priorDynamicsConstraintsCovarianceInverse.rows()
               && covariance.columns() == m_pimpl->priorDynamicsConstraintsCovarianceInverse.columns());
        BerdySparseMAPSolver::BerdySparseMAPSolverPimpl::invertSparseMatrix(covariance, m_pimpl->priorDynamicsConstraintsCovarianceInverse);
    }

    void BerdySparseMAPSolver::setDynamicsRegularizationPriorCovariance(const iDynTree::SparseMatrix<iDynTree::ColumnMajor>& covariance)
    {
        assert(m_pimpl);
        assert(covariance.rows() == m_pimpl->priorDynamicsRegularizationCovarianceInverse.rows()
               && covariance.columns() == m_pimpl->priorDynamicsRegularizationCovarianceInverse.columns());
        BerdySparseMAPSolver::BerdySparseMAPSolverPimpl::invertSparseMatrix(covariance, m_pimpl->priorDynamicsRegularizationCovarianceInverse);
    }

    void BerdySparseMAPSolver::setDynamicsRegularizationPriorExpectedValue(const iDynTree::VectorDynSize& expectedValue)
    {
        assert(m_pimpl);
        assert(expectedValue.size() == m_pimpl->priorDynamicsRegularizationExpectedValue.size());
        m_pimpl->priorDynamicsRegularizationExpectedValue = expectedValue;
    }

    void BerdySparseMAPSolver::setMeasurementsPriorCovariance(const iDynTree::SparseMatrix<iDynTree::ColumnMajor>& covariance)
    {
        assert(m_pimpl);
        assert(covariance.rows() == m_pimpl->priorMeasurementsCovarianceInverse.rows()
               && covariance.columns() == m_pimpl->priorMeasurementsCovarianceInverse.columns());
        BerdySparseMAPSolver::BerdySparseMAPSolverPimpl::invertSparseMatrix(covariance, m_pimpl->priorMeasurementsCovarianceInverse);
    }

    const iDynTree::SparseMatrix<iDynTree::ColumnMajor>& BerdySparseMAPSolver::dynamicsConstraintsPriorCovarianceInverse() const
    {
        assert(m_pimpl);
        return m_pimpl->priorDynamicsConstraintsCovarianceInverse;
    }

    iDynTree::SparseMatrix<iDynTree::ColumnMajor>& BerdySparseMAPSolver::dynamicsConstraintsPriorCovarianceInverse()
    {
        assert(m_pimpl);
        return m_pimpl->priorDynamicsConstraintsCovarianceInverse;
    }

    const iDynTree::SparseMatrix<iDynTree::ColumnMajor>& BerdySparseMAPSolver::dynamicsRegularizationPriorCovarianceInverse() const
    {
        assert(m_pimpl);
        return m_pimpl->priorDynamicsRegularizationCovarianceInverse;
    }

    iDynTree::SparseMatrix<iDynTree::ColumnMajor>& BerdySparseMAPSolver::dynamicsRegularizationPriorCovarianceInverse()
    {
        assert(m_pimpl);
        return m_pimpl->priorDynamicsRegularizationCovarianceInverse;
    }

    const iDynTree::VectorDynSize& BerdySparseMAPSolver::dynamicsRegularizationPriorExpectedValue() const
    {
        assert(m_pimpl);
        return m_pimpl->priorDynamicsRegularizationExpectedValue;
    }

    iDynTree::VectorDynSize& BerdySparseMAPSolver::dynamicsRegularizationPriorExpectedValue()
    {
        assert(m_pimpl);
        return m_pimpl->priorDynamicsRegularizationExpectedValue;
    }

    const iDynTree::SparseMatrix<iDynTree::ColumnMajor>& BerdySparseMAPSolver::measurementsPriorCovarianceInverse() const
    {
        assert(m_pimpl);
        return m_pimpl->priorMeasurementsCovarianceInverse;
    }

    iDynTree::SparseMatrix<iDynTree::ColumnMajor>& BerdySparseMAPSolver::measurementsPriorCovarianceInverse()
    {
        assert(m_pimpl);
        return m_pimpl->priorMeasurementsCovarianceInverse;
    }

    bool BerdySparseMAPSolver::isValid()
    {
        assert(m_pimpl);
        return m_pimpl->valid;
    }

    bool BerdySparseMAPSolver::initialize()
    {
        assert(m_pimpl);
        bool init = m_pimpl->initialize();
        return init && m_pimpl->valid;
    }

    void BerdySparseMAPSolver::updateEstimateInformationFixedBase(const iDynTree::JointPosDoubleArray& jointsConfiguration,
                                                                  const iDynTree::JointDOFsDoubleArray& jointsVelocity,
                                                                  const FrameIndex fixedFrame,
                                                                  const Vector3& gravity,
                                                                  const iDynTree::VectorDynSize& measurements)
    {
        assert(m_pimpl);

        m_pimpl->jointsConfiguration = jointsConfiguration;
        m_pimpl->jointsVelocity = jointsVelocity;
        m_pimpl->measurements = measurements;

        m_pimpl->berdy.updateKinematicsFromFixedBase(m_pimpl->jointsConfiguration, m_pimpl->jointsVelocity, fixedFrame, gravity);
    }

    void BerdySparseMAPSolver::updateEstimateInformationFloatingBase(const iDynTree::JointPosDoubleArray& jointsConfiguration,
                                                                     const iDynTree::JointDOFsDoubleArray& jointsVelocity,
                                                                     const FrameIndex floatingFrame,
                                                                     const Vector3& bodyAngularVelocityOfSpecifiedFrame,
                                                                     const iDynTree::VectorDynSize& measurements)
    {
        assert(m_pimpl);

        m_pimpl->jointsConfiguration = jointsConfiguration;
        m_pimpl->jointsVelocity = jointsVelocity;
        m_pimpl->measurements = measurements;

        m_pimpl->berdy.updateKinematicsFromFloatingBase(m_pimpl->jointsConfiguration, m_pimpl->jointsVelocity,
                                                        floatingFrame, bodyAngularVelocityOfSpecifiedFrame);
    }

    bool BerdySparseMAPSolver::doEstimate()
    {
        assert(m_pimpl);
        using iDynTree::toEigen;
#ifdef EIGEN_RUNTIME_NO_MALLOC
        Eigen::internal::set_is_malloc_allowed(false);
#endif
        bool computePermutation = false;
        m_pimpl->computeMAP(computePermutation);

#ifdef EIGEN_RUNTIME_NO_MALLOC
        Eigen::internal::set_is_malloc_allowed(true);
#endif
        return true;
    }

    void BerdySparseMAPSolver::getLastEstimate(iDynTree::VectorDynSize& lastEstimate) const
    {
        assert(m_pimpl);
        lastEstimate = m_pimpl->expectedDynamicsAPosteriori;
    }

    const iDynTree::VectorDynSize& BerdySparseMAPSolver::getLastEstimate() const
    {
        assert(m_pimpl);
        return m_pimpl->expectedDynamicsAPosteriori;
    }

    void BerdySparseMAPSolver::BerdySparseMAPSolverPimpl::computeMAP(bool computePermutation)
    {
        /*
         * Get berdy matrices
         */
        berdy.getBerdyMatrices(dynamicsConstraintsMatrix,
                               dynamicsConstraintsBias,
                               measurementsMatrix,
                               measurementsBias);


        // Compute the maximum a posteriori probability
        // See Latella et al., "Whole-Body Human Inverse Dynamics with
        // Distributed Micro-Accelerometers, Gyros and Force Sensing" in Sensors, 2016

        // Intermediate quantities

        if(berdy.getOptions().berdyVariant!=BERDY_FLOATING_BASE_NON_COLLOCATED_EXT_WRENCHES)
        {
            // Covariance matrix of the prior of the dynamics: var[p(d)], Eq. 10a
            //TODO: find a way to map to iDynTree::SparseMatrix
            covarianceDynamicsPriorInverse
            //    toEigen(m_intermediateQuantities.covarianceDynamicsPriorInverse)
            //Better to assign the "sum" before, and adding the product part later
            = toEigen(priorDynamicsRegularizationCovarianceInverse);

            covarianceDynamicsPriorInverse += toEigen(dynamicsConstraintsMatrix).transpose() * toEigen(priorDynamicsConstraintsCovarianceInverse) * toEigen(dynamicsConstraintsMatrix);

            // decompose m_covarianceDynamicsPriorInverse
            if (computePermutation) {
                covarianceDynamicsPriorInverseDecomposition.analyzePattern(covarianceDynamicsPriorInverse);
            }
            //    m_intermediateQuantities.covarianceDynamicsPriorInverseDecomposition.factorize(toEigen(m_intermediateQuantities.covarianceDynamicsPriorInverse));
            covarianceDynamicsPriorInverseDecomposition.factorize(covarianceDynamicsPriorInverse);

            // Expected value of the prior of the dynamics: E[p(d)], Eq. 10b
            toEigen(expectedDynamicsPriorRHS) = toEigen(priorDynamicsRegularizationCovarianceInverse) * toEigen(priorDynamicsRegularizationExpectedValue) - toEigen(dynamicsConstraintsMatrix).transpose() * toEigen(priorDynamicsConstraintsCovarianceInverse) * toEigen(dynamicsConstraintsBias);

            toEigen(expectedDynamicsPrior) =
            covarianceDynamicsPriorInverseDecomposition.solve(toEigen(expectedDynamicsPriorRHS));
        }
        else
        {
            // Modified eq. 10a and 10b without the dynamics constraints
            covarianceDynamicsPriorInverse = toEigen(priorDynamicsRegularizationCovarianceInverse);
            expectedDynamicsPrior = priorDynamicsRegularizationExpectedValue;
        }
        
        // Final result: covariance matrix of the whole-body dynamics, Eq. 11a
        //TODO: find a way to map to iDynTree::SparseMatrix
        covarianceDynamicsAPosterioriInverse = toEigen(measurementsMatrix).transpose() * toEigen(priorMeasurementsCovarianceInverse) * toEigen(measurementsMatrix);
        covarianceDynamicsAPosterioriInverse += covarianceDynamicsPriorInverse;


        // decompose m_covarianceDynamicsAPosterioriInverse
        if (computePermutation) {
            //        m_intermediateQuantities.covarianceDynamicsAPosterioriInverseDecomposition.analyzePattern(toEigen(m_covarianceDynamicsAPosterioriInverse));
            covarianceDynamicsAPosterioriInverseDecomposition.analyzePattern(covarianceDynamicsAPosterioriInverse);
        }
        //    m_intermediateQuantities.covarianceDynamicsAPosterioriInverseDecomposition.factorize(toEigen(m_covarianceDynamicsAPosterioriInverse));
        covarianceDynamicsAPosterioriInverseDecomposition.factorize(covarianceDynamicsAPosterioriInverse);

        // Final result: expected value of the whole-body dynamics, Eq. 11b
        toEigen(expectedDynamicsAPosterioriRHS) = toEigen(measurementsMatrix).transpose() * toEigen(priorMeasurementsCovarianceInverse) * (toEigen(measurements) - toEigen(measurementsBias));
        toEigen(expectedDynamicsAPosterioriRHS) += covarianceDynamicsPriorInverse * toEigen(expectedDynamicsPrior);
        
        toEigen(expectedDynamicsAPosteriori) =
        covarianceDynamicsAPosterioriInverseDecomposition.solve(toEigen(expectedDynamicsAPosterioriRHS));
    }

    bool BerdySparseMAPSolver::BerdySparseMAPSolverPimpl::initialize()
    {
        valid = false;
        if (!berdy.isValid()) return false;
        // Resize internal variables
        size_t numberOfDynVariables = berdy.getNrOfDynamicVariables();
        size_t numberOfDynEquations = berdy.getNrOfDynamicEquations();
        size_t numberOfMeasurements = berdy.getNrOfSensorsMeasurements();

        jointsConfiguration.resize(berdy.model());
        jointsConfiguration.zero();
        jointsVelocity.resize(berdy.model());
        jointsVelocity.zero();

        measurements.resize(numberOfMeasurements);
        measurements.zero();

        expectedDynamicsAPosteriori.resize(numberOfDynVariables);
        expectedDynamicsAPosteriori.zero();
        covarianceDynamicsAPosterioriInverse.resize(numberOfDynVariables, numberOfDynVariables);

        covarianceDynamicsPriorInverse.resize(numberOfDynVariables, numberOfDynVariables);
        expectedDynamicsPrior.resize(numberOfDynVariables);
        expectedDynamicsPriorRHS.resize(numberOfDynVariables);
        expectedDynamicsAPosterioriRHS.resize(numberOfDynVariables);
        expectedDynamicsAPosterioriRHS.zero();

        // Resize priors and set them to identity.
        // If a prior is specified in config file they will be cleared after
        iDynTree::Triplets identityTriplets;
        priorDynamicsConstraintsCovarianceInverse.resize(numberOfDynEquations, numberOfDynEquations);
        identityTriplets.reserve(numberOfDynEquations);
        identityTriplets.setDiagonalMatrix(0, 0, 1.0, numberOfDynEquations);
        priorDynamicsConstraintsCovarianceInverse.setFromTriplets(identityTriplets);

        priorDynamicsRegularizationCovarianceInverse.resize(numberOfDynVariables, numberOfDynVariables);
        identityTriplets.clear();
        identityTriplets.reserve(numberOfDynVariables);
        identityTriplets.setDiagonalMatrix(0, 0, 1.0, numberOfDynVariables);
        priorDynamicsRegularizationCovarianceInverse.setFromTriplets(identityTriplets);

        priorMeasurementsCovarianceInverse.resize(numberOfMeasurements, numberOfMeasurements);
        identityTriplets.clear();
        identityTriplets.reserve(numberOfMeasurements);
        identityTriplets.setDiagonalMatrix(0, 0, 1.0, numberOfMeasurements);
        priorMeasurementsCovarianceInverse.setFromTriplets(identityTriplets);

        priorDynamicsRegularizationExpectedValue.resize(numberOfDynVariables);
        priorDynamicsRegularizationExpectedValue.zero();

        Vector3 initialGravity;
        initialGravity.zero();
        initialGravity(2) = -9.81;

        berdy.updateKinematicsFromFixedBase(jointsConfiguration, jointsVelocity, berdy.dynamicTraversal().getBaseLink()->getIndex(), initialGravity);
        computeMAP(true);

        valid = true;
        return true;
    }


    bool BerdySparseMAPSolver::BerdySparseMAPSolverPimpl::invertSparseMatrix(const iDynTree::SparseMatrix<iDynTree::ColumnMajor>&in,
                                                                             iDynTree::SparseMatrix<iDynTree::ColumnMajor>& inverted)
    {
        Eigen::SimplicialLDLT<Eigen::SparseMatrix<double, Eigen::ColMajor> > inversion;
        inversion.compute(iDynTree::toEigen(in));
        if (inversion.info() != Eigen::Success) return false;
        
        Eigen::SparseMatrix<double> identity(in.rows(), in.columns());
        identity.setIdentity();
        Eigen::SparseMatrix<double> result = inversion.solve(identity);
        iDynTree::Triplets invertedElements;
        invertedElements.reserve(result.nonZeros());

        for (Eigen::Index k = 0; k < result.outerSize(); ++k) {
            for (Eigen::SparseMatrix<double>::InnerIterator it(result, k); it; ++it)
            {
                invertedElements.pushTriplet(iDynTree::Triplet(it.row(), it.col(), it.value()));
            }
        }
        inverted.setFromTriplets(invertedElements);
        return true;
    }
}
