//
// Copyright (c) 2016-2018 CNRS
//

#ifndef __pinocchio_dynamics_hpp__
#define __pinocchio_dynamics_hpp__

#include "pinocchio/multibody/model.hpp"
#include "pinocchio/multibody/data.hpp"
#include "pinocchio/algorithm/compute-all-terms.hpp"
#include "pinocchio/algorithm/cholesky.hpp"
#include "pinocchio/algorithm/crba.hpp"
#include "pinocchio/algorithm/check.hpp"

#include <Eigen/Cholesky>

namespace pinocchio
{

  ///
  /// \brief Compute the forward dynamics with contact constraints.
  /// \note It computes the following problem:

  ///        \f$ \begin{eqnarray} \underset{\ddot{q}}{\min} & & \| \ddot{q} - \ddot{q}_{\text{free}} \|_{M(q)} \\
  ///           \text{s.t.} & & J (q) \ddot{q} + \gamma (q, \dot{q}) = 0 \end{eqnarray} \f$

  ///       where \f$ \ddot{q}_{\text{free}} \f$ is the free acceleration (i.e. without constraints),
  ///       \f$ M \f$ is the mass matrix, \f$ J \f$ the constraint Jacobian and \f$ \gamma \f$ is the constraint drift.
  ///  By default, the constraint Jacobian is assumed to be full rank, and undamped Cholesky inverse is performed.
  ///
  /// \tparam JointCollection Collection of Joint types.
  /// \tparam ConfigVectorType Type of the joint configuration vector.
  /// \tparam TangentVectorType1 Type of the joint velocity vector.
  /// \tparam TangentVectorType2 Type of the joint torque vector.
  /// \tparam ConstraintMatrixType Type of the constraint matrix.
  /// \tparam DriftVectorType Type of the drift vector.

  ///
  /// \param[in] model The model structure of the rigid body system.
  /// \param[in] data The data structure of the rigid body system.
  /// \param[in] q The joint configuration (vector dim model.nq).
  /// \param[in] v The joint velocity (vector dim model.nv).
  /// \param[in] tau The joint torque vector (dim model.nv).
  /// \param[in] J The Jacobian of the constraints (dim nb_constraints*model.nv).
  /// \param[in] gamma The drift of the constraints (dim nb_constraints).
  /// \param[in] inv_damping Damping factor for cholesky decomposition of JMinvJt. Set to zero if constraints are full rank.
  /// \param[in] updateKinematics If true, the algorithm calls first pinocchio::computeAllTerms. Otherwise, it uses the current dynamic values stored in data. \\
  ///            \note A hint: 1e-12 as the damping factor gave good result in the particular case of redundancy in contact constraints on the two feet.
  ///
  /// \return A reference to the joint acceleration stored in data.ddq. The Lagrange Multipliers linked to the contact forces are available throw data.lambda_c vector.
  ///
  template<typename Scalar, int Options, template<typename,int> class JointCollectionTpl, typename ConfigVectorType, typename TangentVectorType1, typename TangentVectorType2,
  typename ConstraintMatrixType, typename DriftVectorType>
  inline const typename DataTpl<Scalar,Options,JointCollectionTpl>::TangentVectorType &
  forwardDynamics(const ModelTpl<Scalar,Options,JointCollectionTpl> & model,
                  DataTpl<Scalar,Options,JointCollectionTpl> & data,
                  const Eigen::MatrixBase<ConfigVectorType> & q,
                  const Eigen::MatrixBase<TangentVectorType1> & v,
                  const Eigen::MatrixBase<TangentVectorType2> & tau,
                  const Eigen::MatrixBase<ConstraintMatrixType> & J,
                  const Eigen::MatrixBase<DriftVectorType> & gamma,
                  const Scalar inv_damping = 0.,
                  const bool updateKinematics = true
                  )
  {
    assert(q.size() == model.nq);
    assert(v.size() == model.nv);
    assert(tau.size() == model.nv);
    assert(J.cols() == model.nv);
    assert(J.rows() == gamma.size());
    assert(model.check(data) && "data is not consistent with model.");

    typedef DataTpl<Scalar,Options,JointCollectionTpl> Data;

    typename Data::TangentVectorType & a = data.ddq;
    typename Data::VectorXs & lambda_c = data.lambda_c;

    if (updateKinematics)
      computeAllTerms(model, data, q, v);

    // Compute the UDUt decomposition of data.M
    cholesky::decompose(model, data);

    // Compute the dynamic drift (control - nle)
    data.torque_residual = tau - data.nle;
    cholesky::solve(model, data, data.torque_residual);

    data.sDUiJt = J.transpose();
    // Compute U^-1 * J.T
    cholesky::Uiv(model, data, data.sDUiJt);
    for(Eigen::DenseIndex k=0;k<model.nv;++k)
      data.sDUiJt.row(k) /= sqrt(data.D[k]);

    data.JMinvJt.noalias() = data.sDUiJt.transpose() * data.sDUiJt;

    data.JMinvJt.diagonal().array() += inv_damping;
    data.llt_JMinvJt.compute(data.JMinvJt);

    // Compute the Lagrange Multipliers
    lambda_c.noalias() = -J*data.torque_residual;
    lambda_c -= gamma;
    data.llt_JMinvJt.solveInPlace(lambda_c);

    // Compute the joint acceleration
    a.noalias() = J.transpose() * lambda_c;
    cholesky::solve(model, data, a);
    a += data.torque_residual;

    return a;
  }

  ///
  /// \brief Compute the impulse dynamics with contact constraints.
  /// \note It computes the following problem:

