#include "quadrotor_simulator/Quadrotor.h"
#include <iostream>
#include <boost/bind.hpp>
#include "odeint-v2/boost/numeric/odeint.hpp"
#include <Eigen/Geometry>

namespace odeint = boost::numeric::odeint;

namespace QuadrotorSimulator
{

Quadrotor::Quadrotor(void)
{
  g_ = 9.81;
  mass_ = 0.075;
  double Ixx = 2.64e-3, Iyy = 2.64e-3, Izz = 4.96e-3;
  prop_radius_ = 0.099;
  J_ = Eigen::Vector3d(Ixx, Iyy, Izz).asDiagonal();

  kf_ = 5.55e-8;
  // km_ = 2.5e-9; // from Nate
  // km = (Cq/Ct)*Dia*kf
  // Cq/Ct for 8 inch props from UIUC prop db ~ 0.07
  km_ = 0.07*(2*prop_radius_)*kf_;

  arm_length_ = 0.17;
  motor_time_constant_ = 1.0/20;
  min_rpm_ = 1500;
  max_rpm_ = 7500;

  state_.x = Eigen::Vector3d::Zero();
  state_.v = Eigen::Vector3d::Zero();
  state_.R = Eigen::Matrix3d::Identity();
  state_.omega = Eigen::Vector3d::Zero();
  state_.motor_rpm = Eigen::Array4d::Zero();
  updateInternalState();

  input_ = Eigen::Array4d::Zero();
}

void Quadrotor::step(double dt)
{
  odeint::integrate(boost::ref(*this), internal_state_, 0.0, dt, dt);

  for(int i = 0; i < 3; i++)
  {
    state_.x(i) = internal_state_[0+i];
    state_.v(i) = internal_state_[3+i];
    state_.R(i,0) = internal_state_[6+i];
    state_.R(i,1) = internal_state_[9+i];
    state_.R(i,2) = internal_state_[12+i];
    state_.omega(i) = internal_state_[15+i];
  }
  state_.motor_rpm(0) = internal_state_[18];
  state_.motor_rpm(1) = internal_state_[19];
  state_.motor_rpm(2) = internal_state_[20];
  state_.motor_rpm(3) = internal_state_[21];

  // Re-orthonormalize R (polar decomposition)
  Eigen::LLT<Eigen::Matrix3d> llt(state_.R.transpose()*state_.R);
  Eigen::Matrix3d P = llt.matrixL();
  Eigen::Matrix3d R = state_.R*P.inverse();
  state_.R = R;

  // Don't go below zero, simulate floor
  if(state_.x(2) < 0.0 && state_.v(2) < 0)
  {
    state_.x(2) = 0;
    state_.v(2) = 0;
  }
  updateInternalState();
}

void Quadrotor::operator()(const Quadrotor::InternalState &x, Quadrotor::InternalState &dxdt, const double /* t */)
{
  State cur_state;
  for(int i = 0; i < 3; i++)
  {
    cur_state.x(i) = x[0+i];
    cur_state.v(i) = x[3+i];
    cur_state.R(i,0) = x[6+i];
    cur_state.R(i,1) = x[9+i];
    cur_state.R(i,2) = x[12+i];
    cur_state.omega(i) = x[15+i];
  }
  for(int i = 0; i < 4; i++)
  {
    cur_state.motor_rpm(i) = x[18 + i];
  }

  // Re-orthonormalize R (polar decomposition)
  Eigen::LLT<Eigen::Matrix3d> llt(cur_state.R.transpose()*cur_state.R);
  Eigen::Matrix3d P = llt.matrixL();
  Eigen::Matrix3d R = cur_state.R*P.inverse();

  Eigen::Vector3d x_dot, v_dot, omega_dot;
  Eigen::Matrix3d R_dot;
  Eigen::Array4d motor_rpm_dot;
  Eigen::Array4d motor_rpm_sq;
  Eigen::Matrix3d omega_hat(Eigen::Matrix3d::Zero());

  omega_hat(2,1) = cur_state.omega(0);
  omega_hat(1,2) = -cur_state.omega(0);
  omega_hat(0,2) = cur_state.omega(1);
  omega_hat(2,0) = -cur_state.omega(1);
  omega_hat(1,0) = cur_state.omega(2);
  omega_hat(0,1) = -cur_state.omega(2);

  motor_rpm_sq = cur_state.motor_rpm.square();

  double thrust = kf_*motor_rpm_sq.sum();
  Eigen::Vector3d moments;
  moments(0) = kf_*(motor_rpm_sq(2) - motor_rpm_sq(3)) * arm_length_;
  moments(1) = kf_*(motor_rpm_sq(1) - motor_rpm_sq(0)) * arm_length_;
  moments(2) = km_*(motor_rpm_sq(0) + motor_rpm_sq(1) - motor_rpm_sq(2) - motor_rpm_sq(3));

  x_dot = cur_state.v;
  v_dot = -Eigen::Vector3d(0,0,g_) + thrust*R.col(2)/mass_ + external_force_/mass_;
  R_dot = R*omega_hat;
  omega_dot = J_.inverse()*(moments - cur_state.omega.cross(J_*cur_state.omega) + external_moment_);
  motor_rpm_dot = (input_ - cur_state.motor_rpm)/motor_time_constant_;

  for(int i = 0; i < 3; i++)
  {
    dxdt[0+i] = x_dot(i);
    dxdt[3+i] = v_dot(i);
    dxdt[6+i] = R_dot(i,0);
    dxdt[9+i] = R_dot(i,1);
    dxdt[12+i] = R_dot(i,2);
    dxdt[15+i] = omega_dot(i);
  }
  for(int i = 0; i < 4; i++)
  {
    dxdt[18 + i] = motor_rpm_dot(i);
  }
}

void Quadrotor::setInput(double u1, double u2, double u3, double u4)
{
  input_(0) = u1;
  input_(1) = u2;
  input_(2) = u3;
  input_(3) = u4;
  for(int i = 0; i < 4; i++)
  {
    if(input_(i) > max_rpm_)
      input_(i) = max_rpm_;
    else if(input_(i) < min_rpm_)
      input_(i) = min_rpm_;
  }
}

const Quadrotor::State &Quadrotor::getState(void) const
{
  return state_;
}
void Quadrotor::setState(const Quadrotor::State &state)
{
  state_.x = state.x;
  state_.v = state.v;
  state_.R = state.R;
  state_.omega = state.omega;
  state_.motor_rpm = state.motor_rpm;

  updateInternalState();
}

double Quadrotor::getMass(void) const
{
  return mass_;
}
void Quadrotor::setMass(double mass)
{
  mass_ = mass;
}

double Quadrotor::getGravity(void) const
{
  return g_;
}
void Quadrotor::setGravity(double g)
{
  g_ = g;
}

const Eigen::Matrix3d &Quadrotor::getInertia(void) const
{
  return J_;
}
void Quadrotor::setInertia(const Eigen::Matrix3d &inertia)
{
  if(inertia != inertia.transpose())
  {
    std::cerr << "Inertia matrix not symmetric, not setting" << std::endl;
    return;
  }
  J_ = inertia;
}

double Quadrotor::getArmLength(void) const
{
  return arm_length_;
}
void Quadrotor::setArmLength(double d)
{
  if(d <= 0)
  {
    std::cerr << "Arm length <= 0, not setting" << std::endl;
    return;
  }

  arm_length_ = d;
}

double Quadrotor::getPropRadius(void) const
{
  return prop_radius_;
}
void Quadrotor::setPropRadius(double r)
{
  if(r <= 0)
  {
    std::cerr << "Prop radius <= 0, not setting" << std::endl;
    return;
  }
  prop_radius_ = r;
}

double Quadrotor::getPropellerThrustCoefficient(void) const
{
  return kf_;
}
void Quadrotor::setPropellerThrustCoefficient(double kf)
{
  if(kf <= 0)
  {
    std::cerr << "Thrust coefficient <= 0, not setting" << std::endl;
    return;
  }

  kf_ = kf;
}

double Quadrotor::getPropellerMomentCoefficient(void) const
{
  return km_;
}
void Quadrotor::setPropellerMomentCoefficient(double km)
{
  if(km <= 0)
  {
    std::cerr << "Moment coefficient <= 0, not setting" << std::endl;
    return;
  }

  km_ = km;
}

double Quadrotor::getMotorTimeConstant(void) const
{
  return motor_time_constant_;
}
void Quadrotor::setMotorTimeConstant(double k)
{
  if(k <= 0)
  {
    std::cerr << "Motor time constant <= 0, not setting" << std::endl;
    return;
  }

  motor_time_constant_ = k;
}

const Eigen::Vector3d &Quadrotor::getExternalForce(void) const
{
  return external_force_;
}
void Quadrotor::setExternalForce(const Eigen::Vector3d &force)
{
  external_force_ = force;
}

const Eigen::Vector3d &Quadrotor::getExternalMoment(void) const
{
  return external_moment_;
}
void Quadrotor::setExternalMoment(const Eigen::Vector3d &moment)
{
  external_moment_ = moment;
}

double Quadrotor::getMaxRPM(void) const
{
  return max_rpm_;
}
void Quadrotor::setMaxRPM(double max_rpm)
{
  if(max_rpm <= 0)
  {
    std::cerr << "Max rpm <= 0, not setting" << std::endl;
    return;
  }
  max_rpm_ = max_rpm;
}

double Quadrotor::getMinRPM(void) const
{
  return min_rpm_;
}
void Quadrotor::setMinRPM(double min_rpm)
{
  if(min_rpm < 0)
  {
    std::cerr << "Min rpm < 0, not setting" << std::endl;
    return;
  }
  min_rpm_ = min_rpm;
}

void Quadrotor::updateInternalState(void)
{
  for(int i = 0; i < 3; i++)
  {
    internal_state_[0+i] = state_.x(i);
    internal_state_[3+i] = state_.v(i);
    internal_state_[6+i] = state_.R(i,0);
    internal_state_[9+i] = state_.R(i,1);
    internal_state_[12+i] = state_.R(i,2);
    internal_state_[15+i] = state_.omega(i);
  }
  internal_state_[18] = state_.motor_rpm(0);
  internal_state_[19] = state_.motor_rpm(1);
  internal_state_[20] = state_.motor_rpm(2);
  internal_state_[21] = state_.motor_rpm(3);
}

}


// Std includes
#include <cmath>
#include <iostream>
#include <memory>
// Thirdparties includes
#include <Eigen/Dense>
// Lib includes
#include "s0s/runge_kutta_fehlberg.h"
#include "sl0/point.h"
#include "sl0/chain/static.h"
// Simple includes
#include "flow.h"

using TypeScalar = double;
// State
template<int StateSize>
using TypeState = Eigen::Matrix<TypeScalar, StateSize, 1>;
using TypeStateDynamic = Eigen::Matrix<TypeScalar, Eigen::Dynamic, 1>;
// Space
constexpr unsigned int DIM = 3;
using TypeVector = Eigen::Matrix<TypeScalar, DIM, 1>;
// Ref and View
template<typename ...Args>
using TypeRef = Eigen::Ref<Args...>;
template<typename ...Args>
using TypeView = Eigen::Map<Args...>;
// Group Parameters
using TypeStepPoint = sl0::StepPoint<TypeState, DIM, TypeRef, TypeView, Flow>;
// Solver
using TypeSolver = s0s::SolverRungeKuttaFehlberg;

