#include <chrono>
#include <cmath>
#include <thread>

#include <ct/optcon/optcon.h>
#include <Eigen/Eigenvalues>

#include "gflags/gflags.h"
#include "third_party/matplotlib-cpp/matplotlibcpp.h"
#include "y2018/control_loops/python/arm_bounds.h"
#include "y2018/control_loops/python/dlqr.h"

DEFINE_double(boundary_scalar, 1500.0, "Test command-line flag");
DEFINE_double(velocity_boundary_scalar, 10.0, "Test command-line flag");
DEFINE_double(boundary_rate, 20.0, "Sigmoid rate");
DEFINE_bool(linear, false, "If true, linear, else see sigmoid.");
DEFINE_bool(sigmoid, false, "If true, sigmoid, else exponential.");
DEFINE_double(round_corner, 0.0, "Corner radius of the constraint box.");
DEFINE_double(convergance, 1e-12, "Residual before finishing the solver.");
DEFINE_double(position_allowance, 5.0,
              "Distance to Velocity at which we have 0 penalty conversion.");
DEFINE_double(bounds_offset, 0.02, "Offset the quadratic boundary in by this");
DEFINE_double(linear_bounds_offset, 0.00,
              "Offset the linear boundary in by this");
DEFINE_double(yrange, 1.0,
              "+- y max for saturating out the state for the cost function.");
DEFINE_bool(debug_print, false, "Print the debugging print from the solver.");
DEFINE_bool(print_starting_summary, true,
            "Print the summary on the pre-solution.");
DEFINE_bool(print_summary, false, "Print the summary on each iteration.");
DEFINE_bool(quadratic, true, "If true, quadratic bounds penalty.");

DEFINE_bool(reset_every_cycle, false,
            "If true, reset the initial guess every cycle.");

DEFINE_double(seconds, 1.5, "The number of seconds to simulate.");

DEFINE_double(theta0, 1.0, "Starting theta0");
DEFINE_double(theta1, 0.9, "Starting theta1");

DEFINE_double(goal_theta0, -0.5, "Starting theta0");
DEFINE_double(goal_theta1, -0.5, "Starting theta1");

DEFINE_double(qpos1, 0.2, "qpos1");
DEFINE_double(qvel1, 4.0, "qvel1");
DEFINE_double(qpos2, 0.2, "qpos2");
DEFINE_double(qvel2, 4.0, "qvel2");

DEFINE_double(u_over_linear, 0.0, "Linear penalty for too much U.");
DEFINE_double(u_over_quadratic, 4.0, "Quadratic penalty for too much U.");

DEFINE_double(time_horizon, 0.75, "MPC time horizon");

DEFINE_bool(only_print_eigenvalues, false,
            "If true, stop after computing the final eigenvalues");

DEFINE_bool(plot_xy, false, "If true, plot the xy trajectory of the end of the arm.");
DEFINE_bool(plot_cost, false, "If true, plot the cost function.");
DEFINE_bool(plot_state_cost, false,
            "If true, plot the state portion of the cost function.");
DEFINE_bool(plot_states, false, "If true, plot the states.");
DEFINE_bool(plot_u, false, "If true, plot the control signal.");

static constexpr double kDt = 0.00505;

namespace y2018 {
namespace control_loops {

::Eigen::Matrix<double, 4, 4> NumericalJacobianX(
    ::Eigen::Matrix<double, 4, 1> (*fn)(
        ::Eigen::Ref<::Eigen::Matrix<double, 4, 1>> X,
        ::Eigen::Ref<::Eigen::Matrix<double, 2, 1>> U, double dt),
    ::Eigen::Matrix<double, 4, 1> X, ::Eigen::Matrix<double, 2, 1> U, double dt,
    const double kEpsilon = 1e-4) {
  constexpr int num_states = 4;
  ::Eigen::Matrix<double, 4, 4> answer = ::Eigen::Matrix<double, 4, 4>::Zero();

  // It's more expensive, but +- epsilon will be more reliable
  for (int i = 0; i < num_states; ++i) {
    ::Eigen::Matrix<double, 4, 1> dX_plus = X;
    dX_plus(i, 0) += kEpsilon;
    ::Eigen::Matrix<double, 4, 1> dX_minus = X;
    dX_minus(i, 0) -= kEpsilon;
    answer.block<4, 1>(0, i) =
        (fn(dX_plus, U, dt) - fn(dX_minus, U, dt)) / kEpsilon / 2.0;
  }
  return answer;
}

::Eigen::Matrix<double, 4, 2> NumericalJacobianU(
    ::Eigen::Matrix<double, 4, 1> (*fn)(
        ::Eigen::Ref<::Eigen::Matrix<double, 4, 1>> X,
        ::Eigen::Ref<::Eigen::Matrix<double, 2, 1>> U, double dt),
    ::Eigen::Matrix<double, 4, 1> X, ::Eigen::Matrix<double, 2, 1> U, double dt,
    const double kEpsilon = 1e-4) {
  constexpr int num_states = 4;
  constexpr int num_inputs = 2;
  ::Eigen::Matrix<double, num_states, num_inputs> answer =
      ::Eigen::Matrix<double, num_states, num_inputs>::Zero();

  // It's more expensive, but +- epsilon will be more reliable
  for (int i = 0; i < num_inputs; ++i) {
    ::Eigen::Matrix<double, 2, 1> dU_plus = U;
    dU_plus(i, 0) += kEpsilon;
    ::Eigen::Matrix<double, 2, 1> dU_minus = U;
    dU_minus(i, 0) -= kEpsilon;
    answer.block<4, 1>(0, i) =
        (fn(X, dU_plus, dt) - fn(X, dU_minus, dt)) / kEpsilon / 2.0;
  }
  return answer;
}

// This code is for analysis and simulation of a double jointed arm.  It is an
// attempt to see if a MPC could work for arm control under constraints.

// Describes a double jointed arm.
// A large chunk of this code comes from demos.  Most of the raw pointer,
// shared_ptr, and non-const &'s come from the library's conventions.
template <typename SCALAR>
class MySecondOrderSystem : public ::ct::core::ControlledSystem<4, 2, SCALAR> {
 public:
  static const size_t STATE_DIM = 4;
  static const size_t CONTROL_DIM = 2;

  MySecondOrderSystem(::std::shared_ptr<::ct::core::Controller<4, 2, SCALAR>>
                          controller = nullptr)
      : ::ct::core::ControlledSystem<4, 2, SCALAR>(
            controller, ::ct::core::SYSTEM_TYPE::GENERAL) {}

  MySecondOrderSystem(const MySecondOrderSystem &arg)
      : ::ct::core::ControlledSystem<4, 2, SCALAR>(arg) {}

  // Deep copy
  MySecondOrderSystem *clone() const override {
    return new MySecondOrderSystem(*this);
  }
  virtual ~MySecondOrderSystem() {}

  static constexpr SCALAR l1 = 46.25 * 0.0254;
  static constexpr SCALAR l2 = 41.80 * 0.0254;

  static constexpr SCALAR m1 = 9.34 / 2.2;
  static constexpr SCALAR m2 = 9.77 / 2.2;

  static constexpr SCALAR J1 = 2957.05 * 0.0002932545454545454;
  static constexpr SCALAR J2 = 2824.70 * 0.0002932545454545454;

  static constexpr SCALAR r1 = 21.64 * 0.0254;
  static constexpr SCALAR r2 = 26.70 * 0.0254;

  static constexpr SCALAR G1 = 140.0;
  static constexpr SCALAR G2 = 90.0;

  static constexpr SCALAR stall_torque = 1.41;
  static constexpr SCALAR free_speed = (5840.0 / 60.0) * 2.0 * M_PI;
  static constexpr SCALAR stall_current = 89.0;
  static constexpr SCALAR R = 12.0 / stall_current;

  static constexpr SCALAR Kv = free_speed / 12.0;
  static constexpr SCALAR Kt = stall_torque / stall_current;

  // Evaluate the system dynamics.
  //
  // Args:
  //   state: current state (position, velocity)
  //   t: current time (gets ignored)
  //   control: control action
  //   derivative: (velocity, acceleration)
  virtual void computeControlledDynamics(
      const ::ct::core::StateVector<4, SCALAR> &state, const SCALAR & /*t*/,
      const ::ct::core::ControlVector<2, SCALAR> &control,
      ::ct::core::StateVector<4, SCALAR> &derivative) override {
    derivative = Dynamics(state, control);
  }

  static ::Eigen::Matrix<double, 4, 1> Dynamics(
      const ::ct::core::StateVector<4, SCALAR> &X,
      const ::ct::core::ControlVector<2, SCALAR> &U) {
    ::ct::core::StateVector<4, SCALAR> derivative;
    const SCALAR alpha = J1 + r1 * r1 * m1 + l1 * l1 * m2;
    const SCALAR beta = l1 * r2 * m2;
    const SCALAR gamma = J2 + r2 * r2 * m2;

    const SCALAR s = sin(X(0) - X(2));
    const SCALAR c = cos(X(0) - X(2));

    // K1 * d^2 theta/dt^2 + K2 * d theta/dt = torque
    ::Eigen::Matrix<SCALAR, 2, 2> K1;
    K1(0, 0) = alpha;
    K1(1, 0) = K1(0, 1) = c * beta;
    K1(1, 1) = gamma;

    ::Eigen::Matrix<SCALAR, 2, 2> K2 = ::Eigen::Matrix<SCALAR, 2, 2>::Zero();
    K2(0, 1) = s * beta * X(3);
    K2(1, 0) = -s * beta * X(1);

    const SCALAR kNumDistalMotors = 2.0;
    ::Eigen::Matrix<SCALAR, 2, 1> torque;
    torque(0, 0) = G1 * (U(0) * Kt / R - X(1) * G1 * Kt / (Kv * R));
    torque(1, 0) = G2 * (U(1) * kNumDistalMotors * Kt / R -
                         X(3) * G2 * Kt * kNumDistalMotors / (Kv * R));

    ::Eigen::Matrix<SCALAR, 2, 1> velocity;
    velocity(0, 0) = X(0);
    velocity(1, 0) = X(2);

    const ::Eigen::Matrix<SCALAR, 2, 1> accel =
        K1.inverse() * (torque - K2 * velocity);

    derivative(0) = X(1);
    derivative(1) = accel(0);
    derivative(2) = X(3);
    derivative(3) = accel(1);

    return derivative;
  }

