#pragma once

#include <Eigen/Core>

#include <ipc/broad_phase/collision_candidate.hpp>
#include <ipc/utils/eigen_ext.hpp>

namespace ipc {

struct CollisionConstraint {
public:
    virtual ~CollisionConstraint() {}

    virtual std::vector<long> vertex_indices(
        const Eigen::MatrixXi& E, const Eigen::MatrixXi& F) const = 0;

    virtual double compute_distance(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const = 0;

    virtual VectorMax12d compute_distance_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const = 0;

    virtual MatrixMax12d compute_distance_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const = 0;

    virtual double compute_potential(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat) const;

    virtual VectorMax12d compute_potential_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat) const;

    virtual MatrixMax12d compute_potential_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat,
        const bool project_hessian_to_psd) const;

    double minimum_distance = 0;
};

///////////////////////////////////////////////////////////////////////////////

struct VertexVertexConstraint : VertexVertexCandidate, CollisionConstraint {
    VertexVertexConstraint(long vertex0_index, long vertex1_index);
    VertexVertexConstraint(const VertexVertexCandidate& candidate);

    std::vector<long> vertex_indices(
        const Eigen::MatrixXi& E, const Eigen::MatrixXi& F) const override
    {
        return { { vertex0_index, vertex1_index } };
    }

    double compute_distance(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    VectorMax12d compute_distance_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    MatrixMax12d compute_distance_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    double compute_potential(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat) const override;

    VectorMax12d compute_potential_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat) const override;

    MatrixMax12d compute_potential_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat,
        const bool project_hessian_to_psd) const override;

    unsigned int multiplicity = 1;
};

///////////////////////////////////////////////////////////////////////////////

struct EdgeVertexConstraint : EdgeVertexCandidate, CollisionConstraint {
    EdgeVertexConstraint(long edge_index, long vertex_index);
    EdgeVertexConstraint(const EdgeVertexCandidate& candidate);

    std::vector<long> vertex_indices(
        const Eigen::MatrixXi& E, const Eigen::MatrixXi& F) const override
    {
        return { { vertex_index, E(edge_index, 0), E(edge_index, 1) } };
    }

    double compute_distance(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    VectorMax12d compute_distance_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    MatrixMax12d compute_distance_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    double compute_potential(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat) const override;

    VectorMax12d compute_potential_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat) const override;

    MatrixMax12d compute_potential_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat,
        const bool project_hessian_to_psd) const override;

    unsigned int multiplicity = 1;
};

///////////////////////////////////////////////////////////////////////////////

struct EdgeEdgeConstraint : EdgeEdgeCandidate, CollisionConstraint {
    EdgeEdgeConstraint(long edge0_index, long edge1_index, double eps_x);
    EdgeEdgeConstraint(const EdgeEdgeCandidate& candidate, double eps_x);

    std::vector<long> vertex_indices(
        const Eigen::MatrixXi& E, const Eigen::MatrixXi& F) const override
    {
        return { { E(edge0_index, 0), E(edge0_index, 1), //
                   E(edge1_index, 0), E(edge1_index, 1) } };
    }

    double compute_distance(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    VectorMax12d compute_distance_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    MatrixMax12d compute_distance_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    double compute_potential(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat) const override;

    VectorMax12d compute_potential_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat) const override;

    MatrixMax12d compute_potential_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F,
        const double dhat,
        const bool project_hessian_to_psd) const override;

    double eps_x;
};

///////////////////////////////////////////////////////////////////////////////

struct FaceVertexConstraint : FaceVertexCandidate, CollisionConstraint {
    FaceVertexConstraint(long face_index, long vertex_index);
    FaceVertexConstraint(const FaceVertexCandidate& candidate);

    std::vector<long> vertex_indices(
        const Eigen::MatrixXi& E, const Eigen::MatrixXi& F) const override
    {
        return { { vertex_index, //
                   F(face_index, 0), F(face_index, 1), F(face_index, 2) } };
    }

    double compute_distance(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    VectorMax12d compute_distance_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    MatrixMax12d compute_distance_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;
};

///////////////////////////////////////////////////////////////////////////////

struct PlaneVertexConstraint : CollisionConstraint {
    PlaneVertexConstraint(
        const VectorMax3d& plane_origin,
        const VectorMax3d& plane_normal,
        const long vertex_index);

    std::vector<long> vertex_indices(
        const Eigen::MatrixXi& E, const Eigen::MatrixXi& F) const override
    {
        return { vertex_index };
    }

