//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_ELLIPTIC_FUNCTIONS_GENERIC_ELLIPJ_HPP_INCLUDED
#define NT2_ELLIPTIC_FUNCTIONS_GENERIC_ELLIPJ_HPP_INCLUDED

#include <nt2/elliptic/functions/ellipj.hpp>
#include <boost/fusion/tuple.hpp>
#include <nt2/include/constants/eps.hpp>
#include <nt2/include/functions/simd/am.hpp>
#include <nt2/include/functions/simd/if_zero_else.hpp>
#include <nt2/include/functions/simd/is_eqz.hpp>
#include <nt2/include/functions/simd/multiplies.hpp>
#include <nt2/include/functions/simd/oneminus.hpp>
#include <nt2/include/functions/simd/sincos.hpp>
#include <nt2/include/functions/simd/sqr.hpp>
#include <nt2/include/functions/simd/sqrt.hpp>
#include <nt2/sdk/meta/scalar_of.hpp>

namespace nt2 { namespace ext
{
  BOOST_DISPATCH_IMPLEMENT  (ellipj_, tag::cpu_,
                             (A0)(A1),
                             (generic_<unspecified_<A0> >)
                             (generic_<floating_<A1> >)

                            )
  {
    typedef boost::fusion::tuple<A0,A0,A0>        result_type;

    inline result_type operator()(A0 const& a0,A1 const & a1) const
    {
      typedef typename nt2::meta::scalar_of<A0>::type sA0;
      A0 s, c, d;
      nt2::ellipj(a0, a1, nt2::Eps<sA0>(), s, c, d);
      return result_type(s, c, d);
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (ellipj_, tag::cpu_,
                             (A0)(A1)(A2),
                             (generic_<unspecified_<A0> >)
                             (generic_<floating_<A1> >)
                             (scalar_<floating_<A2> >)

                            )
  {
    typedef boost::fusion::tuple<A0,A0,A0>        result_type;

    inline result_type operator()(A0 const& a0,A1 const & a1,A2 const & a2) const
    {
      A0 s, c, d;
      nt2::ellipj(a0, a1, a2, s, c, d);
      return result_type(s, c, d);
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (  ellipj_, tag::cpu_,
                               (A0)(A1)(A2),
                               (generic_<unspecified_<A0> >)
                               (generic_<floating_<A1> >)
                               (scalar_<floating_<A2> >)
                               (generic_<unspecified_<A0> >)
                            )
  {
    typedef boost::fusion::tuple<A0,A0>        result_type;

    inline result_type operator()(A0 const& a0,A1 const & a1,A2 const & a2,A0 & a3) const
    {
      A0 s, c;
      nt2::ellipj(a0,a1,a2,s,c,a3);
      return result_type(s, c);
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (  ellipj_, tag::cpu_,
                               (A0)(A1)(A2),
                               (generic_<unspecified_<A0> >)
                               (generic_<floating_<A1> >)
                               (unspecified_<A2>)               //passing _()
                               (generic_<unspecified_<A0> >)
                            )
  {
    typedef boost::fusion::tuple<A0,A0>        result_type;

    inline result_type operator()(A0 const& a0,A1 const & a1,A2 const &,A0 & a3) const
    {
      typedef typename nt2::meta::scalar_of<A0>::type sA0;
      A0 s, c;
      nt2::ellipj(a0,a1,Eps<sA0>(),s,c,a3);
      return result_type(s, c);
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (  ellipj_, tag::cpu_,
                               (A0)(A1)(A2),
                               (generic_<unspecified_<A0> >)
                               (generic_<floating_<A1> >)
                               (scalar_<floating_<A2> >)
                               (generic_<unspecified_<A0> >)
                               (generic_<unspecified_<A0> >)
                            )
  {
    typedef A0 result_type;
    inline result_type operator()(A0 const& a0,A1 const & a1,A2 const & a2,A0 & a3,A0 & a4) const
    {
      A0 s;
      nt2::ellipj(a0,a1,a2,s,a3,a4);
      return s;
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (  ellipj_, tag::cpu_,
                               (A0)(A1)(A2),
                               (generic_<unspecified_<A0> >)
                               (generic_<floating_<A1> >)
                               (unspecified_<A2>)         //passing _()
                               (generic_<unspecified_<A0> >)
                               (generic_<unspecified_<A0> >)
                            )
  {
    typedef A0 result_type;
    inline result_type operator()(A0 const& a0,A1 const & a1,A2 const &,A0 & a3,A0 & a4) const
    {
      typedef typename nt2::meta::scalar_of<A0>::type sA0;
      A0 s;
      nt2::ellipj(a0,a1,Eps<sA0>(),s,a3,a4);
      return s;
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (  ellipj_, tag::cpu_,
                               (A0)(A1)(A2),
                               (generic_<floating_<A0> >)
                               (generic_<floating_<A1> >)
                               (scalar_<floating_<A2> >)
                               (generic_<floating_<A0> >)
                               (generic_<floating_<A0> >)
                               (generic_<floating_<A0> >)
                            )
  {
    typedef void result_type;
    inline result_type operator()(A0 const& a0, A1 const & a1, A2 const & a2,
                                  A0 & a3, A0 & a4, A0 & a5) const
    {
      A0 a = am(a0, a1, a2, 'm');
      a3 = nt2::sincos(a, a4);
      a5 = nt2::sqrt(oneminus(if_zero_else(is_eqz(a1), a1*nt2::sqr(a3))));
    }
  };

  BOOST_DISPATCH_IMPLEMENT  (  ellipj_, tag::cpu_,
                               (A0)(A1)(A2),
                               (generic_<floating_<A0> >)
                               (generic_<floating_<A1> >)
                               (unspecified_<A2>)       //passing _()
                               (generic_<floating_<A0> >)
                               (generic_<floating_<A0> >)
                               (generic_<floating_<A0> >)
                            )
  {
    typedef void result_type;
    inline result_type operator()(A0 const& a0, A1 const & a1, A2 const &,
                                  A0 & a3, A0 & a4, A0 & a5) const
    {
      typedef typename nt2::meta::scalar_of<A0>::type sA0;
      A0 a = am(a0, a1, Eps<sA0>(), 'm');
      a3 = nt2::sincos(a, a4);
      a5 = nt2::sqrt(oneminus(if_zero_else(is_eqz(a1), a1*nt2::sqr(a3))));
    }
  };

} }
#endif


#include "icp-ceres.h"

#include <Eigen/Dense>
#include <math.h>
#include <unordered_map>
#include <vector>

//#include "Visualize.h"

#include <ceres/local_parameterization.h>
#include <ceres/autodiff_local_parameterization.h>
#include <ceres/types.h>
#include <ceres/rotation.h>
#include <ceres/ceres.h>

#include <ceres/loss_function.h>

#define useLocalParam

namespace ICP_Ceres {

/*
Ceres Solving FAQ extracted from http://ceres-solver.org/solving_faqs.html:

1. For small (a few hundred parameters) or dense problems use DENSE_QR.

2. For general sparse problems (i.e., the Jacobian matrix has a substantial number of zeros) use SPARSE_NORMAL_CHOLESKY.
This requires that you have SuiteSparse or CXSparse installed.

3. For bundle adjustment problems with up to a hundred or so cameras, use DENSE_SCHUR.

4. For larger bundle adjustment problems with sparse Schur Complement/Reduced camera matrices use SPARSE_SCHUR.
This requires that you build Ceres with support for SuiteSparse, CXSparse or Eigen’s sparse linear algebra libraries.
If you do not have access to these libraries for whatever reason, ITERATIVE_SCHUR with SCHUR_JACOBI is an excellent alternative.

5. For large bundle adjustment problems (a few thousand cameras or more) use the ITERATIVE_SCHUR solver.
There are a number of preconditioner choices here. SCHUR_JACOBI offers an excellent balance of speed and accuracy.
This is also the recommended option if you are solving medium sized problems for which DENSE_SCHUR is too slow but SuiteSparse is not available.
Note: If you are solving small to medium sized problems, consider setting Solver::Options::use_explicit_schur_complement to true, it can result in a substantial performance boost.
If you are not satisfied with SCHUR_JACOBI‘s performance try CLUSTER_JACOBI and CLUSTER_TRIDIAGONAL in that order. They require that you have SuiteSparse installed.
Both of these preconditioners use a clustering algorithm. Use SINGLE_LINKAGE before CANONICAL_VIEWS.
*/
ceres::Solver::Options getOptions(){
    // Set a few options
    ceres::Solver::Options options;
    //options.use_nonmonotonic_steps = true;
    //options.preconditioner_type = ceres::IDENTITY;
    options.linear_solver_type = ceres::DENSE_QR;
    options.max_num_iterations = 50;

//    options.preconditioner_type = ceres::SCHUR_JACOBI;
//    options.linear_solver_type = ceres::DENSE_SCHUR;
//    options.use_explicit_schur_complement=true;
//    options.max_num_iterations = 100;

    cout << "Ceres Solver getOptions()" << endl;
    cout << "Ceres preconditioner type: " << options.preconditioner_type << endl;
    cout << "Ceres linear algebra type: " << options.sparse_linear_algebra_library_type << endl;
    cout << "Ceres linear solver type: " << options.linear_solver_type << endl;

    return options;
}

ceres::Solver::Options getOptionsMedium(){
    // Set a few options
    ceres::Solver::Options options;

    #ifdef _WIN32
        options.sparse_linear_algebra_library_type = ceres::EIGEN_SPARSE;
        options.linear_solver_type = ceres::ITERATIVE_SCHUR;
        options.preconditioner_type = ceres::SCHUR_JACOBI;
    #else
        //options.sparse_linear_algebra_library_type = ceres::SUITE_SPARSE;
        options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;
    #endif // _WIN32

    //If you are solving small to medium sized problems, consider setting Solver::Options::use_explicit_schur_complement to true, it can result in a substantial performance boost.
    options.use_explicit_schur_complement=true;
    options.max_num_iterations = 50;

    cout << "Ceres Solver getOptionsMedium()" << endl;
    cout << "Ceres preconditioner type: " << options.preconditioner_type << endl;
    cout << "Ceres linear algebra type: " << options.sparse_linear_algebra_library_type << endl;
    cout << "Ceres linear solver type: " << options.linear_solver_type << endl;

    return options;
}

void solve(ceres::Problem &problem, bool smallProblem=false){
    ceres::Solver::Summary summary;
    ceres::Solve(smallProblem ? getOptions() : getOptionsMedium(), &problem, &summary);
    if(!smallProblem) std::cout << "Final report:\n" << summary.FullReport();
}

void isoToAngleAxis(const Isometry3d& pose, double* cam){
//    Matrix<const double,3,3> rot(pose.linear());
//    cout<<"rotation : "<<pose.linear().data()<<endl;
//    auto begin = pose.linear().data();
    RotationMatrixToAngleAxis(ColumnMajorAdapter4x3(pose.linear().data()), cam);
    Vector3d t(pose.translation());
    cam[3]=t.x();
    cam[4]=t.y();
    cam[5]=t.z();
}

Isometry3d axisAngleToIso(const double* cam){
    Isometry3d poseFinal = Isometry3d::Identity();
    Matrix3d rot;
    ceres::AngleAxisToRotationMatrix(cam,rot.data());
    poseFinal.linear() = rot;
    poseFinal.translation() = Vector3d(cam[3],cam[4],cam[5]);
    return poseFinal;//.cast<float>();
}

Isometry3d eigenQuaternionToIso(const Eigen::Quaterniond& q, const Vector3d& t){
    Isometry3d poseFinal = Isometry3d::Identity();
    poseFinal.linear() = q.toRotationMatrix();
    poseFinal.translation() = t;
    return poseFinal;//.cast<float>();
}

Sophus::SE3d isoToSophus(const Isometry3d& pose){
    return Sophus::SE3d(pose);
}

Isometry3d sophusToIso(Sophus::SE3d soph){
    //    return Isometry3d(soph.matrix());
    Isometry3d poseFinal = Isometry3d::Identity();
    poseFinal.linear() = soph.rotationMatrix();
    poseFinal.translation() = soph.translation();
    return poseFinal;
}

Isometry3d pointToPoint_CeresAngleAxis(vector<Vector3d>&src,vector<Vector3d>&dst){

    double cam[6] = {0,0,0,0,0,0};

    ceres::Problem problem;

    for (int i = 0; i < src.size(); ++i) {
        // first viewpoint : dstcloud, fixed
        // second viewpoint: srcCloud, moves
        ceres::CostFunction* cost_function = ICPCostFunctions::PointToPointError_CeresAngleAxis::Create(dst[i],src[i]);
        problem.AddResidualBlock(cost_function, NULL, cam);
    }

    solve(problem);

    return axisAngleToIso(cam);
}

Isometry3d pointToPoint_EigenQuaternion(vector<Vector3d>&src,vector<Vector3d>&dst){
    Eigen::Quaterniond q = Eigen::Quaterniond::Identity();
    Eigen::Vector3d t(0,0,0);

    ceres::Problem problem;

    for (int i = 0; i < src.size(); ++i) {
        // first viewpoint : dstcloud, fixed
        // second viewpoint: srcCloud, moves
        ceres::CostFunction* cost_function = ICPCostFunctions::PointToPointError_EigenQuaternion::Create(dst[i],src[i]);
        problem.AddResidualBlock(cost_function, NULL, q.coeffs().data(), t.data());
    }

#ifdef useLocalParam
    ceres::LocalParameterization *quaternion_parameterization = new eigen_quaternion::EigenQuaternionParameterization;
    problem.SetParameterization(q.coeffs().data(),quaternion_parameterization);
#endif

    solve(problem);

    return eigenQuaternionToIso(q,t);
}

Isometry3d pointToPlane_CeresAngleAxis(vector<Vector3d> &src,vector<Vector3d> &dst,vector<Vector3d> &nor){

    double cam[6] = {0,0,0,0,0,0};

    ceres::Problem problem;

    for (int i = 0; i < src.size(); ++i) {
        // first viewpoint : dstcloud, fixed
        // second viewpoint: srcCloud, moves
        // nor is normal of dst
        ceres::CostFunction* cost_function = ICPCostFunctions::PointToPlaneError_CeresAngleAxis::Create(dst[i],src[i],nor[i]);
        problem.AddResidualBlock(cost_function, NULL, cam);
    }

    solve(problem);

    return axisAngleToIso(cam);
}

Isometry3d pointToPlane_EigenQuaternion(vector<Vector3d>&src,vector<Vector3d>&dst,vector<Vector3d> &nor){
    Eigen::Quaterniond q = Eigen::Quaterniond::Identity();
    Eigen::Vector3d t(0,0,0);

    ceres::Problem problem;

    for (int i = 0; i < src.size(); ++i) {
        // first viewpoint : dstcloud, fixed
        // second viewpoint: srcCloud, moves
        ceres::CostFunction* cost_function = ICPCostFunctions::PointToPlaneError_EigenQuaternion::Create(dst[i],src[i],nor[i]);
        problem.AddResidualBlock(cost_function, NULL, q.coeffs().data(), t.data());
    }

#ifdef useLocalParam
    ceres::LocalParameterization *quaternion_parameterization = new eigen_quaternion::EigenQuaternionParameterization;
    problem.SetParameterization(q.coeffs().data(),quaternion_parameterization);
#endif

    solve(problem);

    return eigenQuaternionToIso(q,t);
}

void ceresOptimizer(vector< std::shared_ptr<Frame> >& frames, bool pointToPlane, bool robust){

    ceres::Problem problem;

//    double* cameras = new double[frames.size()*7]; //4 quaternion, 3 translation

    vector<Eigen::Quaterniond> qs(frames.size());
    vector<Eigen::Vector3d> ts(frames.size());

    //extract initial camera poses

    for(int i=0; i<frames.size(); i++){
      Isometry3d originalPose = frames[i]->pose;
      Eigen::Quaterniond q;// = Eigen::Map<Eigen::Quaterniond>(cameras+i*7);
      Eigen::Vector3d t;// = Eigen::Map<Eigen::Vector3d>(cameras+i*7+4);

      q=Eigen::Quaterniond(originalPose.linear());
      t=Eigen::Vector3d(originalPose.translation());

      qs[i]=q;
      ts[i]=t;

      if (i==0){
          frames[i]->fixed=true;
      }
    }

    cout<<"ok ceres"<<endl;

//    Visualize::spin(1);

    //add edges
    for(int src_id=0; src_id<frames.size(); src_id++){

        Frame& srcCloud = *frames[src_id];
        if(srcCloud.fixed) continue;

        Eigen::Quaterniond& srcQ = qs[src_id];
        Eigen::Vector3d& srcT = ts[src_id];

        for (int j = 0; j < srcCloud.neighbours.size(); ++j) {

            OutgoingEdge& dstEdge = srcCloud.neighbours[j];
            Frame& dstCloud = *frames.at(dstEdge.neighbourIdx);

            int dst_id=dstEdge.neighbourIdx;

            Eigen::Quaterniond& dstQ = qs[dst_id];
            Eigen::Vector3d& dstT = ts[dst_id]; //dstCloud

            for(auto corr : dstEdge.correspondances){

                // first viewpoint : dstcloud, fixed
                // second viewpoint: srcCloud, moves

                ceres::CostFunction* cost_function;

                if(pointToPlane){
                    cost_function = ICPCostFunctions::PointToPlaneErrorGlobal::Create(dstCloud.pts[corr.second],srcCloud.pts[corr.first],dstCloud.nor[corr.second]);
                }else{
                    cost_function = ICPCostFunctions::PointToPointErrorGlobal::Create(dstCloud.pts[corr.second],srcCloud.pts[corr.first]);
                }

                ceres::LossFunction* loss = NULL;
                if(robust) loss = new ceres::SoftLOneLoss(dstEdge.weight);

//                cout<<dstEdge.weight<<endl;
//                problem.AddResidualBlock(cost_function,  srcQ.coeffs().data(),srcT.data(),dstQ.coeffs().data(), dstT.data());
                problem.AddResidualBlock(cost_function, loss, srcQ.coeffs().data(),srcT.data(),dstQ.coeffs().data(), dstT.data());

//                problem.AddResidualBlock(cost_function, NULL, &cameras[src_id*7],&cameras[src_id*7+4],&cameras[dst_id*7],&cameras[dst_id*7+4]);

            }
        }
    }

#ifdef useLocalParam
    eigen_quaternion::EigenQuaternionParameterization *quaternion_parameterization = new eigen_quaternion::EigenQuaternionParameterization;
#endif

    for (int i = 0; i < frames.size(); ++i) {
        #ifdef useLocalParam
//        problem.SetParameterization(&cameras[i*7],quaternion_parameterization);
        problem.SetParameterization(qs[i].coeffs().data(),quaternion_parameterization);
        #endif

        if(frames[i]->fixed){
            std::cout<<i<<" fixed"<<endl;
//            problem.SetParameterBlockConstant(&cameras[i*7]);
//            problem.SetParameterBlockConstant(&cameras[i*7+4]);
            problem.SetParameterBlockConstant(qs[i].coeffs().data());
            problem.SetParameterBlockConstant(ts[i].data());
        }
    }

    solve(problem);

    //update camera poses
    for (int i = 0; i < frames.size(); ++i) {
//        poseFinal.linear() = Eigen::Map<Eigen::Quaterniond>(cameras+i*7).toRotationMatrix();
//        poseFinal.translation() = Eigen::Map<Eigen::Vector3d>(cameras+i*7+4);
        frames[i]->pose=eigenQuaternionToIso(qs[i],ts[i]);
    }
}

void ceresOptimizer_ceresAngleAxis(vector< std::shared_ptr<Frame> >& frames, bool pointToPlane,bool robust){

    ceres::Problem problem;

    double* cameras = new double[frames.size()*6];

    //extract initial camera poses
    for(int i=0; i<frames.size(); i++){
      isoToAngleAxis(frames[i]->pose,&cameras[i*6]);

//      cout<<"pose "<<i<<endl;
//      cout<<frames[i]->pose.matrix()<<endl;
//      cout<<"pose test "<<i<<endl;
//      cout<<axisAngleToIso(&cameras[i*6]).matrix()<<endl;

//      Visualize::spin();

      if (i==0){
          frames[i]->fixed=true;
      }
    }

    //add edges
    for(int src_id=0; src_id<frames.size(); src_id++){

        Frame& srcCloud = *frames[src_id];
        if(srcCloud.fixed) continue;

        for (int j = 0; j < srcCloud.neighbours.size(); ++j) {

            OutgoingEdge& dstEdge = srcCloud.neighbours[j];
            Frame& dstCloud = *frames.at(dstEdge.neighbourIdx);

            int dst_id=dstEdge.neighbourIdx;

            for(auto corr : dstEdge.correspondances){

                // first viewpoint : dstcloud, fixed
                // second viewpoint: srcCloud, moves

                ceres::CostFunction* cost_function;

                if(pointToPlane){
                    cost_function = ICPCostFunctions::PointToPlaneErrorGlobal_CeresAngleAxis::Create(dstCloud.pts[corr.second],srcCloud.pts[corr.first],dstCloud.nor[corr.second]);
                }else{
                    cost_function = ICPCostFunctions::PointToPointErrorGlobal_CeresAngleAxis::Create(dstCloud.pts[corr.second],srcCloud.pts[corr.first]);
                }

                ceres::LossFunction* loss = NULL;
                if(robust) loss = new ceres::SoftLOneLoss(dstEdge.weight);

                problem.AddResidualBlock(cost_function, loss, &cameras[src_id*6],&cameras[dst_id*6]);

            }
        }
    }

    for (int i = 0; i < frames.size(); ++i) {
        if(frames[i]->fixed){
            std::cout<<i<<" fixed"<<endl;
            problem.SetParameterBlockConstant(&cameras[i*6]);
        }
    }

    solve(problem);

    //update camera poses
    for (int i = 0; i < frames.size(); ++i) {
        frames[i]->pose=axisAngleToIso(&cameras[i*6]);
    }
}

void ceresOptimizer_sophusSE3(vector< std::shared_ptr<Frame> >& frames, bool pointToPlane,bool robust,bool automaticDiff){

    ceres::Problem problem;

    std::vector<Sophus::SE3d> cameras;

    //extract initial camera poses
    for(int i=0; i<frames.size(); i++){
      Sophus::SE3d soph = isoToSophus(frames[i]->pose);

//      cout<<"pose "<<i<<endl;
//      cout<<frames[i]->pose.matrix()<<endl;
//      cout<<"pose test "<<i<<endl;
//      cout<<sophusToIso(soph).matrix()<<endl;

//      Visualize::spin();
      cameras.push_back(soph);

      if (i==0){
          frames[i]->fixed=true;
      }
    }

    //add edges
    for(int src_id=0; src_id<frames.size(); src_id++){

        Frame& srcCloud = *frames[src_id];
        if(srcCloud.fixed) continue;

        for (int j = 0; j < srcCloud.neighbours.size(); ++j) {

            OutgoingEdge& dstEdge = srcCloud.neighbours[j];
            Frame& dstCloud = *frames.at(dstEdge.neighbourIdx);

            int dst_id=dstEdge.neighbourIdx;

            for(auto corr : dstEdge.correspondances){

                // first viewpoint : dstcloud, fixed
                // second viewpoint: srcCloud, moves

                ceres::CostFunction* cost_function;

                if(pointToPlane){
                    cost_function = ICPCostFunctions::PointToPlaneErrorGlobal_SophusSE3::Create(dstCloud.pts[corr.second],srcCloud.pts[corr.first],dstCloud.nor[corr.second]);
                }else{
                    cost_function = ICPCostFunctions::PointToPointErrorGlobal_SophusSE3::Create(dstCloud.pts[corr.second],srcCloud.pts[corr.first]);
                }

                ceres::LossFunction* loss = NULL;
                if(robust) loss = new ceres::SoftLOneLoss(dstEdge.weight);

                problem.AddResidualBlock(cost_function, loss, cameras[src_id].data(),cameras[dst_id].data());

            }
        }
    }
#ifdef useLocalParam
    ceres::LocalParameterization* param = sophus_se3::getParameterization(automaticDiff);
#endif
    for (int i = 0; i < frames.size(); ++i) {
        #ifdef useLocalParam
        problem.SetParameterization(cameras[i].data(),param);
        #endif
        if(frames[i]->fixed){
            std::cout<<i<<" fixed"<<endl;
            problem.SetParameterBlockConstant(cameras[i].data());
        }
    }

    solve(problem);

    //update camera poses
    for (int i = 0; i < frames.size(); ++i) {
        frames[i]->pose=sophusToIso(cameras[i]);
    }
}

//Isometry3d pointToPoint_SophusSE3(vector<Vector3d> &src,vector<Vector3d> &dst){
//    Sophus::SE3d soph = isoToSophus(Isometry3d::Identity());
//    Sophus::SE3d soph2 = isoToSophus(Isometry3d::Identity());

//    ceres::Problem problem;

//    for (int i = 0; i < src.size(); ++i) {
//        // first viewpoint : dstcloud, fixed
//        // second viewpoint: srcCloud, moves
//        ceres::CostFunction* cost_function = ICPCostFunctions::PointToPointErrorGlobal_SophusSE3::Create(dst[i],src[i]);
//        problem.AddResidualBlock(cost_function, NULL,soph.data(),soph2.data());
//    }

//    ceres::LocalParameterization* param = sophus_se3::getParameterization(false);

//    problem.SetParameterization(soph.data(),param);
//    problem.SetParameterization(soph2.data(),param);
//    problem.SetParameterBlockConstant(soph2.data());

//    solve(problem);

//    return sophusToIso(soph);
//}

//Isometry3d pointToPlane_SophusSE3(vector<Vector3d> &src,vector<Vector3d> &dst,vector<Vector3d> &nor){
//    Sophus::SE3d soph = isoToSophus(Isometry3d::Identity());
//    Sophus::SE3d soph2 = isoToSophus(Isometry3d::Identity());

//    ceres::Problem problem;

//    for (int i = 0; i < src.size(); ++i) {
//        // first viewpoint : dstcloud, fixed
//        // second viewpoint: srcCloud, moves
//        ceres::CostFunction* cost_function = ICPCostFunctions::PointToPlaneErrorGlobal_SophusSE3::Create(dst[i],src[i],nor[i]);
//        problem.AddResidualBlock(cost_function, NULL,soph.data(),soph2.data());
//    }

//    ceres::LocalParameterization* param = sophus_se3::getParameterization(false);

//    problem.SetParameterization(soph.data(),param);
//    problem.SetParameterization(soph2.data(),param);
//    problem.SetParameterBlockConstant(soph2.data());

//    solve(problem);

//    return sophusToIso(soph);
//}

Isometry3d pointToPoint_SophusSE3(vector<Vector3d> &src,vector<Vector3d> &dst,bool autodiff){
    Sophus::SE3d soph = isoToSophus(Isometry3d::Identity());

    ceres::Problem problem;

    for (int i = 0; i < src.size(); ++i) {
        // first viewpoint : dstcloud, fixed
        // second viewpoint: srcCloud, moves
        ceres::CostFunction* cost_function = ICPCostFunctions::PointToPointError_SophusSE3::Create(dst[i],src[i]);
        problem.AddResidualBlock(cost_function, NULL,soph.data());
    }
#ifdef useLocalParam
    ceres::LocalParameterization* param = sophus_se3::getParameterization(autodiff);
    problem.SetParameterization(soph.data(),param);
#endif

    solve(problem);

    return sophusToIso(soph);
}

Isometry3d pointToPlane_SophusSE3(vector<Vector3d> &src,vector<Vector3d> &dst,vector<Vector3d> &nor, bool autodiff){
    Sophus::SE3d soph = isoToSophus(Isometry3d::Identity());

    ceres::Problem problem;

    for (int i = 0; i < src.size(); ++i) {
        // first viewpoint : dstcloud, fixed
        // second viewpoint: srcCloud, moves
        ceres::CostFunction* cost_function = ICPCostFunctions::PointToPlaneError_SophusSE3::Create(dst[i],src[i],nor[i]);
        problem.AddResidualBlock(cost_function, NULL,soph.data());
    }
#ifdef useLocalParam
    ceres::LocalParameterization* param = sophus_se3::getParameterization(autodiff);
    problem.SetParameterization(soph.data(),param);
#endif

    solve(problem);

    return sophusToIso(soph);
}

} //end namespace


/*
 * Copyright 2018 James Dyer
 *
 * 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.
 *
 * This is being developed for the TANGO Project: http://tango-project.eu
 */

#include <sstream>
#include <NTL/RR.h>
#include <json/json.h>
#include "HE1NEncrypter.h"
#include "Random.h"

/**
 * Default destructor
 */
HE1NEncrypter::~HE1NEncrypter()
{
	delete kappamod;
}

HE1NEncrypter::HE1NEncrypter(int lambda, int eta, int nu){
	kappa = NTL::RandomPrime_ZZ(nu, 20);
	generateParameters(lambda,eta);
	NTL::ZZ_p::init(modulus);
	NTL::ZZ_p tmp = to_ZZ_p(kappa);
	kappamod = new NTL::ZZ_p(tmp);
}

HE1NEncrypter::HE1NEncrypter(int n, int d, int rho, int rhoprime)
{
	NTL::RR two(2);
	NTL::RR exp1(rho*d);
	NTL::RR nplusone(n+1);
	NTL::RR exp2(d);
	NTL::RR exp3(rho);
	NTL::RR kloBound = pow(two,exp1)*pow(nplusone,exp2);
	NTL::ZZ kappaLowerBound;
	conv(kappaLowerBound,kloBound);
	long nu = NumBits(kappaLowerBound)+1;
	kappa = NTL::RandomPrime_ZZ(nu, 20);
	NTL::RR kappareal = to_RR(kappa);
	NTL::RR ploBound = pow(pow(two,exp3)+pow(kappareal,two),exp2)*pow(nplusone,exp2);
	NTL::ZZ pLowerBound;
	conv(pLowerBound,ploBound);
	long lambda = NumBits(pLowerBound) + 1;
	long eta = ((lambda * lambda / rhoprime) - lambda);
	generateParameters(lambda,eta);
	NTL::ZZ_p::init(modulus);
	NTL::ZZ_p tmp = to_ZZ_p(kappa);
	kappamod = new NTL::ZZ_p(tmp);
};

HE1NEncrypter::HE1NEncrypter(std::string& secrets, std::string& parameters){
	Json::Value rootP;
	Json::Reader reader;
	bool parsingSuccessful = reader.parse(parameters,rootP);
	if (parsingSuccessful){
		modulus = NTL::conv<NTL::ZZ>(rootP["modulus"].asCString());
		NTL::ZZ_p::init(modulus);
		Json::Value rootS;
		bool parsingSuccessful = reader.parse(secrets,rootS);
		if (parsingSuccessful){
			p = NTL::conv<NTL::ZZ>(rootS["p"].asCString());
			NTL::ZZ_p ptmp = to_ZZ_p(p);
			pmod = new NTL::ZZ_p(ptmp);
			q = modulus/p;
			kappa = NTL::conv<NTL::ZZ>(rootS["kappa"].asCString());
			NTL::ZZ_p ktmp = to_ZZ_p(kappa);
			kappamod = new NTL::ZZ_p(ktmp);
		}
	}
}

NTL::vec_ZZ HE1NEncrypter::getKey()
{
	NTL::vec_ZZ ret;
	ret.append(p);
	ret.append(kappa);
    return ret;
};

NTL::ZZ_p HE1NEncrypter::encrypt(NTL::ZZ& plaintext)
{
    NTL::ZZ_p r = to_ZZ_p(rng->nextBigInteger(ONE, q));
    NTL::ZZ_p s = to_ZZ_p(rng->nextBigInteger(kappa));
    NTL::ZZ_p ptext = to_ZZ_p(plaintext);
    return ptext+r*(*pmod)+s*(*kappamod);
};

std::string HE1NEncrypter::writeSecretsToJSON()
{
	Json::Value root;
	std::stringstream pStr, kappaStr;
	pStr << p;
	kappaStr << kappa;
	root["p"] = pStr.str();
	root["kappa"] = kappaStr.str();
	Json::FastWriter writer;
	return writer.write(root);
};


//
// Created by brad on 1/4/16.
//

#include "MFCCs.h"
#include "armadillo"
#include <boost/property_tree/xml_parser.hpp>
#include <boost/property_tree/ptree.hpp>
//#include <log4cxx/logger.h>
#include <iostream>
//#include <boost/foreach.hpp>

using namespace std;

//initialize everything to default params
MFCCs::MFCCs(){// :
//  logger(log4cxx::Logger::getLogger("ade.asr.MFCCs")) {
  normVolume = false;
  numDiffs = 2;
  smoothFactor = .25;
}

MFCCs::~MFCCs() {

}

vector<vector<double>> MFCCs::getFeatures(std::vector<double> PCMData) {
  vector<vector<double>> out;

  //TODO: is it worth getting rid of this redundancy?
  //calculating minimum possible size for PCMData
  int fs = 16000;
  int Tw = 25;
  int Ts = 10;

  int numWindows = (int) round((1E-3) * Tw * fs);  //frame duration in samples
  int numShift = (int) round((1E-3) * Ts * fs);

  int minSegmentSize = -numShift + numWindows;

  if ((((int) PCMData.size() - numWindows) / numShift + 1) > 0) {

    //normalize volume
    if (normVolume) PCMData = normalizeVolume(PCMData);

    //caluclate MFCCS
    //will always smooth, can set smooth PCT to 0 for no smoothing
    vector<vector<double>> mfccs = smooth(calculateMFCCs(PCMData));

    if (numDiffs > 0) {

      vector<vector<double>> firstDiffs;
      vector<vector<double>> secondDiffs;

      if (numDiffs >= 1) {
        firstDiffs = smooth(getDiffs(mfccs));
      }
      if (numDiffs >= 2) {
        secondDiffs = smooth(getDiffs(firstDiffs));
      }
      if (numDiffs >= 3) {
        //LOG4CXX_WARN(logger, "we only support 2 diffs, and that's probably all you want");
        cout<< "we only support 2 diffs, and that's probably all you want";
      }

      //combine into single feature vector
      for (int i = 0; i < mfccs.size(); i++) {
        vector<double> combinedFeature;
        combinedFeature.insert(combinedFeature.end(), mfccs[i].begin(), mfccs[i].end());
        if (numDiffs >= 1) {
          combinedFeature.insert(combinedFeature.end(), firstDiffs[i].begin(),
                                 firstDiffs[i].end());
        }
        if (numDiffs >= 2) {
          combinedFeature.insert(combinedFeature.end(), secondDiffs[i].begin(),
                                 secondDiffs[i].end());
        }
        out.push_back(combinedFeature);
      }
    }
    else {
      out = mfccs;
    }
  }
  else {
    //LOG4CXX_WARN(logger,
    //             "No MFCCs extracted because input segment is too small.\n segment size:"<<
    //             PCMData.size()<<" min segment size: "<<minSegmentSize<<
    //             ". Try changing filter params");
  }
  //LOG4CXX_DEBUG(logger, "MFCCs generated: "<<out.size());
 // cout<<"MFCCs generated: "<<out.size();
  return out;
}

vector<double> MFCCs::normalizeVolume(vector<double> raw) {
  vector<double> norm(raw.size(), 0);
  double maxPCM = 0;
  for (double r: raw) {
    if (r > maxPCM)maxPCM = r;
  }
  for (int i = 0; i < raw.size(); i++) {
    norm[i] = raw[i] / maxPCM;
  }
  return norm;
}

vector<vector<double>> MFCCs::smooth(vector<vector<double>> rough) {
  vector<double> padding(rough[0].size(), 0.0);
  vector<vector<double>> smoothed;
  smoothed.push_back(padding);

  //double smoothFactor = .1;

  for (int i = 0; i < rough.size(); i++) {
    vector<double> smoothedPoint(rough[i].size(), 0);
    for (int k = 0; k < smoothedPoint.size(); k++) {
      smoothedPoint[k] +=
        smoothed[i][k] * (1.0 - smoothFactor) + rough[i][k] * smoothFactor;
    }
    smoothed.push_back(smoothedPoint);
  }
  smoothed.erase(smoothed.begin());
  return smoothed;
}

vector<vector<double>> MFCCs::getDiffs(vector<vector<double>> original) {
  vector<vector<double>> diffs(original.size(),
                               vector<double>(original[0].size(), 0));
  //insert copies of first and last time step
  vector<double> first = original[0];
  original.insert(original.begin(), first);
  vector<double> last = original[original.size() - 1];
  original.insert(original.end(), last);

  for (int i = 1; i < original.size() - 1; i++) {
    for (int j = 0; j < original[i].size(); j++) {
      diffs[i - 1][j] = (original[i + 1][j] - original[i - 1][j]) / 2.0;
    }
  }
  return diffs;
}

//cepstral lifter routine
inline vector<double> MFCCs::ceplifter(double n, double l) {
  vector<double> lift;
  for (int i = 0; i < n; i++) {
    lift.push_back(1 + 0.5 * l * sin(M_PI * i / l));
  }
  return lift;
}

inline int MFCCs::NextPowerOf2(int val) {
  val--;
  val = (val>>1) | val;
  val = (val>>2) | val;
  val = (val>>4) | val;
  val = (val>>8) | val;
  val = (val>>16) | val;
  return ++val;
}

//htk pre-emphasis function
vector<double> MFCCs::preEmphasis(double coeff, vector<double> x) {
  vector<double> y(x.size(), 0);
  y[0] = x[0];
  for (int n = 1; n < x.size(); n++) {
    y[n] = x[n] - coeff * x[n - 1];
  }
  return y;
}

//hamming window a la matlab
vector<double> MFCCs::window(int windowSize) {
  vector<double> windowed(windowSize, 0);
  for (int n = 0; n < windowSize; n++) {
    windowed[n] = 0.54 - .46 * cos((2 * M_PI) * ((double) n / windowSize - 1));
  }
  return windowed;
}

//based on mike's adaptation of  Kamil Wojcicki's "HTK MFCC MATLAB CODE"
//http://www.mathworks.com/matlabcentral/fileexchange/32849-htk-mfcc-matlab/content/mfcc/mfcc.m
vector<vector<double>> MFCCs::calculateMFCCs(vector<double> signal) {
  vector<vector<double>> out;
  //ofstream fout;
  //fout.open("mfccstuff.txt");

  //params hard coded for now
  int fs = 16000;
  int Tw = 25;
  int Ts = 10;
  double alpha = .97;
  int numFilterBankChannels = 26;
  int numCC = 13;//12 + 1;// but don't use the first
  int cepstralSineLifter = 22;

  double fMin = 0; // filter coefficients start at this frequency (Hz)
  double fLow = 20; //minFreq;//       % 20 lower cutoff frequency (Hz) for the filterbank
  double fHigh = 4400; //maxFreq;//      % 4400 upper cutoff frequency (Hz) for the filterbank
  double fMax =
    0.5 * fs; //     % filter coefficients end at this frequency (Hz)

  int numWindows = (int) round(
    (1E-3) * Tw * fs);  // % frame duration (in samples)
  int numShift = (int) round((1E-3) * Ts * fs);   //% frame shift (in samples)

  double power = log2(NextPowerOf2(numWindows));
  int nfft = pow(2, power);     // %length of FFT analysis
  int uniqueFFT = (nfft / 2) + 1;    //% length of the unique part of the FFT

  //preemphasize signal (high pass filter)

  //setting up frame buckets for fft
  //int numFrames = int((filteredSpeech.size() - numWindows) / (numShift) + 1);
  //vector<vector<double>> frames;
  //for (int i = 0; i < numFrames; i++) {
  //  vector<double> nextFrame(filteredSpeech.begin() + (i * numShift),
  //                           filteredSpeech.begin() + (i * numShift) + numWindows);
  //  frames.push_back(nextFrame);
  //}
  int numFrames = ((int) (signal.size()) - numWindows) / (numShift) + 1;
  vector<vector<double>> frames;
  for (int i = 0; i < numFrames; i++) {
    vector<double> nextFrame(signal.begin() + (i * numShift),
                             signal.begin() + (i * numShift) + numWindows);
    //frames.push_back(filter(filterVec,1,nextFrame));
    frames.push_back(preEmphasis(alpha, nextFrame));
  }

  //create matrix for application of Hamming Window
  vector<double> windowed = window(numWindows);
  arma::vec aWindowed(windowed);
  arma::mat windowMat = arma::diagmat(aWindowed);

  //convert to armadillo for fft and matrix multiplication
  arma::mat mFrames(frames[0].size(), frames.size());
  for (int i = 0; i < frames[0].size(); i++) {
    for (int j = 0; j < frames.size(); j++) {
      mFrames(i, j) = frames[j][i];
    }
  }

  //Apply hamming window to frame buckets
  arma::mat wFrames = windowMat * mFrames;

  //get energy
  double silenceFloor = 50; //50db defaults from HTK
  double energyScale = .1; //default vlaue from HTK
  vector<double> energy(wFrames.n_cols, 0);

  for (int i = 0; i < wFrames.n_cols; i++) {
    double e = 0;
    for (int j = 0; j < wFrames.n_rows; j++) {
      e += pow(wFrames(j, i), 2);
    }
    energy[i] = log(e);
  }

  //normalize energy
  double maxEnergy = 0;
  for (double e: energy) {
    if (e > maxEnergy)maxEnergy = e;
  }
  double minEnergy = maxEnergy - (silenceFloor * log(10.0)) / 10.0;
  for (int i = 0; i < energy.size(); i++) {
    if (energy[i] < minEnergy) energy[i] = minEnergy;
    energy[i] = 1.0 - (maxEnergy - energy[i]) * energyScale;
  }

  // Magnitude spectrum computation
  // could implement HTK's c++ fft here to get results closer to theirs, not sure if its worth it
  arma::mat MAG = arma::abs(arma::fft(wFrames, nfft));

  //calculate frequency range in Hz, based on sized of unique part fo fft
  vector<double> fRange(uniqueFFT, 0);
  for (int i = 0; i < uniqueFFT; i++) {
    fRange[i] = i * (fMax / (uniqueFFT - 1));
  }

  double hz2melf_low = 2595 * log10(1 + fLow / 700);
  double hz2melf_high = 2595 * log10(1 + fHigh / 700);

  vector<double> melChannels;
  for (int i = 0; i <= numFilterBankChannels + 1; i++) {
    double temp = hz2melf_low + i * ((hz2melf_high - hz2melf_low) /
                                     (numFilterBankChannels + 1));
    double x = 700 * (pow(10, (temp / 2595)) - 1);
    melChannels.push_back(x);
  }

  // create filter bank matrix to put magnitude spectrum into Mel Frequency Bands
  // super contrived, based on doing it in MATLAB but probably doesn't need to be a matrix;
  arma::mat H = arma::zeros(numFilterBankChannels, uniqueFFT);
  for (int m = 0; m < numFilterBankChannels; m++) {
    for (int k = 0; k < uniqueFFT; k++) {
      if (fRange[k] >= melChannels[m] && fRange[k] < melChannels[m + 1]) {
        H(m, k) =
          (fRange[k] - melChannels[m]) / (melChannels[m + 1] - melChannels[m]);
      }
      if (fRange[k] >= melChannels[m + 1] && fRange[k] <= melChannels[m + 2]) {
        H(m, k) = (melChannels[m + 2] - fRange[k]) /
                  (melChannels[m + 2] - melChannels[m + 1]);
      }
    }
  }

  // GET FREQUENCY BAND ENERGY
  // Filterbank application to unique part of the magnitude spectrum
  arma::mat FBE =
    H * MAG(arma::span(0, uniqueFFT - 1), arma::span(0, frames.size() - 1));

  // DCT Matrix Computation
  arma::mat DCT = arma::zeros(numCC, numFilterBankChannels);
  for (int i = 0; i < numCC; i++) {
    for (int j = 0; j < numFilterBankChannels; j++) {
      DCT(i, j) = sqrt(2.0 / numFilterBankChannels) *
                  cos((i) * (M_PI * (j + 1 - .5) / numFilterBankChannels));
    }
  }

  // Conversion of logFBEs to cepstral coefficients through DCT
  arma::mat CC = DCT * arma::log(FBE);

  // Cepstral lifter computation
  arma::vec lifter(ceplifter(numCC, cepstralSineLifter));
  arma::mat lifterMat = arma::diagmat(lifter);

  // Cepstral liftering gives liftered cepstral coefficients
  arma::mat MFCCs = lifterMat * CC;

  //convert back to std::vector
  typedef vector<double> stdvec;
  for (int i = 0; i < MFCCs.n_cols; i++) {
    //vector<double> features(MFCCs.begin_col(i) + 1, MFCCs.end_col(i));
    //features.push_back(energy[i]);

    //this order for compatability with Mike's code
    vector<double> features;
    features.push_back(energy[i]);
    features.insert(features.end(), MFCCs.begin_col(i) + 1, MFCCs.end_col(i));

    out.push_back(features);
  }

  //for (int i = 0; i < out.size(); i++) {
  //  for (double j: out[i]) {
  //    fout << setw(10) << j << " ";
  //  }
  //  fout << endl;
  //}
  //
  //fout.close();
  return out;
}


// 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
//
// 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_TO_NLP_HPP_
#define SMOOTH__FEEDBACK__OCP_TO_NLP_HPP_

/**
 * @file
 * @brief Formulate optimal control problem as a nonlinear program
 */

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

#include "collocation/mesh.hpp"
#include "collocation/mesh_function.hpp"
#include "nlp.hpp"
#include "ocp.hpp"

namespace smooth::feedback {

// \cond
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);
}

/**
 * @brief NLP representing an OCP.
 *
 * @note This class allocates matrices internally and returns references to those.
 */
template<FlatOCPType Ocp, MeshType Mesh, diff::Type DT = diff::Type::Default>
class OCPNLP
{
private:
  static constexpr auto Nx = std::decay_t<Ocp>::Nx;
  static constexpr auto Nu = std::decay_t<Ocp>::Nu;
  static constexpr auto Nq = std::decay_t<Ocp>::Nq;

  Ocp ocp_;        // optimal control problem
  Mesh mesh_;      // discretization mesh
  std::size_t N_;  // number of collocation points in mesh

  std::size_t tfvar_B, qvar_B, xvar_B, uvar_B, n_;  // variable start indices
  std::size_t tfvar_L, qvar_L, xvar_L, uvar_L;      // variable lengths

  std::size_t dcon_B, qcon_B, crcon_B, cecon_B, m_;  // constraint start indices
  std::size_t dcon_L, qcon_L, crcon_L, cecon_L;      // constraint lengths

  // scaling
  double w_scaling_{1};

  // variable and constraint bounds
  Eigen::VectorXd xl_, xu_, gl_, gu_;

  // allocated return arguments
  Eigen::VectorXd g_;
  Eigen::SparseMatrix<double> df_dx_, dg_dx_, d2f_dx2_, d2g_dx2_;

  // allocated computation
  MeshValue<0> dyn_out0_, int_out0_, cr_out0_;
  MeshValue<1> dyn_out1_, int_out1_, cr_out1_;
  MeshValue<2> dyn_out2_, int_out2_, cr_out2_;

public:
  /// @brief Constructor
  template<typename OcpArg, typename MeshArg>
  OCPNLP(OcpArg && ocp, MeshArg && mesh)
      : ocp_(std::forward<OcpArg>(ocp)), mesh_(std::forward<MeshArg>(mesh)), N_(mesh_.N_colloc())
  {
    const auto [var_beg, var_len, con_beg, con_len] = detail::ocp_nlp_structure(ocp_, mesh_);

    tfvar_B = var_beg[0];
    qvar_B  = var_beg[1];
    xvar_B  = var_beg[2];
    uvar_B  = var_beg[3];
    n_      = var_beg[4];

    tfvar_L = var_len[0];
    qvar_L  = var_len[1];
    xvar_L  = var_len[2];
    uvar_L  = var_len[3];

    dcon_B  = con_beg[0];
    qcon_B  = con_beg[1];
    crcon_B = con_beg[2];
    cecon_B = con_beg[3];
    m_      = con_beg[4];

    dcon_L  = con_len[0];
    qcon_L  = con_len[1];
    crcon_L = con_len[2];
    cecon_L = con_len[3];

    // Mesh weight
    double max_weight = 1e-6;
    for (auto w : mesh_.all_weights()) { max_weight = std::max(max_weight, w); }
    w_scaling_ = 1. / max_weight;

    // VARIABLE BOUNDS

    xl_.setConstant(n_, -std::numeric_limits<double>::infinity());
    xl_.segment(tfvar_B, tfvar_L).setZero();  // tf lower bounded by zero
    xu_.setConstant(n_, std::numeric_limits<double>::infinity());

    // CONSTRAINT BOUNDS

    gl_.resize(m_);
    gu_.resize(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 (scaled by quadrature weights)
    gl_.segment(crcon_B, crcon_L) = ocp.crl.replicate(N_, 1);
    gu_.segment(crcon_B, crcon_L) = ocp.cru.replicate(N_, 1);
    for (const auto & [i, w] : utils::zip(std::views::iota(0u, N_), mesh_.all_weights())) {
      gl_.segment(crcon_B + i * ocp_.Ncr, ocp_.Ncr) *= w_scaling_ * w;
      gu_.segment(crcon_B + i * ocp_.Ncr, ocp_.Ncr) *= w_scaling_ * w;
    }

    // end constraints
    gl_.segment(cecon_B, cecon_L) = ocp.cel;
    gu_.segment(cecon_B, cecon_L) = ocp.ceu;

    // allocate output args
    g_.setZero(m_);

    df_dx_.resize(1, n_);
    df_dx_.reserve(Eigen::VectorXi::Constant(n_, 1));

    d2f_dx2_.resize(n_, n_);
    dg_dx_.resize(m_, n_);
    d2g_dx2_.resize(n_, n_);

    /// @todo Allocate nnz's to speed up first call?
  }

  std::size_t n() const { return n_; }
  std::size_t m() const { return m_; }
  const Eigen::VectorXd & xl() const { return xl_; }
  const Eigen::VectorXd & xu() const { return xu_; }

  double f(const Eigen::Ref<const Eigen::VectorXd> x) const
  {
    assert(static_cast<std::size_t>(x.size()) == n_);

    const auto x0var_B = xvar_B;
    const auto xfvar_B = xvar_B + xvar_L - Nx;

    const double tf                    = x(tfvar_B);
    const Eigen::Vector<double, Nx> x0 = x.segment(x0var_B, Nx);
    const Eigen::Vector<double, Nx> xf = x.segment(xfvar_B, Nx);
    const Eigen::Vector<double, Nq> q  = x.segment(qvar_B, qvar_L);

    return ocp_.theta(tf, x0, xf, q);
  }

  const Eigen::SparseMatrix<double> & df_dx(const Eigen::Ref<const Eigen::VectorXd> x)
  {
    assert(static_cast<std::size_t>(x.size()) == n_);

    const auto x0var_B = xvar_B;
    const auto xfvar_B = xvar_B + xvar_L - Nx;

    const double tf                    = x(tfvar_B);
    const Eigen::Vector<double, Nx> x0 = x.segment(x0var_B, Nx);
    const Eigen::Vector<double, Nx> xf = x.segment(xfvar_B, Nx);
    const Eigen::Vector<double, Nq> q  = x.segment(qvar_B, qvar_L);

    const auto & [fval, dfval] = diff::dr<1, DT>(ocp_.theta, wrt(tf, x0, xf, q));

    set_zero(df_dx_);

    block_add(df_dx_, 0, tfvar_B, dfval.middleCols(0, 1));           // df / dtf
    block_add(df_dx_, 0, x0var_B, dfval.middleCols(1, Nx));          // df / dx0
    block_add(df_dx_, 0, xfvar_B, dfval.middleCols(1 + Nx, Nx));     // df / dxf
    block_add(df_dx_, 0, qvar_B, dfval.middleCols(1 + 2 * Nx, Nq));  // df / dq

    df_dx_.makeCompressed();
    return df_dx_;
  }

  const Eigen::SparseMatrix<double> & d2f_dx2(Eigen::Ref<const Eigen::VectorXd> x)
  {
    assert(static_cast<std::size_t>(x.size()) == n_);

    const auto x0var_B = xvar_B;
    const auto xfvar_B = xvar_B + xvar_L - Nx;

    const double tf                    = x(tfvar_B);
    const Eigen::Vector<double, Nx> x0 = x.segment(x0var_B, Nx);
    const Eigen::Vector<double, Nx> xf = x.segment(xfvar_B, Nx);
    const Eigen::Vector<double, Nq> q  = x.segment(qvar_B, qvar_L);

    const auto & [fval, dfval, d2fval] = diff::dr<2, DT>(ocp_.theta, wrt(tf, x0, xf, q));

    set_zero(d2f_dx2_);

    // clang-format off
    block_add(d2f_dx2_, tfvar_B, tfvar_B, d2fval.block(         0,          0,  1,  1), 1, true);  // tftf
    block_add(d2f_dx2_, tfvar_B, x0var_B, d2fval.block(         0,          1,  1, Nx), 1, true);  // tfx0
    block_add(d2f_dx2_, tfvar_B, xfvar_B, d2fval.block(         0,     1 + Nx,  1, Nx), 1, true);  // tfxf
    block_add(d2f_dx2_, tfvar_B,  qvar_B, d2fval.block(         0, 1 + 2 * Nx,  1, Nq), 1, true);  // tfq

    block_add(d2f_dx2_, x0var_B, x0var_B, d2fval.block(         1,          1, Nx, Nx), 1, true);  // x0x0
    block_add(d2f_dx2_, x0var_B, xfvar_B, d2fval.block(         1,     1 + Nx, Nx, Nx), 1, true);  // x0xf
    block_add(d2f_dx2_, x0var_B,  qvar_B, d2fval.block(         1, 1 + 2 * Nx, Nx, Nq), 1, true);  // x0q

    block_add(d2f_dx2_, xfvar_B, xfvar_B, d2fval.block(    1 + Nx,     1 + Nx, Nx, Nx), 1, true);  // xfxf
    block_add(d2f_dx2_, xfvar_B,  qvar_B, d2fval.block(    1 + Nx, 1 + 2 * Nx, Nx, Nq), 1, true);  // xfq

    block_add(d2f_dx2_,  qvar_B,  qvar_B, d2fval.block(1 + 2 * Nx, 1 + 2 * Nx, Nq, Nq), 1, true);  // qq
    // clang-format on

    d2f_dx2_.makeCompressed();
    return d2f_dx2_;
  }

  const Eigen::VectorXd & g(const Eigen::Ref<const Eigen::VectorXd> x)
  {
    assert(static_cast<std::size_t>(x.size()) == n_);

    const auto x0var_B = xvar_B;
    const auto xfvar_B = xvar_B + xvar_L - Nx;

    const double t0                    = 0;
    const double tf                    = x(tfvar_B);
    const Eigen::Vector<double, Nx> x0 = x.segment(x0var_B, Nx);
    const Eigen::Vector<double, Nx> xf = x.segment(xfvar_B, Nx);
    const Eigen::Vector<double, Nq> q  = x.segment(qvar_B, qvar_L);

    const Eigen::Map<const Eigen::Matrix<double, Nx, -1>> X(x.data() + xvar_B, Nx, N_ + 1);
    const Eigen::Map<const Eigen::Matrix<double, Nu, -1>> U(x.data() + uvar_B, Nu, N_);

    mesh_dyn<0>(dyn_out0_, mesh_, ocp_.f, t0, tf, X.colwise(), U.colwise());
    mesh_integrate<0>(int_out0_, mesh_, ocp_.g, t0, tf, X.colwise(), U.colwise());
    mesh_eval<0>(cr_out0_, mesh_, ocp_.cr, t0, tf, X.colwise(), U.colwise(), true);

    g_.segment(dcon_B, dcon_L)   = w_scaling_ * dyn_out0_.F;
    g_.segment(qcon_B, qcon_L)   = w_scaling_ * (int_out0_.F - q);
    g_.segment(crcon_B, crcon_L) = w_scaling_ * cr_out0_.F;
    g_.segment(cecon_B, cecon_L) = ocp_.ce(tf, x0, xf, q);

    return g_;
  }
  const Eigen::VectorXd & gl() const { return gl_; }
  const Eigen::VectorXd & gu() const { return gu_; }
  const Eigen::SparseMatrix<double> & dg_dx(const Eigen::Ref<const Eigen::VectorXd> x)
  {
    assert(static_cast<std::size_t>(x.size()) == n_);

    const auto x0var_B = xvar_B;
    const auto xfvar_B = xvar_B + xvar_L - Nx;

    const double t0                    = 0;
    const double tf                    = x(tfvar_B);
    const Eigen::Vector<double, Nx> x0 = x.segment(x0var_B, Nx);
    const Eigen::Vector<double, Nx> xf = x.segment(xfvar_B, Nx);
    const Eigen::Vector<double, Nq> q  = x.segment(qvar_B, qvar_L);

    const Eigen::Map<const Eigen::Matrix<double, Nx, -1>> X(x.data() + xvar_B, Nx, N_ + 1);
    const Eigen::Map<const Eigen::Matrix<double, Nu, -1>> U(x.data() + uvar_B, Nu, N_);

    mesh_dyn<1, DT>(dyn_out1_, mesh_, ocp_.f, t0, tf, X.colwise(), U.colwise());
    mesh_integrate<1, DT>(int_out1_, mesh_, ocp_.g, t0, tf, X.colwise(), U.colwise());
    mesh_eval<1, DT>(cr_out1_, mesh_, ocp_.cr, t0, tf, X.colwise(), U.colwise(), true);
    const auto & [ceval, dceval] = diff::dr<1, DT>(ocp_.ce, wrt(tf, x0, xf, q));

    dyn_out1_.dF.makeCompressed();
    int_out1_.dF.makeCompressed();
    cr_out1_.dF.makeCompressed();

    set_zero(dg_dx_);

    // dynamics constraint
    block_add(dg_dx_, dcon_B, tfvar_B, dyn_out1_.dF.middleCols(1, 1), w_scaling_);
    block_add(dg_dx_, dcon_B, xvar_B, dyn_out1_.dF.middleCols(2, xvar_L), w_scaling_);
    block_add(dg_dx_, dcon_B, uvar_B, dyn_out1_.dF.middleCols(2 + xvar_L, uvar_L), w_scaling_);

    // integral constraint
    block_add(dg_dx_, qcon_B, tfvar_B, int_out1_.dF.middleCols(1, 1), w_scaling_);
    block_add(dg_dx_, qcon_B, xvar_B, int_out1_.dF.middleCols(2, xvar_L), w_scaling_);
    block_add(dg_dx_, qcon_B, uvar_B, int_out1_.dF.middleCols(2 + xvar_L, uvar_L), w_scaling_);
    block_add_identity(dg_dx_, qcon_B, qvar_B, qvar_L, -w_scaling_);

    // running constraint
    block_add(dg_dx_, crcon_B, tfvar_B, cr_out1_.dF.middleCols(1, 1), w_scaling_);
    block_add(dg_dx_, crcon_B, xvar_B, cr_out1_.dF.middleCols(2, xvar_L), w_scaling_);
    block_add(dg_dx_, crcon_B, uvar_B, cr_out1_.dF.middleCols(2 + xvar_L, uvar_L), w_scaling_);

    // end constraint
    block_add(dg_dx_, cecon_B, tfvar_B, dceval.middleCols(0, 1));
    block_add(dg_dx_, cecon_B, xvar_B, dceval.middleCols(1, Nx));
    block_add(dg_dx_, cecon_B, xvar_B + xvar_L - Nx, dceval.middleCols(1 + Nx, Nx));
    block_add(dg_dx_, cecon_B, qvar_B, dceval.middleCols(1 + 2 * Nx, Nq));

    dg_dx_.makeCompressed();
    return dg_dx_;
  }

  const Eigen::SparseMatrix<double> &
  d2g_dx2(const Eigen::Ref<const Eigen::VectorXd> x, const Eigen::Ref<const Eigen::VectorXd> lambda)
  {
    assert(static_cast<std::size_t>(x.size()) == n_);
    assert(static_cast<std::size_t>(lambda.size()) == m_);

    const auto x0var_B = xvar_B;
    const auto xfvar_B = xvar_B + xvar_L - Nx;

    const double t0                    = 0;
    const double tf                    = x(tfvar_B);
    const Eigen::Vector<double, Nx> x0 = x.segment(x0var_B, Nx);
    const Eigen::Vector<double, Nx> xf = x.segment(xfvar_B, Nx);
    const Eigen::Vector<double, Nq> q  = x.segment(qvar_B, qvar_L);

    const Eigen::Map<const Eigen::Matrix<double, Nx, -1>> X(x.data() + xvar_B, Nx, N_ + 1);
    const Eigen::Map<const Eigen::Matrix<double, Nu, -1>> U(x.data() + uvar_B, Nu, N_);

    dyn_out2_.lambda = lambda.segment(dcon_B, dcon_L);
    int_out2_.lambda = lambda.segment(qcon_B, qcon_L);
    cr_out2_.lambda  = lambda.segment(crcon_B, crcon_L);
    mesh_dyn<2, DT>(dyn_out2_, mesh_, ocp_.f, t0, tf, X.colwise(), U.colwise());
    mesh_integrate<2, DT>(int_out2_, mesh_, ocp_.g, t0, tf, X.colwise(), U.colwise());
    mesh_eval<2, DT>(cr_out2_, mesh_, ocp_.cr, t0, tf, X.colwise(), U.colwise(), true);
    const auto & [ceval, dceval, d2ceval] = diff::dr<2, DT>(ocp_.ce, wrt(tf, x0, xf, q));

    dyn_out2_.dF.makeCompressed();
    int_out2_.dF.makeCompressed();
    cr_out2_.dF.makeCompressed();

    dyn_out2_.d2F.makeCompressed();
    int_out2_.d2F.makeCompressed();
    cr_out2_.d2F.makeCompressed();

    set_zero(d2g_dx2_);

    // clang-format off
    block_add(d2g_dx2_, tfvar_B, tfvar_B, dyn_out2_.d2F.block(1, 1, 1, 1), w_scaling_, true);      // tftf
    block_add(d2g_dx2_, tfvar_B, tfvar_B, int_out2_.d2F.block(1, 1, 1, 1), w_scaling_, true);      // tftf
    block_add(d2g_dx2_, tfvar_B, tfvar_B,  cr_out2_.d2F.block(1, 1, 1, 1), w_scaling_, true);      // tftf

    block_add(d2g_dx2_, tfvar_B, xvar_B, dyn_out2_.d2F.block(1, 2, 1, xvar_L), w_scaling_, true);  // tfx
    block_add(d2g_dx2_, tfvar_B, xvar_B, int_out2_.d2F.block(1, 2, 1, xvar_L), w_scaling_, true);  // tfx
    block_add(d2g_dx2_, tfvar_B, xvar_B,  cr_out2_.d2F.block(1, 2, 1, xvar_L), w_scaling_, true);  // tfx

    block_add(d2g_dx2_, tfvar_B, uvar_B, dyn_out2_.d2F.block(1, 2 + xvar_L, 1, uvar_L), w_scaling_, true);  // tfu
    block_add(d2g_dx2_, tfvar_B, uvar_B, int_out2_.d2F.block(1, 2 + xvar_L, 1, uvar_L), w_scaling_, true);  // tfu
    block_add(d2g_dx2_, tfvar_B, uvar_B,  cr_out2_.d2F.block(1, 2 + xvar_L, 1, uvar_L), w_scaling_, true);  // tfu

    block_add(d2g_dx2_, xvar_B, xvar_B, dyn_out2_.d2F.block(2,          2, xvar_L, xvar_L), w_scaling_, true);  // xx
    block_add(d2g_dx2_, xvar_B, xvar_B, int_out2_.d2F.block(2,          2, xvar_L, xvar_L), w_scaling_, true);  // xx
    block_add(d2g_dx2_, xvar_B, xvar_B,  cr_out2_.d2F.block(2,          2, xvar_L, xvar_L), w_scaling_, true);  // xx

    block_add(d2g_dx2_, xvar_B, uvar_B, dyn_out2_.d2F.block(2, 2 + xvar_L, xvar_L, uvar_L), w_scaling_, true);  // xu
    block_add(d2g_dx2_, xvar_B, uvar_B, int_out2_.d2F.block(2, 2 + xvar_L, xvar_L, uvar_L), w_scaling_, true);  // xu
    block_add(d2g_dx2_, xvar_B, uvar_B,  cr_out2_.d2F.block(2, 2 + xvar_L, xvar_L, uvar_L), w_scaling_, true);  // xu

    block_add(d2g_dx2_, uvar_B, uvar_B, dyn_out2_.d2F.block(2 + xvar_L, 2 + xvar_L, uvar_L, uvar_L), w_scaling_, true);  // uu
    block_add(d2g_dx2_, uvar_B, uvar_B, int_out2_.d2F.block(2 + xvar_L, 2 + xvar_L, uvar_L, uvar_L), w_scaling_, true);  // uu
    block_add(d2g_dx2_, uvar_B, uvar_B,  cr_out2_.d2F.block(2 + xvar_L, 2 + xvar_L, uvar_L, uvar_L), w_scaling_, true);  // uu
    // clang-format on

    for (auto j = 0u; j < ocp_.Nce; ++j) {
      const auto b0 = (1 + 2 * Nx + ocp_.Nq) * j;
      // clang-format off
      block_add(d2g_dx2_, tfvar_B, tfvar_B, d2ceval.block(         0, b0 +          0,  1,  1), lambda(cecon_B + j), true);  // tftf
      block_add(d2g_dx2_, tfvar_B, x0var_B, d2ceval.block(         0, b0 +          1,  1, Nx), lambda(cecon_B + j), true);  // tfx0
      block_add(d2g_dx2_, tfvar_B, xfvar_B, d2ceval.block(         0, b0 +     1 + Nx,  1, Nx), lambda(cecon_B + j), true);  // tfxf
      block_add(d2g_dx2_, tfvar_B,  qvar_B, d2ceval.block(         0, b0 + 1 + 2 * Nx,  1, Nq), lambda(cecon_B + j), true);  // tfq

      block_add(d2g_dx2_, x0var_B, x0var_B, d2ceval.block(         1, b0 +          1, Nx, Nx), lambda(cecon_B + j), true);  // x0x0
      block_add(d2g_dx2_, x0var_B, xfvar_B, d2ceval.block(         1, b0 +     1 + Nx, Nx, Nx), lambda(cecon_B + j), true);  // x0xf
      block_add(d2g_dx2_, x0var_B,  qvar_B, d2ceval.block(         1, b0 + 1 + 2 * Nx, Nx, Nq), lambda(cecon_B + j), true);  // x0q

      block_add(d2g_dx2_, xfvar_B, xfvar_B, d2ceval.block(    1 + Nx, b0 +     1 + Nx, Nx, Nx), lambda(cecon_B + j), true);  // xfxf
      block_add(d2g_dx2_, xfvar_B,  qvar_B, d2ceval.block(    1 + Nx, b0 + 1 + 2 * Nx, Nx, Nq), lambda(cecon_B + j), true);  // xfq

      block_add(d2g_dx2_,  qvar_B,  qvar_B, d2ceval.block(1 + 2 * Nx, b0 + 1 + 2 * Nx, Nq, Nq), lambda(cecon_B + j), true);  // qq
      // clang-format on
    }

    d2g_dx2_.makeCompressed();
    return d2g_dx2_;
  }
};

}  // namespace detail
// \endcond

/**
 * @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()
 */
template<diff::Type DT = diff::Type::Default>
auto ocp_to_nlp(FlatOCPType auto && ocp, MeshType auto && mesh)
  -> detail::OCPNLP<std::decay_t<decltype(ocp)>, std::decay_t<decltype(mesh)>, DT>
{
  return detail::OCPNLP<std::decay_t<decltype(ocp)>, std::decay_t<decltype(mesh)>, DT>(
    std::forward<decltype(ocp)>(ocp), std::forward<decltype(mesh)>(mesh));
}

/**
 * @brief Convert nonlinear program solution to ocp solution
 */
auto nlpsol_to_ocpsol(
  const FlatOCPType auto & ocp, const MeshType auto & mesh, const NLPSolution & nlp_sol)
{
  using ocp_t = std::decay_t<decltype(ocp)>;

  static constexpr auto Nx  = ocp_t::Nx;
  static constexpr auto Nu  = ocp_t::Nu;
  static constexpr auto Nq  = ocp_t::Nq;
  static constexpr auto Ncr = ocp_t::Ncr;

  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::Vector<double, Nq> 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::Vector<double, Nx> {
    return mesh.template eval<Eigen::Vector<double, Nx>>(
      (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::Vector<double, Nu> {
    return mesh.template eval<Eigen::Vector<double, Nu>>(
      (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::Vector<double, Nx> {
    return mesh.template eval<Eigen::Vector<double, Nx>>(
      (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::Vector<double, Ncr> {
    return mesh.template eval<Eigen::Vector<double, Ncr>>(
      (t - t0) / (tf - t0), Lcr.colwise(), 0, false);
  };

  return OCPSolution<typename ocp_t::X, typename ocp_t::U, ocp_t::Nq, ocp_t::Ncr, ocp_t::Nce>{
    .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 solution
 *
 * @note Allocates memory for return type.
 */
NLPSolution
ocpsol_to_nlpsol(const FlatOCPType auto & ocp, const MeshType auto & mesh, const auto & ocpsol)
{
  const auto 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 = ocpsol.tf;

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

  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 (const auto & [i, tau] : utils::zip(std::views::iota(0u), mesh.all_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));
    }
  }

  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_TO_NLP_HPP_


//
// Created by eliane on 06/03/19.
//

#include "VinaLikeCovalentReversible.h"

#include <exception>

#include <Structures/Atom.h>

#include <Engines/Internals/InternalsUtilityFunctions.h>

#include <boost/log/trivial.hpp>

#include "VinaLikeCommon.h"
#include "VinaExtendedCommon.h"

namespace SmolDock::Score {

    const std::array<std::string, VinaLikeCovalentReversible::numCoefficients>
    VinaLikeCovalentReversible::coefficientsNames =  {"Gauss1", "Gauss2", "RepulsionExceptCovalent", "Hydrophobic","Hydrogen", "CovalentReversible"};

    template<bool OnlyIntermolecular, bool useNonDefaultCoefficients> // default : false
    double VinaLikeCovalentReversibleIntermolecularScoringFunction(const iConformer &ligand_, iTransform &transform,
                                                 const iProtein &protein,
                                                 std::array<double, VinaLikeCovalentReversible_numCoefficients> nonDefaultCoeffs) {

        BOOST_ASSERT(!ligand_.x.empty());
        BOOST_ASSERT(!protein.x.empty());

        BOOST_ASSERT(transform.bondRotationsAngles.size() == ligand_.num_rotatable_bond);

        if(std::abs(transform.rota.norm() - 1) > 0.1) {
            transform.rota.normalize();
        }

        double score_raw = 0;

        iConformer ligand = ligand_;
        applyBondRotationInPlace(ligand, transform);

        Eigen::Vector3d ProtCenterPosition = {protein.center_x, protein.center_y, protein.center_z};

        if constexpr(!OnlyIntermolecular) // C++17
        {
            for (unsigned int idxLig = 0; idxLig < ligand.x.size(); idxLig++) {
                for (unsigned int idxLig2 = idxLig; idxLig2 < ligand.x.size(); idxLig2++) {

                    if(idxLig == idxLig2)
                        continue;

                    Eigen::Vector3d LigDistance = {ligand.x[idxLig] - ligand.x[idxLig2],
                                                   ligand.y[idxLig] - ligand.y[idxLig2],
                                                   ligand.z[idxLig] - ligand.z[idxLig2]};

                    double distance_raw = LigDistance.norm();
                    const double distance = distanceFromRawDistance(distance_raw, ligand.atomicRadius[idxLig],
                                                                    ligand.atomicRadius[idxLig2]);

//                score_raw += VinaClassic::coeff_gauss1      * vinaGaussComponent(distance, 0.0, 0.5);
//                score_raw += VinaClassic::coeff_gauss2      * vinaGaussComponent(distance, 3.0, 2.0);
                    score_raw += VinaClassic::coeff_repulsion   * vinaRepulsionComponent(distance, 0.0);
                }
            }
        }

        for (unsigned int idxLig = 0; idxLig < ligand.x.size(); idxLig++) {
            for (unsigned int idxProt = 0; idxProt < protein.x.size(); idxProt++) {

                Eigen::Vector3d LigPosition = {ligand.x[idxLig], ligand.y[idxLig], ligand.z[idxLig]};
                applyRigidTransformInPlace(LigPosition, transform);

                Eigen::Vector3d ProtPosition = {protein.x[idxProt], protein.y[idxProt], protein.z[idxProt]};
                Eigen::Vector3d distToCenterVector = LigPosition - ProtCenterPosition;

                double distanceToProteinCenter = distToCenterVector.norm();

                if (distanceToProteinCenter > (protein.radius - 1)) {
                    score_raw += std::pow((distanceToProteinCenter - protein.radius), 4) + 10;
                    continue;
                }

                Eigen::Vector3d distVect = ProtPosition - LigPosition;

                double rawDist = distVect.norm();

                if (rawDist >= VinaClassic::interaction_cutoff)
                    continue;

                const double distance = distanceFromRawDistance(rawDist, ligand.atomicRadius[idxLig],
                                                                protein.atomicRadius[idxProt]);

                const unsigned int atom1AtomicNumber = ligand.type[idxLig];
                const unsigned int atom1AtomVariant = ligand.variant[idxLig];
                const unsigned int atom2AtomicNumber = protein.type[idxProt];
                const unsigned int atom2AtomVariant = protein.variant[idxProt];

                if constexpr(useNonDefaultCoefficients)
                {
                    score_raw += nonDefaultCoeffs[0]      * vinaGaussComponent(distance, 0.0, 0.5);
                    score_raw += nonDefaultCoeffs[1]     * vinaGaussComponent(distance, 3.0, 2.0);
                    score_raw += nonDefaultCoeffs[2]   * vinaRepulsionComponent(distance, 0.0);
                    score_raw += nonDefaultCoeffs[3] * vinaHydrophobicComponent(distance,
                                                                               atom1AtomicNumber, atom1AtomVariant,
                                                                               atom2AtomicNumber, atom2AtomVariant);

                    score_raw += nonDefaultCoeffs[4]    * vinaHydrogenComponent(distance,
                                                                               atom1AtomicNumber, atom1AtomVariant,
                                                                               atom2AtomicNumber, atom2AtomVariant);
                    score_raw += nonDefaultCoeffs[5] * VinaExtended::covalentReversibleComponent(distance,
                                                                                atom1AtomicNumber, atom1AtomVariant,
                                                                                atom2AtomicNumber, atom2AtomVariant);
                }else {
                    score_raw += VinaClassic::coeff_gauss1      * vinaGaussComponent(distance, 0.0, 0.5);
                    score_raw += VinaClassic::coeff_gauss2      * vinaGaussComponent(distance, 3.0, 2.0);
                    score_raw += VinaClassic::coeff_repulsion   * VinaExtended::RepulsionExceptForCovalentComponent(distance, 0.0,
                                                                                atom1AtomicNumber, atom1AtomVariant,
                                                                                atom2AtomicNumber, atom2AtomVariant);
                    score_raw += VinaClassic::coeff_hydrophobic * vinaHydrophobicComponent(distance,
                                                                                atom1AtomicNumber, atom1AtomVariant,
                                                                                atom2AtomicNumber, atom2AtomVariant);

                    score_raw += VinaClassic::coeff_hydrogen    * vinaHydrogenComponent(distance,
                                                                                atom1AtomicNumber, atom1AtomVariant,
                                                                                atom2AtomicNumber, atom2AtomVariant);

                    score_raw += VinaExtended::coeff_CovalentReversible * VinaExtended::covalentReversibleComponent(distance,
                                                                                atom1AtomicNumber, atom1AtomVariant,
                                                                                atom2AtomicNumber, atom2AtomVariant);
                }

            } // for
        } // for

        double final_score = score_raw / (1 + (VinaClassic::coeff_entropic * ligand.num_rotatable_bond));
        return final_score;
    }

    VinaLikeCovalentReversible::VinaLikeCovalentReversible(const iConformer &startingConformation_,
                       const iProtein &p,
                       const iTransform &initialTransform_,
                       double differential_epsilon_,
                        bool useNonDefaultCoefficient) :
            useNonDefaultCoefficient(useNonDefaultCoefficient),
            startingConformation(startingConformation_),
            prot(p),
            initialTransform(initialTransform_),
            differential_epsilon(differential_epsilon_) {
        this->numberOfRotatableBonds = this->startingConformation.num_rotatable_bond;
        this->numberOfParamInState = 7 + (this->numberOfRotatableBonds);

        if (this->initialTransform.bondRotationsAngles.size() != this->numberOfRotatableBonds) {
            BOOST_LOG_TRIVIAL(error)
                    << "Discrepency between the number of rotatable bonds in the iConformer and iTransform ("
                    << this->numberOfRotatableBonds << " != " << this->initialTransform.bondRotationsAngles.size()
                    << ")";
            std::terminate();
        }

        this->nonDefaultCoefficients = {VinaClassic::coeff_gauss1, VinaClassic::coeff_gauss2,
                                        VinaClassic::coeff_repulsion, VinaClassic::coeff_hydrophobic,
                                        VinaClassic::coeff_hydrogen, VinaExtended::coeff_CovalentReversible };

    }

    double VinaLikeCovalentReversible::Evaluate(const arma::mat &x) {
        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);

        tr.doHousekeeping();

        double score_ = VinaLikeCovalentReversibleIntermolecularScoringFunction(this->startingConformation, tr, this->prot);

        return score_;
    }

    double VinaLikeCovalentReversible::EvaluateWithGradient(const arma::mat &x, arma::mat &grad) {

        BOOST_ASSERT(!x.has_nan());
        BOOST_ASSERT(!grad.has_nan());
        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);
        BOOST_ASSERT(grad.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);
        tr.doHousekeeping();

        double score_ = VinaLikeCovalentReversibleIntermolecularScoringFunction(this->startingConformation, tr, this->prot);

        // Translation
        {
            iTransform transform_dx = tr;
            transform_dx.transl.x() += this->differential_epsilon;
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeCovalentReversibleIntermolecularScoringFunction<false,true>(this->startingConformation, transform_dx, this->prot, this->nonDefaultCoefficients)
                                                                    : VinaLikeCovalentReversibleIntermolecularScoringFunction<false,false>(this->startingConformation, transform_dx, this->prot);
            grad[0] = gradScore - score_;
        }

        {
            iTransform transform_dy = tr;
            transform_dy.transl.y() += this->differential_epsilon;
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeCovalentReversibleIntermolecularScoringFunction<false,true>(this->startingConformation, transform_dy, this->prot, this->nonDefaultCoefficients)
                                                                    : VinaLikeCovalentReversibleIntermolecularScoringFunction<false,false>(this->startingConformation, transform_dy, this->prot);
            grad[1] = gradScore - score_;
        }

        {
            iTransform transform_dz = tr;
            transform_dz.transl.z() += this->differential_epsilon;
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeCovalentReversibleIntermolecularScoringFunction<false,true>(this->startingConformation, transform_dz, this->prot, this->nonDefaultCoefficients)
                                                                    : VinaLikeCovalentReversibleIntermolecularScoringFunction<false,false>(this->startingConformation, transform_dz, this->prot);
            grad[2] = gradScore - score_;
        }

        // Rotation

        {
            iTransform transform_dqs = tr;
            transform_dqs.rota.w() += this->differential_epsilon;
            transform_dqs.doHousekeeping();
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeCovalentReversibleIntermolecularScoringFunction<false,true>(this->startingConformation, transform_dqs, this->prot, this->nonDefaultCoefficients)
                                                                    : VinaLikeCovalentReversibleIntermolecularScoringFunction<false,false>(this->startingConformation, transform_dqs, this->prot);
            grad[3] = gradScore - score_;
        }

        {
            iTransform transform_dqx = tr;
            transform_dqx.rota.x() += this->differential_epsilon;
            transform_dqx.doHousekeeping();
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeCovalentReversibleIntermolecularScoringFunction<false,true>(this->startingConformation, transform_dqx, this->prot, this->nonDefaultCoefficients)
                                                                    : VinaLikeCovalentReversibleIntermolecularScoringFunction<false,false>(this->startingConformation, transform_dqx, this->prot);
            grad[4] = gradScore - score_;

        }

        {
            iTransform transform_dqy = tr;
            transform_dqy.rota.x() += this->differential_epsilon;
            transform_dqy.doHousekeeping();
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeCovalentReversibleIntermolecularScoringFunction<false,true>(this->startingConformation, transform_dqy, this->prot, this->nonDefaultCoefficients)
                                                                    : VinaLikeCovalentReversibleIntermolecularScoringFunction<false,false>(this->startingConformation, transform_dqy, this->prot);
            grad[5] = gradScore - score_;
        }

        {
            iTransform transform_dqz = tr;
            transform_dqz.rota.x() += this->differential_epsilon;
            transform_dqz.doHousekeeping();
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeCovalentReversibleIntermolecularScoringFunction<false,true>(this->startingConformation, transform_dqz, this->prot, this->nonDefaultCoefficients)
                                                                    : VinaLikeCovalentReversibleIntermolecularScoringFunction<false,false>(this->startingConformation, transform_dqz, this->prot);
            grad[6] = gradScore - score_;
        }

        for (unsigned int i = 0; i < this->numberOfRotatableBonds; i++) {
            iTransform transform_dbondrot = tr;
            transform_dbondrot.bondRotationsAngles[i] += this->differential_epsilon;
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeCovalentReversibleIntermolecularScoringFunction<false,true>(this->startingConformation, transform_dbondrot, this->prot, this->nonDefaultCoefficients)
                                                                    : VinaLikeCovalentReversibleIntermolecularScoringFunction<false,false>(this->startingConformation, transform_dbondrot, this->prot);
            grad[7 + i] = gradScore - score_;

        }

        /*
        BOOST_LOG_TRIVIAL(debug) << "Transform: " << x.t();
        BOOST_LOG_TRIVIAL(debug) << "Score: " << score_;
        BOOST_LOG_TRIVIAL(debug) << "Gradient";
        BOOST_LOG_TRIVIAL(debug) << "     ds: " << grad[3];
        BOOST_LOG_TRIVIAL(debug) << "     du: " << grad[4] << "   dx: " << grad[0];
        BOOST_LOG_TRIVIAL(debug) << "     dv: " << grad[5] << "   dy: " << grad[1];
        BOOST_LOG_TRIVIAL(debug) << "     dt: " << grad[6] << "   dx: " << grad[2];
        //*/

        BOOST_ASSERT(score_ == score_); // catches NaN
        return score_;
    }

    double VinaLikeCovalentReversible::getDifferentialEpsilon() const {
        return this->differential_epsilon;
    }

    arma::mat VinaLikeCovalentReversible::getStartingConditions() const {
        return this->externalToInternalRepr(this->initialTransform);
    }

    iConformer VinaLikeCovalentReversible::getConformerForParamMatrix(const arma::mat &x) {
        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);
        tr.doHousekeeping();

        iConformer ret = this->startingConformation;
        applyBondRotationInPlace(ret, tr);
        applyRigidTransformInPlace(ret, tr);

        return ret;
    }

    unsigned int VinaLikeCovalentReversible::getParamVectorDimension() const {
        return this->numberOfParamInState;
    }

    std::vector<std::tuple<std::string, double>> VinaLikeCovalentReversible::EvaluateSubcomponents(const arma::mat &x) {

        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);
        tr.doHousekeeping();

        return VinaLikeCovalentReversibleIntermolecularComponents(this->startingConformation, tr, this->prot);
    }

    double VinaLikeCovalentReversible::EvaluateOnlyIntermolecular(const arma::mat &x) {
        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);

        tr.doHousekeeping();

        // Template parameter controls whether onlyIntermolecular interaction are taken into account. Here we want true
        double score_ = VinaLikeCovalentReversibleIntermolecularScoringFunction<true>(this->startingConformation, tr, this->prot);

        return score_;
    }

    unsigned int VinaLikeCovalentReversible::getCoefficientsVectorWidth() {
        return this->numCoefficients;
    }

    std::vector<std::string> VinaLikeCovalentReversible::getCoefficientsNames() {
        return std::vector<std::string>(this->coefficientsNames.begin(), this->coefficientsNames.end());
    }

    std::vector<double> VinaLikeCovalentReversible::getCurrentCoefficients() {
        if(this->useNonDefaultCoefficient)
        {
            return std::vector<double>(this->nonDefaultCoefficients.begin(),this->nonDefaultCoefficients.end());
        }
        return {VinaClassic::coeff_gauss1, VinaClassic::coeff_gauss2,
                VinaClassic::coeff_repulsion, VinaClassic::coeff_hydrophobic,
                VinaClassic::coeff_hydrogen, VinaExtended::coeff_CovalentReversible };

    }

    bool VinaLikeCovalentReversible::setNonDefaultCoefficients(std::vector<double> coeffs) {
        if(coeffs.size() != this->numCoefficients)
        {
            BOOST_LOG_TRIVIAL(error) << "Trying to set " << this->numCoefficients <<" coefficients with vector of " << coeffs.size() << " values.";
            return false;
        }
        if(this->useNonDefaultCoefficient == false)
        {
            BOOST_LOG_TRIVIAL(error) << "Trying to set non default coefficient, but this scoring function was constructed with default coefficients only.";
            BOOST_LOG_TRIVIAL(error) << "Check the parameters passed to the scoring function constructor.";
            return false;
        }
        for (unsigned int j = 0; j < this->nonDefaultCoefficients.size(); ++j) {
            this->nonDefaultCoefficients[j] = coeffs[j];
        }
        return true;
    }

    template<bool useNonDefaultCoefficients> // default : false
    std::vector<std::tuple<std::string, double>> VinaLikeCovalentReversibleIntermolecularComponents(const iConformer &conformer, iTransform &transform,
                                                                                const iProtein &protein,
                                                                                std::array<double, VinaLikeCovalentReversible_numCoefficients> nonDefaultCoeffs)
    {
        BOOST_ASSERT(!conformer.x.empty());
        BOOST_ASSERT(!protein.x.empty());
        BOOST_ASSERT(transform.bondRotationsAngles.size() == conformer.num_rotatable_bond);

        if(std::abs(transform.rota.norm() - 1) > 0.1) {
            transform.rota.normalize();
        }

        std::vector<std::tuple<std::string, double>> ret;

        double gauss1_total = 0.0;
        double gauss2_total = 0.0;
        double repulsion_total = 0.0;
        double repulsion_VinaClassic_total = 0.0;
        double hydrogen_total = 0.0;
        double hydrophobic_total = 0.0;
        double covrev_total = 0.0;
        double score_raw = 0.0;

        double intramolecular_repuls_total = 0.0;
        double intramolecular_score = 0.0;

        iConformer ligand = conformer;
        applyBondRotationInPlace(ligand, transform);

        Eigen::Vector3d ProtCenterPosition = {protein.center_x, protein.center_y, protein.center_z};

        for (unsigned int idxLig = 0; idxLig < ligand.x.size(); idxLig++) {
            for (unsigned int idxLig2 = idxLig; idxLig2 < ligand.x.size(); idxLig2++) {

                if(idxLig == idxLig2)
                    continue;

                Eigen::Vector3d LigDistance = {ligand.x[idxLig] - ligand.x[idxLig2],
                                               ligand.y[idxLig] - ligand.y[idxLig2],
                                               ligand.z[idxLig] - ligand.z[idxLig2]};

                double distance_raw = LigDistance.norm();
                const double distance = distanceFromRawDistance(distance_raw, ligand.atomicRadius[idxLig],
                                                                ligand.atomicRadius[idxLig2]);

//                score_raw += VinaClassic::coeff_gauss1      * vinaGaussComponent(distance, 0.0, 0.5);
//                score_raw += VinaClassic::coeff_gauss2      * vinaGaussComponent(distance, 3.0, 2.0);
                intramolecular_repuls_total +=  vinaRepulsionComponent(distance, 0.0);

            }
        }

        intramolecular_score = VinaClassic::coeff_repulsion   * intramolecular_repuls_total;

        for (unsigned int idxLig = 0; idxLig < ligand.x.size(); idxLig++) {
            for (unsigned int idxProt = 0; idxProt < protein.x.size(); idxProt++) {

                Eigen::Vector3d LigPosition = {ligand.x[idxLig], ligand.y[idxLig], ligand.z[idxLig]};
                applyRigidTransformInPlace(LigPosition, transform);

                Eigen::Vector3d ProtPosition = {protein.x[idxProt], protein.y[idxProt], protein.z[idxProt]};
                Eigen::Vector3d distToCenterVector = LigPosition - ProtCenterPosition;

                double distanceToProteinCenter = distToCenterVector.norm();

                if (distanceToProteinCenter > (protein.radius - 1)) {
                    score_raw += std::pow((distanceToProteinCenter - protein.radius), 4) + 10;
                    continue;
                }

                Eigen::Vector3d distVect = ProtPosition - LigPosition;

                double rawDist = distVect.norm();

                if (rawDist >= VinaClassic::interaction_cutoff)
                    continue;

                const double distance = distanceFromRawDistance(rawDist, ligand.atomicRadius[idxLig],
                                                          protein.atomicRadius[idxProt]);

                const unsigned int atom1AtomicNumber = ligand.type[idxLig];
                const unsigned int atom1AtomVariant = ligand.variant[idxLig];
                const unsigned int atom2AtomicNumber = protein.type[idxProt];
                const unsigned int atom2AtomVariant = protein.variant[idxProt];

                gauss1_total += vinaGaussComponent(distance, 0.0, 0.5);
                gauss2_total += vinaGaussComponent(distance, 3.0, 2.0);
                repulsion_total += VinaExtended::RepulsionExceptForCovalentComponent(distance, 0.0,
                                                                                atom1AtomicNumber, atom1AtomVariant,
                                                                                atom2AtomicNumber, atom2AtomVariant);
                hydrophobic_total += vinaHydrophobicComponent(distance,
                                                           atom1AtomicNumber, atom1AtomVariant,
                                                           atom2AtomicNumber, atom2AtomVariant);
                hydrogen_total += vinaHydrogenComponent(distance,
                                                           atom1AtomicNumber, atom1AtomVariant,
                                                           atom2AtomicNumber, atom2AtomVariant);
                covrev_total += VinaExtended::covalentReversibleComponent(distance,
                                                                          atom1AtomicNumber, atom1AtomVariant,
                                                                          atom2AtomicNumber, atom2AtomVariant);

                // Not used in the calculation of the score
                repulsion_VinaClassic_total += vinaRepulsionComponent(distance, 0.0);

            } // for
        } // for

        double score_sum = 0.0;
        double score_NoCovRev = 0.0;
        if constexpr(useNonDefaultCoefficients)
        {
            ret.emplace_back(std::make_tuple("NonDefaultCoeffs", 1.0));
            score_sum =   nonDefaultCoeffs[0] * gauss1_total
                                 + nonDefaultCoeffs[1] * gauss2_total
                                 + nonDefaultCoeffs[2] * repulsion_total
                                 + nonDefaultCoeffs[3] * hydrophobic_total
                                 + nonDefaultCoeffs[4] * hydrogen_total
                                 + nonDefaultCoeffs[5] * covrev_total;

            score_NoCovRev =   nonDefaultCoeffs[0] * gauss1_total
                               + nonDefaultCoeffs[1] * gauss2_total
                               + nonDefaultCoeffs[2] * repulsion_total
                               + nonDefaultCoeffs[3] * hydrophobic_total
                               + nonDefaultCoeffs[4] * hydrogen_total;
        }else {
            score_sum =   VinaClassic::coeff_gauss1 * gauss1_total
                                 + VinaClassic::coeff_gauss2 * gauss2_total
                                 + VinaClassic::coeff_repulsion * repulsion_total
                                 + VinaClassic::coeff_hydrophobic * hydrophobic_total
                                 + VinaClassic::coeff_hydrogen * hydrogen_total
                                 + VinaExtended::coeff_CovalentReversible * covrev_total;

            score_NoCovRev =   VinaClassic::coeff_gauss1 * gauss1_total
                                      + VinaClassic::coeff_gauss2 * gauss2_total
                                      + VinaClassic::coeff_repulsion * repulsion_total
                                      + VinaClassic::coeff_hydrophobic * hydrophobic_total
                                      + VinaClassic::coeff_hydrogen * hydrogen_total;
        }

        double final_score = score_sum / (1 + (VinaClassic::coeff_entropic * ligand.num_rotatable_bond));

        double final_score_nocovrev = score_NoCovRev / (1 + (VinaClassic::coeff_entropic * ligand.num_rotatable_bond));

        ret.emplace_back(std::make_tuple("Gauss1", gauss1_total));
        ret.emplace_back(std::make_tuple("Gauss2", gauss2_total));
        ret.emplace_back(std::make_tuple("Repulsion_NotUsedInScore", repulsion_VinaClassic_total));
        ret.emplace_back(std::make_tuple("RepulsionExceptCovalent", repulsion_total));
        ret.emplace_back(std::make_tuple("Hydrophobic", hydrophobic_total));
        ret.emplace_back(std::make_tuple("Hydrogen", hydrogen_total ));
        ret.emplace_back(std::make_tuple("CovalentReversible", covrev_total));
        ret.emplace_back(std::make_tuple("numRot", ligand.num_rotatable_bond));
        ret.emplace_back(std::make_tuple("Intra_Repuls", intramolecular_repuls_total));
        ret.emplace_back(std::make_tuple("Intra_Score", intramolecular_score));
        ret.emplace_back(std::make_tuple("ScoreRaw_NoCovRev", score_NoCovRev));
        ret.emplace_back(std::make_tuple("Score_NoCovRev", final_score_nocovrev));
        ret.emplace_back(std::make_tuple("Score_Raw", score_sum));
        ret.emplace_back(std::make_tuple("Score", final_score));
        return ret;
    }

}