  // Runge-Kutta.
  static ::Eigen::Matrix<double, 4, 1> DiscreteDynamics(
      ::Eigen::Ref<::Eigen::Matrix<double, 4, 1>> X,
      ::Eigen::Ref<::Eigen::Matrix<double, 2, 1>> U, double dt) {
    const double half_dt = dt * 0.5;
    ::Eigen::Matrix<double, 4, 1> k1 = Dynamics(X, U);
    ::Eigen::Matrix<double, 4, 1> k2 = Dynamics(X + half_dt * k1, U);
    ::Eigen::Matrix<double, 4, 1> k3 = Dynamics(X + half_dt * k2, U);
    ::Eigen::Matrix<double, 4, 1> k4 = Dynamics(X + dt * k3, U);
    return X + dt / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4);
  }
};

template <size_t STATE_DIM, size_t CONTROL_DIM, typename SCALAR_EVAL = double,
          typename SCALAR = SCALAR_EVAL>
class ObstacleAwareQuadraticCost
    : public ::ct::optcon::TermBase<STATE_DIM, CONTROL_DIM, SCALAR_EVAL,
                                    SCALAR> {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, 1> state_vector_t;
  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, STATE_DIM> state_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, CONTROL_DIM> control_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, STATE_DIM>
      control_state_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, STATE_DIM>
      state_matrix_double_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, CONTROL_DIM>
      control_matrix_double_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, STATE_DIM>
      control_state_matrix_double_t;

  ObstacleAwareQuadraticCost(const ::Eigen::Matrix<double, 2, 2> &R,
                             const ::Eigen::Matrix<double, 4, 4> &Q)
      : R_(R), Q_(Q) {}

  ObstacleAwareQuadraticCost(const ObstacleAwareQuadraticCost &arg)
      : R_(arg.R_), Q_(arg.Q_) {}
      static constexpr double kEpsilon = 1.0e-5;

  virtual ~ObstacleAwareQuadraticCost() {}

  ObstacleAwareQuadraticCost<STATE_DIM, CONTROL_DIM, SCALAR_EVAL, SCALAR>
      *clone() const override {
    return new ObstacleAwareQuadraticCost(*this);
  }

  double SaturateX(double x, double yrange) {
    return 2.0 * ((1.0 / (1.0 + ::std::exp(-x * 2.0 / yrange)) - 0.5)) * yrange;
  }

  SCALAR distance(const Eigen::Matrix<SCALAR, STATE_DIM, 1> &x,
                  const Eigen::Matrix<SCALAR, CONTROL_DIM, 1> & /*u*/) {
    constexpr double kCornerNewUpper0 = 0.35;
    // constexpr double kCornerUpper1 = 3.13;
    // Push it up a bit further (non-real) until we have an actual path cost.
    constexpr double kCornerNewUpper1 = 3.39;
    constexpr double kCornerNewUpper0_far = 10.0;

    // Push it up a bit further (non-real) until we have an actual path cost.
    // constexpr double kCornerUpper0 = 0.315;
    constexpr double kCornerUpper0 = 0.310;
    // constexpr double kCornerUpper1 = 3.13;
    constexpr double kCornerUpper1 = 3.25;
    constexpr double kCornerUpper0_far = 10.0;

    constexpr double kCornerLower0 = 0.023;
    constexpr double kCornerLower1 = 1.57;
    constexpr double kCornerLower0_far = 10.0;

    const Segment new_upper_segment(
        Point(kCornerNewUpper0, kCornerNewUpper1),
        Point(kCornerNewUpper0_far, kCornerNewUpper1));
    const Segment upper_segment(Point(kCornerUpper0, kCornerUpper1),
                                Point(kCornerUpper0_far, kCornerUpper1));
    const Segment lower_segment(Point(kCornerLower0, kCornerLower1),
                                Point(kCornerLower0_far, kCornerLower1));

    Point current_point(x(0, 0), x(2, 0));

    SCALAR result = 0.0;
    if (intersects(new_upper_segment,
                   Segment(current_point,
                           Point(FLAGS_goal_theta0, FLAGS_goal_theta1)))) {
      result += hypot(current_point.x() - kCornerNewUpper0,
                      current_point.y() - kCornerNewUpper1);
      current_point = Point(kCornerNewUpper0, kCornerNewUpper1);
    }

    if (intersects(upper_segment,
                   Segment(current_point,
                           Point(FLAGS_goal_theta0, FLAGS_goal_theta1)))) {
      result += hypot(current_point.x() - kCornerUpper0,
                      current_point.y() - kCornerUpper1);
      current_point = Point(kCornerUpper0, kCornerUpper1);
    }

    if (intersects(lower_segment,
                   Segment(current_point,
                           Point(FLAGS_goal_theta0, FLAGS_goal_theta1)))) {
      result += hypot(current_point.x() - kCornerLower0,
                      current_point.y() - kCornerLower1);
      current_point = Point(kCornerLower0, kCornerLower1);
    }
    result += hypot(current_point.x() - FLAGS_goal_theta0,
                    current_point.y() - FLAGS_goal_theta1);
    return result;
  }

  virtual SCALAR evaluate(const Eigen::Matrix<SCALAR, STATE_DIM, 1> &x,
                          const Eigen::Matrix<SCALAR, CONTROL_DIM, 1> &u,
                          const SCALAR & /*t*/) override {
    // Positive means violation.
    Eigen::Matrix<SCALAR, STATE_DIM, 1> saturated_x = x;
    SCALAR d = distance(x, u);
    saturated_x(0, 0) = d;
    saturated_x(2, 0) = 0.0;

    saturated_x(0, 0) = SaturateX(saturated_x(0, 0), FLAGS_yrange);
    saturated_x(2, 0) = 0.0;

    //SCALAR saturation_scalar = saturated_x(0, 0) / d;
    //saturated_x(1, 0) *= saturation_scalar;
    //saturated_x(3, 0) *= saturation_scalar;

    SCALAR result = (saturated_x.transpose() * Q_ * saturated_x +
                     u.transpose() * R_ * u)(0, 0);

    if (::std::abs(u(0, 0)) > 11.0) {
      result += (::std::abs(u(0, 0)) - 11.0) * FLAGS_u_over_linear;
      result += (::std::abs(u(0, 0)) - 11.0) * (::std::abs(u(0, 0)) - 11.0) *
                FLAGS_u_over_quadratic;
    }
    if (::std::abs(u(1, 0)) > 11.0) {
      result += (::std::abs(u(1, 0)) - 11.0) * FLAGS_u_over_linear;
      result += (::std::abs(u(1, 0)) - 11.0) * (::std::abs(u(1, 0)) - 11.0) *
                FLAGS_u_over_quadratic;
    }
    return result;
  }

  ct::core::StateVector<STATE_DIM, SCALAR_EVAL> stateDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    SCALAR epsilon = SCALAR(kEpsilon);

    ct::core::StateVector<STATE_DIM, SCALAR_EVAL> result =
        ct::core::StateVector<STATE_DIM, SCALAR_EVAL>::Zero();

    // Perterb x for both position axis and return the result.
    for (size_t i = 0; i < STATE_DIM; i += 1) {
      ct::core::StateVector<STATE_DIM, SCALAR_EVAL> plus_perterbed_x = x;
      ct::core::StateVector<STATE_DIM, SCALAR_EVAL> minus_perterbed_x = x;
      plus_perterbed_x[i] += epsilon;
      minus_perterbed_x[i] -= epsilon;
      result[i] = (evaluate(plus_perterbed_x, u, t) -
                   evaluate(minus_perterbed_x, u, t)) /
                  (epsilon * 2.0);
    }
    return result;
  }

  // Compute second order derivative of this cost term w.r.t. the state
  state_matrix_t stateSecondDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    state_matrix_t result = state_matrix_t::Zero();

    SCALAR epsilon = SCALAR(kEpsilon);

    // Perterb x a second time.
    for (size_t i = 0; i < STATE_DIM; i += 1) {
      ct::core::StateVector<STATE_DIM, SCALAR_EVAL> plus_perterbed_x = x;
      ct::core::StateVector<STATE_DIM, SCALAR_EVAL> minus_perterbed_x = x;
      plus_perterbed_x[i] += epsilon;
      minus_perterbed_x[i] -= epsilon;
      state_vector_t delta = (stateDerivative(plus_perterbed_x, u, t) -
                              stateDerivative(minus_perterbed_x, u, t)) /
                             (epsilon * 2.0);

      result.col(i) = delta;
    }
    //::std::cout << "Q_numeric " << result << " endQ" << ::std::endl;
    return result;
  }

  ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> controlDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> result =
        ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL>::Zero();

    SCALAR epsilon = SCALAR(kEpsilon);

    // Perterb x a second time.
    for (size_t i = 0; i < CONTROL_DIM; i += 1) {
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> plus_perterbed_u = u;
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> minus_perterbed_u = u;
      plus_perterbed_u[i] += epsilon;
      minus_perterbed_u[i] -= epsilon;
      SCALAR delta = (evaluate(x, plus_perterbed_u, t) -
                      evaluate(x, minus_perterbed_u, t)) /
                     (epsilon * 2.0);

      result[i] = delta;
    }
    //::std::cout << "cd " << result(0, 0) << " " << result(1, 0) << " endcd"
                //<< ::std::endl;

    return result;
  }

  control_state_matrix_t stateControlDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> & /*x*/,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> & /*u*/,
      const SCALAR_EVAL & /*t*/) override {
    // No coupling here, so let's not bother to calculate it.
    control_state_matrix_t result = control_state_matrix_t::Zero();
    return result;
  }

  control_matrix_t controlSecondDerivative(
      const ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    control_matrix_t result = control_matrix_t::Zero();

    SCALAR epsilon = SCALAR(kEpsilon);

    //static int j = 0;
    //::std::this_thread::sleep_for(::std::chrono::milliseconds(j % 10));
    //int k = ++j;
    // Perterb x a second time.
    for (size_t i = 0; i < CONTROL_DIM; i += 1) {
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> plus_perterbed_u = u;
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> minus_perterbed_u = u;
      plus_perterbed_u[i] += epsilon;
      minus_perterbed_u[i] -= epsilon;
      ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> delta =
          (controlDerivative(x, plus_perterbed_u, t) -
           controlDerivative(x, minus_perterbed_u, t)) /
          (epsilon * 2.0);

      //::std::cout << "delta: " << delta(0, 0) << " " << delta(1, 0) << " k "
                  //<< k << ::std::endl;
      result.col(i) = delta;
    }
    //::std::cout << "R_numeric " << result << " endR 0.013888888888888888    k:" << k
                //<< ::std::endl;
    //::std::cout << "x " << x << " u " << u << "    k " << k << ::std::endl;

    return result;
  }

 private:
  const ::Eigen::Matrix<double, 2, 2> R_;
  const ::Eigen::Matrix<double, 4, 4> Q_;
};

template <size_t STATE_DIM, size_t CONTROL_DIM, typename SCALAR_EVAL = double,
          typename SCALAR = SCALAR_EVAL>
class MyTermStateBarrier : public ::ct::optcon::TermBase<STATE_DIM, CONTROL_DIM,
                                                         SCALAR_EVAL, SCALAR> {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, 1> state_vector_t;
  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, STATE_DIM> state_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, CONTROL_DIM> control_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, STATE_DIM>
      control_state_matrix_t;
  typedef Eigen::Matrix<SCALAR_EVAL, STATE_DIM, STATE_DIM>
      state_matrix_double_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, CONTROL_DIM>
      control_matrix_double_t;
  typedef Eigen::Matrix<SCALAR_EVAL, CONTROL_DIM, STATE_DIM>
      control_state_matrix_double_t;

  MyTermStateBarrier(BoundsCheck *bounds_check) : bounds_check_(bounds_check) {}

  MyTermStateBarrier(const MyTermStateBarrier &arg)
      : bounds_check_(arg.bounds_check_) {}

  static constexpr double kEpsilon = 5.0e-6;

  virtual ~MyTermStateBarrier() {}

  MyTermStateBarrier<STATE_DIM, CONTROL_DIM, SCALAR_EVAL, SCALAR> *clone()
      const override {
    return new MyTermStateBarrier(*this);
  }

  SCALAR distance(const Eigen::Matrix<SCALAR, STATE_DIM, 1> &x,
                  const Eigen::Matrix<SCALAR, CONTROL_DIM, 1> & /*u*/,
                  const SCALAR & /*t*/, Eigen::Matrix<SCALAR, 2, 1> *norm) {
    return bounds_check_->min_distance(Point(x(0, 0), x(2, 0)), norm);
  }

  virtual SCALAR evaluate(const Eigen::Matrix<SCALAR, STATE_DIM, 1> &x,
                          const Eigen::Matrix<SCALAR, CONTROL_DIM, 1> & u,
                          const SCALAR & t) override {
    Eigen::Matrix<SCALAR, 2, 1> norm = Eigen::Matrix<SCALAR, 2, 1>::Zero();
    SCALAR min_distance = distance(x, u, t, &norm);

    // Velocity component (+) towards the wall.
    SCALAR velocity_penalty = -(x(1, 0) * norm(0, 0) + x(3, 0) * norm(1, 0));
    if (min_distance + FLAGS_bounds_offset < 0.0) {
      velocity_penalty = 0.0;
    }

    SCALAR result;
    //if (FLAGS_quadratic) {
    result = FLAGS_boundary_scalar *
                 ::std::max(0.0, min_distance + FLAGS_bounds_offset) *
                 ::std::max(0.0, min_distance + FLAGS_bounds_offset) +
             FLAGS_boundary_rate *
                 ::std::max(0.0, min_distance + FLAGS_linear_bounds_offset) +
             FLAGS_velocity_boundary_scalar *
                 ::std::max(0.0, min_distance + FLAGS_linear_bounds_offset) *
                 ::std::max(0.0, velocity_penalty) *
                 ::std::max(0.0, velocity_penalty);
    /*
} else if (FLAGS_linear) {
result =
FLAGS_boundary_scalar * ::std::max(0.0, min_distance) +
FLAGS_velocity_boundary_scalar * ::std::max(0.0, -velocity_penalty);
} else if (FLAGS_sigmoid) {
result = FLAGS_boundary_scalar /
    (1.0 + ::std::exp(-min_distance * FLAGS_boundary_rate)) +
FLAGS_velocity_boundary_scalar /
    (1.0 + ::std::exp(-velocity_penalty * FLAGS_boundary_rate));
} else {
// Values of 4 and 15 work semi resonably.
result = FLAGS_boundary_scalar *
    ::std::exp(min_distance * FLAGS_boundary_rate) +
FLAGS_velocity_boundary_scalar *
    ::std::exp(velocity_penalty * FLAGS_boundary_rate);
}
if (result < 0.0) {
printf("Result negative %f\n", result);
}
*/
    return result;
  }

  ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> stateDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    SCALAR epsilon = SCALAR(kEpsilon);

    ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> result =
        ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL>::Zero();

    // Perturb x for both position axis and return the result.
    for (size_t i = 0; i < STATE_DIM; i += 2) {
      ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> plus_perterbed_x = x;
      ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> minus_perterbed_x = x;
      plus_perterbed_x[i] += epsilon;
      minus_perterbed_x[i] -= epsilon;
      result[i] = (evaluate(plus_perterbed_x, u, t) -
                   evaluate(minus_perterbed_x, u, t)) /
                  (epsilon * 2.0);
    }
    return result;
  }

  // Compute second order derivative of this cost term w.r.t. the state
  state_matrix_t stateSecondDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> &x,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> &u,
      const SCALAR_EVAL &t) override {
    state_matrix_t result = state_matrix_t::Zero();

    SCALAR epsilon = SCALAR(kEpsilon);

    // Perturb x a second time.
    for (size_t i = 0; i < STATE_DIM; i += 1) {
      ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> plus_perterbed_x = x;
      ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> minus_perterbed_x = x;
      plus_perterbed_x[i] += epsilon;
      minus_perterbed_x[i] -= epsilon;
      state_vector_t delta = (stateDerivative(plus_perterbed_x, u, t) -
                              stateDerivative(minus_perterbed_x, u, t)) /
                             (epsilon * 2.0);

      result.col(i) = delta;
    }
    return result;
  }

  ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> controlDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> & /*x*/,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> & /*u*/,
      const SCALAR_EVAL & /*t*/) override {
    return ::ct::core::StateVector<CONTROL_DIM, SCALAR_EVAL>::Zero();
  }

  control_state_matrix_t stateControlDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> & /*x*/,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> & /*u*/,
      const SCALAR_EVAL & /*t*/) override {
    control_state_matrix_t result = control_state_matrix_t::Zero();

    return result;
  }

  control_matrix_t controlSecondDerivative(
      const ::ct::core::StateVector<STATE_DIM, SCALAR_EVAL> & /*x*/,
      const ::ct::core::ControlVector<CONTROL_DIM, SCALAR_EVAL> & /*u*/,
      const SCALAR_EVAL & /*t*/) override {
    control_matrix_t result = control_matrix_t::Zero();
    return result;
  }

  /*
    // TODO(austin): Implement this for the automatic differentiation.
    virtual ::ct::core::ADCGScalar evaluateCppadCg(
        const ::ct::core::StateVector<STATE_DIM, ::ct::core::ADCGScalar> &x,
        const ::ct::core::ControlVector<CONTROL_DIM, ::ct::core::ADCGScalar> &u,
        ::ct::core::ADCGScalar t) override {
      ::ct::core::ADCGScalar c = ::ct::core::ADCGScalar(0.0);
      for (size_t i = 0; i < STATE_DIM; i++)
        c += barriers_[i].computeActivation(x(i));
      return c;
    }
  */

  BoundsCheck *bounds_check_;
};