  ///        \f$ \begin{eqnarray} \underset{\dot{q}^{+}}{\min} & & \| \dot{q}^{+} - \dot{q}^{-} \|_{M(q)} \\
  ///           \text{s.t.} & & J (q) \dot{q}^{+} = - \epsilon J (q) \dot{q}^{-}  \end{eqnarray} \f$

  ///       where \f$ \dot{q}^{-} \f$ is the generalized velocity before impact,
  ///       \f$ M \f$ is the joint space mass matrix, \f$ J \f$ the constraint Jacobian and \f$ \epsilon \f$ is the coefficient of restitution (1 for a fully elastic impact or 0 for a rigid impact).
  ///
  /// \tparam JointCollection Collection of Joint types.
  /// \tparam ConfigVectorType Type of the joint configuration vector.
  /// \tparam TangentVectorType Type of the joint velocity vector.
  /// \tparam ConstraintMatrixType Type of the constraint matrix.
  ///
  /// \param[in] model The model structure of the rigid body system.
  /// \param[in] data The data structure of the rigid body system.
  /// \param[in] q The joint configuration (vector dim model.nq).
  /// \param[in] v_before The joint velocity before impact (vector dim model.nv).
  /// \param[in] J The Jacobian of the constraints (dim nb_constraints*model.nv).
  /// \param[in] r_coeff The coefficient of restitution. Must be in [0;1].
  /// \param[in] updateKinematics If true, the algorithm calls first pinocchio::crba. Otherwise, it uses the current mass matrix value stored in data.
  ///
  /// \return A reference to the generalized velocity after impact stored in data.dq_after. The Lagrange Multipliers linked to the contact impulsed are available throw data.impulse_c vector.
  ///
  template<typename Scalar, int Options, template<typename,int> class JointCollectionTpl, typename ConfigVectorType, typename TangentVectorType, typename ConstraintMatrixType>
  inline const typename DataTpl<Scalar,Options,JointCollectionTpl>::TangentVectorType &
  impulseDynamics(const ModelTpl<Scalar,Options,JointCollectionTpl> & model,
                  DataTpl<Scalar,Options,JointCollectionTpl> & data,
                  const Eigen::MatrixBase<ConfigVectorType> & q,
                  const Eigen::MatrixBase<TangentVectorType> & v_before,
                  const Eigen::MatrixBase<ConstraintMatrixType> & J,
                  const Scalar r_coeff = 0,
                  const bool updateKinematics = true
                  )
  {
    assert(q.size() == model.nq);
    assert(v_before.size() == model.nv);
    assert(J.cols() == model.nv);
    assert(model.check(data) && "data is not consistent with model.");

    typedef DataTpl<Scalar,Options,JointCollectionTpl> Data;

    typename Data::VectorXs & impulse_c = data.impulse_c;
    typename Data::TangentVectorType & dq_after = data.dq_after;

    // Compute the mass matrix
    if (updateKinematics)
      crba(model, data, q);

    // Compute the UDUt decomposition of data.M
    cholesky::decompose(model, data);

    data.sDUiJt = J.transpose();
    // Compute U^-1 * J.T
    cholesky::Uiv(model, data, data.sDUiJt);
    for(int k=0;k<model.nv;++k) data.sDUiJt.row(k) /= sqrt(data.D[k]);

    data.JMinvJt.noalias() = data.sDUiJt.transpose() * data.sDUiJt;
    data.llt_JMinvJt.compute(data.JMinvJt);

    // Compute the Lagrange Multipliers related to the contact impulses
    impulse_c.noalias() = (-r_coeff - 1.) * (J * v_before);
    data.llt_JMinvJt.solveInPlace(impulse_c);

    // Compute the joint velocity after impacts
    dq_after.noalias() = J.transpose() * impulse_c;
    cholesky::solve(model, data, dq_after);
    dq_after += v_before;

    return dq_after;
  }
} // namespace pinocchio

#endif // ifndef __pinocchio_dynamics_hpp__


#include <ceres/ceres.h>
#include "ceres/rotation.h"
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <algorithm>
#include <glog/logging.h>
#include <iterator>
#include <unordered_set>
#include <cmath>
#include <math.h>
#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/Geometry>

#define EPS T(0.00001)

using namespace pybind11::literals;
namespace py = pybind11;
using namespace Eigen;

struct FittingResidual {
    FittingResidual(double *rt_in, double* time_in, double coeff_3_in): rt(rt_in), time(time_in), coeff_3(coeff_3_in){}
    double *rt;
    double *time;
    double coeff_3;

    template<typename T>
    bool operator()(const T* const spline, T *residual) const {
        // parameter to be optimized
        // residual size 2
        residual[0] = spline[0]*time[0]*time[0]*time[0] + spline[1]*time[0]*time[0] +\
                      spline[2]*time[0] + spline[3] - rt[3];
        residual[1] = spline[4]*time[0]*time[0]*time[0] + spline[5]*time[0]*time[0] +\
                      spline[6]*time[0] + spline[7]  - rt[4];
        auto vx = 3.0*spline[0]*time[0]*time[0] + 2*time[0]*spline[1] + spline[2];
        auto vy = 3.0*spline[4]*time[0]*time[0] + 2*time[0]*spline[5] + spline[6];
        auto v_norm = sqrt(vx*vx + vy*vy);
        vx /= v_norm, vy /= v_norm;
        residual[2] = 1.0*(vx - rt[0]);
        residual[3] = 1.0*(vy - rt[1]);
        residual[4] = coeff_3*spline[0];
        residual[5] = coeff_3*spline[4];
        residual[6] = 0.0*spline[1];
        residual[7] = 0.0*spline[5];
        // auto grad_square = (spline[0]*time[0]+spline[1])*(spline[0]*time[0]+spline[1]) +\
        //          (spline[3]*time[0]+spline[4])*(spline[3]*time[0]+spline[4]);
        // auto cross = (spline[0]*time[0]+spline[1])*spline[3]-(spline[3]*time[0]+spline[4])*spline[0];
        return true;
    }
};