const unsigned int np = 10;

int main () {
    // Parameters
    TypeVector x0 = TypeVector({0.0, 0.0, 0.0});
    TypeScalar t0 = 0.0;
    TypeScalar dt = 1e-3;
    unsigned int nt = std::round(1.0 / dt);
    double l = 1.0;
    // Create chain
    std::shared_ptr<TypeStepPoint> sStepPoint = std::make_shared<TypeStepPoint>(std::make_shared<Flow>());
    sl0::ChainStatic<TypeState, DIM, TypeRef, np, TypeView, TypeStepPoint, TypeSolver> chain(sStepPoint, l, 4);
    // Init
    for(std::size_t i = 0; i < np; i++) {
        sStepPoint->x(chain.sStep->memberState(chain.state, i)) = x0;
        sStepPoint->x(chain.sStep->memberState(chain.state, i))[0] += i * l/double(np-1);
    }
    std::cout << "Init Length : " << "\n" << chain.sStep->actualLength(chain.state) << "\n";
    std::cout << "Init State : " << "\n" << chain.state << "\n";
    chain.t = t0;
    // Computation
    std::cout << "Computing" << "\n";
    for(std::size_t i = 0; i < nt; i++) {
        chain.update(dt);
    }
    // out
    std::cout << "\n";
    std::cout << "Chain advected following an exponential flow, exp(" << chain.t << ") = " << "\n";
    std::cout << "\n";
    std::cout << "Final State : " << "\n" << chain.state << "\n";
    std::cout << "Final Length : " << "\n" << chain.sStep->actualLength(chain.state) << "\n";
    std::cout << std::endl;
}


#include "TensorCoreLinesEvaluator.hh"

#include <Eigen/Geometry>

#include <boost/algorithm/cxx11/any_of.hpp>
#include <boost/range/algorithm/min_element.hpp>
#include <boost/range/algorithm/max_element.hpp>

#include <vector>
#include <tuple>
#include <utility>

using namespace cpp_utils;

namespace tl
{

template <typename T, std::size_t... Degrees>
using TPBT = TensorProductBezierTriangle<T, double, Degrees...>;

std::pair<std::array<TPBT<double, 1, 2>, 3>, std::array<TPBT<double, 0, 3>, 3>>
tensorCoreLinesCoeffs(const TensorInterp& t,
                      const std::array<TensorInterp, 3>& dt,
                      const Triangle& r)
{
    using Coords1 = TPBT<double, 1, 2>::Coords;
    using Coords2 = TPBT<double, 0, 3>::Coords;

    // (T * r) x r
    auto eval_ev =
            [&](const Coords1& coords, const TensorInterp& t, int i) -> double {
        auto rv = r(coords.tail<3>());
        auto tv = t(coords.head<3>());
        return (tv * rv).cross(rv)[i];
    };

    // ((\nabla T * r) * r) x r
    auto eval_deriv_ev = [&](const Coords2& coords,
                             const std::array<TensorInterp, 3>& dt,
                             int i) -> double {
        auto rv = r(coords.tail<3>());
        auto txv = dt[0](coords.head<3>());
        auto tyv = dt[1](coords.head<3>());
        auto tzv = dt[2](coords.head<3>());

        return ((txv * rv[0] + tyv * rv[1] + tzv * rv[2]) * rv).cross(rv)[i];
    };

    auto eval1 = [&](const Coords1& coords) { return eval_ev(coords, t, 0); };
    auto eval2 = [&](const Coords1& coords) { return eval_ev(coords, t, 1); };
    auto eval3 = [&](const Coords1& coords) { return eval_ev(coords, t, 2); };

    auto eval4 = [&](const Coords2& coords) {
        return eval_deriv_ev(coords, dt, 0);
    };
    auto eval5 = [&](const Coords2& coords) {
        return eval_deriv_ev(coords, dt, 1);
    };
    auto eval6 = [&](const Coords2& coords) {
        return eval_deriv_ev(coords, dt, 2);
    };

    return {{TPBT<double, 1, 2>{eval1},
             TPBT<double, 1, 2>{eval2},
             TPBT<double, 1, 2>{eval3}},
            {TPBT<double, 0, 3>{eval4},
             TPBT<double, 0, 3>{eval5},
             TPBT<double, 0, 3>{eval6}}};
}

using TSHE = TensorCoreLinesEvaluator;

TSHE::TensorCoreLinesEvaluator(const DoubleTri& tri,
                               const TensorInterp& t,
                               const std::array<TensorInterp, 3>& dt,
                               const Options& opts)
        : _tri(tri),
          _opts(opts)
{
    std::tie(_target_funcs_t, _target_funcs_dt) =
            tensorCoreLinesCoeffs(t, dt, tri.dir_tri);
}

std::array<TSHE, 4> TSHE::split() const
{
    if(_last_split_dir)
    {
        return split<0>();
    }
    return split<1>();
}

Result TSHE::eval()
{
    // Check if any of the error components can not become zero in the
    // current subdivision triangles
    auto has_nonzero =
            boost::algorithm::any_of(_target_funcs_t, [](const auto& c) {
                    return sameSign(c.coefficients()) != 0;
            })
            || boost::algorithm::any_of(_target_funcs_dt, [](const auto& c) {
                   return sameSign(c.coefficients()) != 0;
               });

    // Discard triangles if no roots can occur inside
    if(has_nonzero)
    {
        return Result::Discard;
    }

    // Compute upper bound for target functions
    auto max_error = std::max(abs_max_upper_bound(_target_funcs_t),
                              abs_max_upper_bound(_target_funcs_dt));

    if(max_error < _opts.tolerance)
    {
        return Result::Accept;
    }

    return Result::Split;
}

double TSHE::error() const
{
    return std::max(upper_bound_norm(_target_funcs_t),
                    upper_bound_norm(_target_funcs_dt));
}

double distance(const TSHE& t1, const TSHE& t2)
{
    return distance(t1.tris(), t2.tris());
}

bool operator==(const TSHE& t1, const TSHE& t2)
{
    return t1._tri == t2._tri && t1._target_funcs_t == t2._target_funcs_t
           && t1._target_funcs_dt == t2._target_funcs_dt
           && t1._last_split_dir == t2._last_split_dir
           && t1._split_level == t2._split_level && t1._opts == t2._opts;
}

bool operator!=(const TSHE& t1, const TSHE& t2)
{
    return !(t1 == t2);
}

}


#include "mLibInclude.h"
#include "../../shared/OptUtils.h"
#define GLOG_NO_ABBREVIATED_SEVERITIES
#include <Eigen33b1/Eigen>
#include <Eigen33b1/IterativeLinearSolvers>
#include <cuda_runtime.h>

#include <time.h>

#define LEAST_SQ_CONJ_GRADIENT 0
#define SPARSE_QR 1

#define SOLVER LEAST_SQ_CONJ_GRADIENT

typedef Eigen::Triplet<float> Tripf;
typedef Eigen::SparseMatrix<float> SpMatrixf;
#include "EigenSolverPoissonImageEditing.h"
#include <Eigen33b1/OrderingMethods>

#if SOLVER == LEAST_SQ_CONJ_GRADIENT
typedef Eigen::LeastSquaresConjugateGradient<SpMatrixf > AxEqBSolver;
#elif SOLVER == SPARSE_QR
typedef Eigen::SparseQR<SpMatrixf, Eigen::COLAMDOrdering<int> > AxEqBSolver;
#endif

struct vec2iHash {
    size_t operator()(const vec2i& v) const {
        return std::hash < int > {}(v.x) ^ std::hash < int > {}(v.y);
    }
};

float4 sampleImage(float4* image, vec2i p, int W) {
    return image[p.y*W + p.x];
}

void setPixel(float4* image, vec2i p, int W, float r, float g, float b) {
    image[p.y*W + p.x].x = r;
    image[p.y*W + p.x].y = g;
    image[p.y*W + p.x].z = b;
}