int Main() {
  // PRELIMINIARIES
  BoundsCheck arm_space = MakeClippedArmSpace();

  constexpr size_t state_dim = MySecondOrderSystem<double>::STATE_DIM;
  constexpr size_t control_dim = MySecondOrderSystem<double>::CONTROL_DIM;

  ::std::shared_ptr<ct::core::ControlledSystem<state_dim, control_dim>>
  oscillator_dynamics(new MySecondOrderSystem<double>());

  ::std::shared_ptr<ct::core::SystemLinearizer<state_dim, control_dim>>
      ad_linearizer(new ::ct::core::SystemLinearizer<state_dim, control_dim>(
          oscillator_dynamics));

  const double kQPos1 = FLAGS_qpos1;
  const double kQVel1 = FLAGS_qvel1;
  const double kQPos2 = FLAGS_qpos2;
  const double kQVel2 = FLAGS_qvel2;

  ::Eigen::Matrix<double, 4, 4> Q_step;
  Q_step << 1.0 / (kQPos1 * kQPos1), 0.0, 0.0, 0.0, 0.0,
      1.0 / (kQVel1 * kQVel1), 0.0, 0.0, 0.0, 0.0, 1.0 / (kQPos2 * kQPos2), 0.0,
      0.0, 0.0, 0.0, 1.0 / (kQVel2 * kQVel2);
  ::Eigen::Matrix<double, 2, 2> R_step;
  R_step << 1.0 / (12.0 * 12.0), 0.0, 0.0, 1.0 / (12.0 * 12.0);
  ::std::shared_ptr<::ct::optcon::TermQuadratic<state_dim, control_dim>>
      quadratic_intermediate_cost(
          new ::ct::optcon::TermQuadratic<state_dim, control_dim>(Q_step,
                                                                  R_step));
  // TODO(austin): Move back to this with the new Q and R
  ::std::shared_ptr<ObstacleAwareQuadraticCost<state_dim, control_dim>>
      intermediate_cost(new ObstacleAwareQuadraticCost<4, 2>(R_step, Q_step));

  ::Eigen::Matrix<double, 4, 4> final_A =
      NumericalJacobianX(MySecondOrderSystem<double>::DiscreteDynamics,
                         Eigen::Matrix<double, 4, 1>::Zero(),
                         Eigen::Matrix<double, 2, 1>::Zero(), kDt);

  ::Eigen::Matrix<double, 4, 2> final_B =
      NumericalJacobianU(MySecondOrderSystem<double>::DiscreteDynamics,
                         Eigen::Matrix<double, 4, 1>::Zero(),
                         Eigen::Matrix<double, 2, 1>::Zero(), kDt);

  ::Eigen::Matrix<double, 4, 4> S_lqr;
  ::Eigen::Matrix<double, 2, 4> K_lqr;
  ::frc971::controls::dlqr(final_A, final_B, Q_step, R_step, &K_lqr, &S_lqr);
  ::std::cout << "A -> " << ::std::endl << final_A << ::std::endl;
  ::std::cout << "B -> " << ::std::endl << final_B << ::std::endl;
  ::std::cout << "K -> " << ::std::endl << K_lqr << ::std::endl;
  ::std::cout << "S -> " << ::std::endl << S_lqr << ::std::endl;
  ::std::cout << "Q -> " << ::std::endl << Q_step << ::std::endl;
  ::std::cout << "R -> " << ::std::endl << R_step << ::std::endl;
  ::std::cout << "Eigenvalues: " << (final_A - final_B * K_lqr).eigenvalues()
              << ::std::endl;

  ::Eigen::Matrix<double, 4, 4> Q_final = 0.5 * S_lqr;
  ::Eigen::Matrix<double, 2, 2> R_final = ::Eigen::Matrix<double, 2, 2>::Zero();
  ::std::shared_ptr<ct::optcon::TermQuadratic<state_dim, control_dim>>
      final_cost(new ::ct::optcon::TermQuadratic<state_dim, control_dim>(
          Q_final, R_final));
  if (FLAGS_only_print_eigenvalues) {
    return 0;
  }

  ::std::shared_ptr<MyTermStateBarrier<state_dim, control_dim>> bounds_cost(
      new MyTermStateBarrier<4, 2>(&arm_space));

  // TODO(austin): Cost function needs constraints.
  ::std::shared_ptr<::ct::optcon::CostFunctionQuadratic<state_dim, control_dim>>
      cost_function(
          new ::ct::optcon::CostFunctionAnalytical<state_dim, control_dim>());
  //cost_function->addIntermediateTerm(quadratic_intermediate_cost);
  cost_function->addIntermediateTerm(intermediate_cost);
  cost_function->addIntermediateTerm(bounds_cost);
  cost_function->addFinalTerm(final_cost);

  // STEP 1-D: set up the box constraints for the control input
  // input box constraint boundaries with sparsities in constraint toolbox
  // format
  Eigen::VectorXd u_lb(control_dim);
  Eigen::VectorXd u_ub(control_dim);
  u_ub.setConstant(12.0);
  u_lb = -u_ub;
  //::std::cout << "uub " << u_ub << ::std::endl;
  //::std::cout << "ulb " << u_lb << ::std::endl;

  // constraint terms
  std::shared_ptr<::ct::optcon::ControlInputConstraint<state_dim, control_dim>>
      controlConstraint(
          new ::ct::optcon::ControlInputConstraint<state_dim, control_dim>(
              u_lb, u_ub));
  controlConstraint->setName("ControlInputConstraint");
  // create constraint container
  std::shared_ptr<
      ::ct::optcon::ConstraintContainerAnalytical<state_dim, control_dim>>
      box_constraints(
          new ::ct::optcon::ConstraintContainerAnalytical<state_dim,
                                                          control_dim>());
  // add and initialize constraint terms
  box_constraints->addIntermediateConstraint(controlConstraint, true);
  box_constraints->initialize();

  // Starting point.
  ::ct::core::StateVector<state_dim> x0;
  x0 << FLAGS_theta0, 0.0, FLAGS_theta1, 0.0;

  const ::ct::core::Time kTimeHorizon = FLAGS_time_horizon;
  ::ct::optcon::OptConProblem<state_dim, control_dim> opt_con_problem(
      kTimeHorizon, x0, oscillator_dynamics, cost_function, ad_linearizer);
  ::ct::optcon::NLOptConSettings ilqr_settings;
  ilqr_settings.nThreads = 4;
  ilqr_settings.dt = kDt;  // the control discretization in [sec]
  ilqr_settings.integrator = ::ct::core::IntegrationType::RK4;
  ilqr_settings.debugPrint = FLAGS_debug_print;
  ilqr_settings.discretization =
      ::ct::optcon::NLOptConSettings::APPROXIMATION::FORWARD_EULER;
  // ilqr_settings.discretization =
  //   NLOptConSettings::APPROXIMATION::MATRIX_EXPONENTIAL;
  ilqr_settings.max_iterations = 40;
  ilqr_settings.min_cost_improvement = FLAGS_convergance;
  ilqr_settings.nlocp_algorithm =
      //::ct::optcon::NLOptConSettings::NLOCP_ALGORITHM::ILQR;
      ::ct::optcon::NLOptConSettings::NLOCP_ALGORITHM::GNMS;
  // the LQ-problems are solved using a custom Gauss-Newton Riccati solver
  ilqr_settings.lqocp_solver =
      ::ct::optcon::NLOptConSettings::LQOCP_SOLVER::GNRICCATI_SOLVER;
  //ilqr_settings.lqocp_solver =
      //::ct::optcon::NLOptConSettings::LQOCP_SOLVER::HPIPM_SOLVER;
  ilqr_settings.printSummary = FLAGS_print_starting_summary;
  if (ilqr_settings.lqocp_solver ==
      ::ct::optcon::NLOptConSettings::LQOCP_SOLVER::HPIPM_SOLVER) {
    //opt_con_problem.setBoxConstraints(box_constraints);
  }

  const size_t num_steps = ilqr_settings.computeK(kTimeHorizon);
  printf("Using %d steps\n", static_cast<int>(num_steps));

  // Vector of feeback matricies.
  ::ct::core::FeedbackArray<state_dim, control_dim> u0_fb(
      num_steps, ::ct::core::FeedbackMatrix<state_dim, control_dim>::Zero());
  ::ct::core::ControlVectorArray<control_dim> u0_ff(
      num_steps, ::ct::core::ControlVector<control_dim>::Zero());
  ::ct::core::StateVectorArray<state_dim> x_ref_init(num_steps + 1, x0);
  ::ct::core::StateFeedbackController<state_dim, control_dim>
      initial_controller(x_ref_init, u0_ff, u0_fb, ilqr_settings.dt);

  // STEP 2-C: create an NLOptConSolver instance
  ::ct::optcon::NLOptConSolver<state_dim, control_dim> iLQR(opt_con_problem,
                                                            ilqr_settings);
  // Seed it with the initial guess
  iLQR.setInitialGuess(initial_controller);
  // we solve the optimal control problem and retrieve the solution
  iLQR.solve();
  ::ct::core::StateFeedbackController<state_dim, control_dim> initial_solution =
      iLQR.getSolution();
  // MPC-EXAMPLE
  // we store the initial solution obtained from solving the initial optimal
  // control problem, and re-use it to initialize the MPC solver in the
  // following.

  // STEP 1: first, we set up an MPC instance for the iLQR solver and configure
  // it. Since the MPC class is wrapped around normal Optimal Control Solvers,
  // we need to different kind of settings, those for the optimal control
  // solver, and those specific to MPC:

  // 1) settings for the iLQR instance used in MPC. Of course, we use the same
  // settings as for solving the initial problem ...
  ::ct::optcon::NLOptConSettings ilqr_settings_mpc = ilqr_settings;
  ilqr_settings_mpc.max_iterations = 40;
  // and we limited the printouts, too.
  ilqr_settings_mpc.printSummary = FLAGS_print_summary;
  // 2) settings specific to model predictive control. For a more detailed
  // description of those, visit ct/optcon/mpc/MpcSettings.h
  ::ct::optcon::mpc_settings mpc_settings;
  mpc_settings.stateForwardIntegration_ = true;
  mpc_settings.postTruncation_ = false;
  mpc_settings.measureDelay_ = false;
  mpc_settings.fixedDelayUs_ = 5000 * 0;  // Ignore the delay for now.
  mpc_settings.delayMeasurementMultiplier_ = 1.0;
  // mpc_settings.mpc_mode = ::ct::optcon::MPC_MODE::FIXED_FINAL_TIME;
  mpc_settings.mpc_mode = ::ct::optcon::MPC_MODE::CONSTANT_RECEDING_HORIZON;
  mpc_settings.coldStart_ = false;

  // STEP 2 : Create the iLQR-MPC object, based on the optimal control problem
  // and the selected settings.
  ::ct::optcon::MPC<::ct::optcon::NLOptConSolver<state_dim, control_dim>>
      ilqr_mpc(opt_con_problem, ilqr_settings_mpc, mpc_settings);
  // initialize it using the previously computed initial controller
  ilqr_mpc.setInitialGuess(initial_solution);
  // STEP 3: running MPC
  // Here, we run the MPC loop. Note that the general underlying idea is that
  // you receive a state-estimate together with a time-stamp from your robot or
  // system. MPC needs to receive both that time information and the state from
  // your control system. Here, "simulate" the time measurement using
  // ::std::chrono and wrap everything into a for-loop.
  // The basic idea of operation is that after receiving time and state
  // information, one executes the finishIteration() method of MPC.
  ///
  auto start_time = ::std::chrono::high_resolution_clock::now();
  // limit the maximum number of runs in this example
  size_t maxNumRuns = FLAGS_seconds / kDt;
  ::std::cout << "Starting to run MPC" << ::std::endl;

  ::std::vector<double> time_array;
  ::std::vector<double> theta1_array;
  ::std::vector<double> omega1_array;
  ::std::vector<double> theta2_array;
  ::std::vector<double> omega2_array;

  ::std::vector<double> u0_array;
  ::std::vector<double> u1_array;

  ::std::vector<double> x_array;
  ::std::vector<double> y_array;

  // TODO(austin): Plot x, y of the end of the arm.

  for (size_t i = 0; i < maxNumRuns; i++) {
    ::std::cout << "Solving iteration " << i << ::std::endl;
    // Time which has passed since start of MPC
    auto current_time = ::std::chrono::high_resolution_clock::now();
    ::ct::core::Time t =
        1e-6 *
        ::std::chrono::duration_cast<::std::chrono::microseconds>(current_time -
                                                                  start_time)
            .count();
    {
      if (FLAGS_reset_every_cycle) {
        ::ct::core::FeedbackArray<state_dim, control_dim> u0_fb(
            num_steps,
            ::ct::core::FeedbackMatrix<state_dim, control_dim>::Zero());
        ::ct::core::ControlVectorArray<control_dim> u0_ff(
            num_steps, ::ct::core::ControlVector<control_dim>::Zero());
        ::ct::core::StateVectorArray<state_dim> x_ref_init(num_steps + 1, x0);
        ::ct::core::StateFeedbackController<state_dim, control_dim>
            resolved_controller(x_ref_init, u0_ff, u0_fb, ilqr_settings.dt);

        iLQR.setInitialGuess(initial_controller);
        // we solve the optimal control problem and retrieve the solution
        iLQR.solve();
        resolved_controller = iLQR.getSolution();
        ilqr_mpc.setInitialGuess(resolved_controller);
      }
    }

    // prepare mpc iteration
    ilqr_mpc.prepareIteration(t);
    // new optimal policy
    ::std::shared_ptr<ct::core::StateFeedbackController<state_dim, control_dim>>
        newPolicy(
            new ::ct::core::StateFeedbackController<state_dim, control_dim>());
    // timestamp of the new optimal policy
    ::ct::core::Time ts_newPolicy;
    current_time = ::std::chrono::high_resolution_clock::now();
    t = 1e-6 *
        ::std::chrono::duration_cast<::std::chrono::microseconds>(current_time -
                                                                  start_time)
            .count();
    // TODO(austin): This is only iterating once...  I need to fix that...
    //  NLOptConSolver::solve() runs for upto N iterations.  This call runs
    //  runIteration() effectively once.  (nlocAlgorithm_ is iLQR)
    bool success = ilqr_mpc.finishIteration(x0, t, *newPolicy, ts_newPolicy);
    // we break the loop in case the time horizon is reached or solve() failed
    if (ilqr_mpc.timeHorizonReached() | !success) break;

    ::std::cout << "Solved  for time " << newPolicy->time()[0] << " state "
                << x0.transpose() << " next time " << newPolicy->time()[1]
                << ::std::endl;
    ::std::cout << "  Solution: Uff " << newPolicy->uff()[0].transpose()
                << " x_ref_ " << newPolicy->x_ref()[0].transpose()
                << ::std::endl;

    time_array.push_back(ilqr_settings.dt * i);
    theta1_array.push_back(x0(0));
    omega1_array.push_back(x0(1));
    theta2_array.push_back(x0(2));
    omega2_array.push_back(x0(3));

    u0_array.push_back(newPolicy->uff()[0](0, 0));
    u1_array.push_back(newPolicy->uff()[0](1, 0));

    ::std::cout << "xref[1] " << newPolicy->x_ref()[1].transpose()
                << ::std::endl;
    ilqr_mpc.doForwardIntegration(0.0, ilqr_settings.dt, x0, newPolicy);
    ::std::cout << "Next X:  " << x0.transpose() << ::std::endl;

    x_array.push_back(MySecondOrderSystem<double>::l1 * sin(x0(0)) +
                      MySecondOrderSystem<double>::r2 * sin(x0(2)));
    y_array.push_back(MySecondOrderSystem<double>::l1 * cos(x0(0)) +
                      MySecondOrderSystem<double>::r2 * cos(x0(2)));

    // TODO(austin): Re-use the policy. Maybe?  Or maybe mpc already does that.
  }
  // The summary contains some statistical data about time delays, etc.
  ilqr_mpc.printMpcSummary();

  if (FLAGS_plot_states) {
    // Now plot our simulation.
    matplotlibcpp::plot(time_array, theta1_array, {{"label", "theta1"}});
    matplotlibcpp::plot(time_array, omega1_array, {{"label", "omega1"}});
    matplotlibcpp::plot(time_array, theta2_array, {{"label", "theta2"}});
    matplotlibcpp::plot(time_array, omega2_array, {{"label", "omega2"}});
    matplotlibcpp::legend();
  }

  if (FLAGS_plot_xy) {
    matplotlibcpp::figure();
    matplotlibcpp::plot(x_array, y_array, {{"label", "xy trajectory"}});
    matplotlibcpp::legend();
  }

  if (FLAGS_plot_u) {
    matplotlibcpp::figure();
    matplotlibcpp::plot(time_array, u0_array, {{"label", "u0"}});
    matplotlibcpp::plot(time_array, u1_array, {{"label", "u1"}});
    matplotlibcpp::legend();
  }

  ::std::vector<::std::vector<double>> cost_x;
  ::std::vector<::std::vector<double>> cost_y;
  ::std::vector<::std::vector<double>> cost_z;
  ::std::vector<::std::vector<double>> cost_state_z;

  for (double x_coordinate = -0.5; x_coordinate < 1.2; x_coordinate += 0.05) {
    ::std::vector<double> cost_x_row;
    ::std::vector<double> cost_y_row;
    ::std::vector<double> cost_z_row;
    ::std::vector<double> cost_state_z_row;

    for (double y_coordinate = -1.0; y_coordinate < 6.0; y_coordinate += 0.05) {
      cost_x_row.push_back(x_coordinate);
      cost_y_row.push_back(y_coordinate);
      Eigen::Matrix<double, 4, 1> state_matrix;
      state_matrix << x_coordinate, 0.0, y_coordinate, 0.0;
      Eigen::Matrix<double, 2, 1> u_matrix =
          Eigen::Matrix<double, 2, 1>::Zero();
      cost_state_z_row.push_back(
          intermediate_cost->distance(state_matrix, u_matrix));
      cost_z_row.push_back(
          ::std::min(bounds_cost->evaluate(state_matrix, u_matrix, 0.0), 50.0));
    }
    cost_x.push_back(cost_x_row);
    cost_y.push_back(cost_y_row);
    cost_z.push_back(cost_z_row);
    cost_state_z.push_back(cost_state_z_row);
  }

  if (FLAGS_plot_cost) {
    matplotlibcpp::plot_surface(cost_x, cost_y, cost_z);
  }

  if (FLAGS_plot_state_cost) {
    matplotlibcpp::plot_surface(cost_x, cost_y, cost_state_z);
  }

  matplotlibcpp::figure();
  matplotlibcpp::plot(theta1_array, theta2_array, {{"label", "trajectory"}});
  ::std::vector<double> bounds_x;
  ::std::vector<double> bounds_y;
  for (const Point p : arm_space.points()) {
    bounds_x.push_back(p.x());
    bounds_y.push_back(p.y());
  }
  matplotlibcpp::plot(bounds_x, bounds_y, {{"label", "allowed region"}});
  matplotlibcpp::legend();

  matplotlibcpp::show();

  return 0;
}

}  // namespace control_loops
}  // namespace y2018