py::object spline_fitting(py::buffer rt_, py::buffer time_, py::buffer spline_, py::buffer cost_,
                int num_obj, int num_iter, double coeff_3){
    auto buf0 = rt_.request(); // num_obj*6
    double *rt = static_cast<double *>(buf0.ptr);
    auto buf1 = time_.request();
    double *time = static_cast<double *>(buf1.ptr);
    auto buf2 = spline_.request();
    double *spline = static_cast<double *>(buf2.ptr); // 6
    auto buf3 = cost_.request();
    double *cost = static_cast<double *>(buf3.ptr);
    double *rt_i, *time_i;
    ceres::Problem problem;
    for (int i = 0; i <num_obj; i++){
        rt_i = rt + 6*i;
        time_i = time + i;
        ceres::CostFunction *cost_function;
        cost_function = new ceres::AutoDiffCostFunction <FittingResidual, 8, 8>(
                      new FittingResidual(rt_i, time_i, coeff_3)
                );
        problem.AddResidualBlock(cost_function, NULL, spline);

    }
    ceres::Solver::Options options;
    options.max_num_iterations = num_iter;
    options.minimizer_progress_to_stdout = false;
    options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;  //ceres::SPARSE_SCHUR;  //ceres::DENSE_SCHUR;
    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);

    // std::cout<<summary.FullReport()<<"\n";
    cost[0] = summary.final_cost;
    // return optimized values
    py::array_t <double> x_spline = py::array_t <double>(
            py::buffer_info(
                    spline,
                    sizeof(double), //itemsize
                    py::format_descriptor <double>::format(),
                    1, // ndim
                    std::vector<size_t>{(unsigned long) (8)}, // shape
                    std::vector<size_t>{sizeof(double)} // strides
            )
    );
    py::array_t <double> x_cost = py::array_t <double>(
            py::buffer_info(
                    cost,
                    sizeof(double), //itemsize
                    py::format_descriptor <double>::format(),
                    1, // ndim
                    std::vector<size_t>{(unsigned long) (1)}, // shape
                    std::vector<size_t>{sizeof(double)} // strides
            )
    );
    py::list outputs;
    outputs.append(x_spline);
    outputs.append(x_cost);
    return outputs;
}

PYBIND11_PLUGIN(ceres_spline) {
        py::module m("ceres_spline", "Python bindings to the Ceres-Solver minimizer.");
        // google::InitGoogleLogging("ceres_spline");
        m.def("spline_fitting", &spline_fitting, "Fitting a spline to a car trajectory");
        return m.ptr();
}


#include "InverseProblem_Adjoint_YoungsModulus.h"

#include <Eigen/Dense>
#include <tinyformat.h>
#include <cinder/Log.h>
#include <numeric>
#include <LBFGS.h>
#include <fstream>

#include "GradientDescent.h"
#include "AdjointUtilities.h"

using namespace std;
using namespace Eigen;

#define DEBUG_SAVE_MATRICES 0

namespace ar {

    InverseProblem_Adjoint_YoungsModulus::InverseProblem_Adjoint_YoungsModulus()
        : initialYoungsModulus(50)
        , youngsModulusPrior(0.1)
        , numIterations(10)
    {
    }