void solveAxEqb(AxEqBSolver& solver, const Eigen::VectorXf& b, Eigen::VectorXf& x) {

#if SOLVER==LEAST_SQ_CONJ_GRADIENT
    x = solver.solveWithGuess(b, x);
    //std::cout << "#iterations:     " << solver.iterations() << std::endl;
    //std::cout << "estimated error: " << solver.error() << std::endl;
#else
    x = solver.solve(b);
#endif
}
double EigenSolverPoissonImageEditing::solve(const NamedParameters& solverParameters, const NamedParameters& problemParameters, bool profileSolve, std::vector<SolverIteration>& iters)
{
    int numUnknowns = 0;
    std::unordered_map<vec2i, int, vec2iHash> pixelLocationsToIndex;
    std::vector<vec2i> pixelLocations;
    size_t pixelCount = m_dims[0] * m_dims[1];
    std::vector<float4> h_unknownFloat(pixelCount);
    std::vector<float4> h_target(pixelCount);
    std::vector<float>  h_mask(pixelCount);

    findAndCopyArrayToCPU("X", h_unknownFloat, problemParameters);
    findAndCopyArrayToCPU("T", h_target, problemParameters);
    findAndCopyArrayToCPU("M", h_mask, problemParameters);

    for (int y = 0; y < (int)m_dims[1]; ++y) {
        for (int x = 0; x < (int)m_dims[0]; ++x) {
            if (h_mask[y*m_dims[0] + x] == 0.0f) {
                ++numUnknowns;
                vec2i p(x, y);
                pixelLocationsToIndex[p] =(int)pixelLocations.size();
                pixelLocations.push_back(p);
            }
        }
    }
    printf("# Unknowns: %d\n", numUnknowns);
    int numResiduals = (int)pixelLocations.size() * 4;

    Eigen::VectorXf x_r(numUnknowns), b_r(numResiduals);
    Eigen::VectorXf x_g(numUnknowns), b_g(numResiduals);
    Eigen::VectorXf x_b(numUnknowns), b_b(numResiduals);
    Eigen::VectorXf x_a(numUnknowns), b_a(numResiduals);

    b_r.setZero();
    b_g.setZero();
    b_b.setZero();
    b_a.setZero();

    for (int i = 0; i < pixelLocations.size(); ++i) {
        vec2i p = pixelLocations[i];
        float4 color = sampleImage(h_unknownFloat.data(), p, m_dims[0]);
        x_r[i] = color.x;
        //printf("%f\n", color.x);
        x_g[i] = color.y;
        x_b[i] = color.z;
        x_a[i] = color.w;
    }
    SpMatrixf A(numResiduals, numUnknowns);
    A.setZero();
    printf("Constructing Matrix\n");
    std::vector<Tripf> entriesA;

    std::vector<vec2i> offsets;
    offsets.push_back(vec2i(-1, 0));
    offsets.push_back(vec2i(1, 0));
    offsets.push_back(vec2i(0, -1));
    offsets.push_back(vec2i(0, 1));

    for (int i = 0; i < pixelLocations.size(); ++i) {
        vec2i p = pixelLocations[i];
        int numInternalNeighbors = 0;
        float4 g_p = sampleImage(h_target.data(), p, m_dims[0]);
        int j = 0;

        for (vec2i off : offsets) {
            vec2i q = p + off;
            if (q.x >= 0 && q.y >= 0 && q.x < (int)m_dims[0] && q.y < (int)m_dims[1]) {
                auto it = pixelLocationsToIndex.find(q);
                int row = 4 * i + j;
                if (it == pixelLocationsToIndex.end()) {
                    float4 f_q = sampleImage(h_unknownFloat.data(), q, m_dims[0]);
                    b_r[row] += f_q.x;
                    b_g[row] += f_q.y;
                    b_b[row] += f_q.z;
                    b_a[row] += f_q.w;
                }
                else {
                    entriesA.push_back(Tripf(row, it->second, -1.0f));
                }
                entriesA.push_back(Tripf(row, i, 1.0f));

                float4 g_q = sampleImage(h_target.data(), q, m_dims[0]);
                b_r[row] += (g_p.x - g_q.x);
                b_g[row] += (g_p.y - g_q.y);
                b_b[row] += (g_p.z - g_q.z);
                b_a[row] += (g_p.w - g_q.w);
            }
            ++j;

        }
    }

    printf("Entries Set\n");
    A.setFromTriplets(entriesA.begin(), entriesA.end());
    printf("Sparse Matrix Constructed\n");
    A.makeCompressed();
    printf("Matrix Compressed\n");
    {
        float totalCost = 0.0f;

        float cost_r = (A*x_r - b_r).squaredNorm();
        float cost_g = (A*x_g - b_g).squaredNorm();
        float cost_b = (A*x_b - b_b).squaredNorm();
        float cost_a = (A*x_a - b_a).squaredNorm();
        totalCost = cost_r + cost_g + cost_b + cost_a;
        printf("Initial Cost: %f : (%f, %f, %f, %f)\n", totalCost, cost_r, cost_g, cost_b, cost_a);

    }

    AxEqBSolver solver;
    solver.setMaxIterations(97);
    printf("Solvers Initialized\n");

    clock_t start = clock(), diff;

    solver.compute(A);
    //printf("solver.compute(A)\n");
    solveAxEqb(solver, b_r, x_r);
    //printf("Red solve done\n");
    solveAxEqb(solver, b_g, x_g);
    //printf("Green solve done\n");
    solveAxEqb(solver, b_b, x_b);
    //printf("Blue solve done\n");
    solveAxEqb(solver, b_a, x_a);

    diff = clock() - start;
    printf("Time taken %f ms\n", diff*1000.0 / double(CLOCKS_PER_SEC));

    float totalCost = 0.0f;

    float cost_r = (A*x_r - b_r).squaredNorm();
    float cost_g = (A*x_g - b_g).squaredNorm();
    float cost_b = (A*x_b - b_b).squaredNorm();
    float cost_a = (A*x_a - b_a).squaredNorm();
    totalCost = cost_r + cost_g + cost_b + cost_a;
    printf("Final Cost: %f : (%f, %f, %f, %f)\n", totalCost, cost_r, cost_g, cost_b, cost_a);

    for (int i = 0; i < pixelLocations.size(); ++i) {
        setPixel(h_unknownFloat.data(), pixelLocations[i], m_dims[0], x_r[i], x_g[i], x_b[i]);
    }
    findAndCopyToArrayFromCPU("X", h_unknownFloat, problemParameters);;
    return (double)totalCost;

}

/* Proper Poisson Image Editing

for (int i = 0; i < pixelLocations.size(); ++i) {
vec2i p = pixelLocations[i];
int row = i;
int numInternalNeighbors = 0;
float4 g_p = sampleImage(h_target, p, m_width);
for (int off_y = -1; off_y <= 1; off_y += 2) {
for (int off_x = -1; off_x <= 1; off_x += 2) {
vec2i q(p.x + off_x, p.y + off_y);
auto it = pixelLocationsToIndex.find(q);
if (it != pixelLocationsToIndex.end()) {
++numInternalNeighbors;
entriesA.push_back(Tripf(row, it->second, -1.0f));
} else {

float4 f_star_q = sampleImage(h_target, q, m_width);
b_r[i] += f_star_q.x;
b_g[i] += f_star_q.y;
b_b[i] += f_star_q.z;
}
float4 g_q = sampleImage(h_target, q, m_width);
b_r[i] += (g_p.x - g_q.x);
b_g[i] += (g_p.y - g_q.y);
b_b[i] += (g_p.z - g_q.z);
}
}
entriesA.push_back(Tripf(row, row, (float)numInternalNeighbors));
}

*/


// Copyright (c) 2021 CNES
//
// All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.
#pragma once
#include <pybind11/eigen.h>

#include <Eigen/Core>
#include <optional>
#include <tuple>
#include <unordered_map>

#include "pyinterp/detail/broadcast.hpp"
#include "pyinterp/detail/math.hpp"
#include "pyinterp/eigen.hpp"
#include "pyinterp/geodetic/box.hpp"
#include "pyinterp/geodetic/point.hpp"
#include "pyinterp/geodetic/polygon.hpp"

namespace pyinterp::geohash::int64 {

/// Returns the precision in longitude/latitude and degrees for the given
/// precision
[[nodiscard]] inline auto error_with_precision(const uint32_t precision)
    -> std::tuple<double, double> {
  auto lat_bits = static_cast<int32_t>(precision >> 1U);
  auto lng_bits = static_cast<int32_t>(precision - lat_bits);

  return std::make_tuple(360 * detail::math::power2(-lng_bits),
                         180 * detail::math::power2(-lat_bits));
}

// Encode a point into geohash with the given precision
[[nodiscard]] auto encode(const geodetic::Point& point, uint32_t precision)
    -> uint64_t;

// Encode points into geohash with the given precision
[[nodiscard]] inline auto encode(const Eigen::Ref<const Eigen::VectorXd>& lon,
                                 const Eigen::Ref<const Eigen::VectorXd>& lat,
                                 uint32_t precision) -> Vector<uint64_t> {
  detail::check_eigen_shape("lon", lon, "lat", lat);
  auto size = lon.size();
  auto result = Vector<uint64_t>(size);
  for (Eigen::Index ix = 0; ix < size; ++ix) {
    result(ix) = encode({lon[ix], lat[ix]}, precision);
  }
  return result;
}

// Returns the region encoded by the integer geohash with the specified
// precision.
[[nodiscard]] auto bounding_box(uint64_t hash, uint32_t precision)
    -> geodetic::Box;

// Decode a hash into a geographic point with the given precision.
// If round is true, the coordinates of the points will be rounded to the
// accuracy defined by the GeoHash.
[[nodiscard]] inline auto decode(const uint64_t hash, const uint32_t precision,
                                 const bool round) -> geodetic::Point {
  auto bbox = bounding_box(hash, precision);
  return round ? bbox.round() : bbox.centroid();
}

// Decode hashes into a geographic points with the given bit depth.
// If round is true, the coordinates of the points will be rounded to the
// accuracy defined by the GeoHash.
[[nodiscard]] inline auto decode(const Eigen::Ref<const Vector<uint64_t>>& hash,
                                 const uint32_t precision, const bool center)
    -> std::tuple<Eigen::VectorXd, Eigen::VectorXd> {
  auto lon = Eigen::VectorXd(hash.size());
  auto lat = Eigen::VectorXd(hash.size());
  auto point = geodetic::Point();
  for (Eigen::Index ix = 0; ix < hash.size(); ++ix) {
    point = decode(hash(ix), precision, center);
    lon[ix] = point.lon();
    lat[ix] = point.lat();
  }
  return std::make_tuple(lon, lat);
}

// Returns all neighbors hash clockwise from north around northwest at the given
// precision.
// 7 0 1
// 6 x 2
// 5 4 3
[[nodiscard]] auto neighbors(uint64_t hash, uint32_t precision)
    -> Eigen::Matrix<uint64_t, 8, 1>;

// Returns the property of the grid covering the given box: geohash of the
// minimum corner point, number of boxes in longitudes and latitudes.
[[nodiscard]] auto grid_properties(const geodetic::Box& box, uint32_t precision)
    -> std::tuple<uint64_t, size_t, size_t>;

// Returns all the GeoHash codes within the box.
[[nodiscard]] auto bounding_boxes(const std::optional<geodetic::Box>& box,
                                  uint32_t precision) -> Vector<uint64_t>;

// Returns all the GeoHash codes within the polygon.
[[nodiscard]] inline auto bounding_boxes(const geodetic::Polygon& polygon,
                                         uint32_t chars) -> Vector<uint64_t> {
  auto box = geodetic::Box();
  boost::geometry::envelope<geodetic::Polygon, geodetic::Box>(polygon, box);
  return bounding_boxes(box, chars);
}

// Returns the area covered by the GeoHash
[[nodiscard]] inline auto area(uint64_t hash, uint32_t precision,
                               const std::optional<geodetic::System>& wgs)
    -> double {
  return bounding_box(hash, precision).area(wgs);
}

// Returns the start and end indexes of the different GeoHash boxes.
[[nodiscard]] auto where(
    const pybind11::EigenDRef<const Eigen::Matrix<uint64_t, -1, -1>>& hash)
    -> std::unordered_map<uint64_t, std::tuple<std::tuple<int64_t, int64_t>,
                                     std::tuple<int64_t, int64_t>>>;

}  // namespace pyinterp::geohash::int64


/*
 * dynamical-system.hpp
 *
 *  Created on: 19 mai 2014
 *      Author: alexis
 */

#ifndef DYNAMICAL_SYSTEM_HPP_
#define DYNAMICAL_SYSTEM_HPP_

#include <vector>

#include <state-observation/api.h>
#include <state-observation/dynamical-system/dynamical-system-functor-base.hpp>
#include <state-observation/noise/noise-base.hpp>
#include <state-observation/sensors-simulation/accelerometer-gyrometer.hpp>
#include <state-observation/tools/hrp2.hpp>
#include <state-observation/tools/rigid-body-kinematics.hpp>

#include <Eigen/Cholesky>

namespace stateObservation
{

namespace flexibilityEstimation
{

/**
 * \class  DynamicalSystem
 * \brief  This class describes the dynamics of a robot's flexibility
 *         this dynamics is the simplest possible system, the flexibility
 *         is expressed as a rotation against the contact positions with no
 *         other hypothesis than that the contact points are at constant position
 *
 */
class STATE_OBSERVATION_DLLAPI IMUElasticLocalFrameDynamicalSystem : public stateObservation::DynamicalSystemFunctorBase
{
public:
  struct input
  {
    /// indices of the different components of a vector of the input state
    static constexpr unsigned posCom = 0;
    static constexpr unsigned velCom = 3;
    static constexpr unsigned accCom = 6;
    static constexpr unsigned inertia = 9;
    static constexpr unsigned angMoment = 15;
    static constexpr unsigned dotInertia = 18;
    static constexpr unsigned dotAngMoment = 24;
    static constexpr unsigned posIMU = 27;
    static constexpr unsigned oriIMU = 30;
    static constexpr unsigned linVelIMU = 33;
    static constexpr unsigned angVelIMU = 36;
    static constexpr unsigned linAccIMU = 39;
    static constexpr unsigned additionalForces = 42;
    static constexpr unsigned contacts = 48;