int main(int argc, char **argv) {
  gflags::ParseCommandLineFlags(&argc, &argv, false);
  return ::y2018::control_loops::Main();
}


#include <fast_tf/fast_tf.hpp>
#include <benchmark/benchmark.h>

#include <tf2/LinearMath/Quaternion.h>
#include <tf2/LinearMath/Vector3.h>

#include <Eigen/Geometry>

using benchmark::State;

static void
eigen_isometry_prod(State& _state) {
  const Eigen::Isometry3d tf1(Eigen::Translation3d(1, 2, 3) *
                              Eigen::Quaterniond(-0.002, -0.678, 0.226, 0.699));

  const Eigen::Isometry3d tf2(Eigen::Translation3d(3, 4, 1) *
                              Eigen::Quaterniond(-0.003, -0.945, 0.315, 0.091));

  for (auto _ : _state)
    benchmark::DoNotOptimize(tf1 * tf2);
}

struct tf_transform {
  tf2::Vector3 v;
  tf2::Quaternion q;
};

// without this the compiler will still remove the call to the product. this can
// be seen when running the benchmark below with valgrind --tool=callgrind.
__attribute__((noinline)) tf_transform
tf2_product(const tf_transform& _l, const tf_transform& _r) {
  return {tf2::quatRotate(_l.q, _r.v) + _l.v, _l.q * _r.q};
}

static void
tf_transform_accum_prod(State& _state) {
  tf2::Quaternion q1(-0.002, -0.678, 0.226, 0.699);
  tf2::Vector3 v1(1, 2, 3);

  tf2::Quaternion q2(-0.003, -0.945, 0.315, 0.091);
  tf2::Vector3 v2(3, 4, 1);

  tf_transform tf1{v1, q1};
  tf_transform tf2{v2, q2};
  for (auto _ : _state)
    benchmark::DoNotOptimize(tf2_product(tf1, tf2));
}

// benchmarks comparing the tf product of transform with the eigen-based.
BENCHMARK(eigen_isometry_prod);
BENCHMARK(tf_transform_accum_prod);


#ifndef KNAPSACK_HPP
#define KNAPSACK_HPP

#include <vector>
#include <array>
#include <cmath>

#include <hpx/config.hpp>
#include <boost/serialization/access.hpp>

#include <hpx/util/tuple.hpp>

#include "util/NodeGenerator.hpp"

/* A representation of a knapsack current solution */
struct KPSolution {
  std::vector<int> items;
  int profit;
  int weight;

  template <class Archive>
  void serialize(Archive & ar, const unsigned int version) {
    ar & items;
    ar & profit;
    ar & weight;
  }
};

template <unsigned N>
struct KPSpace {
  std::array<int, N> profits;
  std::array<int, N> weights;
  int numItems;
  int capacity;

  template <class Archive>
  void serialize(Archive & ar, const unsigned int version) {
    ar & profits;
    ar & weights;
    ar & numItems;
    ar & capacity;
  }
};

struct KPNode {
  KPSolution sol;
  std::vector<int> rem;

  int getObj() const {
    return sol.profit;
  }

  template <class Archive>
  void serialize(Archive & ar, const unsigned int version) {
    ar & sol;
    ar & rem;
  }
};

template <unsigned numItems>
struct GenNode : YewPar::NodeGenerator<KPNode, KPSpace<numItems> > {
  std::vector<int> items;
  int pos;

  std::reference_wrapper<const KPSpace<numItems> > space;
  std::reference_wrapper<const KPNode> n;

  GenNode (const KPSpace<numItems> & space, const KPNode & n) :
      pos(0), space(std::cref(space)), n(std::cref(n)) {
    this->numChildren = n.rem.size();
  }

  KPNode next() override {
    auto i = n.get().rem[pos];
    auto newSol = n.get().sol;
    newSol.items.push_back(i);
    newSol.profit += space.get().profits[i];
    newSol.weight += space.get().weights[i];

    ++pos;

    std::vector<int> newRem;
    std::copy_if(n.get().rem.begin() + pos, n.get().rem.end(), std::back_inserter(newRem),
                 [&](const int i) {
                   return newSol.weight + space.get().weights[i] <= space.get().capacity;
                 });

    return { newSol, std::move(newRem) };
  }
};

template <unsigned numItems>
int upperBound(const KPSpace<numItems> & space, const KPNode & n) {
  auto sol = n.sol;

  double profit = sol.profit;
  auto weight  = sol.weight;

  for (auto i = sol.items.back() + 1; i < space.numItems; i++) {
    // If there is enough space for a full item we take it all
    if (space.weights[i] + weight <= space.capacity) {
      profit += space.profits[i];
      weight += space.weights[i];
    } else {
      // Only space for some fraction of the last item
      profit = profit + (space.capacity - weight) * ((double) space.profits[i] / (double) space.weights[i]);
      break;
    }
  }

  return std::ceil(profit);
}

#endif


/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          http://www.mrpt.org/                          |
   |                                                                        |
   | Copyright (c) 2005-2017, Individual contributors, see AUTHORS file     |
   | See: http://www.mrpt.org/Authors - All rights reserved.                |
   | Released under BSD License. See details in http://www.mrpt.org/License |
   +------------------------------------------------------------------------+ */

#include "vision-precomp.h"  // Precompiled headers
#include <iostream>
#include <vector>
#include <cmath>

#include <mrpt/utils/types_math.h>  // Eigen must be included first via MRPT to enable the plugin system
#include <Eigen/Dense>
#include <Eigen/SVD>
#include <Eigen/StdVector>
#include <unsupported/Eigen/Polynomials>

#include "rpnp.h"

mrpt::vision::pnp::rpnp::rpnp(
	Eigen::MatrixXd obj_pts_, Eigen::MatrixXd img_pts_, Eigen::MatrixXd cam_,
	int n0)
{
	obj_pts = obj_pts_;
	img_pts = img_pts_;
	cam_intrinsic = cam_;
	n = n0;

	// Store obj_pts as 3XN and img_projections as 2XN matrices
	P = obj_pts.transpose();
	Q = img_pts.transpose();

	for (int i = 0; i < n; i++) Q.col(i) = Q.col(i) / Q.col(i).norm();

	R.setZero();
	t.setZero();
}