    InverseProblemOutput InverseProblem_Adjoint_YoungsModulus::solveGrid(int deformedTimestep,
        BackgroundWorker* worker, IntermediateResultCallback_t callback)
    {
        CI_LOG_I("Solve for Young's Modulus at timestep " << deformedTimestep);
        // Create simulation
        worker->setStatus("Adjoint - Young's Modulus Grid: create simulation");
        SoftBodyGrid2D simulation;
        simulation.setGridResolution(input->gridResolution_);
        simulation.setSDF(input->gridReferenceSdf_);
        simulation.setExplicitDiffusion(true);
        simulation.setHardDirichletBoundaries(false);
        simulation.resetBoundaries();
        for (const auto& b : input->gridDirichletBoundaries_)
            simulation.addDirichletBoundary(b.first.first, b.first.second, b.second);
        for (const auto& b : input->gridNeumannBoundaries_)
            simulation.addNeumannBoundary(b.first.first, b.first.second, b.second);
        if (worker->isInterrupted()) return InverseProblemOutput();

        // Set parameters with everything that can't be reconstructed here
        simulation.setGravity(input->settings_.gravity_);
        simulation.setMass(input->settings_.mass_);
        simulation.setDamping(input->settings_.dampingAlpha_, input->settings_.dampingBeta_);
        simulation.setRotationCorrection(SoftBodySimulation::RotationCorrection::None);
        if (worker->isInterrupted()) return InverseProblemOutput();

        //cost function:
        //F(youngsModulus) = 0.5*|| disp(youngsModulus) - dispRef ||^2 + priorWeight / youngsModulus

        GridUtils2D::grid_t outputSdf(input->gridResolution_, input->gridResolution_);

		//TODO: once the gradient is fixed, switch to LBFGS
#if 1
		//Gradient Descent
		//define gradient of the cost function
		real finalCost = 0;
        const auto gradient = [&simulation, &outputSdf, &finalCost, this, worker, deformedTimestep](const Vector1& youngsModulusV)
        {
            real youngsModulus = youngsModulusV.x();
            real cost;
            real gradient = gradientGrid(deformedTimestep, simulation, youngsModulus, outputSdf, cost, worker);
			finalCost = cost;
            return Vector1(gradient);
        };
		//Optimize
        GradientDescent<Vector1> gd(Vector1(initialYoungsModulus), gradient, 1e-10, 10);
		gd.setMinStepsize(0.0001);
		gd.setMaxStepsize(0.5);
        int oi;
        for (oi = 0; oi < numIterations; ++oi) {
            worker->setStatus(tfm::format("Adjoint - Young's Modulus Grid: optimization %d/%d", (oi + 1), numIterations));
            if (gd.step()) break;

            {
				InverseProblemOutput output;
				output.youngsModulus_ = gd.getCurrentSolution()[0];
				output.resultGridSdf_ = outputSdf;
				output.finalCost_ = finalCost;
				callback(output);
            }

            if (worker->isInterrupted()) return InverseProblemOutput();
        }
        real finalValue = gd.getCurrentSolution()[0];
#else
		//LBFGS
		LBFGSpp::LBFGSParam<real> params;
		params.epsilon = 1e-10;
		params.max_iterations = numIterations;
		LBFGSpp::LBFGSSolver<real> lbfgs(params);
		//define gradient
		LBFGSpp::LBFGSSolver<real>::ObjectiveFunction_t fun([&simulation, &outputSdf, this, worker, deformedTimestep](const VectorX& x, VectorX& gradient) -> real {
			real youngsModulus = x[0];
			real cost;
			real grad = gradientGrid(deformedTimestep, simulation, youngsModulus, outputSdf, cost, worker);
			gradient[0] = grad;
			return cost;
		});
		LBFGSpp::LBFGSSolver<real>::CallbackFunction_t lbfgsCallback([worker, callback, &outputSdf, this](const VectorX& x, const real& v, int k) -> bool {
			worker->setStatus(tfm::format("Adjoint - Young's Modulus: optimization %d/%d", k, numIterations));

			InverseProblemOutput output;
			output.youngsModulus_ = x[0];
			output.resultGridSdf_ = outputSdf;
			output.finalCost_ = v;
			callback(output);

			return !worker->isInterrupted();
		});
		//optimize
		real finalCost = 0;
		VectorX finalValueV = VectorX::Constant(1, initialYoungsModulus);
		int oi = lbfgs.minimize(fun, finalValueV, finalCost, lbfgsCallback);
		if (worker->isInterrupted()) return InverseProblemOutput();
		real finalValue = finalValueV[0];
#endif
        CI_LOG_I("Optimization: final Youngs' Modulus is " << finalValue << " after " << oi << " optimization steps, reference value is " << input->settings_.youngsModulus_);

        InverseProblemOutput output;
        output.youngsModulus_ = finalValue;
        output.resultGridSdf_ = outputSdf;
		output.finalCost_ = finalCost;
        return output;
    }

    InverseProblemOutput InverseProblem_Adjoint_YoungsModulus::solveMesh(int deformedTimestep,
        BackgroundWorker* worker, IntermediateResultCallback_t callback)
    {
        CI_LOG_I("Solve for Young's Modulus at timestep " << deformedTimestep);
        // Create simulation
        worker->setStatus("Adjoint - Young's Modulus: create simulation");
        SoftBodyMesh2D simulation;
        simulation.setMesh(input->meshReferencePositions_, input->meshReferenceIndices_);
        simulation.resetBoundaries();
        for (const auto& b : input->meshDirichletBoundaries_)
            simulation.addDirichletBoundary(b.first, b.second);
        for (const auto& b : input->meshNeumannBoundaries_)
            simulation.addNeumannBoundary(b.first, b.second);
        simulation.reorderNodes();
        if (worker->isInterrupted()) return InverseProblemOutput();

        // Set parameters with everything that can't be reconstructed here
        simulation.setGravity(input->settings_.gravity_);
        simulation.setMass(input->settings_.mass_);
        simulation.setDamping(input->settings_.dampingAlpha_, input->settings_.dampingBeta_);
        simulation.setRotationCorrection(SoftBodySimulation::RotationCorrection::None);
        if (worker->isInterrupted()) return InverseProblemOutput();

        //cost function:
        //F(youngsModulus) = 0.5*|| disp(youngsModulus) - dispRef ||^2 + priorWeight / youngsModulus

		VectorX outputU;
#if 0
		//Gradient Descent
        //define gradient of the cost function
        const auto gradient = [&simulation, &outputU, this, worker, deformedTimestep](const Vector1& youngsModulusV)
        {
            real youngsModulus = youngsModulusV.x();
            real cost;
            real gradient = gradientMesh(deformedTimestep, simulation, youngsModulus, outputU, cost, worker);
            return Vector1(gradient);
        };

        //run optimization
        GradientDescent<Vector1> gd(Vector1(initialYoungsModulus), gradient, 1e-10, 10);
        int oi;
        for (oi = 0; oi < numIterations; ++oi) {
            worker->setStatus(tfm::format("Adjoint - Young's Modulus: optimization %d/%d", (oi + 1), numIterations));
            if (gd.step()) break;
            if (worker->isInterrupted()) return InverseProblemOutput();
        }
        real finalValue = gd.getCurrentSolution()[0];
#else
		//LBFGS
		LBFGSpp::LBFGSParam<real> params;
		params.epsilon = 1e-15;
		params.max_iterations = numIterations;
		LBFGSpp::LBFGSSolver<real> lbfgs(params);
		//define gradient
		LBFGSpp::LBFGSSolver<real>::ObjectiveFunction_t fun([&simulation, &outputU, this, worker, deformedTimestep](const VectorX& x, VectorX& gradient) -> real {
			real youngsModulus = x[0];
			real cost;
			real grad = gradientMesh(deformedTimestep, simulation, youngsModulus, outputU, cost, worker);
			gradient[0] = grad;
			return cost;
		});
		LBFGSpp::LBFGSSolver<real>::CallbackFunction_t lbfgsCallback([worker, &simulation, &outputU, callback, this](const VectorX& x, const VectorX& g, const real& v, int k) -> bool {
			worker->setStatus(tfm::format("Adjoint - Young's Modulus: optimization %d/%d", k, numIterations));

			InverseProblemOutput output;
			output.youngsModulus_ = x[0];
			output.resultMeshDisp_ = SoftBodyMesh2D::Vector2List(simulation.getNumNodes());
			for (int i = 0; i<simulation.getNumNodes(); ++i)
			{
				if (simulation.getNodeToFreeMap()[i] < 0)
					output.resultMeshDisp_->at(i).setZero();
				else
					output.resultMeshDisp_->at(i) = outputU.segment<2>(2 * simulation.getNodeToFreeMap()[i]);
			}
			output.finalCost_ = v;
			callback(output);

			return !worker->isInterrupted();
		});
		//optimize
		real finalCost = 0;
		VectorX finalValueV = VectorX::Constant(1, initialYoungsModulus);
		int oi = lbfgs.minimize(fun, finalValueV, finalCost, lbfgsCallback);
		if (worker->isInterrupted()) return InverseProblemOutput();
		real finalValue = finalValueV[0];
#endif
        CI_LOG_I("Optimization: final Youngs' Modulus is " << finalValue << " after " << oi << " optimization steps, reference value is " << input->settings_.youngsModulus_);

        InverseProblemOutput output;
        output.youngsModulus_ = finalValue;
        output.resultMeshDisp_ = SoftBodyMesh2D::Vector2List(simulation.getNumNodes());
        for (int i=0; i<simulation.getNumNodes(); ++i)
        {
            if (simulation.getNodeToFreeMap()[i] < 0)
                output.resultMeshDisp_->at(i).setZero();
            else
                output.resultMeshDisp_->at(i) = outputU.segment<2>(2 * simulation.getNodeToFreeMap()[i]);
        }
		output.finalCost_ = finalCost;
        return output;
    }