    double compute_distance(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    VectorMax12d compute_distance_gradient(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    MatrixMax12d compute_distance_hessian(
        const Eigen::MatrixXd& V,
        const Eigen::MatrixXi& E,
        const Eigen::MatrixXi& F) const override;

    VectorMax3d plane_origin;
    VectorMax3d plane_normal;
    long vertex_index;
};

///////////////////////////////////////////////////////////////////////////////

struct Constraints {
    std::vector<VertexVertexConstraint> vv_constraints;
    std::vector<EdgeVertexConstraint> ev_constraints;
    std::vector<EdgeEdgeConstraint> ee_constraints;
    std::vector<FaceVertexConstraint> fv_constraints;
    std::vector<PlaneVertexConstraint> pv_constraints;

    size_t size() const;

    size_t num_constraints() const;

    bool empty() const;

    void clear();

    CollisionConstraint& operator[](size_t idx);
    const CollisionConstraint& operator[](size_t idx) const;
};

} // namespace ipc


///////////////////////////////////////////////////////////////
//  Copyright 2020 John Maddock. Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at https://www.boost.org/LICENSE_1_0.txt

#ifdef _MSC_VER
#define _SCL_SECURE_NO_WARNINGS
#endif

#define BOOST_MP_GCD_DEBUG

#include "test.hpp"
#include <boost/multiprecision/cpp_int.hpp>
#include <boost/random.hpp>
#include <map>
#include <tuple>
//
// clang in c++14 mode only has a problem with this file: it's an order of instantiation
// issue caused by us using cpp_int within the gcd algorithm as an error check.
// Just exclude that combination from testing for now as it's purely a testing issue
// and we have other compilers that cover this sanity check...
//
#if !(defined(__clang__) && (__cplusplus > 201300))

using boost::multiprecision::cpp_int;

std::tuple<std::vector<cpp_int>, std::vector<cpp_int>, std::vector<cpp_int> >& get_test_vector(std::size_t bits)
{
   static std::map<std::size_t, std::tuple<std::vector<cpp_int>, std::vector<cpp_int>, std::vector<cpp_int> > > data;

   std::tuple<std::vector<cpp_int>, std::vector<cpp_int>, std::vector<cpp_int> >& result = data[bits];

   if (std::get<0>(result).size() == 0)
   {
      boost::random::mt19937                     mt;
      boost::random::uniform_int_distribution<cpp_int> ui(cpp_int(1) << (bits - 1), cpp_int(1) << bits);

      std::vector<cpp_int>& a = std::get<0>(result);
      std::vector<cpp_int>& b = std::get<1>(result);
      std::vector<cpp_int>& c = std::get<2>(result);

      for (unsigned i = 0; i < 1000; ++i)
      {
         a.push_back(ui(mt));
         b.push_back(ui(mt));
         if (b.back() > a.back())
            b.back().swap(a.back());
         c.push_back(0);
      }
   }
   return result;
}

std::vector<cpp_int>& get_test_vector_a(std::size_t bits)
{
   return std::get<0>(get_test_vector(bits));
}
std::vector<cpp_int>& get_test_vector_b(std::size_t bits)
{
   return std::get<1>(get_test_vector(bits));
}
std::vector<cpp_int>& get_test_vector_c(std::size_t bits)
{
   return std::get<2>(get_test_vector(bits));
}

cpp_int gcd_euler(cpp_int u, cpp_int v)
{
   while (v)
   {
      cpp_int t(v);
      v = u % v;
      u = t;
   }
   return u;
}

namespace boost {
namespace multiprecision {
namespace backends {

std::size_t total_lehmer_gcd_calls      = 0;
std::size_t total_lehmer_gcd_bits_saved = 0;
std::size_t total_lehmer_gcd_cycles     = 0;

}}}

int main()
{
   using boost::multiprecision::backends::total_lehmer_gcd_calls;
   using boost::multiprecision::backends::total_lehmer_gcd_bits_saved;
   using boost::multiprecision::backends::total_lehmer_gcd_cycles;

   unsigned bits = 2048;

   std::vector<cpp_int>& a = get_test_vector_a(bits);
   std::vector<cpp_int>& b = get_test_vector_b(bits);
   std::vector<cpp_int>& c = get_test_vector_c(bits);

   for (unsigned i = 0; i < a.size(); ++i)
   {
      c[i] = gcd(a[i], b[i]);
      cpp_int t = gcd_euler(a[i], b[i]);
      BOOST_CHECK_EQUAL(t, c[i]);
   }

   float average_bits_saved_per_lehmer_call = static_cast<float>(total_lehmer_gcd_bits_saved) / total_lehmer_gcd_calls;
   float average_number_of_euler_cycles_saved = static_cast<float>(total_lehmer_gcd_cycles) / total_lehmer_gcd_calls;

   std::cout << "Average number of bits saved per Lehmer call:           " << average_bits_saved_per_lehmer_call << std::endl;
   std::cout << "Average number of Euclid cycles saved per Lehmer call:  " << average_number_of_euler_cycles_saved << std::endl;

#ifndef BOOST_HAS_INT128
   BOOST_CHECK_GT(average_bits_saved_per_lehmer_call, 30);
#else
   BOOST_CHECK_GT(average_bits_saved_per_lehmer_call, 62);
#endif

   return boost::report_errors();
}

#else

int main() { return 0; }

#endif


/*
 * Filename: similarity.cpp
 *
 * Copyright 2020 Tecnalia
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "similarity.h"

#include <Eigen/LU>

#include <numeric>

namespace manipulability_metrics
{
double volumeIntersection(const Ellipsoid& desired_ellipsoid, const Eigen::Matrix<double, 6, Eigen::Dynamic>& jacobian)
{
  // Force evaluation by casting to avoid multiple (lazy) computations when iterating
  auto jjti = static_cast<Eigen::Matrix<double, 6, 6>>((jacobian * jacobian.transpose()).inverse());

  return std::accumulate(cbegin(desired_ellipsoid), cend(desired_ellipsoid), 1.0, [&](auto volume, const auto& ax) {
    double nu = 1.0 / sqrt(ax.unit.transpose() * jjti * ax.unit);
    return volume * std::min(nu, ax.len);
  });
}

double inverseShapeDiscrepancy(const Ellipsoid& desired_ellipsoid,
                               const Eigen::Matrix<double, 6, Eigen::Dynamic>& jacobian)
{
  // Force evaluation by casting to avoid multiple (lazy) computations when iterating
  auto jjti = static_cast<Eigen::Matrix<double, 6, 6>>((jacobian * jacobian.transpose()).inverse());

  return 1.0 / std::accumulate(cbegin(desired_ellipsoid), cend(desired_ellipsoid),
                               std::numeric_limits<double>::epsilon(), [&](auto sq_diff, const auto& ax) {
                                 double nu = 1.0 / sqrt(ax.unit.transpose() * jjti * ax.unit);
                                 return sq_diff + pow(nu - ax.len, 2.0);
                               });
}

double dualVolumeIntersection(const Ellipsoid& desired_ellipsoid,
                              const Eigen::Matrix<double, 6, Eigen::Dynamic>& left_jacobian,
                              const Eigen::Matrix<double, 6, Eigen::Dynamic>& right_jacobian)
{
  const auto left_jjti =
      static_cast<Eigen::Matrix<double, 6, 6>>((left_jacobian * left_jacobian.transpose()).inverse());
  const auto right_jjti =
      static_cast<Eigen::Matrix<double, 6, 6>>((right_jacobian * right_jacobian.transpose()).inverse());

  double volume = 1.0;
  for (auto i = 0; i < 6; ++i)
  {
    double left_nu = 1.0 / sqrt(desired_ellipsoid[i].unit.transpose() * left_jjti * desired_ellipsoid[i].unit);
    double right_nu = 1.0 / sqrt(desired_ellipsoid[i].unit.transpose() * right_jjti * desired_ellipsoid[i].unit);
    volume *= std::min({ desired_ellipsoid[i].len, left_nu, right_nu });
  }

  return volume;
}

double dualInverseShapeDiscrepancy(const Ellipsoid& desired_ellipsoid,
                                   const Eigen::Matrix<double, 6, Eigen::Dynamic>& left_jacobian,
                                   const Eigen::Matrix<double, 6, Eigen::Dynamic>& right_jacobian)
{
  const auto left_jjti =
      static_cast<Eigen::Matrix<double, 6, 6>>((left_jacobian * left_jacobian.transpose()).inverse());
  const auto right_jjti =
      static_cast<Eigen::Matrix<double, 6, 6>>((right_jacobian * right_jacobian.transpose()).inverse());

  double sq_diff = 0.0;
  for (auto i = 0; i < 6; ++i)
  {
    double left_nu = 1.0 / sqrt(desired_ellipsoid[i].unit.transpose() * left_jjti * desired_ellipsoid[i].unit);
    double right_nu = 1.0 / sqrt(desired_ellipsoid[i].unit.transpose() * right_jjti * desired_ellipsoid[i].unit);
    sq_diff += pow(desired_ellipsoid[i].len - std::min(left_nu, right_nu), 2.0);
  }

  return 1.0 / (sq_diff + std::numeric_limits<double>::epsilon());
}
}  // namespace manipulability_metrics


/*
 * Copyright 2012-2019 CNRS-UM LIRMM, CNRS-AIST JRL
 */

// includes
// std
#include <fstream>
#include <iostream>
#include <tuple>

// boost
#define BOOST_TEST_MODULE QPMultiRobotTest
#include <boost/math/constants/constants.hpp>
#include <boost/test/unit_test.hpp>

// RBDyn
#include <RBDyn/EulerIntegration.h>
#include <RBDyn/FK.h>
#include <RBDyn/FV.h>
#include <RBDyn/ID.h>
#include <RBDyn/MultiBody.h>
#include <RBDyn/MultiBodyConfig.h>
#include <RBDyn/MultiBodyGraph.h>

// Tasks
#include "Tasks/Bounds.h"
#include "Tasks/QPConstr.h"
#include "Tasks/QPContactConstr.h"
#include "Tasks/QPMotionConstr.h"
#include "Tasks/QPSolver.h"
#include "Tasks/QPTasks.h"

// Arms
#include "arms.h"

// Test contact between two robot.
// We set two identical robot at the same positio
// then we link the end effector and add a task
// to make it move on the second robot.
// The first robot must have the same motion.
BOOST_AUTO_TEST_CASE(TwoArmContactTest)
{
  using namespace Eigen;
  using namespace sva;
  using namespace rbd;
  using namespace tasks;
  namespace cst = boost::math::constants;

  MultiBody mb1, mb2;
  MultiBodyConfig mbc1Init, mbc2Init;

  std::tie(mb1, mbc1Init) = makeZXZArm();
  std::tie(mb2, mbc2Init) = makeZXZArm();

  forwardKinematics(mb1, mbc1Init);
  forwardVelocity(mb1, mbc1Init);
  forwardKinematics(mb2, mbc2Init);
  forwardVelocity(mb2, mbc2Init);

  sva::PTransformd X_0_b1(mbc1Init.bodyPosW.back());
  sva::PTransformd X_0_b2(mbc2Init.bodyPosW.back());
  sva::PTransformd X_b1_b2(X_0_b2 * X_0_b1.inv());

  std::vector<MultiBody> mbs = {mb1, mb2};
  std::vector<MultiBodyConfig> mbcs = {mbc1Init, mbc2Init};

  // Test ContactAccConstr constraint
  // Also test PositionTask on the second robot

  qp::QPSolver solver;

  std::vector<qp::UnilateralContact> contVec = {qp::UnilateralContact(0, 1, "b3", "b3", {Vector3d::Zero()},
                                                                      RotX(cst::pi<double>() / 2.), X_b1_b2, 3,
                                                                      std::tan(cst::pi<double>() / 4.))};

#if defined __i386__ || defined __aarch64__
  Matrix3d oriD = RotZ(cst::pi<double>() / 4.);
  if(solver.solver() == "QLD")
  {
    oriD = RotZ(0.0);
  }
#else
  Matrix3d oriD = RotZ(cst::pi<double>() / 4.);
#endif
  Vector3d posD(oriD * mbc2Init.bodyPosW.back().translation());
  qp::PositionTask posTask(mbs, 1, "b3", posD);
  qp::SetPointTask posTaskSp(mbs, 1, &posTask, 1000., 1.);

  qp::ContactAccConstr contCstrAcc;

  contCstrAcc.addToSolver(solver);
  solver.addTask(&posTaskSp);

  solver.nrVars(mbs, contVec, {});
  solver.updateConstrSize();

  // 3 dof + 3 dof + 3 lambda
  BOOST_CHECK_EQUAL(solver.nrVars(), 3 + 3 + 3);

  for(int i = 0; i < 1000; ++i)
  {
    BOOST_REQUIRE(solver.solve(mbs, mbcs));
    for(std::size_t r = 0; r < mbs.size(); ++r)
    {
      eulerIntegration(mbs[r], mbcs[r], 0.001);

      forwardKinematics(mbs[r], mbcs[r]);
      forwardVelocity(mbs[r], mbcs[r]);
    }
    // check that the link hold
    sva::PTransformd X_0_b1_post(mbcs[0].bodyPosW.back());
    sva::PTransformd X_0_b2_post(mbcs[1].bodyPosW.back());
    sva::PTransformd X_b1_b2_post(X_0_b2 * X_0_b1.inv());
    BOOST_CHECK_SMALL((X_b1_b2.matrix() - X_b1_b2_post.matrix()).norm(), 1e-5);
  }
  // check that the task is well minimized
  BOOST_CHECK_SMALL(posTask.eval().norm(), 1e-5);

  contCstrAcc.removeFromSolver(solver);
  solver.removeTask(&posTaskSp);

  // Test ContactSpeedConstr constraint
  // Also test OrientationTask on the second robot

  mbcs = {mbc1Init, mbc2Init};
  qp::OrientationTask oriTask(mbs, 1, "b3", oriD);
  qp::SetPointTask oriTaskSp(mbs, 1, &oriTask, 1000., 1.);

  qp::ContactSpeedConstr contCstrSpeed(0.001);

  contCstrSpeed.addToSolver(solver);
  solver.addTask(&oriTaskSp);

  solver.nrVars(mbs, contVec, {});
  solver.updateConstrSize();

  for(int i = 0; i < 1000; ++i)
  {
    BOOST_REQUIRE(solver.solve(mbs, mbcs));
    for(std::size_t r = 0; r < mbs.size(); ++r)
    {
      eulerIntegration(mbs[r], mbcs[r], 0.001);

      forwardKinematics(mbs[r], mbcs[r]);
      forwardVelocity(mbs[r], mbcs[r]);
    }
    // check that the link hold
    sva::PTransformd X_0_b1_post(mbcs[0].bodyPosW.back());
    sva::PTransformd X_0_b2_post(mbcs[1].bodyPosW.back());
    sva::PTransformd X_b1_b2_post(X_0_b2 * X_0_b1.inv());
    BOOST_CHECK_SMALL((X_b1_b2.matrix() - X_b1_b2_post.matrix()).norm(), 1e-5);
  }
  // check that the task is well minimized
  BOOST_CHECK_SMALL(oriTask.eval().norm(), 1e-5);
}

// Test Motion constraint
// We setup two arm, one with a fixed base and the second
// with a freebase put on the body b3 of the first robot.
// First we launch an impossible motion to check the dynamics
// After we try with an unilateral contact
// Then we try with a bilateral contact.
BOOST_AUTO_TEST_CASE(TwoArmDDynamicContactTest)
{
  using namespace Eigen;
  using namespace sva;
  using namespace rbd;
  using namespace tasks;
  namespace cst = boost::math::constants;

  MultiBody mb1, mb2;
  MultiBodyConfig mbc1Init, mbc2Init;

  std::tie(mb1, mbc1Init) = makeZXZArm();

  forwardKinematics(mb1, mbc1Init);
  forwardVelocity(mb1, mbc1Init);

  std::tie(mb2, mbc2Init) = makeZXZArm(false);
  Vector3d mb2InitPos = mbc1Init.bodyPosW.back().translation();
  Quaterniond mb2InitOri(RotY(cst::pi<double>() / 2.));
  mbc2Init.q[0] = {mb2InitOri.w(), mb2InitOri.x(),     mb2InitOri.y(), mb2InitOri.z(),
                   mb2InitPos.x(), mb2InitPos.y() + 1, mb2InitPos.z()};
  forwardKinematics(mb2, mbc2Init);
  forwardVelocity(mb2, mbc2Init);

  sva::PTransformd X_0_b1(mbc1Init.bodyPosW.back());
  sva::PTransformd X_0_b2(mbc2Init.bodyPosW.front());
  sva::PTransformd X_b1_b2(X_0_b2 * X_0_b1.inv());

  std::vector<MultiBody> mbs = {mb1, mb2};
  std::vector<MultiBodyConfig> mbcs = {mbc1Init, mbc2Init};

  // Test ContactAccConstr constraint
  // Also test PositionTask on the second robot

  qp::QPSolver solver;

  std::vector<Eigen::Vector3d> points = {
      Vector3d(0.1, 0., 0.1),
      Vector3d(0.1, 0., -0.1),
      Vector3d(-0.1, 0., -0.1),
      Vector3d(-0.1, 0., 0.1),
  };

  std::vector<Eigen::Vector3d> biPoints = {
      Vector3d(0., 0., 0.),
      Vector3d(0., 0., 0.),
      Vector3d(0., 0., 0.),
      Vector3d(0., 0., 0.),
  };

  const int nrGen = 4;
  std::vector<Eigen::Matrix3d> biFrames = {
      RotX((1. * cst::pi<double>()) / 4.),
      RotX((3. * cst::pi<double>()) / 4.),
      Matrix3d(RotX((1. * cst::pi<double>()) / 4.) * RotY(cst::pi<double>() / 2.)),
      Matrix3d(RotX((3. * cst::pi<double>()) / 4.) * RotY(cst::pi<double>() / 2.)),
  };

  // The fixed robot can pull the other
  std::vector<qp::UnilateralContact> contVecFail = {
      qp::UnilateralContact(0, 1, "b3", "b0", points, RotX(-cst::pi<double>() / 2.), X_b1_b2, nrGen, 0.7)};

  // The fixed robot can push the other
  std::vector<qp::UnilateralContact> contVec = {
      qp::UnilateralContact({0, 1, "b3", "b0"}, points, RotX(cst::pi<double>() / 2.), X_b1_b2, nrGen, 0.7)};

  // The fixed robot has non coplanar force apply on the other
  std::vector<qp::BilateralContact> contVecBi = {
      qp::BilateralContact({0, 1, "b3", "b0"}, biPoints, biFrames, X_b1_b2, nrGen, 1.)};

  qp::PostureTask posture1Task(mbs, 0, mbc1Init.q, 2., 1.);
  qp::PostureTask posture2Task(mbs, 1, mbc2Init.q, 2., 1.);

  qp::ContactSpeedConstr contCstrSpeed(0.001);

  const double Inf = std::numeric_limits<double>::infinity();
  std::vector<std::vector<double>> torqueMin1 = {{}, {-Inf}, {-Inf}, {-Inf}};
  std::vector<std::vector<double>> torqueMax1 = {{}, {Inf}, {Inf}, {Inf}};
  std::vector<std::vector<double>> torqueMin2 = {{0., 0., 0., 0., 0., 0.}, {-Inf}, {-Inf}, {-Inf}};
  std::vector<std::vector<double>> torqueMax2 = {{0., 0., 0., 0., 0., 0.}, {Inf}, {Inf}, {Inf}};
  std::vector<std::vector<double>> torqueDtMin1 = {{}, {-Inf}, {-Inf}, {-Inf}};
  std::vector<std::vector<double>> torqueDtMax1 = {{}, {Inf}, {Inf}, {Inf}};
  std::vector<std::vector<double>> torqueDtMin2 = {{0., 0., 0., 0., 0., 0.}, {-Inf}, {-Inf}, {-Inf}};
  std::vector<std::vector<double>> torqueDtMax2 = {{0., 0., 0., 0., 0., 0.}, {Inf}, {Inf}, {Inf}};
  qp::MotionConstr motion1(mbs, 0, {torqueMin1, torqueMax1}, {torqueDtMin1, torqueDtMax1}, 0.005);
  qp::MotionConstr motion2(mbs, 1, {torqueMin2, torqueMax2}, {torqueDtMin2, torqueDtMax2}, 0.005);
  qp::PositiveLambda plCstr;

  motion1.addToSolver(solver);
  motion2.addToSolver(solver);
  plCstr.addToSolver(solver);

  contCstrSpeed.addToSolver(solver);
  solver.addTask(&posture1Task);
  solver.addTask(&posture2Task);

  // check the impossible motion
  solver.nrVars(mbs, contVecFail, {});
  solver.updateConstrSize();

  // 3 dof + 9 dof + 4*nrGen lambda
  BOOST_CHECK_EQUAL(solver.nrVars(), 3 + 9 + 4 * nrGen);
  BOOST_REQUIRE(!solver.solve(mbs, mbcs));

  // check the unilateral motion
  mbcs = {mbc1Init, mbc2Init};
  solver.nrVars(mbs, contVec, {});
  solver.updateConstrSize();

  for(int i = 0; i < 1000; ++i)
  {
    BOOST_REQUIRE(solver.solve(mbs, mbcs));
    for(std::size_t r = 0; r < mbs.size(); ++r)
    {
      eulerIntegration(mbs[r], mbcs[r], 0.001);

      forwardKinematics(mbs[r], mbcs[r]);
      forwardVelocity(mbs[r], mbcs[r]);
    }
    // check that the link hold
    sva::PTransformd X_0_b1_post(mbcs[0].bodyPosW.back());
    sva::PTransformd X_0_b2_post(mbcs[1].bodyPosW.front());
    sva::PTransformd X_b1_b2_post(X_0_b2 * X_0_b1.inv());
    BOOST_CHECK_SMALL((X_b1_b2.matrix() - X_b1_b2_post.matrix()).norm(), 1e-5);

    // force in world frame must be the same
    auto f1 = contVec[0].force(solver.lambdaVec(0), contVec[0].r1Cone);
    auto f2 = contVec[0].force(solver.lambdaVec(0), contVec[0].r2Cone);
    BOOST_CHECK_SMALL((f1 + f2).norm(), 1e-5);
  }

  // check the bilateral motion
  mbcs = {mbc1Init, mbc2Init};
  solver.nrVars(mbs, {}, contVecBi);
  solver.updateConstrSize();
  // 3 dof + 9 dof + 4*nrGen lambda
  BOOST_CHECK_EQUAL(solver.nrVars(), 3 + 9 + 4 * nrGen);

  for(int i = 0; i < 1000; ++i)
  {
    // std::cout << i << std::endl;
    BOOST_REQUIRE(solver.solve(mbs, mbcs));
    for(std::size_t r = 0; r < mbs.size(); ++r)
    {
      eulerIntegration(mbs[r], mbcs[r], 0.001);

      forwardKinematics(mbs[r], mbcs[r]);
      forwardVelocity(mbs[r], mbcs[r]);
    }
    // check that the link hold
    sva::PTransformd X_0_b1_post(mbcs[0].bodyPosW.back());
    sva::PTransformd X_0_b2_post(mbcs[1].bodyPosW.front());
    sva::PTransformd X_b1_b2_post(X_0_b2 * X_0_b1.inv());
    BOOST_CHECK_SMALL((X_b1_b2.matrix() - X_b1_b2_post.matrix()).norm(), 1e-5);

    // force in world frame must be the same
    auto f1 = contVec[0].force(solver.lambdaVec(0), contVec[0].r1Cone);
    auto f2 = contVec[0].force(solver.lambdaVec(0), contVec[0].r2Cone);
    BOOST_CHECK_SMALL((f1 + f2).norm(), 1e-5);
  }

  plCstr.removeFromSolver(solver);
  motion2.removeFromSolver(solver);
  motion1.removeFromSolver(solver);
  contCstrSpeed.removeFromSolver(solver);

  solver.removeTask(&posture1Task);
  solver.removeTask(&posture2Task);
}

// Test the MultiCoMTask.
// We try to move the CoM of two arm at a specific position.
BOOST_AUTO_TEST_CASE(TwoArmMultiCoMTest)
{
  using namespace Eigen;
  using namespace sva;
  using namespace rbd;
  using namespace tasks;
  namespace cst = boost::math::constants;

  MultiBody mb1, mb2;
  MultiBodyConfig mbc1Init, mbc2Init;

  std::tie(mb1, mbc1Init) = makeZXZArm(true, sva::PTransformd(Vector3d(-0.5, 0., 0.)));
  forwardKinematics(mb1, mbc1Init);
  forwardVelocity(mb1, mbc1Init);

  std::tie(mb2, mbc2Init) = makeZXZArm(true, sva::PTransformd(Vector3d(0.5, 0., 0.)));
  forwardKinematics(mb2, mbc2Init);
  forwardVelocity(mb2, mbc2Init);

  sva::PTransformd X_0_b1(mbc1Init.bodyPosW.back());
  sva::PTransformd X_0_b2(mbc2Init.bodyPosW.back());
  sva::PTransformd X_b1_b2(X_0_b2 * X_0_b1.inv());

  std::vector<MultiBody> mbs = {mb1, mb2};
  std::vector<MultiBodyConfig> mbcs = {mbc1Init, mbc2Init};

  // Test ContactAccConstr constraint
  // Also test PositionTask on the second robot

  qp::QPSolver solver;

  const int nrGen = 3;
  // The fixed robot can push the other
  std::vector<qp::UnilateralContact> contVec = {qp::UnilateralContact(
      {0, 1, "b3", "b3"}, {Vector3d(0., 0., 0.)}, RotX(cst::pi<double>() / 2.), X_b1_b2, nrGen, 0.7)};

  qp::PostureTask posture1Task(mbs, 0, mbc1Init.q, 2., 1.);
  qp::PostureTask posture2Task(mbs, 1, mbc2Init.q, 2., 1.);
  Vector3d comD((rbd::computeCoM(mb1, mbc1Init) + rbd::computeCoM(mb2, mbc2Init)) / 2. + Vector3d(0., 0., 0.5));
  qp::MultiCoMTask multiCoM(mbs, {0, 1}, comD, 10., 500.);
  // call this method just for test coverage
  multiCoM.updateInertialParameters(mbs);

  qp::ContactSpeedConstr contCstrSpeed(0.001);

  solver.addTask(&posture1Task);
  solver.addTask(&posture2Task);

  solver.nrVars(mbs, contVec, {});

  // Add MultiCoMTask and ContactSpeedConstr after the nrVars call to test
  // addTask and addToSolver with nrVars init
  solver.addTask(mbs, &multiCoM);
  contCstrSpeed.addToSolver(mbs, solver);

  solver.updateConstrSize();

  // 3 dof + 3 dof + 1*nrGen lambda
  BOOST_CHECK_EQUAL(solver.nrVars(), 3 + 3 + 1 * nrGen);

  for(int i = 0; i < 2000; ++i)
  {
    BOOST_REQUIRE(solver.solve(mbs, mbcs));
    for(std::size_t r = 0; r < mbs.size(); ++r)
    {
      eulerIntegration(mbs[r], mbcs[r], 0.001);

      forwardKinematics(mbs[r], mbcs[r]);
      forwardVelocity(mbs[r], mbcs[r]);
    }
    // check that the link hold
    sva::PTransformd X_0_b1_post(mbcs[0].bodyPosW.back());
    sva::PTransformd X_0_b2_post(mbcs[1].bodyPosW.back());
    sva::PTransformd X_b1_b2_post(X_0_b2 * X_0_b1.inv());
    BOOST_CHECK_SMALL((X_b1_b2.matrix() - X_b1_b2_post.matrix()).norm(), 1e-5);
  }
  BOOST_CHECK_SMALL(multiCoM.speed().norm(), 1e-3);

  contCstrSpeed.removeFromSolver(solver);

  solver.removeTask(&posture1Task);
  solver.removeTask(&posture2Task);
  solver.removeTask(&multiCoM);
}

// Test the MultiRobotTransformTask
// We try to set he two end effector at the same frame.
BOOST_AUTO_TEST_CASE(MultiRobotTransformTest)
{
  using namespace Eigen;
  using namespace sva;
  using namespace rbd;
  using namespace tasks;
  namespace cst = boost::math::constants;

  MultiBody mb1, mb2;
  MultiBodyConfig mbc1Init, mbc2Init;

  std::tie(mb1, mbc1Init) =
      makeZXZArm(true, sva::PTransformd(sva::RotZ(-cst::pi<double>() / 4.), Vector3d(-0.5, 0., 0.)));
  forwardKinematics(mb1, mbc1Init);
  forwardVelocity(mb1, mbc1Init);

  std::tie(mb2, mbc2Init) =
      makeZXZArm(false, sva::PTransformd(sva::RotZ(cst::pi<double>() / 2.), Vector3d(0.5, 0., 0.)));
  forwardKinematics(mb2, mbc2Init);
  forwardVelocity(mb2, mbc2Init);

  std::vector<MultiBody> mbs = {mb1, mb2};
  std::vector<MultiBodyConfig> mbcs = {mbc1Init, mbc2Init};

  // Test ContactAccConstr constraint
  // Also test PositionTask on the second robot

  qp::QPSolver solver;

  qp::PostureTask posture1Task(mbs, 0, mbc1Init.q, 0.1, 10.);
  qp::PostureTask posture2Task(mbs, 1, mbc2Init.q, 0.1, 10.);
  qp::MultiRobotTransformTask mrtt(mbs, 0, 1, "b3", "b3", sva::PTransformd(sva::RotZ(-cst::pi<double>() / 8.)),
                                   sva::PTransformd::Identity(), 100., 1000.);
  mrtt.dimWeight((Vector6d() << 0., 0., 1., 1., 1., 0.).finished());

  solver.addTask(&posture1Task);
  solver.addTask(&posture2Task);
  solver.addTask(&mrtt);

  solver.nrVars(mbs, {}, {});
  solver.updateConstrSize();
  // 3 dof + 9 dof
  BOOST_CHECK_EQUAL(solver.nrVars(), 3 + 9);

  for(int i = 0; i < 2000; ++i)
  {
    BOOST_REQUIRE(solver.solve(mbs, mbcs));
    for(std::size_t r = 0; r < mbs.size(); ++r)
    {
      eulerIntegration(mbs[r], mbcs[r], 0.005);

      forwardKinematics(mbs[r], mbcs[r]);
      forwardVelocity(mbs[r], mbcs[r]);
    }
  }
  BOOST_CHECK_SMALL(mrtt.eval().norm(), 1e-3);

  solver.removeTask(&posture1Task);
  solver.removeTask(&posture2Task);
  solver.removeTask(&mrtt);
}

// Test the TorqueTask
BOOST_AUTO_TEST_CASE(TorqueTaskTest)
{
  using namespace Eigen;
  using namespace sva;
  using namespace rbd;
  using namespace tasks;
  namespace cst = boost::math::constants;

  MultiBody mb1, mb2;
  MultiBodyConfig mbc1Init, mbc2Init;

  std::tie(mb1, mbc1Init) =
      makeZXZArm(true, sva::PTransformd(sva::RotZ(-cst::pi<double>() / 4.), Vector3d(-0.5, 0., 0.)));
  forwardKinematics(mb1, mbc1Init);
  forwardVelocity(mb1, mbc1Init);

  std::tie(mb2, mbc2Init) =
      makeZXZArm(false, sva::PTransformd(sva::RotZ(cst::pi<double>() / 2.), Vector3d(0.5, 0., 0.)));
  forwardKinematics(mb2, mbc2Init);
  forwardVelocity(mb2, mbc2Init);

  std::vector<MultiBody> mbs = {mb1, mb2};
  std::vector<MultiBodyConfig> mbcs = {mbc1Init, mbc2Init};

  // Test ContactAccConstr constraint
  // Also test PositionTask on the second robot

  qp::QPSolver solver;

  std::vector<std::vector<double>> lsup;
  std::vector<std::vector<double>> linf;
  std::vector<double> sup;
  std::vector<double> inf;
  std::vector<std::vector<double>> lsupDt;
  std::vector<std::vector<double>> linfDt;
  std::vector<double> supDt;
  std::vector<double> infDt;

  for(const auto j : mb1.joints())
  {
    sup.resize(j.dof());
    inf.resize(j.dof());
    std::fill(sup.begin(), sup.end(), 1e4);
    std::fill(inf.begin(), inf.end(), -1e4);
    lsup.push_back(sup);
    linf.push_back(inf);

    supDt.resize(j.dof());
    infDt.resize(j.dof());
    std::fill(supDt.begin(), supDt.end(), 1e8);
    std::fill(infDt.begin(), infDt.end(), -1e8);
    lsupDt.push_back(supDt);
    linfDt.push_back(infDt);
  }

  TorqueBound tb(lsup, linf);
  TorqueDBound tdb(lsup, linf);

  qp::PostureTask posture1Task(mbs, 0, mbc1Init.q, 0.1, 10.);
  qp::PostureTask posture2Task(mbs, 1, mbc2Init.q, 0.1, 10.);
  qp::TorqueTask tt(mbs, 0, tb, tdb, 0.005, 1);

  solver.addTask(&posture1Task);
  solver.addTask(&posture2Task);
  solver.addTask(&tt);

  solver.nrVars(mbs, {}, {});
  solver.updateConstrSize();
  // 3 dof + 9 dof
  BOOST_CHECK_EQUAL(solver.nrVars(), 3 + 9);

  for(int i = 0; i < 2000; ++i)
  {
    BOOST_REQUIRE(solver.solve(mbs, mbcs));
    for(std::size_t r = 0; r < mbs.size(); ++r)
    {
      eulerIntegration(mbs[r], mbcs[r], 0.005);

      forwardKinematics(mbs[r], mbcs[r]);
      forwardVelocity(mbs[r], mbcs[r]);
    }
  }
  solver.removeTask(&posture1Task);
  solver.removeTask(&posture2Task);
  solver.removeTask(&tt);
}


#include "misc.h"

#include <stdlib.h>
#include <cmath>
#include <random>

#include "../molecules/bead.h"
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

double getDist(Bead* b1, Bead* b2, double box_l[], int npbc) {
  double dist = 0;
  for (int i = 0; i < 3; i++) {
    double di = b1->GetCrd(1, i) - b2->GetCrd(1,i);
    if (i < npbc) {
      // Periodic boundary conditions
      di -= box_l[i] * round(di / box_l[i]);
    }
    dist += (di * di);
  }

  return sqrt(dist);

}

double getDist(Bead& b1, Bead& b2, double box_l[], int npbc) {
  double dist = 0;
  for (int i = 0; i < 3; i++) {
    double di = b1.GetCrd(1, i) - b2.GetCrd(1,i);
    if (i < npbc) {
      // Periodic boundary conditions
      di -= box_l[i] * round(di / box_l[i]);
    }
    dist += (di * di);
  }

  return sqrt(dist);

}

void GetDistVector(Bead& b1, Bead& b2, double box_l[], int npbc,
                   double (&dist)[3]) {
  // Vector pointing from Bead1 to Bead2.
  for (int i = 0; i < 3; i++) {
    double di = b2.GetCrd(1, i) - b1.GetCrd(1, i);

    // Periodic boundary conditions.
    if (i < npbc) {
      di -= box_l[i] * round(di / box_l[i]);
    }
    dist[i] = di;
  }

}

void GetDistVectorConsistent(Bead& b1, Bead& b2, double box_l[], int npbc,
                             double (&dist)[3]) {

  double xyz1[3] = {b1.GetCrd(1, 0), b1.GetCrd(1, 1), b1.GetCrd(1, 2)};
  double xyz2[3] = {b2.GetCrd(1, 0), b2.GetCrd(1, 1), b2.GetCrd(1, 2)};
  for (int i = 0; i < 3; i++) {
    xyz1[i] -= box_l[i] * floor(xyz1[i] / box_l[i]);
    xyz2[i] -= box_l[i] * floor(xyz2[i] / box_l[i]);
    dist[i] = xyz2[i] - xyz1[i];
  }

}

void GetDistVectorC(Bead& b1, Bead& b2, double box_l[], int npbc,
                    double (&dist)[3]) {
  // Vector pointing from Bead1 to Bead2.
  for (int i = 0; i < 3; i++) {
    double di = b2.GetCrd(0, i) - b1.GetCrd(0, i);

    // Periodic boundary conditions.
    if (i < npbc) {
      di -= box_l[i] * round(di / box_l[i]);
    }
    // Finding the shortest distance between periodic images.
    if (abs(di) > box_l[i]/2.0) {
      if (di < 0) {
        di += box_l[i];
      }
      else {
        di -= box_l[i];
      }
    }
    dist[i] = di;
  }

}

void randSphere(double vec[], mt19937& rand_gen) {
  double rand_square = 2;
  double r1, r2;
  while (rand_square > 1) {
    // Random numbers between -1 and 1.
    r1 = 1 - 2 * ((double)rand_gen() / rand_gen.max());
    r2 = 1 - 2 * ((double)rand_gen() / rand_gen.max());
    rand_square = r1 * r1 + r2 * r2;
  }
  double ranh = 2 * sqrt(1 - rand_square);
  vec[0] = r1 * ranh;
  vec[1] = r2 * ranh;
  vec[2] = (1 - 2*rand_square);

}

double gasdev(double mean, double stdev, mt19937& ranGen) {
  static int iset = 0;
  static double gset;
  double fac, rsq, v1, v2;
  if (iset == 0) {
    do {
      v1=2.0*((double)ranGen() / ranGen.max()) - 1.0;
      v2=2.0*((double)ranGen() / ranGen.max()) - 1.0;
      rsq = v1*v1 + v2*v2;
    } while (rsq >= 1.0 || rsq == 0.0 );
      fac = sqrt(-2.0 * log(rsq)/rsq);

      gset = v1*fac;
      iset = 1;
      return v2*fac * stdev + mean;
  }
  else {
    iset = 0;
    return gset * stdev + mean;
  }

  return 0;

}

string YesOrNo(bool input) {
  if (input) {
    return "yes";
  }
  else {
    return "no";
  }

}

int factorial(int n) {
  if (n != 1) {
     return n*factorial(n-1);
  }
  else {
    return 1;
  }

}

/** This is a simple least square fit procedure. */
double Interpolate(double * bin_pos, double * f, int len, double sigma) {
  /*                     INDEX 2
     X = | bin1_pos 1 |  B = | b1 |  Y = | g(bin1_pos) | d
         | bin2_pos 1 |      | b2 |      | g(bin2_pos) | i     INDEX 1
         |    ...     |                  |     ...     | m e
   */
  MatrixXf X(len, 2);
  VectorXf Y(len);
  for (int i = 0; i < len; i++) {
    X(i, 0) = bin_pos[i];
    X(i, 1) = 1.0;
    Y(i) = f[i];
  }
  VectorXf B = (X.transpose() * X).ldlt().solve(X.transpose() * Y);

  /*
  cout << sigma << '\t' << sigma * B(0) + B(1) << endl;
  for (int i = 0; i < len; i++)
  cout << bin_pos[i] << '\t' << f[i] << endl;
  cout << endl;
  */

  return sigma * B(0) + B(1);

}

double Interpolate2(double * bin_pos, double * f, int len, double sigma) {
  MatrixXf X(len, 2);
  VectorXf Y(len);
  for (int i = 0; i < len; i++) {
    X(i, 0) = bin_pos[i];
    X(i, 1) = 1.0;
    Y(i) = f[i];
  }
  VectorXf B = (X.transpose() * X).ldlt().solve(X.transpose() * Y);
  return bin_pos[len-1] * B(0) + B(1);

}


// smooth_feedback: Control theory on Lie groups
// https://github.com/pettni/smooth_feedback
//
// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
//
// Copyright (c) 2021 Petter Nilsson, John B. Mains
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#ifndef SMOOTH__FEEDBACK__OCP_HPP_
#define SMOOTH__FEEDBACK__OCP_HPP_

/**
 * @file
 * @brief Optimal control problem definition and conversion to QP.
 */

#include <Eigen/Core>
#include <smooth/lie_group.hpp>

#include "collocation.hpp"
#include "nlp.hpp"
#include "traits.hpp"

namespace smooth::feedback {

/**
 * @brief Optimal control problem definition
 * @tparam _X state space
 * @tparam _U input space
 *
 * Problem is defined on the interval \f$ t \in [0, t_f] \f$.
 * \f[
 * \begin{cases}
 *  \min              & \theta(t_f, x_0, x_f, q)                                         \\
 *  \text{s.t.}       & x(0) = x_0                                                       \\
 *                    & x(t_f) = x_f                                                     \\
 *                    & \dot x(t) = f(t, x(t), u(t))                                     \\
 *                    & q = \int_{0}^{t_f} g(t, x(t), u(t)) \mathrm{d}t                  \\
 *                    & c_{rl} \leq c_r(t, x(t), u(t)) \leq c_{ru} \quad t \in [0, t_f]  \\
 *                    & c_{el} \leq c_e(t_f, x_0, x_f, q) \leq c_{eu}
 * \end{cases}
 * \f]
 *
 * The optimal control problem depends on arbitrary functions \f$ \theta, f, g, c_r, c_e \f$.
 * The type of those functions are template pararamters in this structure.
 *
 * @note To enable automatic differentiation \f$ \theta, f, g, c_r, c_e \f$ must be templated over the
 * scalar type.
 */
template<LieGroup _X, Manifold _U, typename Theta, typename F, typename G, typename CR, typename CE>
struct OCP
{
  using X = _X;
  using U = _U;