bool mrpt::vision::pnp::rpnp::compute_pose(
	Eigen::Ref<Eigen::Matrix3d> R_, Eigen::Ref<Eigen::Vector3d> t_)
{
	// selecting an edge $P_{ i1 }P_{ i2 }$ by n random sampling
	int i1 = 0, i2 = 1;
	double lmin =
		Q(0, i1) * Q(0, i2) + Q(1, i1) * Q(1, i2) + Q(2, i1) * Q(2, i2);

	Eigen::MatrixXi rij(n, 2);

	R_ = Eigen::MatrixXd::Identity(3, 3);
	t_ = Eigen::Vector3d::Zero();

	for (int i = 0; i < n; i++)
		for (int j = 0; j < 2; j++) rij(i, j) = rand() % n;

	for (int ii = 0; ii < n; ii++)
	{
		int i = rij(ii, 0), j = rij(ii, 1);

		if (i == j) continue;

		double l = Q(0, i) * Q(0, j) + Q(1, i) * Q(1, j) + Q(2, i) * Q(2, j);

		if (l < lmin)
		{
			i1 = i;
			i2 = j;
			lmin = l;
		}
	}

	// calculating the rotation matrix of $O_aX_aY_aZ_a$.
	Eigen::Vector3d p1, p2, p0, x, y, z, dum_vec;

	p1 = P.col(i1);
	p2 = P.col(i2);
	p0 = (p1 + p2) / 2;

	x = p2 - p0;
	x /= x.norm();

	if (std::abs(x(1)) < std::abs(x(2)))
	{
		dum_vec << 0, 1, 0;
		z = x.cross(dum_vec);
		z /= z.norm();
		y = z.cross(x);
		y /= y.norm();
	}
	else
	{
		dum_vec << 0, 0, 1;
		y = dum_vec.cross(x);
		y /= y.norm();
		z = x.cross(y);
		x /= x.norm();
	}

	Eigen::Matrix3d R0;

	R0.col(0) = x;
	R0.col(1) = y;
	R0.col(2) = z;

	for (int i = 0; i < n; i++) P.col(i) = R0.transpose() * (P.col(i) - p0);

	// Dividing the n - point set into(n - 2) 3 - point subsets
	// and setting up the P3P equations

	Eigen::Vector3d v1 = Q.col(i1), v2 = Q.col(i2);
	double cg1 = v1.dot(v2);
	double sg1 = sqrt(1 - cg1 * cg1);
	double D1 = (P.col(i1) - P.col(i2)).norm();
	Eigen::MatrixXd D4(n - 2, 5);

	int j = 0;
	Eigen::Vector3d vi;
	Eigen::VectorXd rowvec(5);
	for (int i = 0; i < n; i++)
	{
		if (i == i1 || i == i2) continue;

		vi = Q.col(i);
		double cg2 = v1.dot(vi);
		double cg3 = v2.dot(vi);
		double sg2 = sqrt(1 - cg2 * cg2);
		double D2 = (P.col(i1) - P.col(i)).norm();
		double D3 = (P.col(i) - P.col(i2)).norm();

		// get the coefficients of the P3P equation from each subset.

		rowvec = getp3p(cg1, cg2, cg3, sg1, sg2, D1, D2, D3);
		D4.row(j) = rowvec;
		j += 1;

		if (j > n - 3) break;
	}

	Eigen::VectorXd D7(8), dumvec(8), dumvec1(5);
	D7.setZero();

	for (int i = 0; i < n - 2; i++)
	{
		dumvec1 = D4.row(i);
		dumvec = getpoly7(dumvec1);
		D7 += dumvec;
	}

	Eigen::PolynomialSolver<double, 7> psolve(D7.reverse());
	Eigen::VectorXcd comp_roots = psolve.roots().transpose();
	Eigen::VectorXd real_comp, imag_comp;
	real_comp = comp_roots.real();
	imag_comp = comp_roots.imag();

	Eigen::VectorXd::Index max_index;

	double max_real = real_comp.cwiseAbs().maxCoeff(&max_index);

	std::vector<double> act_roots_;

	int cnt = 0;

	for (int i = 0; i < imag_comp.size(); i++)
	{
		if (std::abs(imag_comp(i)) / max_real < 0.001)
		{
			act_roots_.push_back(real_comp(i));
			cnt++;
		}
	}

	double* ptr = &act_roots_[0];
	Eigen::Map<Eigen::VectorXd> act_roots(ptr, cnt);

	if (cnt == 0)
	{
		return false;
	}

	Eigen::VectorXd act_roots1(cnt);
	act_roots1 << act_roots.segment(0, cnt);

	std::vector<Eigen::Matrix3d> R_cum(cnt);
	std::vector<Eigen::Vector3d> t_cum(cnt);
	std::vector<double> err_cum(cnt);

	for (int i = 0; i < cnt; i++)
	{
		double root = act_roots(i);

		// Compute the rotation matrix

		double d2 = cg1 + root;

		Eigen::Vector3d unitx, unity, unitz;
		unitx << 1, 0, 0;
		unity << 0, 1, 0;
		unitz << 0, 0, 1;
		x = v2 * d2 - v1;
		x /= x.norm();
		if (std::abs(unity.dot(x)) < std::abs(unitz.dot(x)))
		{
			z = x.cross(unity);
			z /= z.norm();
			y = z.cross(x);
			y / y.norm();
		}
		else
		{
			y = unitz.cross(x);
			y /= y.norm();
			z = x.cross(y);
			z /= z.norm();
		}
		R.col(0) = x;
		R.col(1) = y;
		R.col(2) = z;

		// calculating c, s, tx, ty, tz

		Eigen::MatrixXd D(2 * n, 6);
		D.setZero();

		R0 = R.transpose();
		Eigen::VectorXd r(
			Eigen::Map<Eigen::VectorXd>(R0.data(), R0.cols() * R0.rows()));

		for (int j = 0; j < n; j++)
		{
			double ui = img_pts(j, 0), vi = img_pts(j, 1), xi = P(0, j),
				   yi = P(1, j), zi = P(2, j);
			D.row(2 * j) << -r(1) * yi + ui * (r(7) * yi + r(8) * zi) -
								r(2) * zi,
				-r(2) * yi + ui * (r(8) * yi - r(7) * zi) + r(1) * zi, -1, 0,
				ui, ui * r(6) * xi - r(0) * xi;

			D.row(2 * j + 1)
				<< -r(4) * yi + vi * (r(7) * yi + r(8) * zi) - r(5) * zi,
				-r(5) * yi + vi * (r(8) * yi - r(7) * zi) + r(4) * zi, 0, -1,
				vi, vi * r(6) * xi - r(3) * xi;
		}

		Eigen::MatrixXd DTD = D.transpose() * D;

		Eigen::EigenSolver<Eigen::MatrixXd> es(DTD);

		Eigen::VectorXd Diag = es.pseudoEigenvalueMatrix().diagonal();

		Eigen::MatrixXd V_mat = es.pseudoEigenvectors();

		Eigen::MatrixXd::Index min_index;

		Diag.minCoeff(&min_index);

		Eigen::VectorXd V = V_mat.col(min_index);

		V /= V(5);

		double c = V(0), s = V(1);
		t << V(2), V(3), V(4);

		// calculating the camera pose by 3d alignment
		Eigen::VectorXd xi, yi, zi;
		xi = P.row(0);
		yi = P.row(1);
		zi = P.row(2);

		Eigen::MatrixXd XXcs(3, n), XXc(3, n);
		XXc.setZero();

		XXcs.row(0) = r(0) * xi + (r(1) * c + r(2) * s) * yi +
					  (-r(1) * s + r(2) * c) * zi +
					  t(0) * Eigen::VectorXd::Ones(n);
		XXcs.row(1) = r(3) * xi + (r(4) * c + r(5) * s) * yi +
					  (-r(4) * s + r(5) * c) * zi +
					  t(1) * Eigen::VectorXd::Ones(n);
		XXcs.row(2) = r(6) * xi + (r(7) * c + r(8) * s) * yi +
					  (-r(7) * s + r(8) * c) * zi +
					  t(2) * Eigen::VectorXd::Ones(n);

		for (int ii = 0; ii < n; ii++)
			XXc.col(ii) = Q.col(ii) * XXcs.col(ii).norm();

		Eigen::Matrix3d R2;
		Eigen::Vector3d t2;

		Eigen::MatrixXd XXw = obj_pts.transpose();

		calcampose(XXc, XXw, R2, t2);

		R_cum[i] = R2;
		t_cum[i] = t2;

		for (int k = 0; k < n; k++) XXc.col(k) = R2 * XXw.col(k) + t2;

		Eigen::MatrixXd xxc(2, n);

		xxc.row(0) = XXc.row(0).array() / XXc.row(2).array();
		xxc.row(1) = XXc.row(1).array() / XXc.row(2).array();

		double res = ((xxc.row(0) - img_pts.col(0).transpose()).norm() +
					  (xxc.row(1) - img_pts.col(1).transpose()).norm()) /
					 2;

		err_cum[i] = res;
	}

	int pos_cum =
		std::min_element(err_cum.begin(), err_cum.end()) - err_cum.begin();

	R_ = R_cum[pos_cum];
	t_ = t_cum[pos_cum];

	return true;
}

void mrpt::vision::pnp::rpnp::calcampose(
	Eigen::MatrixXd& XXc, Eigen::MatrixXd& XXw, Eigen::Matrix3d& R2,
	Eigen::Vector3d& t2)
{
	Eigen::MatrixXd X = XXc;
	Eigen::MatrixXd Y = XXw;
	Eigen::MatrixXd K =
		Eigen::MatrixXd::Identity(n, n) - Eigen::MatrixXd::Ones(n, n) * 1 / n;
	Eigen::VectorXd ux, uy;
	uy = X.rowwise().mean();
	ux = Y.rowwise().mean();

	// Need to verify sigmax2
	double sigmax2 =
		(((X * K).array() * (X * K).array()).colwise().sum()).mean();

	Eigen::MatrixXd SXY = Y * K * (X.transpose()) / n;

	Eigen::JacobiSVD<Eigen::MatrixXd> svd(
		SXY, Eigen::ComputeThinU | Eigen::ComputeThinV);

	Eigen::Matrix3d S = Eigen::MatrixXd::Identity(3, 3);
	if (SXY.determinant() < 0) S(2, 2) = -1;

	R2 = svd.matrixV() * S * svd.matrixU().transpose();

	double c2 = (svd.singularValues().asDiagonal() * S).trace() / sigmax2;
	t2 = uy - c2 * R2 * ux;

	Eigen::Vector3d x, y, z;
	x = R2.col(0);
	y = R2.col(1);
	z = R2.col(2);

	if ((x.cross(y) - z).norm() > 0.02) R2.col(2) = -R2.col(2);
}

Eigen::VectorXd mrpt::vision::pnp::rpnp::getpoly7(const Eigen::VectorXd& vin)
{
	Eigen::VectorXd vout(8);
	vout << 4 * pow(vin(0), 2), 7 * vin(1) * vin(0),
		6 * vin(2) * vin(0) + 3 * pow(vin(1), 2),
		5 * vin(3) * vin(0) + 5 * vin(2) * vin(1),
		4 * vin(4) * vin(0) + 4 * vin(3) * vin(1) + 2 * pow(vin(2), 2),
		3 * vin(4) * vin(1) + 3 * vin(3) * vin(2),
		2 * vin(4) * vin(2) + pow(vin(3), 2), vin(4) * vin(3);
	return vout;
}

Eigen::VectorXd mrpt::vision::pnp::rpnp::getp3p(
	double l1, double l2, double A5, double C1, double C2, double D1, double D2,
	double D3)
{
	double A1 = (D2 / D1) * (D2 / D1);
	double A2 = A1 * pow(C1, 2) - pow(C2, 2);
	double A3 = l2 * A5 - l1;
	double A4 = l1 * A5 - l2;
	double A6 = (pow(D3, 2) - pow(D1, 2) - pow(D2, 2)) / (2 * pow(D1, 2));
	double A7 = 1 - pow(l1, 2) - pow(l2, 2) + l1 * l2 * A5 + A6 * pow(C1, 2);

	Eigen::VectorXd vec(5);

	vec << pow(A6, 2) - A1 * pow(A5, 2), 2 * (A3 * A6 - A1 * A4 * A5),
		pow(A3, 2) + 2 * A6 * A7 - A1 * pow(A4, 2) - A2 * pow(A5, 2),
		2 * (A3 * A7 - A2 * A4 * A5), pow(A7, 2) - A2 * pow(A4, 2);

	return vec;
}


#include <iostream>
#include <fstream>
#include <sstream>
#include <map>
#include <list>
#include <string>
#include <set>
#include <vector>
#include <utility>
#include <boost/algorithm/string.hpp>

using namespace std;
using namespace boost;

struct Ing {
    __int64 quantity;
    string name;
};
typedef list<Ing> IngList;
struct RecipeElement {
    __int64 makes;
    IngList ings;
    RecipeElement() : makes(0) {}
};
typedef map<string, RecipeElement> Recipe;
typedef map<string, __int64> Stash;

Ing Parse(string s)
{
    Ing r;
    stringstream ss(s);
    ss >> r.quantity;
    ss >> r.name;
    trim(r.name);

    return r;
}

Recipe Load()
{
    Recipe r;
    ifstream f("Data.txt");

    while (!f.eof()) {
        string s;
        getline(f, s);

        vector<string> parts;
        split(parts, s, is_any_of("=>"), token_compress_on);
        vector<string> ings;
        split(ings, parts[0], is_any_of(","), token_compress_on);

        Ing result = Parse(parts[1]);
        auto& re = r[result.name];
        re.makes = result.quantity;

        for (const string& ing : ings) {
            Ing src = Parse(ing);
            re.ings.push_back(src);
        }
    }

    return r;
}

__int64 GetRequiredOre(const Recipe& r, const string& type, __int64 num, Stash& stash)
{
    __int64 ore = 0;
    __int64 in_stash = stash[type];

    //cout << "Need " << num << " " << type << "> ";

    // Check if we have enough in stash
    if (in_stash > num) {
        stash[type] -= num;
        //cout << "In stash" << endl;
        return 0;
    }

    // Exhaust stash first
    //cout << stash[type] << " from stash, ";
    num -= stash[type];
    stash.erase(type);

    // Get recipe
    auto it = r.find(type);
    if (it == r.end()) throw std::exception("No recipe!");
    auto& re = it->second;

    // Recipe makes X lots per batch, so need to make
    __int64 to_make = (__int64)ceil(1.0 * num / re.makes);
    //cout << "Making " << to_make * re.makes << "< ";
    //for (auto ing : re.ings) cout << ing.name << ": " << ing.quantity * to_make << " ";
    //cout << endl;

    for (auto ing : re.ings) {
        if (ing.name == "ORE"s) {
            ore += ing.quantity * to_make;
        }
        else {
            ore += GetRequiredOre(r, ing.name, ing.quantity * to_make, stash);
        }
    }

    stash[type] += (to_make * re.makes - num);
    return ore;
}

int main()
{
    Recipe r = Load();
    Stash stash;
    __int64 needs = GetRequiredOre(r, "FUEL", 1, stash);
    cout << "FUEL requires " << needs << " ORE" << endl;
    cout << "Left over:" << endl;
    for (auto x : stash) {
        if (x.second > 0) {
            cout << "  " << x.first << " = " << x.second << endl;
        }
    }

    // Part 2 - for 1e12 ore, how much fuel.  Search!
    __int64 target = 1000000000000i64;
    __int64 guess = target / needs + 1;
    __int64 step = 16777216;
    while (true) {
        stash.clear();
        needs = GetRequiredOre(r, "FUEL", guess, stash);
        cout << "Guess " << guess << " gives " << needs << endl;
        if (needs > target) {
            guess -= step;
        }
        else if (step > 1) {
            step /= 2;
            guess += step;
        }
        else {
            break;
        }
    }
    cout << guess << " FUEL from " << needs << " ORE" << endl;

    return 0;
}


#include "ale.h"

#include <nlohmann/json.hpp>

#include <gsl/gsl_interp.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_poly.h>

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

#include <iostream>
#include <Python.h>

#include <string>
#include <iostream>
#include <stdexcept>

using json = nlohmann::json;
using namespace std;

namespace ale {

  // Position Data Functions
  vector<double> getPosition(vector<vector<double>> coords, vector<double> times, double time,
                             interpolation interp) {
    // Check that all of the data sizes are okay
    // TODO is there a cleaner way to do this? We're going to have to do this a lot.
    if (coords.size() != 3) {
      throw invalid_argument("Invalid input positions, expected three vectors.");
    }

    // GSL setup
    vector<double> coordinate = {0.0, 0.0, 0.0};

    coordinate = { interpolate(coords[0], times, time, interp, 0),
                   interpolate(coords[1], times, time, interp, 0),
                   interpolate(coords[2], times, time, interp, 0) };

    return coordinate;
  }