    real InverseProblem_Adjoint_YoungsModulus::gradientMesh(
        int deformedTimestep, const SoftBodyMesh2D& simulation,
        real youngsModulus, VectorX& outputU, real& outputCost, BackgroundWorker* worker) const
    {
        real poissonsRatio = input->settings_.poissonsRatio_;

        //FORWARD

        MatrixX K = MatrixX::Zero(2 * simulation.getNumFreeNodes(), 2 * simulation.getNumFreeNodes());
        VectorX F = VectorX::Zero(2 * simulation.getNumFreeNodes());

        //assemble force vector F and stiffness matrix K
        simulation.assembleForceVector(F);
        if (worker->isInterrupted()) return 0;
        real materialMu, materialLambda;
        SoftBodySimulation::computeMaterialParameters(youngsModulus, poissonsRatio, materialMu, materialLambda);
        Matrix3 C = SoftBodySimulation::computeMaterialMatrix(materialMu, materialLambda);
        simulation.assembleStiffnessMatrix(C, K, &F, SoftBodyMesh2D::Vector2List(simulation.getNumNodes(), Vector2::Zero()));
        if (worker->isInterrupted()) return 0;
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(K, tfm::format("AdjointYoung-GradientMesh-K-k%f.csv", youngsModulus));
#endif

        //assemble derivate of matrix K with respect to the youngsModulus
        MatrixX KDyoung = MatrixX::Zero(2 * simulation.getNumFreeNodes(), 2 * simulation.getNumFreeNodes());
        real muDyoung, lambdaDyoung;
        AdjointUtilities::computeMaterialParameters_D_YoungsModulus(youngsModulus, poissonsRatio, muDyoung, lambdaDyoung);
        Matrix3 CDyoung = SoftBodySimulation::computeMaterialMatrix(muDyoung, lambdaDyoung);
        simulation.assembleStiffnessMatrix(CDyoung, KDyoung, nullptr, SoftBodyMesh2D::Vector2List(simulation.getNumNodes(), Vector2::Zero()));
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(KDyoung, tfm::format("AdjointYoung-GradientMesh-KDyoung-k%f.csv", youngsModulus));
#endif

        //solve for the current displacement
        PartialPivLU<MatrixX> Klu = K.partialPivLu();
        VectorX u = Klu.solve(F);
        outputU = u;

        //compute the partial gradient of the operations with respect to Young's Modulus
        VectorX EDyoung = -KDyoung * u;
        //VectorX EDyoung = Klu.solve(-KDyoung * u);

        // BACKWARD

        //compute derivate of the cost function with respect to the output u (displacement)
        VectorX costDu(2 * simulation.getNumFreeNodes());
        for (int i = 0; i < simulation.getNumNodes(); ++i)
        {
            if (simulation.getNodeToFreeMap()[i] < 0) continue;
            costDu.segment<2>(2 * simulation.getNodeToFreeMap()[i])
                = u.segment<2>(2 * simulation.getNodeToFreeMap()[i])
                - input->meshResultsDisplacement_[deformedTimestep][i];
        }
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(costDu, tfm::format("AdjointYoung-GradientMesh-costDu-k%f.csv", youngsModulus));
#endif

        //solve for the dual
        //In theory, I have to use K', but because I know that K is symmetric,
        //I can use K directly and not the transpose.
        VectorX lambda = Klu.solve(costDu);

        //compute gradient of the prior term
        real priorDyoung = -youngsModulusPrior / (youngsModulus*youngsModulus);

        //evalue the final gradient
        real gradient = lambda.transpose() * EDyoung + priorDyoung;

        //For testing: evaluate cost
        real cost = youngsModulusPrior / youngsModulus;
        for (int i = 0; i < simulation.getNumNodes(); ++i)
        {
            if (simulation.getNodeToFreeMap()[i] < 0) continue;
            cost += (u.segment<2>(2 * simulation.getNodeToFreeMap()[i])
                - input->meshResultsDisplacement_[deformedTimestep][i]).squaredNorm() / 2;
        }
        outputCost = cost;
        CI_LOG_I("Optimization: Youngs Modulus " << youngsModulus << " -> cost " << cost << ", gradient " << gradient);

        return gradient;
    }