    static constexpr unsigned sizeBase = 48;
  };

  struct state
  {
    static constexpr unsigned pos = 0;
    static constexpr unsigned ori = 3;
    static constexpr unsigned linVel = 6;
    static constexpr unsigned angVel = 9;
    static constexpr unsigned fc = 12;
    static constexpr unsigned unmodeledForces = 24;
    static constexpr unsigned comBias = 30;
    static constexpr unsigned drift = 32;

    static constexpr unsigned size = 35;
  };

  struct contactModel
  {
    /// indices of the different components of a vector of the input state
    static constexpr unsigned elasticContact = 1;
    static constexpr unsigned pendulum = 2;
    static constexpr unsigned none = 0;
  };

  typedef Eigen::LLT<Matrix3> LLTMatrix3;

  /// constructor
  explicit IMUElasticLocalFrameDynamicalSystem(double dt);

  /// virtual destructor
  virtual ~IMUElasticLocalFrameDynamicalSystem();

  void test();

  // Get the contact wrench
  void computeContactWrench(const Matrix3 & orientation,
                            const Vector3 & position,
                            const IndexedVectorArray & contactPosV,
                            const IndexedVectorArray & contactOriV,
                            const Vector & fc,
                            const Vector & tc,
                            const Vector3 & fm,
                            const Vector3 & tm,
                            const Vector3 & addForce,
                            const Vector3 & addMoment);

  stateObservation::Vector computeAccelerations(const Vector & x, const Vector & u);

  // computation of the acceleration linear
  virtual void computeAccelerations(const Vector3 & positionCom,
                                    const Vector3 & velocityCom,
                                    const Vector3 & accelerationCom,
                                    const Vector3 & AngMomentum,
                                    const Vector3 & dotAngMomentum,
                                    const Matrix3 & Inertia,
                                    const Matrix3 & dotInertia,
                                    const IndexedVectorArray & contactPos,
                                    const IndexedVectorArray & contactOri,
                                    const Vector3 & position,
                                    const Vector3 & linVelocity,
                                    Vector3 & linearAcceleration,
                                    const Vector3 & oriVector,
                                    const Matrix3 & orientation,
                                    const Vector3 & angularVel,
                                    Vector3 & angularAcceleration,
                                    const Vector & fc,
                                    const Vector & tc,
                                    const Vector3 & fm,
                                    const Vector3 & tm,
                                    const Vector3 & addForces,
                                    const Vector3 & addMoments);

  /// Description of the state dynamics
  virtual stateObservation::Vector stateDynamics(const stateObservation::Vector & x,
                                                 const stateObservation::Vector & u,
                                                 TimeIndex k);

  /// compute the jacobien of the state dynamics at the last computed value
  stateObservation::Matrix stateDynamicsJacobian();

  /// compute the jacobien of the state dynamics at a given state
  stateObservation::Matrix stateDynamicsJacobian(const stateObservation::Vector & x,
                                                 const stateObservation::Vector & u,
                                                 TimeIndex k);

  /// sets the finite differences derivation step vector
  void setFDstep(const stateObservation::Vector & dx);

  /// Description of the sensor's dynamics
  virtual stateObservation::Vector measureDynamics(const stateObservation::Vector & x,
                                                   const stateObservation::Vector & u,
                                                   TimeIndex k);

  /// compute the Jacobien of the measurements dynamics at the last computed value
  stateObservation::Matrix measureDynamicsJacobian();

  /// compute the Jacobien of the measurements dynamics at a given state value
  stateObservation::Matrix measureDynamicsJacobian(const stateObservation::Vector & x,
                                                   const stateObservation::Vector & u,
                                                   TimeIndex k);

  /// Sets a noise which disturbs the state dynamics
  virtual void setProcessNoise(stateObservation::NoiseBase *);

  /// Removes the process noise
  virtual void resetProcessNoise();

  /// Gets the process noise
  virtual stateObservation::NoiseBase * getProcessNoise() const;

  /// Sets a noise which disturbs the measurements
  virtual void setMeasurementNoise(stateObservation::NoiseBase *);

  /// Removes the measurement noise
  virtual void resetMeasurementNoise();

  /// Gets a pointer on the measurement noise
  virtual stateObservation::NoiseBase * getMeasurementNoise() const;

  /// Set the period of the time discretization
  virtual void setSamplingPeriod(double dt);

  /// Gets the state size
  virtual Index getStateSize() const;

  /// Gets the input size
  virtual Index getInputSize() const;

  /// Sets the input size
  virtual void setInputSize(Index i);

  /// Gets the contact number
  /// virtual

  /// Gets the contacts position

  /// Gets the measurement size
  virtual Index getMeasurementSize() const;

  /// Sets the number of contacts
  virtual void setContactsNumber(unsigned);

  virtual void setPe(stateObservation::Vector3 Pe)
  {
    pe = Pe;
  }

  /// Gets the nimber of contacts
  inline unsigned getContactsNumber(void) const
  {
    return nbContacts_;
  }

  virtual void setContactModel(unsigned nb);

  virtual void setPrinted(bool b)
  {
    printed_ = b;
  }

  virtual void computeElastContactForcesAndMoments(const IndexedVectorArray & contactPosArray,
                                                   const IndexedVectorArray & contactOriArray,
                                                   const IndexedVectorArray & contactVelArray,
                                                   const IndexedVectorArray & contactAngVelArray,
                                                   const Vector3 & position,
                                                   const Vector3 & linVelocity,
                                                   const Vector3 & oriVector,
                                                   const Matrix3 & orientation,
                                                   const Vector3 & angVel,
                                                   Vector & fc,
                                                   Vector & tc);

  virtual void computeElastPendulumForcesAndMoments(const IndexedVectorArray & PrArray,
                                                    const Vector3 & position,
                                                    const Vector3 & linVelocity,
                                                    const Vector3 & oriVector,
                                                    const Matrix3 & orientation,
                                                    const Vector3 & angVel,
                                                    Vector & forces,
                                                    Vector & moments);

  void computeForcesAndMoments(const IndexedVectorArray & position1,
                               const IndexedVectorArray & position2,
                               const IndexedVectorArray & velocity1,
                               const IndexedVectorArray & velocity2,
                               const Vector3 & position,
                               const Vector3 & linVelocity,
                               const Vector3 & oriVector,
                               const Matrix3 & orientation,
                               const Vector3 & angVel,
                               Vector & fc,
                               Vector & tc);

  virtual void computeForcesAndMoments(const Vector & x, const Vector & u);

  virtual Vector getForcesAndMoments();

  virtual Vector getForcesAndMoments(const Vector & x, const Vector & u);

  virtual Vector getMomentaDotFromForces(const Vector & x, const Vector & u);
  virtual Vector getMomentaDotFromKinematics(const Vector & x, const Vector & u);

  virtual void iterateDynamicsEuler(const Vector3 & positionCom,
                                    const Vector3 & velocityCom,
                                    const Vector3 & accelerationCom,
                                    const Vector3 & AngMomentum,
                                    const Vector3 & dotAngMomentum,
                                    const Matrix3 & Inertia,
                                    const Matrix3 & dotInertia,
                                    const IndexedVectorArray & contactPos,
                                    const IndexedVectorArray & contactOri,
                                    Vector3 & position,
                                    Vector3 & linVelocity,
                                    Vector & fc1,
                                    Vector3 & oriVector,
                                    Vector3 & angularVel,
                                    Vector & fc2,
                                    const Vector3 & fm,
                                    const Vector3 & tm,
                                    const Vector3 & addForces,
                                    const Vector3 & addMoments,
                                    double dt);

  virtual void iterateDynamicsRK4(const Vector3 & positionCom,
                                  const Vector3 & velocityCom,
                                  const Vector3 & accelerationCom,
                                  const Vector3 & AngMomentum,
                                  const Vector3 & dotAngMomentum,
                                  const Matrix3 & Inertia,
                                  const Matrix3 & dotInertia,
                                  const IndexedVectorArray & contactPos,
                                  const IndexedVectorArray & contactOri,
                                  Vector3 & position,
                                  Vector3 & linVelocity,
                                  Vector & fc1,
                                  Vector3 & oriVector,
                                  Vector3 & angularVel,
                                  Vector & fc2,
                                  const Vector3 & fm,
                                  const Vector3 & tm,
                                  const Vector3 & addForces,
                                  const Vector3 & addMoments,
                                  double dt);

  virtual void setWithForceMeasurements(bool b);
  virtual bool getWithForceMeasurements() const;
  virtual void setWithComBias(bool b);
  virtual bool getWithComBias() const;
  virtual void setWithAbsolutePosition(bool b);
  virtual bool getWithAbsolutePosition() const;
  void setWithUnmodeledForces(bool b);

  virtual void setKfe(const Matrix3 & m);
  virtual void setKfv(const Matrix3 & m);
  virtual void setKte(const Matrix3 & m);
  virtual void setKtv(const Matrix3 & m);

  virtual void setKfeRopes(const Matrix3 & m);
  virtual void setKfvRopes(const Matrix3 & m);
  virtual void setKteRopes(const Matrix3 & m);
  virtual void setKtvRopes(const Matrix3 & m);

  virtual Matrix getKfe() const;
  virtual Matrix getKfv() const;
  virtual Matrix getKte() const;
  virtual Matrix getKtv() const;

  virtual void setRobotMass(double d);

  virtual double getRobotMass() const;

protected:
  bool printed_;

  stateObservation::AccelerometerGyrometer sensor_;

  stateObservation::NoiseBase * processNoise_;

  void updateMeasurementSize_();

  double dt_;

  double robotMass_;
  double robotMassInv_;

  Matrix3 & computeRotation_(const Vector3 & x, int i);

  static constexpr Index stateSize_ = state::size;
  Index inputSize_;
  static constexpr Index measurementSizeBase_ = 6;
  unsigned nbContacts_;
  unsigned contactModel_;

  Vector fc_;
  Vector tc_;

  Vector dx_;

  Vector xk1_;
  Vector xk_;
  Vector uk_;

  Vector xk_fory_;
  Vector yk_;
  Vector uk_fory_;

  Index measurementSize_;

  std::vector<Vector3, Eigen::aligned_allocator<Vector3>> contactPositions_;

  Matrix3 Kfe_, Kte_, Kfv_, Ktv_;
  Matrix3 KfeRopes_, KteRopes_, KfvRopes_, KtvRopes_;

  TimeIndex kcurrent_;

  bool withForceMeasurements_;
  bool withComBias_;
  bool withAbsolutePos_;
  bool withUnmodeledForces_;

  stateObservation::Vector3 pe;

  double marginalStabilityFactor_;
  // a scaling factor a=1-epsilon to avoid the natural marginal stability of
  // the dynamics x_{k+1}=x_k we replace it with x_{k+1}=a*x_k

  unsigned index_;

  struct Optimization
  {

    Vector6 momentaDot;

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    Vector3 positionFlex;
    Vector3 velocityFlex;
    Vector3 accelerationFlex;
    Vector3 orientationFlexV;
    Vector3 angularVelocityFlex;
    Vector3 angularAccelerationFlex;
    Vector3 positionComBias;

    Matrix3 rFlex;
    Matrix3 rFlexT;

    Vector3 drift;
    Vector3 pdrift;
    Matrix3 rdrift;