  /// @brief State dimension \f$ n_{x} \f$
  std::size_t nx;
  /// @brief Input dimension \f$ n_{u} \f$
  std::size_t nu;
  /// @brief Number of integrals \f$ n_{q} \f$
  std::size_t nq;
  /// @brief Number of running constraints \f$ n_{cr} \f$
  std::size_t ncr;
  /// @brief Number of end constraints \f$ n_{ce} \f$
  std::size_t nce;

  /// @brief Objective function \f$ \theta : R \times X \times X \times R^{n_q} \rightarrow R \f$
  Theta theta;

  /// @brief System dynamics \f$ f : R \times X \times U \rightarrow Tangent<X> \f$
  F f;
  /// @brief Integrals \f$ g : R \times X \times U \rightarrow R^{n_q} \f$
  G g;

  /// @brief Running constraint \f$ c_r : R \times X \times U \rightarrow R^{n_{cr}} \f$
  CR cr;
  /// @brief Running constraint lower bound \f$ c_{rl} \in R^{n_{cr}} \f$
  Eigen::VectorXd crl;
  /// @brief Running constraint upper bound \f$ c_{ru} \in R^{n_{cr}} \f$
  Eigen::VectorXd cru;

  /// @brief End constraint \f$ c_e : R \times X \times X \times R^{n_q} \rightarrow R^{n_{ce}} \f$
  CE ce;
  /// @brief End constraint lower bound \f$ c_{el} \in R^{n_{ce}} \f$
  Eigen::VectorXd cel;
  /// @brief End constraint upper bound \f$ c_{eu} \in R^{n_{ce}} \f$
  Eigen::VectorXd ceu;
};

/// @brief Concept that is true for OCP specializations
template<typename T>
concept OCPType = traits::is_specialization_of_v<T, OCP>;

/// @brief OCP defined on flat spaces
template<typename Theta, typename F, typename G, typename CR, typename CE>
using FlatOCP = OCP<Eigen::VectorXd, Eigen::VectorXd, Theta, F, G, CR, CE>;

/// @brief Concept that is true for FlatOCP specializations
template<typename T>
concept FlatOCPType = OCPType<T> &&(
  std::is_same_v<typename T::X, Eigen::VectorXd> && std::is_same_v<typename T::U, Eigen::VectorXd>);

/**
 * @brief Check if an OCP is properly defined.
 */
inline bool check_ocp(const OCPType auto & ocp)
{
  using X = typename std::decay_t<decltype(ocp)>::X;
  using U = typename std::decay_t<decltype(ocp)>::U;

  const double t = 0;
  const X x      = Default<X>(ocp.nx);
  const U u      = Default<U>(ocp.nu);

  if (!(static_cast<std::size_t>(dof(x)) == ocp.nx)) { return false; }
  if (!(static_cast<std::size_t>(dof(u)) == ocp.nu)) { return false; }

  const auto dx = ocp.f.template operator()<double>(t, x, u);
  if (!(static_cast<std::size_t>(dx.size()) == ocp.nx)) { return false; }

  const auto g = ocp.g.template operator()<double>(t, x, u);
  if (!(static_cast<std::size_t>(g.size()) == ocp.nq)) { return false; }

  [[maybe_unused]] const double obj = ocp.theta.template operator()<double>(t, x, x, g);

  const auto cr = ocp.cr.template operator()<double>(t, x, u);
  if (!(static_cast<std::size_t>(cr.size()) == ocp.ncr)) { return false; }
  if (!(static_cast<std::size_t>(ocp.crl.size()) == ocp.ncr)) { return false; }
  if (!(static_cast<std::size_t>(ocp.cru.size()) == ocp.ncr)) { return false; }

  const auto ce = ocp.ce.template operator()<double>(t, x, x, g);
  if (!(static_cast<std::size_t>(ce.size()) == ocp.nce)) { return false; }
  if (!(static_cast<std::size_t>(ocp.cel.size()) == ocp.nce)) { return false; }
  if (!(static_cast<std::size_t>(ocp.ceu.size()) == ocp.nce)) { return false; }

  return true;
}

/**
 * @brief Flatten a LieGroup OCP by defining it in the tangent space around a trajectory.
 *
 * @param ocp OCPType defined on a LieGroup
 * @param xl nominal state trajectory
 * @param ul nominal state trajectory
 *
 * @return FlatOCPType in variables (xe, ue) obtained via variables change x = xl ⊕ xe, u = ul ⊕ ue,
 */
inline auto flatten_ocp(const OCPType auto & ocp, auto && xl_fun, auto && ul_fun)
{
  using Eigen::VectorX;
  using X = typename std::decay_t<decltype(ocp)>::X;
  using U = typename std::decay_t<decltype(ocp)>::U;

  assert(Dof<X> == -1 || ocp.nx == Dof<X>);
  assert(Dof<U> == -1 || ocp.nu == Dof<U>);

  auto f_new = [f = ocp.f, xl_fun = xl_fun, ul_fun = ul_fun]<typename T>(
                 const T & t, const VectorX<T> & xe, const VectorX<T> & ue) -> VectorX<T> {
    using X_T = smooth::CastT<T, X>;
    using U_T = smooth::CastT<T, U>;

    // can not double-differentiate, so we neglect derivative of linearization w.r.t. t
    const double tdbl    = static_cast<double>(t);
    const auto [xl, dxl] = diff::dr(xl_fun, wrt(tdbl));
    const auto ul        = ul_fun(tdbl);

    const X_T x = rplus(xl.template cast<T>(), xe);
    const U_T u = rplus(ul.template cast<T>(), ue);

    return dr_expinv<X_T>(xe) * f.template operator()<T>(t, x, u)
         - dl_expinv<X_T>(xe) * dxl.template cast<T>();
  };

  auto g_new = [g = ocp.g, xl_fun = xl_fun, ul_fun = ul_fun]<typename T>(
                 const T & t, const VectorX<T> & xe, const VectorX<T> & ue) -> VectorX<T> {
    return g.template operator()<T>(t, rplus(xl_fun(t), xe), rplus(ul_fun(t), ue));
  };

  auto cr_new = [cr = ocp.cr, xl_fun = xl_fun, ul_fun = ul_fun]<typename T>(
                  const T & t, const VectorX<T> & xe, const VectorX<T> & ue) -> VectorX<T> {
    return cr.template operator()<T>(t, rplus(xl_fun(t), xe), rplus(ul_fun(t), ue));
  };

  auto theta_new =
    [theta = ocp.theta, xl_fun = xl_fun]<typename T>(
      const T & tf, const VectorX<T> & xe0, const VectorX<T> & xef, const VectorX<T> & q) -> T {
    return theta.template operator()<T>(tf, rplus(xl_fun(T(0.)), xe0), rplus(xl_fun(tf), xef), q);
  };

  auto ce_new = [ce = ocp.ce, xl_fun = xl_fun]<typename T>(
                  const T & tf,
                  const VectorX<T> & xe0,
                  const VectorX<T> & xef,
                  const VectorX<T> & q) -> VectorX<T> {
    return ce.template operator()<T>(tf, rplus(xl_fun(T(0.)), xe0), rplus(xl_fun(tf), xef), q);
  };

  return FlatOCP<
    decltype(theta_new),
    decltype(f_new),
    decltype(g_new),
    decltype(cr_new),
    decltype(ce_new)>{
    .nx    = ocp.nx,
    .nu    = ocp.nu,
    .nq    = ocp.nq,
    .ncr   = ocp.ncr,
    .nce   = ocp.nce,
    .theta = std::move(theta_new),
    .f     = std::move(f_new),
    .g     = std::move(g_new),
    .cr    = std::move(cr_new),
    .crl   = ocp.crl,
    .cru   = ocp.cru,
    .ce    = std::move(ce_new),
    .cel   = ocp.cel,
    .ceu   = ocp.ceu,
  };
}

/**
 * @brief Solution to OCP problem.
 */
template<LieGroup X, Manifold U>
struct OCPSolution
{
  double t0;
  double tf;