    real InverseProblem_Adjoint_YoungsModulus::gradientGrid(int deformedTimestep, const SoftBodyGrid2D& simulation,
        real youngsModulus, GridUtils2D::grid_t& outputSdf, real& outputCost, BackgroundWorker* worker) const
    {
        real poissonsRatio = input->settings_.poissonsRatio_;

        //FORWARD
        int resolution = input->gridResolution_;
        real h = 1.0 / (resolution - 1);
        Vector2 pos(0, 0); //lower left corner of the grid
        Vector2 size(h, h); //size of each cell

		//collect degrees of freedom in the stifness solve
		const Eigen::MatrixXi& posToIndex = simulation.getPosToIndex();
		const SoftBodyGrid2D::indexToPos_t& indexToPos = simulation.getIndexToPos();
		const int dof = simulation.getDoF();

        MatrixX K = MatrixX::Zero(dof * 2, dof * 2);
        VectorX f = VectorX::Zero(dof * 2);
		const VectorX prevU = VectorX::Zero(dof * 2); //previous displacements for rotation correction. Not needed here

        //assemble force vector F and stiffness matrix K
		VectorX collisionForces = VectorX::Zero(2 * dof);
        simulation.assembleForceVector(f, posToIndex, collisionForces);
        if (worker->isInterrupted()) return 0;
        real materialMu, materialLambda;
        SoftBodySimulation::computeMaterialParameters(youngsModulus, poissonsRatio, materialMu, materialLambda);
        Matrix3 C = SoftBodySimulation::computeMaterialMatrix(materialMu, materialLambda);
        simulation.assembleStiffnessMatrix(C, materialMu, materialLambda, K, &f, posToIndex, prevU);
        if (worker->isInterrupted()) return 0;
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(K, tfm::format("AdjointYoung-GradientGrid-K-k%f.csv", youngsModulus));
#endif

        //assemble derivate of matrix K with respect to the youngsModulus
        MatrixX KDyoung = MatrixX::Zero(dof * 2, dof * 2);
        real muDyoung, lambdaDyoung;
		AdjointUtilities::computeMaterialParameters_D_YoungsModulus(youngsModulus, poissonsRatio, muDyoung, lambdaDyoung);
        Matrix3 CDyoung = SoftBodySimulation::computeMaterialMatrix(muDyoung, lambdaDyoung);
        simulation.assembleStiffnessMatrix(CDyoung, muDyoung, lambdaDyoung, KDyoung, nullptr, posToIndex, prevU);
        if (worker->isInterrupted()) return 0;
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(KDyoung, tfm::format("AdjointYoung-GradientGrid-KDyoung-k%f.csv", youngsModulus));
#endif

        //solve for the current displacement
        PartialPivLU<MatrixX> Klu = K.partialPivLu();
        VectorX solution = Klu.solve(f);
        if (worker->isInterrupted()) return 0;

        //map back to a grid
        SoftBodyGrid2D::grid_t uGridX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t uGridY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        for (int i = 0; i < dof; ++i) {
            Vector2i p = indexToPos.at(i);
            uGridX(p.x(), p.y()) = solution[2 * i];
            uGridY(p.x(), p.y()) = solution[2 * i + 1];
        }
        if (worker->isInterrupted()) return 0;

        //perform diffusion step
        GridUtils2D::bgrid_t validCells = simulation.computeValidCells(uGridX, uGridY, posToIndex);
        uGridX = GridUtils2D::fillGridDiffusion(uGridX, validCells);
        uGridY = GridUtils2D::fillGridDiffusion(uGridY, validCells);
        if (worker->isInterrupted()) return 0;

#if 0
        //invert displacements
        SoftBodyGrid2D::grid_t uGridInvX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t uGridInvY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::invertDisplacement(uGridX, uGridY, uGridInvX, uGridInvY, h);
        if (worker->isInterrupted()) return 0;

        //advect levelset
        outputSdf = GridUtils2D::advectGridSemiLagrange(input->gridReferenceSdf_, uGridInvX, uGridInvY, -1 / h);
        if (worker->isInterrupted()) return 0;
#else
        //advect levelset
		GridUtils2D::grid_t advectionWeights(resolution, resolution);
        outputSdf = GridUtils2D::advectGridDirectForward(input->gridReferenceSdf_, uGridX, uGridY, -1 / h, &advectionWeights);
        if (worker->isInterrupted()) return 0;
#endif

        //compute the partial gradient of the operations with respect to Young's Modulus
        VectorX EDyoung = -KDyoung * solution;
        //VectorX EDyoung = Klu.solve(-KDyoung * solution);
        if (worker->isInterrupted()) return 0;

        // BACKWARD

#if 1
        //compute derivate of the cost function with respect to the output u (displacement)
        //SoftBodyGrid2D::grid_t costDu = outputSdf - input->gridResultsSdf_[deformedTimestep];
        SoftBodyGrid2D::grid_t costDu =
            (2 * ((2 / M_PI)*outputSdf.atan() - (2 / M_PI)*input->gridResultsSdf_[deformedTimestep].atan())) /
            (M_PI * (1 + outputSdf.square()));

#if 0
        //Adjoint: advect levelset
        SoftBodyGrid2D::grid_t adjInputSdf = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridInvX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridInvY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::advectGridSemiLagrangeAdjoint(
            adjInputSdf, adjUGridInvX, adjUGridInvY, -1 / h, costDu,
            uGridInvX, uGridInvY, input->gridReferenceSdf_);
        if (worker->isInterrupted()) return 0;

        //Adjoint: invert levelset
        SoftBodyGrid2D::grid_t adjUGridX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::invertDisplacementDirectShepardAdjoint(
            adjUGridX, adjUGridY, adjUGridInvX, adjUGridInvY,
            uGridX, uGridY, uGridInvX, uGridInvY, h);
        if (worker->isInterrupted()) return 0;
#else
        SoftBodyGrid2D::grid_t adjInputSdf = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::advectGridDirectForwardAdjoint(adjInputSdf, adjUGridX, adjUGridY,
            costDu, uGridX, uGridY, input->gridReferenceSdf_, outputSdf, -1 / h, advectionWeights);
        if (worker->isInterrupted()) return 0;
#endif

        //Adjoint: perform diffusion step
        GridUtils2D::fillGridDiffusionAdjoint(adjUGridX, adjUGridX, uGridX, validCells); //note that the adjoint is both input and output
        GridUtils2D::fillGridDiffusionAdjoint(adjUGridY, adjUGridY, uGridY, validCells); //it is modified in-place
        if (worker->isInterrupted()) return 0;

        //Adjoint: map back to a grid
        VectorX adjSolution(2 * dof);
        for (int i = 0; i < dof; ++i) {
            Vector2i p = indexToPos.at(i);
            adjSolution[2 * i] = adjUGridX(p.x(), p.y());
            adjSolution[2 * i + 1] = adjUGridY(p.x(), p.y());
        }
        if (worker->isInterrupted()) return 0;

        //Adjoint: solve for the current displacement
        //since K is symmetric, I can use K directly instead of K' as normally required by the Adjoint Method.
        VectorX lambda = Klu.solve(adjSolution);
        if (worker->isInterrupted()) return 0;

#else
        //Recompute inverted true solution
        SoftBodyGrid2D::grid_t gridResultsInvUx = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t gridResultsInvUy = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::invertDisplacementDirectShepard(
            input->gridResultsUx_[deformedTimestep], input->gridResultsUy_[deformedTimestep],
            gridResultsInvUx, gridResultsInvUy, h);

        //Reduced problem only on the levelset
        SoftBodyGrid2D::grid_t adjUGridInvX = uGridInvX - gridResultsInvUx;
        SoftBodyGrid2D::grid_t adjUGridInvY = uGridInvY - gridResultsInvUy;

        //Adjoint: invert levelset
        SoftBodyGrid2D::grid_t adjUGridX = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        SoftBodyGrid2D::grid_t adjUGridY = SoftBodyGrid2D::grid_t::Zero(resolution, resolution);
        GridUtils2D::invertDisplacementDirectShepardAdjoint(
            adjUGridX, adjUGridY, adjUGridInvX, adjUGridInvY,
            uGridX, uGridY, uGridInvX, uGridInvY, h);
        if (worker->isInterrupted()) return 0;

        //Adjoint: perform diffusion step
        GridUtils2D::fillGridDiffusionAdjoint(adjUGridX, adjUGridX, uGridX, validCells); //note that the adjoint is both input and output
        GridUtils2D::fillGridDiffusionAdjoint(adjUGridY, adjUGridY, uGridY, validCells); //it is modified in-place
        if (worker->isInterrupted()) return 0;

        //Adjoint: map back to a grid
        VectorX adjSolution(2 * dof);
        for (int i = 0; i < dof; ++i) {
            Vector2i p = indexToPos[i];
            adjSolution[2 * i] = adjUGridX(p.x(), p.y());
            adjSolution[2 * i + 1] = adjUGridY(p.x(), p.y());
        }
        if (worker->isInterrupted()) return 0;
#if DEBUG_SAVE_MATRICES==1
        saveAsCSV(adjSolution, tfm::format("AdjointYoung-GradientGrid-costDu-k%f.csv", youngsModulus));
#endif

        //Adjoint: solve for the current displacement
        //since K is symmetric, I can use K directly instead of K' as normally required by the Adjoint Method.
        VectorX lambda = -Klu.solve(adjSolution);
        if (worker->isInterrupted()) return 0;
#endif

        //Done with the adjoint computation
        //Now I can assemble the final gradient

        //compute gradient of the prior term
        real priorDyoung = -youngsModulusPrior / (youngsModulus*youngsModulus);

        //evalue the final gradient
        real gradient = -lambda.transpose() * EDyoung + priorDyoung;

        //For testing: evaluate cost
#if 1
        real cost = youngsModulusPrior / youngsModulus
            + ((2/M_PI)*outputSdf.atan() - (2/M_PI)*input->gridResultsSdf_[deformedTimestep].atan()).matrix().squaredNorm() / 2;
#else
        //reduced problem
        real cost = (uGridInvX - gridResultsInvUx).matrix().squaredNorm() / 2 +
                    (uGridInvY - gridResultsInvUy).matrix().squaredNorm() / 2;
#endif
        outputCost = cost;
        CI_LOG_I("Optimization: Youngs Modulus " << youngsModulus << " -> cost " << cost << ", gradient " << gradient);

        return gradient;
    }