  vector<double> getVelocity(vector<vector<double>> coords, vector<double> times,
                             double time, interpolation interp) {
    // Check that all of the data sizes are okay
    // TODO is there a cleaner way to do this? We're going to have to do this a lot.
    if (coords.size() != 3) {
     throw invalid_argument("Invalid input positions, expected three vectors.");
    }

    // GSL setup
    vector<double> coordinate = {0.0, 0.0, 0.0};

    coordinate = { interpolate(coords[0], times, time, interp, 1),
                   interpolate(coords[1], times, time, interp, 1),
                   interpolate(coords[2], times, time, interp, 1) };

    return coordinate;
  }

  // Postion Function Functions
  // vector<double> coeffs = [[cx_0, cx_1, cx_2 ..., cx_n],
  //                          [cy_0, cy_1, cy_2, ... cy_n],
  //                          [cz_0, cz_1, cz_2, ... cz_n]]
  // The equations evaluated by this function are:
  //                x = cx_n * t^n + cx_n-1 * t^(n-1) + ... + cx_0
  //                y = cy_n * t^n + cy_n-1 * t^(n-1) + ... + cy_0
  //                z = cz_n * t^n + cz_n-1 * t^(n-1) + ... + cz_0
  vector<double> getPosition(vector<vector<double>> coeffs, double time) {

    if (coeffs.size() != 3) {
      throw invalid_argument("Invalid input coeffs, expected three vectors.");
    }

    vector<double> coordinate = {0.0, 0.0, 0.0};
    coordinate[0] = evaluatePolynomial(coeffs[0], time, 0); // X
    coordinate[1] = evaluatePolynomial(coeffs[1], time, 0); // Y
    coordinate[2] = evaluatePolynomial(coeffs[2], time, 0); // Z

    return coordinate;
  }

  // Velocity Function
  // Takes the coefficients from the position equation
  vector<double> getVelocity(vector<vector<double>> coeffs, double time) {

    if (coeffs.size() != 3) {
      throw invalid_argument("Invalid input coeffs, expected three vectors.");
    }

    vector<double> coordinate = {0.0, 0.0, 0.0};
    coordinate[0] = evaluatePolynomial(coeffs[0], time, 1); // X
    coordinate[1] = evaluatePolynomial(coeffs[1], time, 1); // Y
    coordinate[2] = evaluatePolynomial(coeffs[2], time, 1); // Z

    return coordinate;
  }

  // Rotation Data Functions
  vector<double> getRotation(vector<vector<double>> rotations,
                             vector<double> times, double time,  interpolation interp) {
    // Check that all of the data sizes are okay
    // TODO is there a cleaner way to do this? We're going to have to do this a lot.
    if (rotations.size() != 4) {
     throw invalid_argument("Invalid input rotations, expected four vectors.");
    }

    // Alot of copying and reassignment becuase conflicting data types
    // probably should rethink our vector situation to guarentee contiguous
    // memory. Should be easy to switch to a contiguous column-major format
    // if we stick with Eigen.
    for (size_t i = 0; i<rotations[0].size(); i++) {
      Eigen::Quaterniond quat(rotations[0][i], rotations[1][i], rotations[2][i], rotations[3][i]);
      quat.normalize();

      rotations[0][i] = quat.w();
      rotations[1][i] = quat.x();
      rotations[2][i] = quat.y();
      rotations[3][i] = quat.z();
    }

    // GSL setup
    vector<double> coordinate = {0.0, 0.0, 0.0, 0.0};

    coordinate = { interpolate(rotations[0], times, time, interp, 0),
                   interpolate(rotations[1], times, time, interp, 0),
                   interpolate(rotations[2], times, time, interp, 0),
                   interpolate(rotations[3], times, time, interp, 0)};

    // Eigen::Map to ensure the array isn't copied, only the pointer is
    Eigen::Map<Eigen::MatrixXd> quat(coordinate.data(), 4, 1);
    quat.normalize();
    return coordinate;
  }

  vector<double> getAngularVelocity(vector<vector<double>> rotations,
                                    vector<double> times, double time,  interpolation interp) {
    // Check that all of the data sizes are okay
    // TODO is there a cleaner way to do this? We're going to have to do this a lot.
    if (rotations.size() != 4) {
     throw invalid_argument("Invalid input rotations, expected four vectors.");
    }

    double data[] = {0,0,0,0};
    for (size_t i = 0; i<rotations[0].size(); i++) {
      Eigen::Quaterniond quat(rotations[0][i], rotations[1][i], rotations[2][i], rotations[3][i]);
      quat.normalize();
      rotations[0][i] = quat.w();
      rotations[1][i] = quat.x();
      rotations[2][i] = quat.y();
      rotations[3][i] = quat.z();
    }

    // GSL setup

    Eigen::Quaterniond quat(interpolate(rotations[0], times, time, interp, 0),
                            interpolate(rotations[1], times, time, interp, 0),
                            interpolate(rotations[2], times, time, interp, 0),
                            interpolate(rotations[3], times, time, interp, 0));
    quat.normalize();

    Eigen::Quaterniond dQuat(interpolate(rotations[0], times, time, interp, 1),
                             interpolate(rotations[1], times, time, interp, 1),
                             interpolate(rotations[2], times, time, interp, 1),
                             interpolate(rotations[3], times, time, interp, 1));

     Eigen::Quaterniond avQuat = quat.conjugate() * dQuat;

     vector<double> coordinate = {-2 * avQuat.x(), -2 * avQuat.y(), -2 * avQuat.z()};
     return coordinate;
  }

  // Rotation Function Functions
  std::vector<double> getRotation(vector<vector<double>> coeffs, double time) {

    if (coeffs.size() != 3) {
      throw invalid_argument("Invalid input coefficients, expected three vectors.");
    }

    vector<double> rotation = {0.0, 0.0, 0.0};

    rotation[0] = evaluatePolynomial(coeffs[0], time, 0); // X
    rotation[1] = evaluatePolynomial(coeffs[1], time, 0); // Y
    rotation[2] = evaluatePolynomial(coeffs[2], time, 0); // Z

    Eigen::Quaterniond quat;
    quat = Eigen::AngleAxisd(rotation[0] * M_PI / 180, Eigen::Vector3d::UnitZ())
                * Eigen::AngleAxisd(rotation[1] * M_PI / 180, Eigen::Vector3d::UnitX())
                * Eigen::AngleAxisd(rotation[2] * M_PI / 180, Eigen::Vector3d::UnitZ());

    quat.normalize();

    vector<double> rotationQ = {quat.w(), quat.x(), quat.y(), quat.z()};
    return rotationQ;
  }

  vector<double> getAngularVelocity(vector<vector<double>> coeffs, double time) {

    if (coeffs.size() != 3) {
      throw invalid_argument("Invalid input coefficients, expected three vectors.");
    }

    double phi = evaluatePolynomial(coeffs[0], time, 0); // X
    double theta = evaluatePolynomial(coeffs[1], time, 0); // Y
    double psi = evaluatePolynomial(coeffs[2], time, 0); // Z

    double phi_dt = evaluatePolynomial(coeffs[0], time, 1);
    double theta_dt = evaluatePolynomial(coeffs[1], time, 1);
    double psi_dt = evaluatePolynomial(coeffs[2], time, 1);

    Eigen::Quaterniond quat1, quat2;
    quat1 = Eigen::AngleAxisd(phi * M_PI / 180, Eigen::Vector3d::UnitZ());
    quat2 =  Eigen::AngleAxisd(theta * M_PI / 180, Eigen::Vector3d::UnitX());

    Eigen::Vector3d velocity =  phi_dt * Eigen::Vector3d::UnitZ();
    velocity += theta_dt * (quat1 *  Eigen::Vector3d::UnitX());
    velocity += psi_dt * (quat1 * quat2 *  Eigen::Vector3d::UnitZ());

    return {velocity[0], velocity[1], velocity[2]};
  }

  // Polynomial evaluation helper function
  // The equation evaluated by this function is:
  //                x = cx_0 + cx_1 * t^(1) + ... + cx_n * t^n
  // The d parameter is for which derivative of the polynomial to compute.
  // Supported options are
  //   0: no derivative
  //   1: first derivative
  //   2: second derivative
  double evaluatePolynomial(vector<double> coeffs, double time, int d){
    if (coeffs.empty()) {
      throw invalid_argument("Invalid input coeffs, must be non-empty.");
    }

    if (d < 0) {
      throw invalid_argument("Invalid derivative degree, must be non-negative.");
    }

    vector<double> derivatives(d + 1);
    gsl_poly_eval_derivs(coeffs.data(), coeffs.size(), time,
                         derivatives.data(), derivatives.size());

    return derivatives.back();
  }

 double interpolate(vector<double> points, vector<double> times, double time, interpolation interp, int d) {
   size_t numPoints = points.size();
   if (numPoints < 2) {
     throw invalid_argument("At least two points must be input to interpolate over.");
   }
   if (points.size() != times.size()) {
     throw invalid_argument("Invalid gsl_interp_type data, must have the same number of points as times.");
   }
   if (time < times.front() || time > times.back()) {
     throw invalid_argument("Invalid gsl_interp_type time, outside of input times.");
   }

   // convert our interp enum into a GSL one,
   // should be easy to add non GSL interp methods here later
   const gsl_interp_type *interp_methods[] = {gsl_interp_linear, gsl_interp_cspline};

   gsl_interp *interpolator = gsl_interp_alloc(interp_methods[interp], numPoints);
   gsl_interp_init(interpolator, &times[0], &points[0], numPoints);
   gsl_interp_accel *acc = gsl_interp_accel_alloc();

   // GSL evaluate
   double result;
   switch(d) {
     case 0:
       result = gsl_interp_eval(interpolator, &times[0], &points[0], time, acc);
       break;
     case 1:
       result = gsl_interp_eval_deriv(interpolator, &times[0], &points[0], time, acc);
       break;
     case 2:
       result = gsl_interp_eval_deriv2(interpolator, &times[0], &points[0], time, acc);
       break;
     default:
       throw invalid_argument("Invalid derivitive option, must be 0, 1 or 2.");
       break;
   }

   // GSL clean up
   gsl_interp_free(interpolator);
   gsl_interp_accel_free(acc);

   return result;
 }

 std::string getPyTraceback() {
    PyObject* err = PyErr_Occurred();
    if (err != NULL) {
        PyObject *ptype, *pvalue, *ptraceback;
        PyObject *pystr, *module_name, *pyth_module, *pyth_func;
        char *str;
        char *full_backtrace;
        char *error_description;

        PyErr_Fetch(&ptype, &pvalue, &ptraceback);
        pystr = PyObject_Str(pvalue);
        str = PyBytes_AS_STRING(PyUnicode_AsUTF8String(pystr));
        error_description = strdup(str);

        // See if we can get a full traceback
        module_name = PyUnicode_FromString("traceback");
        pyth_module = PyImport_Import(module_name);
        Py_DECREF(module_name);

        if (pyth_module == NULL) {
            throw runtime_error("getPyTraceback - Failed to import Python traceback Library");
        }

        pyth_func = PyObject_GetAttrString(pyth_module, "format_exception");
        PyObject *pyth_val;
        pyth_val = PyObject_CallFunctionObjArgs(pyth_func, ptype, pvalue, ptraceback, NULL);

        pystr = PyObject_Str(pyth_val);
        str = PyBytes_AS_STRING(PyUnicode_AsUTF8String(pystr));
        full_backtrace = strdup(str);
        Py_DECREF(pyth_val);

        std::string join_cmd = "trace = ''.join(list(" + std::string(full_backtrace) + "))";
        PyRun_SimpleString(join_cmd.c_str());

        PyObject *evalModule = PyImport_AddModule( (char*)"__main__" );
        PyObject *evalDict = PyModule_GetDict( evalModule );
        PyObject *evalVal = PyDict_GetItemString( evalDict, "trace" );
        full_backtrace = PyBytes_AS_STRING(PyUnicode_AsUTF8String(evalVal));

        return std::string(error_description) + "\n" + std::string(full_backtrace);
    }

    // no traceback to return
    return "";
 }

 std::string loads(std::string filename, std::string props, std::string formatter, bool verbose) {
     static bool first_run = true;
     if(first_run) {
         // Initialize the Python interpreter but only once.
         first_run = !first_run;
         Py_Initialize();
         atexit(Py_Finalize);
     }

     // Import the file as a Python module.
     PyObject *pModule = PyImport_Import(PyUnicode_FromString("ale"));
     if(!pModule) {
       throw runtime_error("Failed to import ale. Make sure the ale python library is correctly installed.");
     }
     // Create a dictionary for the contents of the module.
     PyObject *pDict = PyModule_GetDict(pModule);

     // Get the add method from the dictionary.
     PyObject *pFunc = PyDict_GetItemString(pDict, "loads");
     if(!pFunc) {
       // import errors do not set a PyError flag, need to use a custom
       // error message instead.
       throw runtime_error("Failed to import ale.loads function from Python."
                           "This Usually indicates an error in the Ale Python Library."
                           "Check if Installed correctly and the function ale.loads exists.");
     }

     // Create a Python tuple to hold the arguments to the method.
     PyObject *pArgs = PyTuple_New(3);
     if(!pArgs) {
       throw runtime_error(getPyTraceback());
     }

     // Set the Python int as the first and second arguments to the method.
     PyObject *pStringFileName = PyUnicode_FromString(filename.c_str());
     PyTuple_SetItem(pArgs, 0, pStringFileName);

     PyObject *pStringProps = PyUnicode_FromString(props.c_str());
     PyTuple_SetItem(pArgs, 1, pStringProps);

     PyObject *pStringFormatter = PyUnicode_FromString(formatter.c_str());
     PyTuple_SetItem(pArgs, 2, pStringFormatter);

     // Call the function with the arguments.
     PyObject* pResult = PyObject_CallObject(pFunc, pArgs);

     if(!pResult) {
        throw invalid_argument("No Valid instrument found for label.");
     }

     PyObject *pResultStr = PyObject_Str(pResult);
     PyObject *temp_bytes = PyUnicode_AsUTF8String(pResultStr); // Owned reference

     if(!temp_bytes){
       throw invalid_argument(getPyTraceback());
     }
     std::string cResult;
     char *temp_str = PyBytes_AS_STRING(temp_bytes); // Borrowed pointer
     cResult = temp_str; // copy into std::string

     Py_DECREF(pResultStr);
     Py_DECREF(pStringFileName);
     Py_DECREF(pStringProps);
     Py_DECREF(pStringFormatter);

     return cResult;
 }