    double cy, sy;

    AngleAxis orientationAA;

    Vector xk1;
    Vector xk;
    Vector xk1dx;

    Vector xdx;

    Vector xk_fory;
    Vector yk;
    Vector ykdy;

    TimeIndex k_fory;

    Matrix3 orinertia;

    LLTMatrix3 invinertia;

    Matrix3 rtotal;
    Vector3 ptotal;
    AngleAxis aatotal;
    Vector3 oritotal;

    Matrix Jx;
    Matrix Jy;

    Matrix3 rimu;
    Vector3 imuAcc;
    Vector3 imuOmega;
    Vector sensorState;

    Vector3 positionCom;
    Vector3 velocityCom;
    Vector3 accelerationCom;
    Vector3 AngMomentum;
    Vector3 dotAngMomentum;

    Vector3 positionControl;
    Vector3 velocityControl;
    Vector3 accelerationControl;
    Vector3 orientationControlV;
    Vector3 angularVelocityControl;

    Matrix3 rControl;

    IndexedVectorArray contactPosV;
    IndexedVectorArray contactOriV;
    IndexedVectorArray contactVelArray;
    IndexedVectorArray contactAngVelArray;

    Matrix3 inertia;
    Matrix3 dotInertia;

    IndexedVectorArray efforts;

    Vector3 f, fi;
    Vector3 t;

    // unmodelled and unmeasured forces
    Vector3 fm;
    Vector3 tm;

    Vector3 linearAcceleration;
    Vector3 angularAcceleration;

    Vector3 vf;
    Vector3 vt;

    Vector3 crosstempV;
    Matrix3 crosstempM;

    // elastic contact forces and moments
    Matrix3 Rci; // rotation of contact i
    Matrix3 Rcit; // transpose of previous
    Vector3 contactPos; //
    Vector3 contactVel;
    Vector3 RciContactPos;
    Vector3 globalContactPos;
    Matrix3 Rt;

    Vector3 forcei;
    Vector3 momenti;

    // additional forces and moments
    Vector3 addForce;
    Vector3 addMoment;

    Matrix3 skewV;
    Matrix3 skewV2;
    Matrix3 skewVR;
    Matrix3 skewV2R;
    Matrix3 RIRT;
    Vector3 wx2Rc;
    Vector3 _2wxRv;
    Vector3 Ra;
    Vector3 Rc;
    Vector3 Rcp;

    // optimization of orientation transformation between vector3 to rotation matrix

    Matrix3 curRotation0;
    Vector3 orientationVector0;
    Matrix3 curRotation1;
    Vector3 orientationVector1;
    Matrix3 curRotation2;
    Vector3 orientationVector2;
    Matrix3 curRotation3;
    Vector3 orientationVector3;

    Optimization()
    : curRotation0(Matrix3::Identity()), orientationVector0(Vector3::Zero()), curRotation1(Matrix3::Identity()),
      orientationVector1(Vector3::Zero()), curRotation2(Matrix3::Identity()), orientationVector2(Vector3::Zero()),
      curRotation3(Matrix3::Identity()), orientationVector3(Vector3::Zero())
    {
    }

    inline Vector3 & orientationVector(int i)
    {
      if(i == 0) return orientationVector0;
      if(i == 1) return orientationVector1;
      if(i == 2) return orientationVector2;

      return orientationVector3;
    }

    inline Matrix3 & curRotation(int i)
    {
      if(i == 0) return curRotation0;
      if(i == 1) return curRotation1;
      if(i == 2) return curRotation2;

      return curRotation3;
    }

  } op_;

public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
};
} // namespace flexibilityEstimation
} // namespace stateObservation

#endif /* DYNAMICAL_SYSTEM_HPP_ */


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

#pragma once

#include <cstddef>
#include <cmath>

#include <algorithm>
#include <functional>
#include <type_traits>
#include <utility>

#include <sentinel/config.hpp>

#include <boost/geometry/core/access.hpp>
#include <boost/geometry/core/coordinate_dimension.hpp>
#include <boost/geometry/core/coordinate_type.hpp>
#include <boost/geometry/core/cs.hpp>
#include <boost/geometry/geometries/register/point.hpp>

namespace sentinel { namespace math {

template<class T, std::size_t N>
struct vector {
    using value_type      = T;
    using size_type       = std::size_t;
    using difference_type = std::ptrdiff_t;
    using reference       = value_type&;
    using const_reference = const value_type&;
    using pointer         = value_type*;
    using const_pointer   = const value_type*;
    using iterator        = pointer;
    using const_iterator  = const_pointer;

    static_assert(N > 0);
    value_type array[N];

    static inline constexpr vector zero = vector(/*ZERO INITIALIZED*/);

    constexpr reference       operator[](size_type pos)       { return array[pos]; }
    constexpr const_reference operator[](size_type pos) const { return array[pos]; }

    constexpr pointer       data() noexcept       { return array; }
    constexpr const_pointer data() const noexcept { return array; }

    static
    constexpr size_type size() noexcept { return N; }

    constexpr iterator begin() noexcept { return array; }
    constexpr iterator end()   noexcept { return begin() + size(); }

    constexpr const_iterator cbegin() const noexcept { return array; }
    constexpr const_iterator cend()   const noexcept { return begin() + size(); }

    constexpr const_iterator begin() const noexcept { return cbegin(); }
    constexpr const_iterator end()   const noexcept { return cend(); }

    constexpr void fill(const T& value) { for (auto& v : *this) v = value; }

    constexpr bool operator==(const vector& other) const {
        return std::mismatch(begin(), end(), other.begin()).first == end();
    }

    constexpr bool operator!=(const vector& other) const {
        return !(*this == other);
    }

    static constexpr vector filled(const T& value)
    {
        auto return_fill_value = [&value] (auto&&...) -> const T& { return value; };
        return [&return_fill_value] <std::size_t... I> (std::index_sequence<I...>) {
            return vector{return_fill_value(I)...};
        } (std::make_index_sequence<N>{});
    }

    template<class BinaryOp>
    constexpr vector& pointwise_transform(const vector& other, BinaryOp op)
    {
        auto other_it = other.begin();
        for (auto it = begin(), end = this->end(); it != end; ++it, ++other_it)
             *it = op(*it, *other_it);
        return *this;
    }

    constexpr vector& operator+=(const vector& other) { return pointwise_transform(other, std::plus{}); }
    constexpr vector& operator-=(const vector& other) { return pointwise_transform(other, std::minus{}); }
    constexpr vector& operator*=(const vector& other) { return pointwise_transform(other, std::multiplies{}); }
    constexpr vector& operator/=(const vector& other) { return pointwise_transform(other, std::divides{}); }

    constexpr vector& operator*=(const T& scalar) { return (*this) *= filled(scalar); }

    constexpr vector operator+(const vector& other) const { vector result = *this; result += other; return result; }
    constexpr vector operator-(const vector& other) const { vector result = *this; result -= other; return result; }
    constexpr vector operator*(const vector& other) const { vector result = *this; result *= other; return result; }
    constexpr vector operator/(const vector& other) const { vector result = *this; result /= other; return result; }

    constexpr vector operator-() const;
};

template<class T, std::size_t N>
constexpr vector<T, N> vector<T, N>::operator-() const
{
    vector res = *this;
    for (auto& v : res)
        v = -v;
    return res;
}

/** \brief Computes the dot product of two vectors.
 */
template<class T, std::size_t N>
constexpr T dot(const vector<T, N>& a, const vector<T, N>& b)
{
    T result = static_cast<T>(0);
    for (std::size_t pos = 0; pos < N; ++pos)
        result += a[pos] * b[pos];
    return result;
}

/** \brief Computes the cross product of two vector.
 */
template<class T>
constexpr vector<T, 3> cross(const vector<T, 3>& a, const vector<T, 3>& b)
{
    return {
        a[1] * b[2] - a[2] * b[1],
        a[2] * b[0] - a[0] * b[2],
        a[0] * b[1] - a[1] * b[0]
    };
}

/** \brief Computes the square of the norm of a vector.
 */
template<class T, std::size_t N>
constexpr float norm2(const vector<T, N>& vec) { return dot(vec, vec); }

/** \brief Computes the norm of a vector.
 */
template<class T, std::size_t N>
constexpr float norm(const vector<T, N>& vec) { return std::sqrt(norm2(vec)); }

/** \brief Computes a normalized vector.
 *
 * If the norm of \a vec is less than or equal to SENTINEL_VECTOR_SMALL_NORM,
 * the returned vector is the `0` vector.
 */
template<class T, std::size_t N>
constexpr vector<T, N> normalized(const vector<T, N>& vec)
{
    const auto length = norm(vec);
    return length > SENTINEL_VECTOR_SMALL_NORM ? (1 / length) * vec
                                               : vector<T, N>::zero;
}

/** \brief Performs scalar-vector multiplication.
 */
template<class T, std::size_t N>
constexpr vector<T, N> operator*(const T& scalar, const vector<T, N>& vec)
{
    return vector<T, N>::filled(scalar) * vec;
}

/** \brief Performs scalar-vector multiplication.
 */
template<class T, std::size_t N>
constexpr vector<T, N> operator*(const vector<T, N>& vec, const T& scalar)
{
    return vec * vector<T, N>::filled(scalar);
}

template<class T, std::size_t N, class BinaryOp>
constexpr vector<T, N> pointwise_compose(const vector<T, N>& a,
                                         const vector<T, N>& b,
                                         BinaryOp op)
{
    vector<T, N> result = a;
    a.pointwise_accumulate(a, std::cref(op));
    return result;
}

using fvec2 = vector<float, 2>;
using fvec3 = vector<float, 3>;

using dvec2 = vector<double, 2>;
using dvec3 = vector<double, 3>;

} } // namespace sentinel::math

BOOST_GEOMETRY_REGISTER_POINT_2D(sentinel::math::fvec2,
                                 float,
                                 boost::geometry::cs::cartesian,
                                 operator[](0),
                                 operator[](1)
                                 );

BOOST_GEOMETRY_REGISTER_POINT_3D(sentinel::math::fvec3,
                                 float,
                                 boost::geometry::cs::cartesian,
                                 operator[](0),
                                 operator[](1),
                                 operator[](2)
                                 );

BOOST_GEOMETRY_REGISTER_POINT_2D(sentinel::math::dvec2,
                                 double,
                                 boost::geometry::cs::cartesian,
                                 operator[](0),
                                 operator[](1)
                                 );

BOOST_GEOMETRY_REGISTER_POINT_3D(sentinel::math::dvec3,
                                 double,
                                 boost::geometry::cs::cartesian,
                                 operator[](0),
                                 operator[](1),
                                 operator[](2)
                                 );


/** @file
 * @brief NPDE WaveABC2D
 * @author Erick Schulz
 * @date 11/12/2019
 * @copyright Developed at ETH Zurich
 */

#include "waveabc2d.h"

#include <iostream>
#include <memory>
#include <string>
// Eigen includes
#include <Eigen/Core>
// Lehrfem++ includes
#include <lf/assemble/assemble.h>
#include <lf/io/io.h>
#include <lf/mesh/hybrid2d/hybrid2d.h>
#include <lf/mesh/utils/utils.h>
#include <lf/uscalfe/uscalfe.h>