  /// @brief Integral values
  Eigen::VectorXd Q;

  /// @brief Callable functions for state and input
  std::function<U(double)> u;
  std::function<X(double)> x;

  /// @brief Multipliers for integral constraints
  Eigen::VectorXd lambda_q;

  /// @brief Multipliers for endpoint constraints
  Eigen::VectorXd lambda_ce;

  /// @brief Multipliers for dynamics equality constraint
  std::function<Eigen::VectorXd(double)> lambda_dyn;

  /// @brief Multipliers for active running constraints
  std::function<Eigen::VectorXd(double)> lambda_cr;
};

/// @brief Solution to OCP problem defined on flat spaces
using FlatOCPSolution = OCPSolution<Eigen::VectorXd, Eigen::VectorXd>;

/**
 * @brief Unflatten a FlatOCPSolution
 *
 * If flat_sol is a solution to flat_ocp = flatten_ocp(ocp, xl_fun, ul_fun),
 * then unflatten_ocpsol(flat_sol, xl_fun, ul_fun) is a solution to ocp.
 */
template<LieGroup X, Manifold U>
OCPSolution<X, U> unflatten_ocpsol(const FlatOCPSolution & flatsol, auto && xl_fun, auto && ul_fun)
{
  auto u_unflat = [ul_fun = std::forward<decltype(ul_fun)>(ul_fun),
                   usol   = flatsol.u](double t) -> U { return rplus(ul_fun(t), usol(t)); };

  auto x_unflat = [xl_fun = std::forward<decltype(xl_fun)>(xl_fun),
                   xsol   = flatsol.x](double t) -> X { return rplus(xl_fun(t), xsol(t)); };

  return {
    .t0         = flatsol.t0,
    .tf         = flatsol.tf,
    .Q          = flatsol.Q,
    .u          = std::move(u_unflat),
    .x          = std::move(x_unflat),
    .lambda_q   = flatsol.lambda_q,
    .lambda_ce  = flatsol.lambda_ce,
    .lambda_dyn = flatsol.lambda_dyn,
    .lambda_cr  = flatsol.lambda_cr,
  };
}

namespace detail {

/// @brief Variable and constraint structure of an OCP NLP
auto ocp_nlp_structure(const FlatOCPType auto & ocp, const MeshType auto & mesh)
{
  std::size_t N = mesh.N_colloc();

  // variable layout
  std::array<std::size_t, 4> var_len{
    1,                 // tf
    ocp.nq,            // integrals
    ocp.nx * (N + 1),  // states
    ocp.nu * N,        // inputs
  };

  // constraint layout
  std::array<std::size_t, 4> con_len{
    ocp.nx * N,   // derivatives
    ocp.nq,       // other integrals
    ocp.ncr * N,  // running constraints
    ocp.nce,      // end constraints
  };

  std::array<std::size_t, 5> var_beg{0};
  std::partial_sum(var_len.begin(), var_len.end(), var_beg.begin() + 1);

  std::array<std::size_t, 5> con_beg{0};
  std::partial_sum(con_len.begin(), con_len.end(), con_beg.begin() + 1);

  return std::make_tuple(var_beg, var_len, con_beg, con_len);
}

}  // namespace detail

/**
 * @brief Formulate an OCP as a NLP using collocation on a Mesh.
 *
 * @param ocp Optimal control problem definition
 * @param mesh collocation point structure
 * @return encoding of ocp as a nonlinear program
 *
 * @see ocpsol_to_nlpsol(), nlpsol_to_ocpsol()
 */
NLP ocp_to_nlp(const FlatOCPType auto & ocp, const MeshType auto & mesh)
{
  const auto [var_beg, var_len, con_beg, con_len] = detail::ocp_nlp_structure(ocp, mesh);

  // OBJECTIVE FUNCTION

  auto f = [var_beg = var_beg, var_len = var_len, ocp = ocp](const Eigen::VectorXd & x) -> double {
    const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
    const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

    assert(std::size_t(x.size()) == n);

    const double t0         = 0;
    const double tf         = x(tfvar_B);
    const Eigen::VectorXd Q = x.segment(qvar_B, qvar_L);
    const Eigen::MatrixXd X = x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);

    return colloc_eval_endpt<false>(1, ocp.nx, ocp.theta, t0, tf, X, Q);
  };