 json load(std::string filename, std::string props, std::string formatter, bool verbose) {
   std::string jsonstr = loads(filename, props, formatter, verbose);
   return json::parse(jsonstr);
 }
}


// 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.
#include <gtest/gtest.h>

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/weighted_kurtosis.hpp>
#include <boost/accumulators/statistics/weighted_mean.hpp>
#include <boost/accumulators/statistics/weighted_median.hpp>
#include <boost/accumulators/statistics/weighted_skewness.hpp>
#include <boost/accumulators/statistics/weighted_sum.hpp>
#include <boost/accumulators/statistics/weighted_variance.hpp>

#include "pyinterp/detail/math/descriptive_statistics.hpp"

namespace math = pyinterp::detail::math;

static double x[20] = {0.00402322, 0.19509434, 0.6425439,  0.66463742,
                       0.76523411, 0.91985221, 0.82729929, 0.21502902,
                       0.48254104, 0.97854649, 0.61394511, 0.00583773,
                       0.06630172, 0.57173946, 0.5881294,  0.30185368,
                       0.18126563, 0.84524097, 0.13754961, 0.17343529};
static double w[20] = {0.45463566, 0.46341234, 0.2072285,  0.02272363,
                       0.76796619, 0.01987153, 0.43634701, 0.1369698,
                       0.65012667, 0.18825124, 0.96310554, 0.31995482,
                       0.28808939, 0.69961506, 0.97369255, 0.98436659,
                       0.05230501, 0.8073624,  0.40509977, 0.6325752};

using Accumulators = boost::accumulators::accumulator_set<
    double,
    boost::accumulators::stats<
        boost::accumulators::tag::count, boost::accumulators::tag::max,
        boost::accumulators::tag::min, boost::accumulators::tag::sum_of_weights,
        boost::accumulators::tag::weighted_kurtosis,
        boost::accumulators::tag::weighted_mean,
        // boost::accumulators::tag::weighted_median(
        //     boost::accumulators::with_p_square_quantile),
        boost::accumulators::tag::weighted_skewness,
        boost::accumulators::tag::weighted_sum,
        boost::accumulators::tag::weighted_variance(boost::accumulators::lazy)>,
    double>;

TEST(math_descriptive_statistics, univariate) {
  auto boost_acc = Accumulators();
  auto acc = math::DescriptiveStatistics<double>();

  for (auto ix = 0; ix < 20; ++ix) {
    boost_acc(x[ix], boost::accumulators::weight = 1);
    acc(x[ix]);
  }

  EXPECT_EQ(boost::accumulators::count(boost_acc), acc.count());
  EXPECT_DOUBLE_EQ(boost::accumulators::min(boost_acc), acc.min());
  EXPECT_DOUBLE_EQ(boost::accumulators::max(boost_acc), acc.max());
  EXPECT_DOUBLE_EQ(boost::accumulators::weighted_mean(boost_acc), acc.mean());
  EXPECT_NEAR(boost::accumulators::weighted_variance(boost_acc), acc.variance(),
              1e-12);
  EXPECT_NEAR(boost::accumulators::weighted_kurtosis(boost_acc), acc.kurtosis(),
              1e-12);
  EXPECT_NEAR(boost::accumulators::weighted_skewness(boost_acc), acc.skewness(),
              1e-12);
  EXPECT_DOUBLE_EQ(boost::accumulators::weighted_sum(boost_acc), acc.sum());
  EXPECT_DOUBLE_EQ(boost::accumulators::sum_of_weights(boost_acc),
                   acc.sum_of_weights());

  auto copy = math::DescriptiveStatistics<double>(
      static_cast<math::Accumulators<double>>(acc));
  EXPECT_EQ(boost::accumulators::count(boost_acc), copy.count());
  EXPECT_DOUBLE_EQ(boost::accumulators::min(boost_acc), copy.min());
  EXPECT_DOUBLE_EQ(boost::accumulators::max(boost_acc), copy.max());
  EXPECT_DOUBLE_EQ(boost::accumulators::weighted_mean(boost_acc), copy.mean());
  EXPECT_NEAR(boost::accumulators::weighted_variance(boost_acc),
              copy.variance(), 1e-12);
  EXPECT_NEAR(boost::accumulators::weighted_kurtosis(boost_acc),
              copy.kurtosis(), 1e-12);
  EXPECT_NEAR(boost::accumulators::weighted_skewness(boost_acc),
              copy.skewness(), 1e-12);
  EXPECT_DOUBLE_EQ(boost::accumulators::weighted_sum(boost_acc), copy.sum());
  EXPECT_DOUBLE_EQ(boost::accumulators::sum_of_weights(boost_acc),
                   copy.sum_of_weights());
}

TEST(math_descriptive_statistics, weighted) {
  auto boost_acc = Accumulators();
  auto acc = math::DescriptiveStatistics<double>();
  auto min = std::numeric_limits<double>::max();
  auto max = std::numeric_limits<double>::min();

  for (auto ix = 0; ix < 20; ++ix) {
    boost_acc(x[ix], boost::accumulators::weight = w[ix]);
    acc(x[ix], w[ix]);
    auto value = x[ix] * w[ix];
    min = std::min(min, value);
    max = std::max(max, value);
  }

  EXPECT_EQ(boost::accumulators::count(boost_acc), acc.count());
  EXPECT_DOUBLE_EQ(min, acc.min());
  EXPECT_DOUBLE_EQ(max, acc.max());
  EXPECT_DOUBLE_EQ(boost::accumulators::weighted_mean(boost_acc), acc.mean());
  EXPECT_NEAR(boost::accumulators::weighted_variance(boost_acc), acc.variance(),
              1e-12);
  EXPECT_NEAR(boost::accumulators::weighted_kurtosis(boost_acc), acc.kurtosis(),
              1e-12);
  EXPECT_NEAR(boost::accumulators::weighted_skewness(boost_acc), acc.skewness(),
              1e-12);
  EXPECT_DOUBLE_EQ(boost::accumulators::weighted_sum(boost_acc), acc.sum());
  EXPECT_DOUBLE_EQ(boost::accumulators::sum_of_weights(boost_acc),
                   acc.sum_of_weights());
}


//==============================================================================
//         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_EXPONENTIAL_FUNCTIONS_SCALAR_IMPL_LOGS_F_LOG_HPP_INCLUDED
#define NT2_EXPONENTIAL_FUNCTIONS_SCALAR_IMPL_LOGS_F_LOG_HPP_INCLUDED

#include <nt2/include/functions/scalar/is_ltz.hpp>
#include <nt2/include/functions/scalar/is_eqz.hpp>
#include <nt2/include/functions/scalar/fma.hpp>
#include <nt2/exponential/functions/scalar/impl/logs/f_kernel.hpp>
#include <nt2/include/constants/zero.hpp>
#include <nt2/include/constants/nan.hpp>
#include <nt2/include/constants/minf.hpp>
#include <nt2/include/constants/mhalf.hpp>
#include <nt2/include/constants/log_2hi.hpp>
#include <nt2/include/constants/log_2lo.hpp>
#include <nt2/include/constants/log2_em1.hpp>
#include <nt2/include/constants/log10_ehi.hpp>
#include <nt2/include/constants/log10_elo.hpp>
#include <nt2/include/constants/log10_2hi.hpp>
#include <nt2/include/constants/log10_2lo.hpp>
#include <nt2/sdk/meta/as_integer.hpp>
#include <nt2/sdk/meta/scalar_of.hpp>
#include <boost/simd/sdk/config.hpp>

#ifndef BOOST_SIMD_NO_NANS
#include <nt2/include/functions/scalar/is_nan.hpp>
#endif
#ifndef BOOST_SIMD_NO_INFINITIES
#include <nt2/include/constants/inf.hpp>
#endif
#ifndef BOOST_SIMD_NO_DENORMALS
#include <nt2/include/functions/scalar/abs.hpp>
#include <nt2/include/constants/smallestposval.hpp>
#include <nt2/include/constants/twotonmb.hpp>
#include <nt2/include/constants/mlogtwo2nmb.hpp>
#include <nt2/include/constants/mlog2two2nmb.hpp>
#include <nt2/include/constants/mlog10two2nmb.hpp>
#endif

namespace nt2 { namespace details
{
  template < class A0,
             class Style ,
             class base_A0 = typename meta::scalar_of<A0>::type>
             struct logarithm{};

  template < class A0 >
  struct logarithm< A0, tag::not_simd_type, float>
  {
    typedef typename meta::as_integer<A0, signed>::type    int_type;
    typedef typename meta::scalar_of<A0>::type                  sA0;
    typedef kernel<A0, tag::not_simd_type, float>          kernel_t;
    static inline A0 log(const A0& a0)
    {
#ifndef BOOST_SIMD_NO_INFINITIES
      if (BOOST_UNLIKELY(a0 == nt2::Inf<A0>())) return a0;
#endif
      if (BOOST_UNLIKELY(nt2::is_eqz(a0))) return nt2::Minf<A0>();
#ifdef BOOST_SIMD_NO_NANS
      if (BOOST_UNLIKELY(nt2::is_ltz(a0))) return nt2::Nan<A0>();
#else
      if (BOOST_UNLIKELY(nt2::is_nan(a0)||nt2::is_ltz(a0))) return nt2::Nan<A0>();
#endif
      A0 z = a0;
#ifndef BOOST_SIMD_NO_DENORMALS
      A0 t = Zero<A0>();
      if(BOOST_UNLIKELY(nt2::abs(z) < Smallestposval<A0>()))
      {
        z *= Twotonmb<A0>();
        t = Mlogtwo2nmb<A0>();
      }
#endif
      A0 x, fe, x2, y;
      kernel_t::log(z, fe, x, x2, y);
      y = nt2::fma(fe, Log_2lo<A0>(), y);
      y = nt2::fma(Mhalf<A0>(), x2, y);
#ifdef BOOST_SIMD_NO_DENORMALS
      return nt2::fma(Log_2hi<A0>(), fe, x+y);
#else
      return nt2::fma(Log_2hi<A0>(), fe, x+y+t);
#endif
    }

    static inline A0 log2(const A0& a0)
    {
#ifndef BOOST_SIMD_NO_INFINITIES
      if (BOOST_UNLIKELY(a0 == nt2::Inf<A0>())) return a0;
#endif
      if (BOOST_UNLIKELY(nt2::is_eqz(a0))) return nt2::Minf<A0>();
#ifdef BOOST_SIMD_NO_NANS
      if (BOOST_UNLIKELY(nt2::is_ltz(a0))) return nt2::Nan<A0>();
#else
      if (BOOST_UNLIKELY(nt2::is_nan(a0)||nt2::is_ltz(a0))) return nt2::Nan<A0>();
#endif
      A0 z = a0;
#ifndef BOOST_SIMD_NO_DENORMALS
      A0 t = Zero<A0>();
      if (BOOST_UNLIKELY(nt2::abs(z) < Smallestposval<A0>()))
      {
        z *= Twotonmb<A0>();
        t = Mlog2two2nmb<A0>();
      }
#endif
      A0 x, fe, x2, y;
      kernel_t::log(z, fe, x, x2, y);
      y = nt2::fma(Mhalf<A0>(),x2, y);
      z = nt2::fma(x,Log2_em1<A0>(),y*Log2_em1<A0>());
#ifdef BOOST_SIMD_NO_DENORMALS
      return ((z+y)+x)+fe;
#else
      return ((z+y)+x)+fe+t;
#endif
    }

    static inline A0 log10(const A0& a0)
    {
#ifndef BOOST_SIMD_NO_INFINITIES
      if (BOOST_UNLIKELY(a0 == nt2::Inf<A0>())) return a0;
#endif
      if (BOOST_UNLIKELY(nt2::is_eqz(a0))) return nt2::Minf<A0>();
#ifdef BOOST_SIMD_NO_NANS
      if (BOOST_UNLIKELY(nt2::is_ltz(a0))) return nt2::Nan<A0>();
#else
      if (BOOST_UNLIKELY(nt2::is_nan(a0)||nt2::is_ltz(a0))) return nt2::Nan<A0>();
#endif
      A0 z = a0;
#ifndef BOOST_SIMD_NO_DENORMALS
      A0 t = Zero<A0>();
      if (BOOST_UNLIKELY(nt2::abs(z) < Smallestposval<A0>()))
      {
        z *= Twotonmb<A0>();
        t = Mlog10two2nmb<A0>();
      }
#endif
      A0 x, fe, x2, y;
      kernel_t::log(z, fe, x, x2, y);
      y = nt2::amul(y, Mhalf<A0>(), x2);
      z = mul(x+y, Log10_elo<A0>());
      z = nt2::amul(z, y, Log10_ehi<A0>());
      z = nt2::amul(z, x, Log10_ehi<A0>());
      z = nt2::amul(z, fe, Log10_2hi<A0>());
#ifdef BOOST_SIMD_NO_DENORMALS
      return nt2::amul(z, fe, Log10_2lo<A0>());
#else
      return nt2::amul(z+t, fe, Log10_2lo<A0>());
#endif
    }
  };
} }

#endif


// Filthy awful hack to test.
#include "gtest/gtest.h"
#include <Eigen/Dense>
#define private public
#include "bond.h"

TEST(ExcitementTest, ZeroExcitement)
{
    // Test the excitement factor of a single bond that isn't stretched
    Eigen::VectorXd rail(2);
    rail << 1.0, 0.0;
    const HarmonicBond bond {1.0, 1.0, 1, 0, rail};

    Eigen::MatrixXd positions(2, 2);
    positions << 0.0, 0.0,
                 1.0, 0.0;

    ASSERT_FLOAT_EQ(bond.get_excitement_factor(positions), 1.0);
}