using namespace WaveABC2D;

int main(int /*argc*/, const char ** /*argv*/) {
  std::cout << "\n" << std::endl;
  std::cout << "PROBLEM - WaveABC2D" << std::endl;

  testConvergenceScalarImplicitTimestepping();

  // Load mesh into a Lehrfem++ object
  std::cout << "Loading mesh..." << std::endl;
  auto mesh_factory = std::make_unique<lf::mesh::hybrid2d::MeshFactory>(2);
  const lf::io::GmshReader reader(
      std::move(mesh_factory), CURRENT_SOURCE_DIR "/../meshes/unitsquare2.msh");
  auto mesh_p = reader.mesh();  // type shared_ptr< const lf::mesh::Mesh>

  // Finite element space
  auto fe_space_p =
      std::make_shared<lf::uscalfe::FeSpaceLagrangeO1<double>>(mesh_p);
  // Obtain local->global index mapping for current finite element space
  const lf::assemble::DofHandler &dofh{fe_space_p->LocGlobMap()};
  // Dimension of finite element space
  const lf::uscalfe::size_type N_dofs(dofh.NumDofs());

  auto mu0 = [](const Eigen::Vector2d &x) -> double {
    return std::sin(x.norm());
  };
  auto nu0 = [](const Eigen::Vector2d &x) -> double { return std::cos(x(1)); };
  auto rho = [](Eigen::Vector2d) -> double { return 1.0; };

  WaveABC2DTimestepper<decltype(rho), decltype(mu0), decltype(nu0)> stepper(
      fe_space_p, rho, 250, 1.0);
  Eigen::VectorXd discrete_solution = stepper.solveWaveABC2D(mu0, nu0);

  double discrete_energy = stepper.energies();

  // Output results to vtk file
  lf::io::VtkWriter vtk_writer(mesh_p, "WaveABC2D_solution.vtk");
  // Write nodal data taking the values of the discrete solution at the vertices
  auto nodal_data = lf::mesh::utils::make_CodimMeshDataSet<double>(mesh_p, 2);
  for (int global_idx = 0; global_idx < N_dofs; global_idx++) {
    nodal_data->operator()(dofh.Entity(global_idx)) =
        discrete_solution[global_idx];
  };
  vtk_writer.WritePointData("WaveABC2D_solution", *nodal_data);

  std::cout << "\nThe WaveABC2D_solution was written to:" << std::endl;
  std::cout << "WaveABC2D_solution.vtk\n" << std::endl;

  std::cout << "The discrete energies E^(k) : " << discrete_energy << std::endl;
  return 0;
}  // main


/*
This file is a part of Raman-Scattering-Code-Conversion.
<https://github.com/Kirbologist/Raman-Scattering-Code-Conversion>

Written by Siwan Li for the UQ School of Maths and Physics.
Based on the SMARTIES MATLAB package by W.R.C. Somerville, B. Auguié, E.C. Le Ru
Copyright (C) 2021-2022 Siwan Li

This source code form is subject to the terms of the MIT License.
If a copy of the MIT License was not distributed with this file,
you can obtain one at <https://opensource.org/licenses/MIT>.

Various small maths functions, helper functions and others used throughout SMARTIES functions that don't fit anywhere.
*/

#ifndef SMARTIES_MATH_HPP
#define SMARTIES_MATH_HPP

#include "core.hpp"
#include <Eigen/LU>
#include <boost/math/special_functions/bessel.hpp>
#include <boost/math/constants/constants.hpp>

using namespace Eigen;
using namespace std;

namespace Smarties {
  /* Returns the mathematical constant pi with the amount of precision specified by template argument */
  template <class Real>
  inline Real mp_pi(void) { return boost::math::constants::pi<Real>(); }

  /* Returns the machine epsilon of the type given in the template argument */
  template <class Real>
  inline Real mp_eps(void) { return numeric_limits<Real>::epsilon(); }

  /* Returns the imaginary unit 'i' with the type given in the template argument */
  template <class Real>
  inline complex<Real> mp_im_unit(void) { return complex<Real>(0.0, 1.0); }

  /*
  View an existing Eigen::Tensor of rank 2 as an Eigen::Map<Eigen::Matrix>
  Rows/Cols are determined from the matrix
  */
  template<typename Scalar>
  auto ArrayMap(Tensor<Scalar, 2>& tensor) {
      return Eigen::Map<ArrayXXr<Scalar>>(tensor.data(), tensor.dimension(0), tensor.dimension(1));
  }

  /*
  Converts an Eigen::Matrix (or expression) to Eigen::Tensor
  with dimensions specified in std::array
  */
  template<typename Derived, typename T, auto rank>
  Tensor<typename Derived::Scalar, rank>
  TensorCast(const EigenBase<Derived>& matrix, const std::array<T, rank>& dims) {
      return Eigen::TensorMap<const Tensor<const typename Derived::Scalar, rank>>
                  (matrix.derived().eval().data(), dims);
  }

  /*
  Converts an Eigen::Matrix (or expression) to Eigen::Tensor
  with dimensions as variadic arguments
  */
  template<typename Derived, typename... Dims>
  auto TensorCast(const EigenBase<Derived>& matrix, const Dims... dims) {
      static_assert(sizeof...(Dims) > 0, "TensorCast: sizeof... (Dims) must be larger than 0");
      return TensorCast(matrix, std::array<Eigen::Index, sizeof...(Dims)>{dims...});
  }

  /*
  Converts an Eigen::Matrix (or expression) to Eigen::Tensor
  with dimensions directly as arguments in a variadic template
  */
  template<typename Derived>
  auto TensorCast(const EigenBase<Derived>& matrix) {
    if constexpr(Derived::ColsAtCompileTime == 1 or Derived::RowsAtCompileTime == 1) {
      return TensorCast(matrix, matrix.size());
    } else {
      return TensorCast(matrix, matrix.rows(), matrix.cols());
    }
  }

  /* Get a readable name/description of the type given in the template argument */
  template <class Real>
  string GetTypeName() {
    string calc_type = typeid(static_cast<Real>(0)).name();
    if (calc_type.find("N5boost14multiprecision6numberINS0_8backends18mpfr_float_backend") != string::npos) {
      size_t prec_begin = calc_type.find("ILj") + 3;
      size_t prec_end = calc_type.find("ELNS");
      string precision = calc_type.substr(prec_begin, prec_end - prec_begin);
      calc_type = "custom_" + precision;
    } else if (calc_type == "f")
      calc_type = "single";
    else if (calc_type == "d")
      calc_type = "double";
    else if (calc_type == "e")
      calc_type = "quad";
    return calc_type;
  }

  /*
  Get a slice of a rank-3 tensor, so that the indices of the slices are given by Eigen::ArithmeticSequence.
  This is just a helper function for sphCheckBesselConvergence, and is not usable for general purpose.
  */
  template <class Real>
  Tensor3c<Real> TensorSlice(Tensor3c<Real>& tensor,
      ArithmeticSequence<long int, long int, long int> slice_dim1,
      ArithmeticSequence<long int, long int, long int> slice_dim2,
      ArithmeticSequence<long int, long int, long int> slice_dim3) {
    long int new_dim1 = slice_dim1.size();
    long int new_dim2 = slice_dim2.size();
    long int new_dim3 = slice_dim3.size();
    Tensor3c<Real> output(new_dim1, new_dim2, new_dim3);
    for (long int i = 0; i < new_dim1; i++) {
      for (long int j = 0; j < new_dim2; j++) {
        for (long int k = 0; k < new_dim3; k++)
          output(i, j, k) = tensor(slice_dim1[i], slice_dim2[j], slice_dim3[k]);
      }
    }
    return output;
  }

  /*
  Initialize a linearly-spaced single-column Eigen::Array by converting from an Eigen::ArithmeticSequence.
  Parameters are the same ones used to construct an Eigen::ArithmeticSequence.
  */
  ArrayXi Seq2Array(long int first, long int last, long int stride);

  /*
  Reduce a complex rank-3 Eigen::Tensor into a complex Eigen::Array by
  reducing the third dimension of the tensor down to only coefficients whose third index is `offset`,
  and reducing the first dimenstion down to only the last `num_rows` rows.
  This was used specifically as a helper function in sphCalculatePQ.
  */
  template <class Real>
  ArrayXXc<Real> Subtensor2ArrMap(const Tensor3c<Real>& tensor,
      const std::array<int, 3>& offsets, const std::array<int, 3>& extents, int rows, int cols) {
    ArrayXXc<Real> output(rows, cols);
    for (int i = 0; i < rows; i++) {
      for (int j = 0; j < cols; j++)
        output(i, j) = tensor(offsets[0] + i, offsets[1], offsets[2] + j);
    }
    return output;
  }

  /* Performs matrix inversion using LU with partial pivoting of columns. This is used in rvhGetTRfromPQ. */
  template <class Real>
  ArrayXXc<Real> InvertLUcol(MatrixXc<Real>& M) {
    PartialPivLU<MatrixXc<Real>> PLU_decomp = M.matrix().lu();
    return PLU_decomp.inverse().array();
  }

  /* Performs the equivalent of MATLAB's logical indexing on a single-column Eigen::Array. */
  template <class Real>
  ArrayXr<Real> LogicalSlice(ArrayXr<Real>& base_array, ArrayXb& bool_array) {
    assert(base_array.size() == bool_array.size());
    int output_size = bool_array.count();
    ArrayXr<Real> output(output_size);
    int base_index = 0;
    for (int i = 0; i < output_size && base_index < bool_array.size(); i++, base_index++) {
      while (!bool_array(base_index))
        base_index++;
      output(i) = base_array(base_index);
    }
    return output;
  }

  /* Performs the equivalent of MATLAB's logical indexing on a single-row Eigen::Array. */
  template <class Real>
  RowArrayXr<Real> LogicalSlice(RowArrayXr<Real>& base_array, RowArrayXb& bool_array) {
    assert(base_array.size() == bool_array.size());
    int output_size = bool_array.count();
    RowArrayXr<Real> output(output_size);
    int base_index = 0;
    for (int i = 0; i < output_size && base_index < bool_array.size(); i++, base_index++) {
      while (!bool_array(base_index))
        base_index++;
      output(i) = base_array(base_index);
    }
    return output;
  }

  /* Gets the indices of all 'true' components of a single-column boolean Eigen::Array. */
  ArrayXi LogicalIndices(ArrayXb& bool_array);

  /* Take a tensor and return the same tensor, but with all its coefficients conjugated. */
  template <class Real>
  Tensor4c<Real> TensorConj(Tensor4c<Real>& base) {
    Tensor4c<Real> output = base;
    for (int i = 0; i < base.dimension(0); i++) {
      for (int j = 0; j < base.dimension(1); j++) {
        for (int k = 0; k < base.dimension(2); k++) {
          for (int l = 0; l < base.dimension(3); l++)
            output(i, j, k, l) = conj(base(i, j, k, l));
        }
      }
    }
    return output;
  }

  /* Calculates the Bessel function of the first kind with many values of `nu` and a fixed value of `x`. */
  template <class Real>
  ArrayXr<Real> ArrBesselJ(ArrayXr<Real>& nu, Real x) {
    ArrayXr<Real> output(nu.size());
    for (int i = 0; i < nu.size(); i++)
      output(i) = boost::math::cyl_bessel_j<Real>(nu(i), x);
    return output;
  }