  // OBJECTIVE JACOBIAN

  auto df_dx = [var_beg = var_beg, var_len = var_len, ocp = ocp](
                 const Eigen::VectorXd & x) -> Eigen::SparseMatrix<double> {
    const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
    const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

    const double t0         = 0;
    const double tf         = x(tfvar_B);
    const Eigen::VectorXd Q = x.segment(qvar_B, qvar_L);
    const Eigen::MatrixXd X = x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);

    const auto [fval, df_dt0, df_dtf, df_dvecX, df_dQ] =
      colloc_eval_endpt<true>(1, ocp.nx, ocp.theta, t0, tf, X, Q);

    return sparse_block_matrix({
      {df_dtf, df_dQ, df_dvecX, Eigen::SparseMatrix<double>(1, uvar_L)},
    });
  };

  // CONSTRAINT FUNCTION

  auto g = [var_beg = var_beg,
            var_len = var_len,
            con_beg = con_beg,
            con_len = con_len,
            mesh    = mesh,
            ocp     = ocp](const Eigen::VectorXd & x) -> Eigen::VectorXd {
    const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
    const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

    const auto [dcon_B, qcon_B, crcon_B, cecon_B, m] = con_beg;
    const auto [dcon_L, qcon_L, crcon_L, cecon_L]    = con_len;

    assert(std::size_t(x.size()) == n);

    const double t0         = 0;
    const double tf         = x(tfvar_B);
    const Eigen::VectorXd Q = x.segment(qvar_B, qvar_L);
    const Eigen::MatrixXd X = x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);
    const Eigen::MatrixXd U = x.segment(uvar_B, uvar_L).reshaped(ocp.nu, uvar_L / ocp.nu);