TEST(ExcitementTest, NegativeExcitement)
{
    // Test the excitement factor of a single bond that is stretched the other way
    Eigen::VectorXd rail(2);
    rail << 1.0, 0.0;
    const HarmonicBond bond {1.0, 1.0, 1, 0, rail};

    Eigen::MatrixXd positions(2, 2);
    positions << 1.0, 0.0,
                 0.0, 0.0;

    ASSERT_FLOAT_EQ(bond.get_excitement_factor(positions), -1.0);
}

TEST(ExcitementTest, DoubleNegativeExcitement)
{
    // Test the excitement factor of a single bond that is stretched the other way
    Eigen::VectorXd rail(2);
    rail << 1.0, 0.0;
    const HarmonicBond bond {1.0, 1.0, 1, 0, rail};

    Eigen::MatrixXd positions(2, 2);
    positions << 2.0, 0.0,
                 0.0, 0.0;

    ASSERT_FLOAT_EQ(bond.get_excitement_factor(positions), -2.0);
}

TEST(ExcitementTest, DoubleExcitement)
{
    // Test the excitement factor of a single bond that is stretched the other way
    Eigen::VectorXd rail(2);
    rail << 1.0, 0.0;
    const HarmonicBond bond {1.0, 1.0, 1, 0, rail};

    Eigen::MatrixXd positions(2, 2);
    positions << 0.0, 0.0,
                 2.0, 0.0;

    ASSERT_FLOAT_EQ(bond.get_excitement_factor(positions), 2.0);
}

TEST(ExcitementTest, OffRailDoubleExcitement)
{
    // Test the excitement factor of a single bond that is off the rail
    Eigen::VectorXd rail(2);
    rail << 1.0, 0.0;
    const HarmonicBond bond {1.0, 1.0, 1, 0, rail};

    Eigen::MatrixXd positions(2, 2);
    positions << 0.0, 0.0,
                 2.0, 1.0;

    ASSERT_FLOAT_EQ(bond.get_excitement_factor(positions), 2.0);
}


#include <functional>
#include <map>
#include <random>
#include <set>
#include <vector>

#include <iostream>

#include "third_party/catch.h"
#include "result.h"
#include <boost/container/flat_set.hpp>
#include <boost/container/flat_map.hpp>

using boost::container::flat_map;
using boost::container::flat_set;

template <typename K, typename V>
flat_map<K, flat_set<V>> build_flat_map_of_flat_sets(std::vector<std::pair<K, V>> buf) {
  std::sort(buf.begin(), buf.end());
  buf.erase(std::unique(buf.begin(), buf.end()), buf.end());

  flat_map<K, flat_set<V>> res;
  for (auto r : srt::group_equals(buf.begin(), buf.end(),
                                  [](const auto& x, const auto& y) { return x.first < y.first; })) {
    res.emplace_hint(res.end(), std::move(r.begin()->first), flat_set<V>{});
    auto& cur_set = (--res.end())->second;
    for (auto& elem : r)
      cur_set.insert(cur_set.end(), std::move(elem.second));
  }

  return res;
}

template <typename C>
std::vector<std::pair<int, int>> to_vector_of_pairs_for_test(const C& c) {
  std::vector<std::pair<int, int>> res;
  for (const auto& pr : c) {
    for (int v : pr.second)
      res.emplace_back(pr.first, v);
  }
  return res;
}

TEST_CASE("build_flat_map_of_flat_sets", "[usage_examples]") {
  const std::vector<std::pair<int, int>> input = [] {
    std::uniform_int_distribution<> dist(0, 100);
    std::mt19937 g;

    std::vector<std::pair<int, int>> res;

    for (int i = 0; i < 100; ++i) {
      res.emplace_back(dist(g), dist(g));
    }

    return res;
  }();

  for (auto f = input.begin(); f != input.end(); ++f)
    for (auto l = f; l != input.end(); ++l) {
      std::map<int, std::set<int>> expected;
      for (auto i = f; i != l; ++i)
        expected[i->first].insert(i->second);

      auto actual = build_flat_map_of_flat_sets(std::vector<std::pair<int, int>>{f, l});

      auto expected_vec = to_vector_of_pairs_for_test(expected);
      auto actual_vec = to_vector_of_pairs_for_test(actual);

      REQUIRE(expected_vec.size() == actual_vec.size());
      REQUIRE(expected_vec == actual_vec);
    }
}


#ifndef PROJECT_GRIDPOINT_HPP
#define PROJECT_GRIDPOINT_HPP

#include <stdint.h>
#include <string>
#include <limits>

#include <Eigen/Dense>

namespace fsd {

using Eigen::Vector2d;
using std::string;
/**
 * @brief The GridPoint class representing a point in a grid by row, column, for
 * fast lookup
 */
class GridPoint {
public:
    /**
     * @brief the type of point this gridpoint is constructed from,
     * must have a (x, y)/coordinate constructor and x(), y()
     * / cordinate getter methods. and probably more ..
     *
    **/
    typedef Vector2d from_point_t;

    typedef int32_t coordinates_t;
    typedef int64_t key_t;
    /**
     * must be doulbe for packing, unpacking
     */
    static_assert(sizeof(coordinates_t) * 2 ==
                  sizeof(key_t), "key_t must be double "
                                 "the size of coordinates_t for"
                                 " packing and stuff");
    /**
     * @brief getter for row
     * @return row
     */
    coordinates_t get_row() const {
        return row;
    }

    /**
     * @brief getter for column
     * @return column
     */
    coordinates_t  get_column() const {
        return column;
    }

    /**
     * @brief hit_cnt how many points mapped to this cell
     */
    uint32_t hit_cnt = 1;

    /**
     * @brief The coordinates struct used for other grid things
     */
    struct coordinates {
        coordinates_t row, column;
        coordinates(coordinates_t row_,
                    coordinates_t column_) :
            row(row_), column(column_) {}
        coordinates() = default;
    };

    static_assert (sizeof (coordinates) == sizeof (key_t), "key size must equal "
                                                           "coordinates");

    /**
     * @brief to_key creates collision free key by packing row and column
     * in one integer
     * @return key
     */
    key_t to_key() const {
        coordinates coords;
        // copy instead of brace initializer to avoid some bugs
        // if the struct is redefined
        coords.row = row;
        coords.column = column;
        return pack_key(coords);
    }

    /**
     * @brief to_coordinates point to coordinates helper
     * @param point
     * @param cell_size
     * @param scale to use more entropy on the coordinates
     * and therefore have less collisions when saving in a hashtable,
     * the coordinates can be scaled the use the whole range
     * of coordinates_t type
     * @return
     */
    static inline coordinates to_coordinates(
            const from_point_t &point, double cell_size, double scale = 1.) {
        double x = point.x() * scale;
        if(x < 0.) {
            x -= cell_size;
        } else {
            // also zero
            x += cell_size;
        }
        double y = point.y() * scale;
        if(y < 0.) {
            y -= cell_size;
        } else {
            // also zero
            y += cell_size;
        }
        // cast is used to truncate towards zero
        coordinates coord;
        coord.row = static_cast<coordinates_t>(x / cell_size);
        coord.column = static_cast<coordinates_t>(y / cell_size);
        return coord;
    }

//#define MIXING8
    /**
     * @brief pack_key helper to pack key, this zigzac
     * packing is used
     * to perform better with std::hash identity hash
     * funtion, the lower bytes are packed,
     * as there is the most entropy
     * @param coords
     * @return
     */
    static inline key_t pack_key(coordinates coords) {
        key_t packed;
#ifdef MIXING8
        uint8_t *packed_ptr = ((uint8_t*)&packed);
        uint8_t *row_ptr = ((uint8_t*)&(coords.row));
        uint8_t *column_ptr = ((uint8_t*)&(coords.column));
        packed_ptr[0] = row_ptr[0];
        packed_ptr[1] = column_ptr[0];
        packed_ptr[2] = row_ptr[1];
        packed_ptr[3] = column_ptr[1];
        packed_ptr[4] = row_ptr[2];
        packed_ptr[5] = column_ptr[2];
        packed_ptr[6] = row_ptr[3];
        packed_ptr[7] = column_ptr[3];
#else
        ((uint32_t*)&packed)[0] = *((uint32_t*)&(coords.row));
        ((uint32_t*)&packed)[1] = *((uint32_t*)&(coords.column));
#endif
        return packed;
    }

    static inline coordinates unpack_key(key_t key) {
        coordinates coords;
#ifdef MIXING8
        uint8_t *packed_ptr = ((uint8_t*)&key);
        uint8_t *row_ptr = ((uint8_t*)&(coords.row));
        uint8_t *column_ptr = ((uint8_t*)&(coords.column));
        row_ptr[0] = packed_ptr[0];
        column_ptr[0] = packed_ptr[1];
        row_ptr[1] = packed_ptr[2];
        column_ptr[1] = packed_ptr[3];
        row_ptr[2] = packed_ptr[4];
        column_ptr[2] = packed_ptr[5];
        row_ptr[3] = packed_ptr[6];
        column_ptr[3] = packed_ptr[7];
#else
        *((uint32_t*)&(coords.row)) = ((uint32_t*)&key)[0];
        *((uint32_t*)&(coords.column)) = ((uint32_t*)&key)[1];
#endif
        return  coords;
    }

    /**
     * @brief to_key fast conversion to key. Allows checking in which
     * cell this point would map without first constructing the point
     * @param point
     * @param cell_size
     * @return key
     */
    static inline key_t to_key(const from_point_t &point,
                               double cell_size, double scale = 1.) {
        auto coords = to_coordinates(point, cell_size, scale);
        return pack_key(coords);
    }

    /**
     * @brief coordinates_would_overflow
     * @param point
     * @param cell_size
     * @return true if coordinates would overflow
     */
    static bool coordinates_would_overflow(
            const from_point_t &point, double cell_size) {
        static_assert (sizeof (coordinates_t) <= 4,
                       "coordinates_would_overflow check unfortunately"
                       "isn't implemented for coordinates_t type bigger than "
                       "4 byte");
        double x = point.x();
        if(x > 0.) {
            x += cell_size;
        } else if(x < 0.){
            x -= cell_size;
        }
        double y = point.y();
        if(y > 0.) {
            y += cell_size;
        } else if(y < 0.) {
            y -= cell_size;
        }

        double row = x / cell_size;
        double column = y / cell_size;
        bool row_overflow = row > std::numeric_limits<coordinates_t>::max()
                || row < std::numeric_limits<coordinates_t>::min();
        bool column_overflow = column > std::numeric_limits<coordinates_t>::max()
                || column < std::numeric_limits<coordinates_t>::min();
        return row_overflow || column_overflow;
    }

    /**
     * @brief to_point to_point creates a point with a datatype which is suitable
     * for doing geometric calculations
     * @param cell_size
     * @param scale scale by which the coordinates where multiplied,
     * so the inverse is taken to get the real coordinates
     * @return
     */
    from_point_t to_point(double cell_size, double scale = 1.) const {
        double tmp_x, tmp_y;
        auto scale_inv = 1. / scale;
        if (row < 0) {
            tmp_x = ((row * cell_size) + .5 * cell_size) * scale_inv;
        } else if(row > 0){
            tmp_x = (((row - 1) * cell_size) + .5 * cell_size) * scale_inv;
        } else {
            tmp_x = 0;
        }
        if (column < 0) {
            tmp_y = ((column * cell_size) + .5 * cell_size) * scale_inv;
        } else if(column > 0){
            tmp_y = (((column - 1) * cell_size) + .5 * cell_size) * scale_inv;
        } else {
            tmp_y = 0;
        }
        return from_point_t(tmp_x, tmp_y);

    }

    /**
     * @brief to_point static method to build a point
     * from key directly
     * @param key
     * @param cell_size
     * @param scale scale by which the coordinates where multiplied,
     * so the inverse is taken to get the real coordinates
     * @return
     */
    static from_point_t to_point(key_t key, double cell_size, double scale = 1.) {
        coordinates coords = unpack_key(key);
        coordinates_t row = coords.row;
        coordinates_t column = coords.column;
        auto scale_inv = 1. / scale;
        double tmp_x, tmp_y;
        if (row < 0) {
            tmp_x = ((row * cell_size) + .5 * cell_size) * scale_inv;
        } else if(row > 0){
            tmp_x =  (((row - 1) * cell_size) + .5 * cell_size) * scale_inv;
        } else {
            tmp_x = 0;
        }

        if (column < 0) {
            tmp_y = ((column * cell_size) + .5 * cell_size) * scale_inv;
        } else if(column > 0){
            tmp_y = (((column - 1) * cell_size) + .5 * cell_size) * scale_inv;
        } else {
            tmp_y = 0;
        }

        return from_point_t(tmp_x, tmp_y);

    }
    /**
     * @brief has_same_grid_position
     * @param other_point
     * @return true if it has same row and column
     */
    bool has_same_grid_position(const GridPoint &other_point) const {
        return row == other_point.row &&
                column == other_point.column;
    }

    /**
     * @brief GridPoint construct gridpoint from vector and cell size
     * @param point
     * @param cell_size
     */
    GridPoint(from_point_t point, double cell_size, double scale = 1.) {
        auto coord = to_coordinates(point, cell_size, scale);
        row = coord.row;
        column = coord.column;
    }

    GridPoint(coordinates_t row, coordinates_t column) :
        row(row), column(column) {

    }

    GridPoint() = default;

    /**
     * @brief operator + adds coordinates of another point by adding row and column
     * @param other_point
     * @return
     */
    GridPoint operator+(const GridPoint &other_point) const {
        GridPoint tmp(other_point);
        tmp.row += row;
        tmp.column += column;
        return tmp;
    }

    /**
     * @brief to string
     * @return
     */
    virtual string to_string() const {
        return "point: row: " + std::to_string(row) +
                ", column: " + std::to_string(column) +
                ", hit_cnt: " + std::to_string(hit_cnt);
    }

    virtual ~GridPoint() {}

protected:
    coordinates_t row = 0;
    coordinates_t column = 0;
};

} // end namespace fsd

#endif //PROJECT_GRIDPOINT_HPP