  /* Calculates the Bessel function of the second kind with many values of `nu` and a fixed value of `x`. */
  template <class Real>
  ArrayXr<Real> ArrBesselY(ArrayXr<Real>& nu, Real x) {
    ArrayXr<Real> output(nu.size());
    for (int i = 0; i < nu.size(); i++)
      output(i) = boost::math::cyl_neumann<Real>(nu(i), x);
    return output;
  }
}

#endif


// 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__EKF_HPP_
#define SMOOTH__FEEDBACK__EKF_HPP_

#include <Eigen/Cholesky>
#include <boost/numeric/odeint.hpp>
#include <smooth/compat/odeint.hpp>
#include <smooth/diff.hpp>
#include <smooth/lie_group.hpp>

namespace smooth::feedback {

/**
 * @brief Extended Kalman filter on Lie groups.
 *
 * The primary methods are predict() and update().
 * - predict(): propagates filter state through a dynamical model.
 * - update(): Bayesian update of filter state from a measurement.
 *
 * Use this class for information fusion (without dynamics) by solely using update().
 *
 * @tparam G \p smooth::LieGroup type.
 * @tparam DiffType \p smooth::diff::Type method for calculating derivatives.
 * @tparam Stpr \p boost::numeric::odeint templated stepper type (\p euler / \p runge_kutta4 /
 * ...). Defaults to \p euler.
 */
template<
  LieGroup G,
  diff::Type DiffType                 = diff::Type::Default,
  template<typename...> typename Stpr = boost::numeric::odeint::euler>
  requires(Dof<G> > 0)
class EKF
{
public:
  //! Scalar type for computations.
  //! Degrees of freedom.
  using CovT = Eigen::Matrix<Scalar<G>, Dof<G>, Dof<G>>;

  /**
   * @brief Reset the state of the EKF.
   *
   * @param g filter value
   * @param P filter covariance
   */
  void reset(const G & g, const CovT & P)
  {
    g_hat_ = g;
    P_     = P;
  }

  /**
   * @brief Access filter state estimate.
   */
  G estimate() const { return g_hat_; }

  /**
   * @brief Access filter covariance.
   */
  CovT covariance() const { return P_; }

  /**
   * @brief Propagate EKF through dynamics \f$ \mathrm{d}^r x_t = f(t, x) \f$ with covariance
   * \f$Q\f$ over a time interval \f$ [0, \tau] \f$.
   *
   * @param f right-hand side \f$ f : \mathbb{R} \times \mathbb{G} \rightarrow \mathbb{R}^{\dim
   * \mathfrak{g}} \f$ of the dynamics. The time type must be the scalar type of G.
   * @param Q process covariance (size \f$ \dim \mathfrak{g} \times \dim \mathfrak{g} \f$)
   * @param tau amount of time to propagate
   * @param dt maximal ODE solver step size (defaults to \p tau, i.e. one step)
   *
   * @note The time \f$ t \f$ argument of \f$ f(t, x) \f$ ranges over the interval \f$t  \in [0,
   * \tau] \f$.
   *
   * @note The covariance \f$ Q \f$ is infinitesimal, i.e. its SI unit is \f$ S^2/T \f$
   * where \f$S\f$ is the unit of state and \f$T\f$ is the unit of time.
   *
   * @note Only the upper triangular part of Q is used.
   */
  template<typename F, typename QDer>
  void predict(
    F && f, const Eigen::MatrixBase<QDer> & Q, Scalar<G> tau, std::optional<Scalar<G>> dt = {})
  {
    const auto state_ode = [&f](const G & g, Tangent<G> & dg, Scalar<G> t) { dg = f(t, g); };

    const auto cov_ode = [this, &f, &Q](const CovT & cov, CovT & dcov, Scalar<G> t) {
      const auto f_x = [&f, &t]<typename _T>(const CastT<_T, G> & x) -> Tangent<CastT<_T, G>> {
        return f(t, x);
      };
      const auto [fv, dr] = diff::dr<1, DiffType>(f_x, wrt(g_hat_));
      const CovT A        = -ad<G>(fv) + dr;
      dcov = (A * cov + cov * A.transpose() + Q).template selfadjointView<Eigen::Upper>();
    };

    Scalar<G> t          = 0;
    const Scalar<G> dt_v = dt.value_or(2 * tau);
    while (t + dt_v < tau) {
      // step covariance first since it depends on g_hat_
      cst_.do_step(cov_ode, P_, t, dt_v);
      sst_.do_step(state_ode, g_hat_, t, dt_v);
      t += dt_v;
    }

    // last step up to time t
    cst_.do_step(cov_ode, P_, t, tau - t);
    sst_.do_step(state_ode, g_hat_, t, tau - t);
  }

  /**
   * @brief Update EKF with a measurement \f$y = h(x) + w\f$ where \f$w \sim \mathcal N(0, R)\f$.
   *
   * @param h measurement function \f$ h : \mathbb{G} \rightarrow \mathbb{Y} \f$
   * @param y measurement value \f$ y \in \mathbb{Y} \f$
   * @param R measurement covariance (size \f$ \dim \mathbb{Y} \times \dim \mathbb{Y} \f$)
   *
   * @note The function h must be differentiable using the desired method.
   *
   * @note Only the upper triangular part of \f$R\f$ is used
   */
  template<typename F, typename RDev, Manifold Y = std::invoke_result_t<F, G>>
  void update(F && h, const Y & y, const Eigen::MatrixBase<RDev> & R)
  {
    const auto [hval, H] = diff::dr<1, DiffType>(h, wrt(g_hat_));

    using Result = std::decay_t<decltype(hval)>;

    static_assert(Manifold<Result>, "h(x) is not a Manifold");

    static constexpr Eigen::Index Ny = Dof<Result>;

    static_assert(Ny > 0, "h(x) must be statically sized");

    const Eigen::Matrix<Scalar<G>, Ny, Ny> S =
      (H * P_.template selfadjointView<Eigen::Upper>() * H.transpose() + R)
        .template triangularView<Eigen::Upper>();

    // solve for Kalman gain
    const Eigen::Matrix<Scalar<G>, Dof<G>, Ny> K =
      S.template selfadjointView<Eigen::Upper>().ldlt().solve(H * P_).transpose();

    // update estimate and covariance
    g_hat_ += K * (y - hval);
    P_ = ((CovT::Identity() - K * H) * P_).template selfadjointView<Eigen::Upper>();
  }

private:
  // filter estimate and covariance
  G g_hat_ = Default<G>();
  CovT P_  = CovT::Identity();