    Eigen::VectorXd ret(m);
    // clang-format off
    ret.segment(dcon_B, dcon_L)   = colloc_dyn<false>(ocp.nx, ocp.f, mesh, t0, tf, X, U);
    ret.segment(qcon_B, qcon_L)   = colloc_int<false>(ocp.nq, ocp.g, mesh, t0, tf, Q, X, U).reshaped();
    ret.segment(crcon_B, crcon_L) = colloc_eval<false>(ocp.ncr, ocp.cr, mesh, t0, tf, X, U).reshaped();
    ret.segment(cecon_B, cecon_L) = colloc_eval_endpt<false>(ocp.nce, ocp.nx, ocp.ce, t0, tf, X, Q).reshaped();
    // clang-format on
    return ret;
  };

  // CONSTRAINT JACOBIAN
  auto dg_dx =
    [var_beg = var_beg, var_len = var_len, mesh = mesh, ocp = ocp](const Eigen::VectorXd & x) {
      const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
      const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

      assert(std::size_t(x.size()) == n);

      const double t0          = 0;
      const double tf          = x(tfvar_B);
      const Eigen::VectorXd Qm = x.segment(qvar_B, qvar_L);
      const Eigen::MatrixXd Xm = x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);
      const Eigen::MatrixXd Um = x.segment(uvar_B, uvar_L).reshaped(ocp.nu, uvar_L / ocp.nu);