    void InverseProblem_Adjoint_YoungsModulus::setupParams(cinder::params::InterfaceGlRef params,
        const std::string& group)
    {
        params->addParam("InverseProblem_Adjoint_YoungsModulus_InitialYoungsModulus", &initialYoungsModulus)
            .group(group).label("Initial Young's Modulus").min(0).step(0.01f);
        params->addParam("InverseProblem_Adjoint_YoungsModulus_YoungsModulusPrior", &youngsModulusPrior)
            .group(group).label("Prior on Young's Modulus").min(0).step(0.01f);
        params->addParam("InverseProblem_Adjoint_YoungsModulus_NumIterations", &numIterations)
            .group(group).label("Iterations").min(1);
    }

    void InverseProblem_Adjoint_YoungsModulus::setParamsVisibility(cinder::params::InterfaceGlRef params,
        bool visible) const
    {
        std::string option = visible ? "visible=true" : "visible=false";
        params->setOptions("InverseProblem_Adjoint_YoungsModulus_InitialYoungsModulus", option);
        params->setOptions("InverseProblem_Adjoint_YoungsModulus_YoungsModulusPrior", option);
        params->setOptions("InverseProblem_Adjoint_YoungsModulus_NumIterations", option);
    }

    std::vector<InverseProblem_Adjoint_YoungsModulus::DataPoint> InverseProblem_Adjoint_YoungsModulus::plotEnergy(int deformedTimestep, BackgroundWorker * worker) const
    {
        //control points
        real trueYoung = input->settings_.youngsModulus_;
        real minYoung = trueYoung / 4;
        real maxYoung = trueYoung + (trueYoung - minYoung);
        int steps = 25; // has to be odd
        std::vector<DataPoint> points(steps);
        for (int i = 0; i < steps; ++i)
            points[i].youngsModulus = minYoung + (maxYoung - minYoung) * i / (steps - 1.0);

        // Create grid simulation
        worker->setStatus("Test Plot Grid: create simulation");
        SoftBodyGrid2D simulationGrid;
        simulationGrid.setGridResolution(input->gridResolution_);
        simulationGrid.setSDF(input->gridReferenceSdf_);
        simulationGrid.setExplicitDiffusion(true);
        simulationGrid.setHardDirichletBoundaries(false);
        simulationGrid.resetBoundaries();
        for (const auto& b : input->gridDirichletBoundaries_)
            simulationGrid.addDirichletBoundary(b.first.first, b.first.second, b.second);
        for (const auto& b : input->gridNeumannBoundaries_)
            simulationGrid.addNeumannBoundary(b.first.first, b.first.second, b.second);
        simulationGrid.setGravity(input->settings_.gravity_);
        simulationGrid.setMass(input->settings_.mass_);
        simulationGrid.setDamping(input->settings_.dampingAlpha_, input->settings_.dampingBeta_);
        simulationGrid.setRotationCorrection(input->settings_.rotationCorrection_);

        //Run grid simulation
        GridUtils2D::grid_t outputGrid;
        for (int i = 0; i < steps; ++i)
        {
            worker->setStatus(tfm::format("Test Plot Grid: simulation %d/%d",i+1,steps));
            real cost;
            real gradient = gradientGrid(deformedTimestep, simulationGrid, points[i].youngsModulus, outputGrid, cost, worker);
            if (worker->isInterrupted()) return {};
            points[i].costGrid = cost;
            points[i].gradientGrid = gradient;
        }

        //Create mesh simulation
        worker->setStatus("Test Plot Mesh: create simulation");
        SoftBodyMesh2D simulationMesh;
        simulationMesh.setMesh(input->meshReferencePositions_, input->meshReferenceIndices_);
        simulationMesh.resetBoundaries();
        for (const auto& b : input->meshDirichletBoundaries_)
            simulationMesh.addDirichletBoundary(b.first, b.second);
        for (const auto& b : input->meshNeumannBoundaries_)
            simulationMesh.addNeumannBoundary(b.first, b.second);
        simulationMesh.reorderNodes();
        simulationMesh.setGravity(input->settings_.gravity_);
        simulationMesh.setMass(input->settings_.mass_);
        simulationMesh.setDamping(input->settings_.dampingAlpha_, input->settings_.dampingBeta_);
        simulationMesh.setRotationCorrection(input->settings_.rotationCorrection_);

        //Run mesh simulation
        VectorX outputMesh;
        for (int i = 0; i < steps; ++i)
        {
            worker->setStatus(tfm::format("Test Plot Mesh: simulation %d/%d", i + 1, steps));
            real cost;
            real gradient = gradientMesh(deformedTimestep, simulationMesh, points[i].youngsModulus, outputMesh, cost, worker);
            if (worker->isInterrupted()) return {};
            points[i].costMesh = cost;
            points[i].gradientMesh = gradient;
        }

        return points;
    }

	void InverseProblem_Adjoint_YoungsModulus::testPlot(int deformedTimestep, BackgroundWorker* worker)
	{
		auto points = plotEnergy(deformedTimestep-1, worker);
		if (worker->isInterrupted()) return;
		fstream file("../plots/Adjoint-YoungsModulus.csv", fstream::out | fstream::trunc);
		file << "Youngs Modulus , Cost Mesh , Gradient Mesh , Cost Grid , Gradient Grid" << endl;
		file << std::fixed;
		for (auto point : points)
		{
			file << point.youngsModulus << " , "
				<< point.costMesh << " , " << point.gradientMesh << " , "
				<< point.costGrid << " , " << point.gradientGrid << endl;
		}
		file.close();
	}

}