  // steppers for numerical ODE solutions
  Stpr<G, Scalar<G>, Tangent<G>, Scalar<G>, boost::numeric::odeint::vector_space_algebra> sst_{};
  Stpr<CovT, Scalar<G>, CovT, Scalar<G>, boost::numeric::odeint::vector_space_algebra> cst_{};
};

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__EKF_HPP_


#include<stdio.h>
#include"mex.h"

#include<ceres/ceres.h>
using namespace std;

#include <Eigen/Core>
#include <Eigen/Geometry>

struct NORMAL_COST
{
  NORMAL_COST(Eigen::Matrix<double, 1, 3> parent_normal, Eigen::Matrix<double, 1, 3> child_normal):_parent_normal(parent_normal),_child_normal(child_normal){}
  template <typename T>
  bool operator()(const T *child_to_parent_q, T *residual)const
  {
    // get the transformation from the child_to_parent_q [qw qx qy qz]
    // Quaternion [w, x, y, z]
    Eigen::Quaternion<T> child_to_parent_q_(child_to_parent_q[0], child_to_parent_q[1], child_to_parent_q[2], child_to_parent_q[3]);
    Eigen::Matrix<T, 3, 3> child_to_parent_rotationalMatrix = child_to_parent_q_.toRotationMatrix();

    // get the vector from the parent_normal and child_normal [nx ny nz]
    Eigen::Matrix<T, 3, 1> parent_normal_v(T(_parent_normal(0, 0)), T(_parent_normal(0, 1)), T(_parent_normal(0, 2)));

    Eigen::Matrix<T, 3, 1> child_normal_v(T(_child_normal(0, 0)), T(_child_normal(0, 1)), T(_child_normal(0, 2)));

    Eigen::Matrix<T, 3, 1> _residual_normal_v = child_to_parent_rotationalMatrix * child_normal_v - parent_normal_v;

    residual[0] = _residual_normal_v(0, 0);
    residual[1] = _residual_normal_v(1, 0);
    residual[2] = _residual_normal_v(2, 0);

    #ifdef DEBUG
      mexPrintf("child_normal_v: [%f, %f, %f]\n", _child_normal(0, 0), _child_normal(0, 1), _child_normal(0, 2));
      mexPrintf("parent_normal_v: [%f, %f, %f]\n", _parent_normal(0, 0), _parent_normal(0, 1), _parent_normal(0, 2));
      mexPrintf("child_to_parent_rotationalMatrix: \n[%f, %f, %f]\n[%f, %f, %f]\n[%f, %f, %f]\n", \
              child_to_parent_rotationalMatrix(0, 0), child_to_parent_rotationalMatrix(0, 1), child_to_parent_rotationalMatrix(0, 2), \
              child_to_parent_rotationalMatrix(1, 0), child_to_parent_rotationalMatrix(1, 1), child_to_parent_rotationalMatrix(1, 2), \
              child_to_parent_rotationalMatrix(2, 0), child_to_parent_rotationalMatrix(2, 1), child_to_parent_rotationalMatrix(2, 2) \
              );
      mexPrintf("residual: [%f, %f, %f]\n", residual[0], residual[1], residual[2]);
    #endif

    return true;
  }
  const Eigen::Matrix<double, 1, 3> _parent_normal, _child_normal;
};

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]){
    // nlhs represent the number of parameters of the output
    // plhs is a array of the mxarray pointers, each pointing to the output
    // nrhs represents the number of parameters of the input
    // prhs is a array of the mxarray pointers, each pointing to the input

    // prhs[0], nx3 double, parent_normal, [nx ny nz]
    // prhs[1], nx3 double, child_normal, [nx ny nz]
    // prhs[2], 1x4 double, child_to_parent_q, [qx qy qz qw]
    // prhs[3], 1x1 double, verbose

    if(nrhs < 4){
        mexErrMsgIdAndTxt( "estimatePitchRollByCoNormalCeresMex:invalidNumInputs", "at least 4 input arguments required");
        return;
    }

    /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // get the parent_normals
    const size_t *dimArrayOfParentNormals = mxGetDimensions(prhs[0]);
    size_t sizeRowsParentNormals = *(dimArrayOfParentNormals + 0);
    size_t sizeColsParentNormals = *(dimArrayOfParentNormals + 1);

    if(sizeColsParentNormals != 3){
        mexErrMsgIdAndTxt( "estimatePitchRollByCoNormalCeresMex:invalidInputs", "the 1st param should be Nx3");
        return;
    }

    double *ptrParentNormals = (double *)(mxGetPr(prhs[0]));
    std::vector<Eigen::Matrix<double, 1, 3>>parent_normals_all;
    for(int i = 0; i < sizeRowsParentNormals; i++){
        Eigen::Matrix<double, 1, 3> current_parent_normal;
        for(int j = 0; j < 3; j++){
            current_parent_normal(0, j) = ptrParentNormals[j * sizeRowsParentNormals + i];
        }
        parent_normals_all.push_back(current_parent_normal);
    }

    /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // get the child_normals
    const size_t *dimArrayOfChildNormals = mxGetDimensions(prhs[1]);
    size_t sizeRowsChildNormals = *(dimArrayOfChildNormals + 0);
    size_t sizeColsChildNormals = *(dimArrayOfChildNormals + 1);
    if(sizeColsChildNormals != 3){
        mexErrMsgIdAndTxt( "estimatePitchRollByCoNormalCeresMex:invalidInputs", "the 2st param should be Nx3");
        return;
    }

    double *ptrChildNormals = (double *)(mxGetPr(prhs[1]));
    std::vector<Eigen::Matrix<double, 1, 3>>child_normals_all;
    for(int i = 0; i < sizeRowsChildNormals; i++){
        Eigen::Matrix<double, 1, 3> current_child_normal;
        for(int j = 0; j < 3; j++){
            current_child_normal(0, j) = ptrChildNormals[j * sizeRowsChildNormals + i];
        }
        child_normals_all.push_back(current_child_normal);
    }

    if(sizeRowsParentNormals != sizeRowsChildNormals){
        mexErrMsgIdAndTxt( "estimatePitchRollByCoNormalCeresMex:invalidInputs", "the 1st param and the 2st param should have the same size");
        return;
    }

    #ifdef DEBUG
        mexPrintf("observation number: %d\n", sizeRowsParentNormals);
    #endif

    /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // get the init_child_to_parent_q
    double child_to_parent_q[4] = {0};
    const size_t *dimArrayOfChildToParentQ = mxGetDimensions(prhs[2]);
    size_t sizeRowsChildToParentQ = *(dimArrayOfChildToParentQ + 0);
    size_t sizeColsChildToParentQ = *(dimArrayOfChildToParentQ + 1);
    if(sizeColsChildToParentQ != 4 && sizeRowsChildToParentQ != 1){
        mexErrMsgIdAndTxt( "estimatePitchRollByCoNormalCeresMex:invalidInputs", "the 3st param should be 1x4");
        return;
    }
    double *ptrChildToParentQ = (double *)(mxGetPr(prhs[2]));
    for(int i = 0; i < 4; i++){
        child_to_parent_q[i] = ptrChildToParentQ[i];
    }

    #ifdef DEBUG
        mexPrintf("init child_to_parent_q: [%.4f,%.4f,%.4f,%.4f]\n", \
                    child_to_parent_q[0], \
                    child_to_parent_q[1], \
                    child_to_parent_q[2], \
                    child_to_parent_q[3]);
    #endif

    /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // get the verbose
    bool verbose = false;
    const size_t *dimArrayOfVerbose = mxGetDimensions(prhs[3]);
    size_t sizeRowsVerbose = *(dimArrayOfVerbose + 0);
    size_t sizeColsVerbose = *(dimArrayOfVerbose + 1);
    if(sizeRowsVerbose != 1 && sizeColsVerbose != 1){
        mexErrMsgIdAndTxt( "estimatePitchRollByCoNormalCeresMex:invalidInputs", "the 4st param should be 1x1");
        return;
    }
    double *ptrVerbose = (double *)(mxGetPr(prhs[3]));
    verbose = bool(ptrVerbose[0]);

    /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    ceres::Problem problem;
    for(int i=0; i<sizeRowsParentNormals; i++)
    {
        problem.AddResidualBlock(
          new ceres::AutoDiffCostFunction<NORMAL_COST, 3, 4>(
              new NORMAL_COST(parent_normals_all[i], child_normals_all[i])
          ),
          NULL,
          child_to_parent_q
        );
    }

    #ifdef DEBUG
        mexPrintf("AddResidualBlock done!\n");
    #endif

    //配置求解器并求解，输出结果
    ceres::Solver::Options options;
    options.linear_solver_type=ceres::DENSE_QR;
    options.minimizer_progress_to_stdout=verbose;
    ceres::Solver::Summary summary;
    ceres::Solve(options,&problem,&summary);

    #ifdef DEBUG
        mexPrintf("estimated child_to_parent_q: [%.4f,%.4f,%.4f,%.4f]\n", \
                    child_to_parent_q[0], \
                    child_to_parent_q[1], \
                    child_to_parent_q[2], \
                    child_to_parent_q[3]);
    #endif

    // the output q will be 1x4
    size_t dimArrayOfParams[2] = { 1, 4 };
    plhs[0] = mxCreateNumericArray(2, dimArrayOfParams, mxDOUBLE_CLASS, mxREAL);
    double *out_params = (double *)mxGetData(plhs[0]);

    for(int i = 0; i < 4; i++){
        out_params[i] = child_to_parent_q[i];
    }
}


#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "vnl/vnl_matrix_ref.h"
#include "vnl/vnl_matrix.h"
#include "itkNumericTraits.h"

#include "tkdCmdParser.h"

#include <iostream>
#include <utility>
#include <fstream>

#include <boost/utility.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/graph_utility.hpp>
#include <boost/graph/betweenness_centrality.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/adjacency_matrix.hpp>
#include <boost/graph/kruskal_min_spanning_tree.hpp>

/**
 * Minimum spanning tree.
 */
class MST
{
public:

	typedef double PixelType;

	typedef boost::property< boost::vertex_index_t, unsigned int, boost::property< boost::vertex_centrality_t, PixelType > >
			VertexPropertyType;
	typedef boost::property< boost::edge_weight_t, PixelType, boost::property< boost::edge_centrality_t, PixelType > > EdgePropertyType;
	typedef boost::adjacency_matrix< boost::undirectedS, VertexPropertyType, EdgePropertyType > GraphMatrixType;
	typedef boost::adjacency_list< boost::vecS, boost::vecS, boost::undirectedS, boost::no_property, boost::property< boost::edge_weight_t,
			PixelType > > GraphListType;

	typedef itk::Image< PixelType, 2 > ImageType;
	typedef itk::ImageFileWriter< ImageType > WriterType;
	typedef vnl_matrix_ref< PixelType > DataMatrixType;
	typedef vnl_vector< PixelType > VectorType;

public:

	/**
	 * Run.
	 */
	void Run( const std::string& inputFileName,
			PixelType lowerThreshold, PixelType upperThreshold, const std::string& outputFileName )
	{
		WriterType::Pointer writer = WriterType::New();

		writer->SetFileName( outputFileName.c_str() );
		writer->SetInput( GetMinimumSpanningTree( inputFileName, lowerThreshold, upperThreshold  ) );

		try
		{
			writer->Update();
		} catch ( itk::ExceptionObject& e )
		{
			std::cerr << "Error writing: " << outputFileName << std::endl;
			std::cerr << e.GetDescription() << std::endl;
		}

	}

protected:

	/**
	 * Return minimum spanning tree as matrix image.
	 */
	ImageType::Pointer GetMinimumSpanningTree( const std::string& filename, PixelType lowerThreshold, PixelType upperThreshold )
	{

		//Example();

		typedef itk::ImageFileReader< ImageType > ReaderType;
		ReaderType::Pointer reader = ReaderType::New();
		reader->SetFileName( filename.c_str() );
		reader->Update();

		ImageType::Pointer image = reader->GetOutput();
		reader = 0;

		PixelType* buffer = image->GetPixelContainer()->GetBufferPointer();
		ImageType::RegionType region = image->GetLargestPossibleRegion();
		ImageType::SizeType size = region.GetSize();
		int rows = size[0];
		int cols = size[1];

		DataMatrixType data( rows, cols, buffer );

		GraphMatrixType g( rows );

		for ( int i = 0; i < rows; ++i )
		{
			for ( int j = i + 1; j < cols; ++j )
			{
				if ( i == j )
				{
					continue;
				}

				PixelType weight = data( i, j );

				if ( weight <= upperThreshold && weight > 0 && weight > lowerThreshold )
				{
					// invert weight ( e.g. correlation coefficients ... )
					boost::add_edge( i, j, EdgePropertyType( 1.0 / weight ), g );
				}
			}
		}

		// [ 2 ]: Calculate minimum spanning tree ...

		typedef boost::graph_traits< GraphMatrixType >::edge_descriptor EdgeType;
		std::vector< EdgeType > spanningTree;

		// get spanning tree ...
		kruskal_minimum_spanning_tree( g, std::back_inserter( spanningTree ) );

		GraphMatrixType spanningTreeGraph();

		// [ 3 ]: Write minimum spanning tree to image ...

		ImageType::Pointer output = ImageType::New();
		output->CopyInformation( image );
		output->SetRegions( image->GetLargestPossibleRegion() );
		output->Allocate();
		output->FillBuffer( 0 );

		boost::graph_traits< GraphMatrixType >::edge_iterator eiter, eiter_end;
		for ( boost::tie( eiter, eiter_end ) = boost::edges( g ); eiter != eiter_end; ++eiter )
		{
			if ( std::find( spanningTree.begin(), spanningTree.end(), *eiter ) != spanningTree.end() )
			{
				PixelType w = boost::get( boost::edge_weight, g, *eiter );

				ImageType::IndexType index;
				index[0] = source( *eiter, g );
				index[1] = target( *eiter, g );
				output->SetPixel( index, w );

				index[1] = source( *eiter, g );
				index[0] = target( *eiter, g );
				output->SetPixel( index, w );
			}
		}
		return output;
	}
};

/**
 * Main.
 */
int main( int argc, char ** argv )
{
	tkd::CmdParser p( "minimum spanning tree", "Calculate minimum spanning tree on undirected graph" );

	std::string inputFileName;
	std::string outputFileName;

	float lowerThreshold = 0;
	float upperThreshold = 1.0;

	p.AddArgument( inputFileName, "input" ) ->AddAlias( "i" ) ->SetDescription( "Input image: 2D undirected adjacency matrix" ) ->SetRequired( true );

	p.AddArgument( upperThreshold, "threshold" ) ->AddAlias( "thu" ) ->AddAlias( "t" ) ->SetDescription(
			"Threshold; only include paths with weight in (0, threshold] (default: 1.0)" );

	p.AddArgument( lowerThreshold, "threshold-lower" ) ->AddAlias( "thl" ) ->SetDescription(
			"Lower threshold; only include paths with weight > threshold (default: 0)" );

	p.AddArgument( outputFileName, "output" ) ->AddAlias( "o" ) ->SetDescription(
			"Output image: minimum spanning tree with weights" ) ->SetRequired( true );

	if ( !p.Parse( argc, argv ) )
	{
		p.PrintUsage( std::cout );
		return -1;
	}

	MST mst;
	mst.Run( inputFileName, lowerThreshold, upperThreshold, outputFileName );

	return EXIT_SUCCESS;
}


#include <Eigen/Dense>
#include <iostream>

using namespace Eigen;
using namespace std;

int main(int, char**)
{
  cout.precision(3);
  Matrix3f m;
m << 1, 2, 3,
     4, 5, 6,
     7, 8, 9;
std::cout << m;

  return 0;
}