      // clang-format off
      const auto [Fval, dF_dt0, dF_dtf, dF_dX, dF_dU]        = colloc_dyn<true>(ocp.nx, ocp.f, mesh, t0, tf, Xm, Um);
      const auto [Gval, dG_dt0, dG_dtf, dG_dQ, dG_dX, dG_dU] = colloc_int<true>(ocp.nq, ocp.g, mesh, t0, tf, Qm, Xm, Um);
      const auto [CRval, dCR_dt0, dCR_dtf, dCR_dX, dCR_dU]   = colloc_eval<true>(ocp.ncr, ocp.cr, mesh, t0, tf, Xm, Um);
      const auto [CEval, dCE_dt0, dCE_dtf, dCE_dX, dCE_dQ]   = colloc_eval_endpt<true>(ocp.nce, ocp.nx, ocp.ce, t0, tf, Xm, Qm);
      // clang-format on

      return sparse_block_matrix({
        // clang-format off
        { dF_dtf,     {},  dF_dX,  dF_dU},
        { dG_dtf,  dG_dQ,  dG_dX,  dG_dU},
        {dCR_dtf,     {}, dCR_dX, dCR_dU},
        {dCE_dtf, dCE_dQ, dCE_dX,     {}},
        // clang-format on
      });
    };

  const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
  const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

  const auto [dcon_B, qcon_B, crcon_B, cecon_B, m] = con_beg;
  const auto [dcon_L, qcon_L, crcon_L, cecon_L]    = con_len;

  // VARIABLE BOUNDS

  Eigen::VectorXd xl = Eigen::VectorXd::Constant(n, -std::numeric_limits<double>::infinity());
  Eigen::VectorXd xu = Eigen::VectorXd::Constant(n, std::numeric_limits<double>::infinity());

  xl.segment(tfvar_B, tfvar_L).setZero();  // tf lower bounded by zero

  // CONSTRAINT BOUNDS

  Eigen::VectorXd gl(m);
  Eigen::VectorXd gu(m);

  // derivative constraints are equalities
  gl.segment(dcon_B, dcon_L).setZero();
  gu.segment(dcon_B, dcon_L).setZero();

  // integral constraints are equalities
  gl.segment(qcon_B, qcon_L).setZero();
  gu.segment(qcon_B, qcon_L).setZero();

  // running constraints
  gl.segment(crcon_B, crcon_L) = ocp.crl.replicate(mesh.N_colloc(), 1);
  gu.segment(crcon_B, crcon_L) = ocp.cru.replicate(mesh.N_colloc(), 1);

  // end constraints
  gl.segment(cecon_B, cecon_L) = ocp.cel;
  gu.segment(cecon_B, cecon_L) = ocp.ceu;

  return {
    .n       = n,
    .m       = m,
    .f       = std::move(f),
    .xl      = std::move(xl),
    .xu      = std::move(xu),
    .g       = std::move(g),
    .gl      = std::move(gl),
    .gu      = std::move(gu),
    .df_dx   = std::move(df_dx),
    .dg_dx   = std::move(dg_dx),
    .d2f_dx2 = {},
    .d2g_dx2 = {},
  };
}

/**
 * @brief Convert nonlinear program solution to ocp solution
 */
FlatOCPSolution nlpsol_to_ocpsol(
  const FlatOCPType auto & ocp, const MeshType auto & mesh, const NLPSolution & nlp_sol)
{
  const std::size_t N                             = mesh.N_colloc();
  const auto [var_beg, var_len, con_beg, con_len] = detail::ocp_nlp_structure(ocp, mesh);

  const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
  const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

  const auto [dcon_B, qcon_B, crcon_B, cecon_B, m] = con_beg;
  const auto [dcon_L, qcon_L, crcon_L, cecon_L]    = con_len;

  const double t0 = 0;
  const double tf = nlp_sol.x(tfvar_B);

  const Eigen::VectorXd Q = nlp_sol.x.segment(qvar_B, qvar_L);

  // state vector has a value at the endpoint

  Eigen::MatrixXd X(ocp.nx, N + 1);
  X = nlp_sol.x.segment(xvar_B, xvar_L).reshaped(ocp.nx, xvar_L / ocp.nx);

  auto xfun = [t0 = t0, tf = tf, mesh = mesh, X = std::move(X)](double t) -> Eigen::VectorXd {
    return mesh.template eval<Eigen::VectorXd>((t - t0) / (tf - t0), X.colwise(), 0, true);
  };

  // for these we repeat last point since there are no values for endpoint

  Eigen::MatrixXd U(ocp.nu, N);
  U = nlp_sol.x.segment(uvar_B, uvar_L).reshaped(ocp.nu, uvar_L / ocp.nu);

  auto ufun = [t0 = t0, tf = tf, mesh = mesh, U = std::move(U)](double t) -> Eigen::VectorXd {
    return mesh.template eval<Eigen::VectorXd>((t - t0) / (tf - t0), U.colwise(), 0, false);
  };

  Eigen::MatrixXd Ldyn(ocp.nx, N);
  Ldyn = nlp_sol.lambda.segment(dcon_B, dcon_L).reshaped(ocp.nx, dcon_L / ocp.nx);

  auto ldfun =
    [t0 = t0, tf = tf, mesh = mesh, Ldyn = std::move(Ldyn)](double t) -> Eigen::VectorXd {
    return mesh.template eval<Eigen::VectorXd>((t - t0) / (tf - t0), Ldyn.colwise(), 0, false);
  };

  Eigen::MatrixXd Lcr(ocp.ncr, N);
  Lcr = nlp_sol.lambda.segment(crcon_B, crcon_L).reshaped(ocp.ncr, crcon_L / ocp.ncr);

  auto lcrfun = [t0 = t0, tf = tf, mesh = mesh, Lcr = std::move(Lcr)](double t) -> Eigen::VectorXd {
    return mesh.template eval<Eigen::VectorXd>((t - t0) / (tf - t0), Lcr.colwise(), 0, false);
  };

  return {
    .t0         = t0,
    .tf         = tf,
    .Q          = std::move(Q),
    .u          = std::move(ufun),
    .x          = std::move(xfun),
    .lambda_q   = nlp_sol.lambda.segment(qcon_B, qcon_L),
    .lambda_ce  = nlp_sol.lambda.segment(cecon_B, cecon_L),
    .lambda_dyn = std::move(ldfun),
    .lambda_cr  = std::move(lcrfun),
  };
}

/**
 * @brief Convert ocp  solution to nonlinear program
 */
NLPSolution ocpsol_to_nlpsol(
  const FlatOCPType auto & ocp, const MeshType auto & mesh, const FlatOCPSolution & ocpsol)
{
  const std::size_t N                             = mesh.N_colloc();
  const auto [var_beg, var_len, con_beg, con_len] = detail::ocp_nlp_structure(ocp, mesh);

  const auto [tfvar_B, qvar_B, xvar_B, uvar_B, n] = var_beg;
  const auto [tfvar_L, qvar_L, xvar_L, uvar_L]    = var_len;

  const auto [dcon_B, qcon_B, crcon_B, cecon_B, m] = con_beg;
  const auto [dcon_L, qcon_L, crcon_L, cecon_L]    = con_len;

  const auto [nodes, weights] = mesh.all_nodes_and_weights();

  const double t0 = 0;
  const double tf = ocpsol.tf;

  Eigen::VectorXd x(n);
  Eigen::VectorXd zl(n), zu(n);
  Eigen::VectorXd lambda(m);

  zl.setZero();
  zu.setZero();

  x(tfvar_B)                = ocpsol.tf;
  x.segment(qvar_B, qvar_L) = ocpsol.Q;

  lambda.segment(qcon_B, qcon_L)   = ocpsol.lambda_q;
  lambda.segment(cecon_B, cecon_L) = ocpsol.lambda_ce;

  for (auto i = 0u; const auto tau : nodes) {
    x.segment(xvar_B + i * ocp.nx, ocp.nx) = ocpsol.x(t0 + tau * (tf - t0));
    if (i < N) {
      x.segment(uvar_B + i * ocp.nu, ocp.nu)         = ocpsol.u(t0 + tau * (tf - t0));
      lambda.segment(dcon_B + i * ocp.nx, ocp.nx)    = ocpsol.lambda_dyn(t0 + tau * (tf - t0));
      lambda.segment(crcon_B + i * ocp.ncr, ocp.ncr) = ocpsol.lambda_cr(t0 + tau * (tf - t0));
    }
    ++i;
  }

  return {
    .status = NLPSolution::Status::Unknown,
    .x      = std::move(x),
    .zl     = Eigen::VectorXd::Zero(n),
    .zu     = Eigen::VectorXd::Zero(n),
    .lambda = std::move(lambda),
  };
}

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__OCP_HPP_


#include "generation-gui.h"

#include <imgui/imgui.h>
#include <Eigen/Dense>
#include <glm/gtc/type_ptr.hpp>
#include <string>

#include "gui/icons-awesome.h"
#include "maths/rbf.h"
#include "components/graphics/material.h"
#include "components/physics/transform.h"

GenerationGui::GenerationGui(Context& ctx, SingletonComponents& scomps)
    : m_ctx(ctx), m_scomps(scomps)
{
    m_controlPointsXYZ.push_back(glm::ivec3(10, 10, 10));
    m_controlPointsWeights.push_back(5);
    m_controlPointsXYZ.push_back(glm::ivec3(0, 0, 0));
    m_controlPointsWeights.push_back(5);
}

GenerationGui::~GenerationGui() {}

void GenerationGui::update() {
    ImGui::Begin(ICON_FA_SEEDLING "  Generation", 0);
        if (ImGui::Button("Add control point")) {
            m_controlPointsXYZ.push_back(glm::ivec3(0, 0, 0));
            m_controlPointsWeights.push_back(1);
        }

        if (m_controlPointsXYZ.size() > 2)
            if (ImGui::Button("Remove control point")) {
                m_controlPointsXYZ.pop_back();
                m_controlPointsWeights.push_back(1);
            }

        ImGui::Spacing();
        ImGui::Separator();
        ImGui::Spacing();

        for (unsigned int i = 0; i < m_controlPointsXYZ.size(); i++) {
            ImGui::SliderFloat(std::to_string(i).c_str(), &m_controlPointsWeights.at(i), 1, 30);
            ImGui::SliderInt3(std::to_string(i).c_str(), (int*) glm::value_ptr(m_controlPointsXYZ.at(i)), 0, 20);
            ImGui::Separator();
        }

        ImGui::Separator();
        ImGui::Spacing();

        if (ImGui::Button("Generate")) {
            Eigen::VectorXd controlPointWeights(m_controlPointsXYZ.size());
            for (unsigned int i = 0; i < m_controlPointsWeights.size(); i++) {
                controlPointWeights[i] = m_controlPointsWeights.at(i);
            }

            std::vector<glm::ivec3> coordWithYtoFind;
            std::vector<met::entity> entityToChange;
            m_ctx.registry.view<comp::Transform>().each([&](met::entity id, comp::Transform& transform){
                coordWithYtoFind.push_back(transform.position);
                entityToChange.push_back(id);
            });
            voxmt::rbfInterpolate(coordWithYtoFind, m_controlPointsXYZ, controlPointWeights, voxmt::RBFType::LINEAR, 0.5f, voxmt::RBFTransformAxis::Y);

            for (size_t i = 0; i < coordWithYtoFind.size(); i++) {
                comp::Transform& trans =  m_ctx.registry.get<comp::Transform>(entityToChange.at(i));
                trans.position = coordWithYtoFind.at(i);
            }

        }
    ImGui::End();
}

void GenerationGui::onEvent(GuiEvent e) {

}


//---------------------------------------------------------------------------//
// Copyright (c) 2014 Roshan <thisisroshansmail@gmail.com>
//
// Distributed under the Boost Software License, Version 1.0
// See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt
//
// See http://kylelutz.github.com/compute for more information.
//---------------------------------------------------------------------------//

#define BOOST_TEST_MODULE TestDiscreteDistribution
#include <boost/test/unit_test.hpp>

#include <boost/compute/system.hpp>
#include <boost/compute/command_queue.hpp>
#include <boost/compute/algorithm/count_if.hpp>
#include <boost/compute/container/vector.hpp>
#include <boost/compute/random/default_random_engine.hpp>
#include <boost/compute/random/discrete_distribution.hpp>
#include <boost/compute/lambda.hpp>

#include "context_setup.hpp"

BOOST_AUTO_TEST_CASE(discrete_distribution_doctest)
{
    using boost::compute::uint_;
    using boost::compute::lambda::_1;

    boost::compute::vector<uint_> vec(100, context);

//! [generate]
// initialize the default random engine
boost::compute::default_random_engine engine(queue);

// initialize weights
int weights[] = {2, 2};

// setup the discrete distribution to produce integers 0 and 1
// with equal weights
boost::compute::discrete_distribution<uint_> distribution(weights, weights+2);

// generate the random values and store them to 'vec'
distribution.generate(vec.begin(), vec.end(), engine, queue);
// ! [generate]

    BOOST_CHECK_EQUAL(
        boost::compute::count_if(
            vec.begin(), vec.end(), _1 > 1, queue
        ),
        size_t(0)
    );
}

BOOST_AUTO_TEST_SUITE_END()


#include <cassert>
#include <cmath>
#include <boost/multiprecision/gmp.hpp>

#include <array>
#include <algorithm>
#include <bitset>
#include <numeric>
#include <sstream>
#include <utility>
#include <vector>

#include <functional>

#include <QtDebug>
#include <QTime>

#include "lrucache.hpp"

using namespace std;
using boost::multiprecision::mpq_rational;

QDebug operator<<(QDebug d, const mpq_rational& r) {
    d.nospace();
    d.noquote();
    stringstream s;
    s << r;
    d << QString::fromStdString(s.str());
    return d.resetFormat();
}

template<class T, std::size_t N>
QDebug operator<<(QDebug d, const array<T, N>& a) {
    d.nospace();
    d.noquote();

    d << "std::array(";
    for(auto i = begin(a); i != end(a); ++i) {
        d << *i;
        if (distance(i, end(a)) > 1) {
            d << ", ";
        }
    }
    d << ")";
    return d.resetFormat();
}

template<class T>
QDebug operator<<(QDebug d, const vector<T>& v) {
    d.nospace();
    d.noquote();
    d << "std::vector(";
    for(auto i = begin(v); i != end(v); ++i) {
        d << *i;
        if (distance(i, end(v)) > 1) {
            d << ", ";
        }
    }
    d << ")";
    return d.resetFormat();
}

template<class S, class T>
QDebug operator<<(QDebug d, const pair<S, T>& p) {
    d.nospace();
    d.noquote();
    d << "std::pair(" << p.first << ", " << p.second << ")";
    return d.resetFormat();
}

// warning: this is only supposed to work for our narrow task of 10 balls for 7 colors. this could go wrong in other cases.
namespace std {
template<size_t L>
struct hash<array<int8_t, L>>
{
    size_t operator()(const array<int8_t, L>& a) const {
        int shift = 0;
        size_t ret = 0;

        foreach (int8_t i, a) {
            ret += i << shift;
            shift += 4;
        }

        return ret;
    }
};
}

template<int NumberOfColors, int NumberOfBallsPerColor>
class Urn {
public:
    Urn() {
        _urn[NumberOfBallsPerColor] = NumberOfColors;
        for (auto i = 0; i < NumberOfBallsPerColor; i += 1) {
            _urn[i] = 0;
        }
    }

    //bool operator==(const Urn<NumberOfColors, NumberOfBallsPerColor>& u) const {
    //    return _urn == u._urn;
    //}

    //bool operator<(const Urn<NumberOfColors, NumberOfBallsPerColor>& u) const {
    //    for (auto i = begin(_urn), j = begin(u._urn); i != end(_urn); ++i, ++j) {
    //        if (*i < *j) return true;
    //    }
    //    return false;
    //}

    bool sumOfElementsIsCorrect() const {
        int8_t sum = 0;
        for (auto i = begin(_urn); i != end(_urn); ++i) {
            sum += *i;
        }
        return sum == NumberOfColors;
    }

    int distinctDrawnColors() const {
        return NumberOfColors - _urn[NumberOfBallsPerColor];
    }

    vector<pair<int /*index*/, mpq_rational /*prob*/>> pickOptions(int remainingBalls) const {
        vector<pair<int, mpq_rational>> options;

        int balls = remainingBalls;

        for (int i = 1; i <= NumberOfBallsPerColor; i += 1) {
            if (_urn[i] > 0) {
                options.emplace_back(make_pair(i, mpq_rational((int)_urn[i]*i, balls)));
            }
        }
        return options;
    }

    void pick(int i) {
        assert(i > 0 && i < (int) _urn.size()); // remove for more speed
        assert(_urn[i] > 0);                    // remove for more speed
        _urn[i] -= 1;
        _urn[i-1] += 1;
    }

    mpq_rational expectedNumberOfDistinctColors(int picks_) const {

        assert(picks_ <= NumberOfColors*NumberOfBallsPerColor);

        function<mpq_rational(Urn<NumberOfColors, NumberOfBallsPerColor>, mpq_rational, int, lru_cache<array<int8_t, NumberOfBallsPerColor + 1>, mpq_rational>*)> recurse = [&recurse, &picks_](Urn<NumberOfColors, NumberOfBallsPerColor> u, mpq_rational p, int picks, lru_cache<array<int8_t, NumberOfBallsPerColor + 1>, mpq_rational>* cache) -> mpq_rational {

            if (picks == 0) {
                return u.distinctDrawnColors();
            }

            if (cache->exists(u._urn)) {
                return cache->get(u._urn);
            }

            mpq_rational e = 0;

            for (const auto& pick : u.pickOptions(NumberOfBallsPerColor*NumberOfColors - (picks_ - picks))) {
                auto uCopy = u;
                uCopy.pick(pick.first);

                e += recurse(uCopy, p*pick.second, picks-1, cache)*pick.second;
            }

            cache->put(u._urn, e);

            return e;
        };

        lru_cache<array<int8_t, NumberOfBallsPerColor + 1>, mpq_rational> cache(100000);

        auto ret = recurse(*this, 1, picks_, &cache);

        return ret;
    }

//private:
public: // for testing
    array<int8_t, NumberOfBallsPerColor + 1> _urn; // _urn[numberofballs] = numberofcolors that have as many balls
};

void test() {
    Urn<7, 10> u710; assert(u710._urn.size() == 11);
    assert(u710._urn[10] == 7);
    for(auto i = 0; i < 10; i += 1) assert(u710._urn[i] == 0);
    assert(u710.sumOfElementsIsCorrect());

    Urn<1, 1> u11;
    assert(u11.expectedNumberOfDistinctColors(1) == 1);

    Urn<1, 2> u12;
    assert(u12.expectedNumberOfDistinctColors(1) == 1);
    assert(u12.expectedNumberOfDistinctColors(2) == 1);

    Urn<2, 3> u22;
    assert(u22.expectedNumberOfDistinctColors(2) == 1*mpq_rational(4, 10) + 2*mpq_rational( 6, 10));
    assert(u22.expectedNumberOfDistinctColors(3) == 1*mpq_rational(4, 40) + 2*mpq_rational(36, 40));
    assert(u22.expectedNumberOfDistinctColors(4) == 2);
    assert(u22.expectedNumberOfDistinctColors(5) == 2);
    assert(u22.expectedNumberOfDistinctColors(6) == 2);

    Urn<3, 2> u32_;
    assert(u32_.expectedNumberOfDistinctColors(2) == 1*mpq_rational( 3, 15) + 2*mpq_rational(12, 15));
    assert(u32_.expectedNumberOfDistinctColors(3) == 2*mpq_rational(18, 30) + 3*mpq_rational(12, 30));
    //assert(u32_.expectedNumberOfDistinctColors(4) == );
    //assert(u32_.expectedNumberOfDistinctColors(5) == );
    assert(u32_.expectedNumberOfDistinctColors(6) == 3);
}

void doit() {
    Urn<7, 10> u;
    qDebug() << u.expectedNumberOfDistinctColors(20);
}

int main() {

    test();

    doit();

    return 0;
}