#include <memory>
#include <chrono>
#include <tuple>
#include <unistd.h>
#include <gflags/gflags.h>
#include <string.h>
#include <assert.h>
#include <Eigen/StdVector>

#include "drake/solvers/snopt_solver.h"
#include "drake/systems/analysis/simulator.h"
#include "drake/systems/framework/diagram.h"
#include "drake/systems/framework/diagram_builder.h"
#include "drake/multibody/parsing/parser.h"
#include <drake/multibody/inverse_kinematics/inverse_kinematics.h>
#include "drake/geometry/geometry_visualization.h"
#include "drake/solvers/solve.h"

#include "common/find_resource.h"
#include "systems/trajectory_optimization/dircon/dircon.h"
#include "multibody/kinematic/world_point_evaluator.h"
#include "solvers/nonlinear_constraint.h"
#include "multibody/kinematic/distance_evaluator.h"
#include "multibody/multibody_utils.h"
#include "multibody/visualization_utils.h"
#include "multibody/kinematic/kinematic_constraints.h"

#include "examples/Spirit/animate_spirit.h"
#include "examples/Spirit/spirit_utils.h"

DEFINE_double(duration, 1, "The stand duration");
DEFINE_double(front2BackToeDistance, 0.35, "Nominal distance between the back and front toes.");
DEFINE_double(side2SideToeDistance, 0.2, "Nominal distance between the back and front toes.");
DEFINE_double(bodyHeight, 0.104, "The spirit body start height (defined in URDF)");
DEFINE_double(lowerHeight,0.15, "The sitting height of the bottom of the robot");
DEFINE_double(upperHeight, 0.35, "The standing height.");
DEFINE_double(inputCost, 3, "The standing height.");
DEFINE_double(velocityCost, 10, "The standing height.");
DEFINE_double(eps, 1e-2, "The wiggle room.");
DEFINE_double(optTol, 1e-4,"Optimization Tolerance");
DEFINE_double(feasTol, 1e-4,"Feasibility Tolerance");
DEFINE_bool(autodiff, false, "Double or autodiff version");
DEFINE_bool(runInitTraj, false, "Animate initial conditions?");
// Parameters which enable dircon-improving features
DEFINE_bool(scale_constraint, true, "Scale the nonlinear constraint values");
DEFINE_bool(scale_variable, false, "Scale the decision variable");

using drake::AutoDiffXd;
using drake::multibody::MultibodyPlant;
using drake::geometry::SceneGraph;
using drake::multibody::Parser;
using drake::trajectories::PiecewisePolynomial;

using Eigen::Vector3d;
using Eigen::VectorXd;
using Eigen::Matrix3d;
using Eigen::MatrixXd;

namespace dairlib {
namespace {

using systems::trajectory_optimization::DirconModeSequence;
using systems::trajectory_optimization::DirconMode;
using systems::trajectory_optimization::Dircon;
using systems::trajectory_optimization::KinematicConstraintType;

using std::vector;
using std::cout;
using std::endl;

template <typename T>
void addConstraints(const MultibodyPlant<T>& plant, Dircon<T>& trajopt){
  // Get position and velocity dictionaries
  int num_knotpoints =10; ///DEBUG
  auto positions_map = multibody::makeNameToPositionsMap(plant);
  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  auto x0 = trajopt.initial_state();
  auto xmid = trajopt.state_vars(0, (num_knotpoints - 1) / 2);
  // auto xf = trajopt.final_state();

  // Initial body positions
  trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, x0(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
  trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, x0(positions_map.at("base_y")));
  trajopt.AddBoundingBoxConstraint(FLAGS_lowerHeight-FLAGS_eps, FLAGS_lowerHeight+FLAGS_eps, x0(positions_map.at("base_z")));

  // Mid body positions
  trajopt.AddBoundingBoxConstraint(-0, 0, xmid(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
  trajopt.AddBoundingBoxConstraint(-0, 0, xmid(positions_map.at("base_y")));
  trajopt.AddBoundingBoxConstraint(FLAGS_upperHeight-FLAGS_eps, FLAGS_upperHeight+FLAGS_eps, xmid(positions_map.at("base_z")));

  return;
}

  /// See runSpiritSquat()
    // std::unique_ptr<MultibodyPlant<T>> plant_ptr,
    // MultibodyPlant<double>* plant_double_ptr,
    // std::unique_ptr<SceneGraph<double>> scene_graph_ptr,
    // double duration,
    // PiecewisePolynomial<double> init_x_traj,
    // PiecewisePolynomial<double> init_u_traj,
    // vector<PiecewisePolynomial<double>> init_l_traj,
    // vector<PiecewisePolynomial<double>> init_lc_traj,
    // vector<PiecewisePolynomial<double>> init_vc_traj
  ///
  /// Given an initial guess for state, control, and forces optimizes a standing behaviour for the spirit robot
template <typename T>
void runSpiritSquat(
    std::unique_ptr<MultibodyPlant<T>> plant_ptr,
    MultibodyPlant<double>* plant_double_ptr,
    std::unique_ptr<SceneGraph<double>> scene_graph_ptr,
    double duration,
    PiecewisePolynomial<double> init_x_traj,
    PiecewisePolynomial<double> init_u_traj,
    vector<PiecewisePolynomial<double>> init_l_traj,
    vector<PiecewisePolynomial<double>> init_lc_traj,
    vector<PiecewisePolynomial<double>> init_vc_traj
    ) {

  drake::systems::DiagramBuilder<double> builder;
  MultibodyPlant<T>& plant = *plant_ptr;
  SceneGraph<double>& scene_graph =
      *builder.AddSystem(std::move(scene_graph_ptr));

  // Get position and velocity dictionaries
  auto positions_map = multibody::makeNameToPositionsMap(plant);
  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  /// For Spirit front left leg->0, back left leg->1, front right leg->2, back right leg->3
  /// Get the frame of each toe and attach a world point to the toe tip (frame is at toe ball center).

  int num_legs = 4;
  double toeRadius = 0.02; // Radius of toe ball
  Vector3d toeOffset(toeRadius,0,0); // vector to "contact point"
  double mu = 1; // Coeff of friction

  int num_knotpoints_per_mode = 10;

  auto sequence = DirconModeSequence<T>(plant);

  dairlib::ModeSequenceHelper msh;

  msh.addMode( // FIRST MODE 1
    (Eigen::Matrix<bool,1,4>() << 1,1,1,1 ).finished(), // contact bools
    num_knotpoints_per_mode,  // number of knot points in the collocation
    Eigen::Vector3d::UnitZ(), // normal
    Eigen::Vector3d::Zero(),  // world offset
    mu //friction
    );

  // auto [modeVector, toeEvals, toeEvalSets] = createSpiritModeSequence(plant, msh.modes , msh.knots , msh.normals , msh.offsets, msh.mus, msh.minTs, msh.maxTs);
  auto [modeVector, toeEvals, toeEvalSets] = createSpiritModeSequence(plant, msh);

  for (auto& mode : modeVector){
    for (int i = 0; i < num_legs; i++ ){
      mode->MakeConstraintRelative(i,0);
      mode->MakeConstraintRelative(i,1);
    }
    mode->SetDynamicsScale(
      {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17}, 1.0 / 150.0);
    mode->SetKinVelocityScale(
      {0, 1, 2, 3}, {0, 1, 2}, 1.0 / 500.0 * 500 * 1 / 1);
    sequence.AddMode(mode.get());
  }

  ///Setup trajectory optimization
  auto trajopt = Dircon<T>(sequence);

  // Set up Trajectory Optimization options
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(),
                           "Print file", "../snopt.out");
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(),
                           "Major iterations limit", 20000);
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(), "Iterations limit", 100000);
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(),
                           "Major optimality tolerance",
                           FLAGS_optTol);  // target optimality
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(), "Major feasibility tolerance", FLAGS_feasTol);
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(), "Verify level",
                           0);  // 0

    // Add duration constraint, currently constrained not bounded
  trajopt.AddDurationBounds(0, duration*2);
  // Initialize the trajectory control state and forces
  for (int j = 0; j < sequence.num_modes(); j++) {
    trajopt.drake::systems::trajectory_optimization::MultipleShooting::
        SetInitialTrajectory(init_u_traj, init_x_traj);
    trajopt.SetInitialForceTrajectory(j, init_l_traj[j], init_lc_traj[j],
                                      init_vc_traj[j]);
  }

  /// Setup all the optimization constraints
  int num_knotpoints = trajopt.N(); // number of knot points total in the collocation
  int n_v = plant.num_velocities();
  // int n_q = plant.num_positions();
  auto u = trajopt.input();
  auto x = trajopt.state();
  auto x0 = trajopt.initial_state();
  auto xmid = trajopt.state_vars(0, (num_knotpoints - 1) / 2);
  auto xf = trajopt.final_state();
  addConstraints(plant,trajopt);

  // // Initial body positions
  // trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, x0(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
  // trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, x0(positions_map.at("base_y")));
  // trajopt.AddBoundingBoxConstraint(FLAGS_lowerHeight-FLAGS_eps, FLAGS_lowerHeight+FLAGS_eps, x0(positions_map.at("base_z")));

  // // Mid body positions
  // trajopt.AddBoundingBoxConstraint(-0, 0, xmid(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
  // trajopt.AddBoundingBoxConstraint(-0, 0, xmid(positions_map.at("base_y")));
  // trajopt.AddBoundingBoxConstraint(FLAGS_upperHeight-FLAGS_eps, FLAGS_upperHeight+FLAGS_eps, xmid(positions_map.at("base_z")));

  // Final Position Constraints
  bool isPeriodic = 1;
  if (isPeriodic){
    trajopt.AddLinearConstraint( x0(positions_map.at("base_x"))==xf(positions_map.at("base_x")) );
    trajopt.AddLinearConstraint( x0(positions_map.at("base_y"))==xf(positions_map.at("base_y")) );
    trajopt.AddLinearConstraint( x0(positions_map.at("base_z"))==xf(positions_map.at("base_z")) );
  }else{
    trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, xf(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
    trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, xf(positions_map.at("base_y")));
    trajopt.AddBoundingBoxConstraint(FLAGS_lowerHeight-FLAGS_eps, FLAGS_lowerHeight+FLAGS_eps, xf(positions_map.at("base_z")));
  }

  /// Decide whether or not to constrain the body orientation at all knotpoints or only at the beginning and end
  bool isConstrainOrientation = false;
  if (isConstrainOrientation){
    // Body pose constraints (keep the body flat) at all the knotpoints including the ends
    for (int i = 0; i < num_knotpoints; i++) {
      auto xi = trajopt.state(i);
      trajopt.AddBoundingBoxConstraint(1, 1, xi(positions_map.at("base_qw")));
      trajopt.AddBoundingBoxConstraint(0, 0, xi(positions_map.at("base_qx")));
      trajopt.AddBoundingBoxConstraint(0, 0, xi(positions_map.at("base_qy")));
      trajopt.AddBoundingBoxConstraint(0, 0, xi(positions_map.at("base_qz")));
    }
  }else{
    // Body pose constraints (keep the body flat) at initial state
    trajopt.AddBoundingBoxConstraint(1, 1, x0(positions_map.at("base_qw")));
    trajopt.AddBoundingBoxConstraint(0, 0, x0(positions_map.at("base_qx")));
    trajopt.AddBoundingBoxConstraint(0, 0, x0(positions_map.at("base_qy")));
    trajopt.AddBoundingBoxConstraint(0, 0, x0(positions_map.at("base_qz")));

    // Body pose constraints (keep the body flat) at mid state
    trajopt.AddBoundingBoxConstraint(1, 1, xmid(positions_map.at("base_qw")));
    trajopt.AddBoundingBoxConstraint(0, 0, xmid(positions_map.at("base_qx")));
    trajopt.AddBoundingBoxConstraint(0, 0, xmid(positions_map.at("base_qy")));
    trajopt.AddBoundingBoxConstraint(0, 0, xmid(positions_map.at("base_qz")));

    // Body pose constraints (keep the body flat) at final state
    trajopt.AddBoundingBoxConstraint(1, 1, xf(positions_map.at("base_qw")));
    trajopt.AddBoundingBoxConstraint(0, 0, xf(positions_map.at("base_qx")));
    trajopt.AddBoundingBoxConstraint(0, 0, xf(positions_map.at("base_qy")));
    trajopt.AddBoundingBoxConstraint(0, 0, xf(positions_map.at("base_qz")));
  }

  /// Start/End velocity constraints of the behavior
  trajopt.AddBoundingBoxConstraint(VectorXd::Zero(n_v), VectorXd::Zero(n_v),
                                   x0.tail(n_v));
  trajopt.AddBoundingBoxConstraint(VectorXd::Zero(n_v), VectorXd::Zero(n_v),
                                   xmid.tail(n_v));
  trajopt.AddBoundingBoxConstraint(VectorXd::Zero(n_v), VectorXd::Zero(n_v),
                                   xf.tail(n_v));

  for (int i = 0; i < num_knotpoints; i++){
    auto xi = trajopt.state(i);
    // legs lined up (front and back hips equal)
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_8") ) == xi( positions_map.at("joint_9") ) );
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_10") ) == xi( positions_map.at("joint_11") )  );
  }
  setSpiritSymmetry(plant, trajopt, "sagittal");
  setSpiritJointLimits(plant, trajopt);
  setSpiritActuationLimits(plant,trajopt);

 double upperLegLength = 0.206; // length of the upper leg link

// Quick check for allowed heights for Spirit
  try{
    if (FLAGS_lowerHeight<FLAGS_bodyHeight/2){
        throw "Body to close to floor" ;
    }
    if (FLAGS_upperHeight>=upperLegLength*2){
        throw "Body can't stand to that height without leaving the ground" ;
    }
  }catch(char const* exc){
      std::cerr<< exc <<std::endl;
  }

  ///Setup the traditional cost function
  const double R = FLAGS_inputCost;  // Cost on input effort
  const MatrixXd Q = FLAGS_velocityCost  * MatrixXd::Identity(n_v, n_v); // Cost on velocity
  trajopt.AddRunningCost( x.tail(n_v).transpose() * Q * x.tail(n_v) );
  trajopt.AddRunningCost( u.transpose()*R*u );

  ///Add regularization costs
  trajopt.AddRunningCost( x0(positions_map.at("base_x")) * 10 * x0(positions_map.at("base_x")) ); // x position cost
  // trajopt.AddRunningCost( (x.segment(12,2)-x.segment(14,2)).transpose() * 1 * (x.segment(12,2)-x.segment(14,2)) );
  // trajopt.AddRunningCost( (x.segment(16,2)-x.segment(18,2)).transpose() * 1 * (x.segment(16,2)-x.segment(18,2)) );

  /// Setup the visualization during the optimization
  int num_ghosts = 3;// Number of ghosts in visualization. NOTE: there are limitations on number of ghosts based on modes and knotpoints
  std::vector<unsigned int> visualizer_poses; // Ghosts for visualizing during optimization
  visualizer_poses.push_back(num_ghosts);

  trajopt.CreateVisualizationCallback(
      dairlib::FindResourceOrThrow("examples/Spirit/spirit_drake.urdf"),
      visualizer_poses, 0.2); // setup which URDF, how many poses, and alpha transparency

  /// Run the optimization using your initial guess
  auto start = std::chrono::high_resolution_clock::now();
  const auto result = Solve(trajopt, trajopt.initial_guess());
  auto finish = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> elapsed = finish - start;
  std::cout << "Solve time:" << elapsed.count() <<std::endl;
  std::cout << "Cost:" << result.get_optimal_cost() <<std::endl;
  std::cout << (result.is_success() ? "Optimization Success" : "Optimization Fail") << std::endl;

  /// Run animation of the final trajectory
  const drake::trajectories::PiecewisePolynomial<double> pp_xtraj =
      trajopt.ReconstructStateTrajectory(result);
  multibody::connectTrajectoryVisualizer(plant_double_ptr,
      &builder, &scene_graph, pp_xtraj);
  auto diagram = builder.Build();
  while (true) {

    drake::systems::Simulator<double> simulator(*diagram);
    simulator.set_target_realtime_rate(0.25);
    simulator.Initialize();
    simulator.AdvanceTo(pp_xtraj.end_time());
    sleep(2);
  }
}
}  // namespace
}  // namespace dairlib

int main(int argc, char* argv[]) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  std::srand(time(0));  // Initialize random number generator.

  auto plant = std::make_unique<MultibodyPlant<double>>(0.0);
  auto plant_vis = std::make_unique<MultibodyPlant<double>>(0.0);
  auto scene_graph = std::make_unique<SceneGraph<double>>();
  Parser parser(plant.get());
  Parser parser_vis(plant_vis.get(), scene_graph.get());
  std::string full_name =
      dairlib::FindResourceOrThrow("examples/Spirit/spirit_drake.urdf");

  parser.AddModelFromFile(full_name);
  parser_vis.AddModelFromFile(full_name);

  plant->mutable_gravity_field().set_gravity_vector(-9.81 *
      Eigen::Vector3d::UnitZ());

  plant->Finalize();
  plant_vis->Finalize();
  Eigen::VectorXd x0 = Eigen::VectorXd::Zero(plant->num_positions() +
                       plant->num_velocities());

  int nu = plant->num_actuators();
  int nx = plant->num_positions() + plant->num_velocities();
  int nq = plant->num_positions();
  int nv = plant->num_velocities();
  int N = 20; // number of timesteps

  std::vector<MatrixXd> init_x;
  std::vector<MatrixXd> init_u;
  std::vector<PiecewisePolynomial<double>> init_l_traj;
  std::vector<PiecewisePolynomial<double>> init_lc_traj;
  std::vector<PiecewisePolynomial<double>> init_vc_traj;

  // Initialize state trajectory
  std::vector<double> init_time;

  VectorXd xInit(nx);
  VectorXd xMid(nx);
  VectorXd xState(nx);
  xInit = Eigen::VectorXd::Zero(plant->num_positions() + plant->num_velocities());
  xMid = Eigen::VectorXd::Zero(plant->num_positions() + plant->num_velocities());
  xState = Eigen::VectorXd::Zero(plant->num_positions() + plant->num_velocities());

  auto positions_map = dairlib::multibody::makeNameToPositionsMap(*plant);
  auto velocities_map = dairlib::multibody::makeNameToVelocitiesMap(*plant);
  auto actuators_map = dairlib::multibody::makeNameToActuatorsMap(*plant);
  int num_joints = 12;

  // Print joint dictionary
  std::cout<<"**********************Joints***********************"<<std::endl;
  for (auto const& element : positions_map)
    std::cout << element.first << " = " << element.second << std::endl;
  for (auto const& element : velocities_map)
    std::cout << element.first << " = " << element.second << std::endl;
  for (auto const& element : actuators_map)
    std::cout << element.first << " = " << element.second << std::endl;
  std::cout<<"***************************************************"<<std::endl;

  dairlib::nominalSpiritStand( *plant, xInit,  FLAGS_lowerHeight); //Update xInit
  dairlib::nominalSpiritStand( *plant, xMid,  FLAGS_upperHeight); //Update xMid

  VectorXd deltaX(nx);
  VectorXd averageV(nx);
  deltaX = xMid-xInit;
  averageV = 2* deltaX / FLAGS_duration;
  xInit.tail(nv-3) = (averageV.head(nq)).tail(nq-4); //Ignoring Orientation make velocity the average

  double time = 0;
  double dt = FLAGS_duration/(N-1);

  // Initial pose
  xState = xInit;

  for (int i = 0; i < N; i++) {
    time=i*dt; // calculate iteration's time
    init_time.push_back(time);

    // Switch the direction of the stand to go back to the initial state (not actually properly periodic initial)
    if ( i > (N-1)/2 ){
        xState.tail(nv) = -xInit.tail(nv);
    }
    // Integrate the positions based on constant velocity  for joints and xyz
    for (int j = 0; j < num_joints; j++){
          xState(positions_map.at("joint_" + std::to_string(j))) =
                        xState(positions_map.at("joint_" + std::to_string(j))) + xState(nq + velocities_map.at("joint_" + std::to_string(j)+"dot" )) * dt;
    }
    xState(positions_map.at("base_x")) =
                        xState(positions_map.at("base_x")) + xState(nq + velocities_map.at("base_vx")) * dt;

    xState(positions_map.at("base_y")) =
                        xState(positions_map.at("base_y")) + xState(nq + velocities_map.at("base_vy")) * dt;

    xState(positions_map.at("base_z")) =
                        xState(positions_map.at("base_z")) + xState(nq + velocities_map.at("base_vz")) * dt;
    // Save timestep state into matrix
    init_x.push_back(xState);
    init_u.push_back(Eigen::VectorXd::Zero(nu));
  }
  // Make matrix into trajectory
  auto init_x_traj = PiecewisePolynomial<double>::ZeroOrderHold(init_time, init_x);
  auto init_u_traj = PiecewisePolynomial<double>::ZeroOrderHold(init_time, init_u);

  // Four contacts so forces are 12 dimensional
  Eigen::VectorXd init_l_vec(12);
  // Initial guess
  init_l_vec << 0, 0, 3*9.81, 0, 0, 3*9.81, 0, 0, 3*9.81, 0, 0, 3*9.81; //gravity and mass distributed

  //Initialize force trajectories
  int num_modes = 1; // MAKE DYNAMIC DEBUG
  for (int j = 0; j < num_modes; j++) {
    std::vector<MatrixXd> init_l_j;
    std::vector<MatrixXd> init_lc_j;
    std::vector<MatrixXd> init_vc_j;
    std::vector<double> init_time_j;
    for (int i = 0; i < N; i++) {
      init_time_j.push_back(i*FLAGS_duration/(N-1));
      init_l_j.push_back(init_l_vec);
      init_lc_j.push_back(init_l_vec);
      init_vc_j.push_back(VectorXd::Zero(12));
    }

    auto init_l_traj_j = PiecewisePolynomial<double>::ZeroOrderHold(init_time_j,init_l_j);
    auto init_lc_traj_j = PiecewisePolynomial<double>::ZeroOrderHold(init_time_j,init_lc_j);
    auto init_vc_traj_j = PiecewisePolynomial<double>::ZeroOrderHold(init_time_j,init_vc_j);

    init_l_traj.push_back(init_l_traj_j);
    init_lc_traj.push_back(init_lc_traj_j);
    init_vc_traj.push_back(init_vc_traj_j);
  }

  // if (FLAGS_autodiff) {
  //   std::unique_ptr<MultibodyPlant<drake::AutoDiffXd>> plant_autodiff =
  //       drake::systems::System<double>::ToAutoDiffXd(*plant);
  //   dairlib::runSpiritSquat<drake::AutoDiffXd>(
  //     std::move(plant_autodiff), plant_vis.get(), std::move(scene_graph),
  //     FLAGS_duration, init_x_traj, init_u_traj, init_l_traj,
  //     init_lc_traj, init_vc_traj);
  // } else
  if (FLAGS_runInitTraj){
    dairlib::runAnimate<double>(
      std::move(plant), plant_vis.get(), std::move(scene_graph), init_x_traj);
  }else {
    dairlib::runSpiritSquat<double>(
      std::move(plant), plant_vis.get(), std::move(scene_graph),
      FLAGS_duration, init_x_traj, init_u_traj, init_l_traj,
      init_lc_traj, init_vc_traj);
  }

}


// tensorrtinputconns.cc ---

// Copyright (C) 2019 Jolibrain http://www.jolibrain.com

// Author: Guillaume Infantes <guillaume.infantes@jolibrain.com>

// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 3
// of the License, or (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

#include "tensorrtinputconns.h"
#include "half.hpp"
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <boost/math/cstdfloat/cstdfloat_types.hpp>

namespace dd
{

  void ImgTensorRTInputFileConn::CVMatToRTBuffer(cv::Mat &img, int i)
  {
    cv::Mat converted;
    int channels = img.channels();
    int offset = channels*_height*_width*i;
    img.convertTo(converted, CV_32F);
    boost::float32_t * fbuf = (boost::float32_t*)(_buf.data());
    boost::float32_t * cvbuf = (boost::float32_t*)converted.data;
    for (int c=0; c<channels; ++c)
      for (int h=0; h<_height; ++h)
	for (int w=0; w<_width; ++w)
	  fbuf[offset++]  = cvbuf[(converted.cols*h+w)*channels+c];
  }

  void ImgTensorRTInputFileConn::applyMeanToRTBuf(int channels, int i)
  {
    int offset = channels*_height*_width*i;
    for (int c=0;c<channels;++c)
      for (int h=0;h<_height;++h)
	for (int w=0;w<_width;++w)
	  {
	    int data_index = (c*_height+h)*_width+w;
	    boost::float32_t * fbuf = static_cast<boost::float32_t*>(_buf.data());
	    fbuf[offset+data_index] -= _mean[c];
	  }
  }

  void ImgTensorRTInputFileConn::transform(const APIData &ad)
  {

    if (ad.has("has_mean_file"))
      _logger->warn("mean file cannot be used with tensorRT backend");

    try
      {
        ImgInputFileConn::transform(ad);
      }
    catch(const std::exception& e)
      {
        throw;
      }

    // ids
    bool set_ids = false;
    if (this->_ids.empty())
      set_ids = true;

    for (int i=0;i<(int)this->_images.size();i++)
      {
	if (set_ids)
	  this->_ids.push_back(this->_uris.at(i));
	_imgs_size.insert(std::pair<std::string,std::pair<int,int>>(this->_ids.at(i),this->_images_size.at(i)));
      }
    _batch_index = 0;
  }

  int ImgTensorRTInputFileConn::process_batch(const unsigned int batch_size)
  {
    if (batch_size < this->_images.size())
      this->_logger->warn("you are giving  more images than max_batch_size, please double check");
    unsigned int i;
    if (_bw)
      _buf.resize(batch_size * height() * width());
    else
      _buf.resize(batch_size * 3 * height() * width());
    for (i=0;
	 i < batch_size  &&  _batch_index< (int)this->_images.size();
	 i++,_batch_index++)
      {
        cv::Mat img = this->_images.at(_batch_index);
        CVMatToRTBuffer(img,i);
        if (_has_mean_scalar)
          {
            applyMeanToRTBuf(img.channels(), i);
          }
      }
    return i;
  }

}


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

#ifndef IMU_FUSION_MADGWICK__IMU_FUSION_MADGWICK_HPP_
#define IMU_FUSION_MADGWICK__IMU_FUSION_MADGWICK_HPP_

#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <Eigen/Geometry>
#include <string>

namespace imu_fusion_madgwick
{

class IMUFusionMadgwick : public rclcpp::Node
{
public:
  IMUFusionMadgwick(const rclcpp::NodeOptions & options = rclcpp::NodeOptions());

  virtual ~IMUFusionMadgwick();

  using Imu = sensor_msgs::msg::Imu;
  using Quaternion = geometry_msgs::msg::Quaternion;

  Quaternion getQuaternion() const;
  Eigen::Quaterniond getEigenQuaternion() const;

  void reset();
  void reset(const Quaternion & quat);
  void reset(const Eigen::Quaterniond & quat);

  /** A wrapper for the Eigen3 based integration implementation
   *
   * @param angular_velocity Measured angular velocity around axes, in rad/sec, e.g. from gyroscope
   * @param interval Time interval to integrate over, in seconds
   * @param linear_acceleration Measurement of the linear acceleration along axes, in meter/seconds^2
   * @param maxGyroError Maximum gyroscope measurement error, in rad/sec. Used to determine correction weight
   */
  void integrate(
    geometry_msgs::msg::Vector3 const & angular_velocity, double interval,
    geometry_msgs::msg::Vector3 const & linear_acceleration,
    double max_gyro_error);

  /** Integrate measurement of angular rate
   *
   * @param angularRate Measured angular rate around axes, in rad/sec, e.g. from gyroscope
   * @param interval Time interval to integrate over, in secons
   */
  void integrate(Eigen::Vector3d const & angularRate, double interval);

  /** Integrate measurement of angular rate, corrected given a reference measurement
   *
   * It uses the technique from paper:
   * "An efficient orientation filter for inertial and inertial/magnetic sensor arrays"
   * Sebastian O.H. Madgwick, April 30, 2010
   *
   * @param angularRate Measured angular rate around axes, in rad/sec, e.g. from gyroscope
   * @param interval Time interval to integrate over, in seconds
   * @param referenceDirMeas Measurement of the reference directory in local reference frame
   * @param maxGyroError Maximum gyroscope measurement error, in rad/sec. Used to determine correction weight
   * @param referenceDir Reference direction that is measured in global reference frame. By default direction of gravity
   */
  void integrate(
    Eigen::Vector3d const & angular_rate, double interval,
    Eigen::Vector3d const & reference_dir_measures,
    double max_gyro_error,
    Eigen::Vector3d const & referenceDir = Eigen::Vector3d{0, 0, 1});
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

private:
  // parameter
  bool use_fixed_dt_;
  float gyro_measuring_error_;

  double dt_;
  rclcpp::Time last_update_time_;

  rclcpp::Subscription<Imu>::SharedPtr sub_;
  rclcpp::Publisher<Imu>::SharedPtr pub_;

  Eigen::Quaterniond orientation_;
};

}  // namespace imu_fusion_madgwick

#endif  // IMU_FUSION_MADGWICK__IMU_FUSION_MADGWICK_HPP_


#include <boost/geometry/strategies/cartesian/buffer_join_round_by_divide.hpp>


#include <iostream>
#include <boost/filesystem.hpp>
#include <xpas/phylo_kmer_db.h>
#include <xpas/serialization.h>
#include <iomanip>

xpas::phylo_kmer_db create_db()
{
    const size_t kmer_size = 3;
    const xpas::phylo_kmer::score_type omega = 1.0;
    const std::string tree;

    xpas::phylo_kmer_db db { kmer_size, omega, xpas::seq_type::name, tree };

    /// branch 0
    db.unsafe_insert(0, { 0, 0.00f });
    db.unsafe_insert(1, { 0, 0.10f });
    db.unsafe_insert(2, { 0, 0.20f });

    /// branch 1
    db.unsafe_insert(1, { 1, 0.11f });
    db.unsafe_insert(2, { 1, 0.21f });
    db.unsafe_insert(3, { 1, 0.31f });

    /// branch 2
    db.unsafe_insert(2, { 2, 0.22f });
    db.unsafe_insert(3, { 2, 0.32f });
    db.unsafe_insert(4, { 2, 0.42f });

    return db;
}

std::ostream& operator<<(std::ostream& out, const xpas::phylo_kmer_db& db)
{
    for (const auto& [key, entries] : db)
    {
        out << key << ":\n";
        for (const auto& [branch, score] : entries)
        {
            out << '\t' << branch << ": " << score << '\n';
        }
    }
    return out;
}

int main()
{
    const auto filename = boost::filesystem::unique_path().string();
    xpas::save(create_db(), filename);

    const auto db = xpas::load(filename);
    std::cout << "Database parameters:" << std::endl
              << "\tSequence type: " << db.sequence_type() << std::endl
              << "\tk: " << db.kmer_size() << std::endl
              << "\tomega: " << db.omega() << std::endl
              << "\tKeep positions: " << (db.positions_loaded() ? "true" : "false") << std::endl << std::endl;
    std::cout << "Loaded a database of " << db.size() << " phylo-kmers. " << std::endl << std::endl;

    std::cout << db;
}


/*******************************************************************************
 * Copyright (c) 2019 UT-Battelle, LLC.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * and Eclipse Distribution License v1.0 which accompanies this
 * distribution. The Eclipse Public License is available at
 * http://www.eclipse.org/legal/epl-v10.html and the Eclipse Distribution
 *License is available at https://eclipse.org/org/documents/edl-v10.php
 *
 * Contributors:
 *   Tyler Kharazi - initial implementation
 *******************************************************************************/
#include "AssignmentErrorKernelDecorator.hpp"
#include "InstructionIterator.hpp"
#include "Utils.hpp"
#include "xacc.hpp"
#include <fstream>
#include <set>
#include <Eigen/Dense>

namespace xacc {
namespace quantum {
void AssignmentErrorKernelDecorator::initialize(
    const HeterogeneousMap &params) {

  if (params.keyExists<bool>("gen-kernel")) {
    gen_kernel = params.get<bool>("gen-kernel");
    std::cout<<"gen_kernel: "<<gen_kernel<<std::endl;
  }
  if (params.keyExists<bool>("multiplex")){
    multiplex = params.get<bool>("multiplex");
    //this is essentially to tell me how to deal with layouts. If multiplex, then we would split layout into two smaller layouts
    //and we would generate two error kernels.
  }
  if (params.keyExists<std::vector<std::size_t>>("layout")){
    layout = params.get<std::vector<std::size_t>>("layout");
    std::cout<<"layout recieved"<<std::endl;
  }
  if (params.keyExists<std::vector<int>>("layout")){
    auto tmp = params.get<std::vector<int>>("layout");
    for (auto& a : tmp) layout.push_back(a);
    std::cout<<"layout recieved"<<std::endl;

  }
} // initialize

void AssignmentErrorKernelDecorator::execute(
    std::shared_ptr<AcceleratorBuffer> buffer,
    const std::shared_ptr<CompositeInstruction> function) {
  int num_bits = buffer->size();
  int size = std::pow(2, num_bits);
  if (!layout.empty()) {
     function->mapBits(layout);
  }
  // get the raw state
  decoratedAccelerator->execute(buffer, function);
  int shots = 0;
  Eigen::VectorXd init_state(size);
  init_state.setZero();
  int i = 0;
  for (auto &x : buffer->getMeasurementCounts()) {
    shots += x.second;
    //std::cout<<"HELLO: " << x.first << ", " << x.second<<std::endl;
    init_state(i) = double(x.second);
    //std::cout<<init_state(i)<<std::endl;
    i++;
  }
  //std::cout << "BEFORE: " << init_state.transpose() << "\n";
  init_state = (double)1/shots*init_state;
  //std::cout<<"num_shots = "<<shots<<std::endl;
  //std::cout<<"INIT STATE:\n"<<init_state<<std::endl;
  if (gen_kernel) {
    if (decoratedAccelerator) {
      generateKernel(buffer);
    }

  } else {
    // generating the list of permutations is O(2^num_bits), we want to minimize
    // the number of times we have to call it.
    if (permutations.empty()) {
      permutations = generatePermutations(num_bits);
    }
  }

  std::cout<<"Error Kernel: \n"<<errorKernel<<std::endl;

  Eigen::VectorXd EM_state = errorKernel * init_state;
  // std::cout<<init_state<<std::endl;
  // std::cout<<EM_state<<std::endl;
  // checking for negative values and performing a "clip and renorm"
  for (int i = 0; i < EM_state.size(); i++) {
    if (EM_state(i) < 0.0) {
      int count = floor(shots * EM_state(i) + 0.5);
      shots += count;
      EM_state(i) = 0.0;
    }
  }

//   std::cout << "Updated EM_state:" << std::endl << EM_state << std::endl;
  // update buffer with new counts and save original counts to extra info

  std::map<std::string, double> origCounts;

  int total = 0;
  i = 0;
  for (auto &x : permutations) {
    origCounts[x] = (double)buffer->getMeasurementCounts()[x];
    int count = floor(shots * EM_state(i) + 0.5);
    //std::cout<<"EM_state = "<<EM_state(i)<<std::endl;
    total += count;
    //std::cout<<"saving " << count <<" counts in slot: "<< x<<std::endl;
    buffer->appendMeasurement(x, count);
    i++;
  }
  // std::cout<<origCounts<<std::endl;
  buffer->addExtraInfo("unmitigated-counts", origCounts);

  return;
} // execute

void AssignmentErrorKernelDecorator::execute(
    const std::shared_ptr<AcceleratorBuffer> buffer,
    const std::vector<std::shared_ptr<CompositeInstruction>> functions) {
  int num_bits = buffer->size();
  if (gen_kernel) {
    if (decoratedAccelerator) {
      generateKernel(buffer);
    }
  } else {
    if (permutations.empty()) {
      permutations = generatePermutations(num_bits);
    }
  }
  // get the raw states
  decoratedAccelerator->execute(buffer, functions);
  int size = std::pow(2, num_bits);

  std::vector<Eigen::VectorXd> init_states;
  // compute number of shots;
  auto buffers = buffer->getChildren();
  int shots = 0;
  for (auto &x : buffers[0]->getMeasurementCounts()) {
    shots += x.second;
  }
//   std::cout << "shots = " << shots << std::endl;

  int i = 0;
  for (auto &b : buffers) {
    Eigen::VectorXd temp(size);
    int j = 0;
    for (auto &x : permutations) {
      //std::cout << b->computeMeasurementProbability(x) << std::endl;
      temp(j) =(double)b->computeMeasurementProbability(x);
      j++;
    }
    // std::cout << "seg fault right here " << std::endl;
    init_states.push_back(temp);
  }
//   std::cout << "unmitigated states: " << std::endl;
  std::vector<Eigen::VectorXd> EM_states;
  i = 0;
  for (auto &x : init_states) {
    //std::cout << x << std::endl << std::endl;
    EM_states.push_back(errorKernel * x);
    i++;
  }
//   std::cout << "mitigated states: " << std::endl;
//   for (auto &x : EM_states) {
//     std::cout << x << std::endl << std::endl;
//   }
  for (int i = 0; i < EM_states.size(); i++) {
    for (int j = 0; j < EM_states[i].size(); j++) {
      if (EM_states[i](j) < 0.0) {
        int count = floor(shots * EM_states[i](j) + 0.5);
        // std::cout << "found negative value, clipping and renorming "
        //           << std::endl;
        // std::cout << "removed " << abs(count) << " shots from total shots"
        //           << std::endl;
        shots += count;
        EM_states[i](j) = 0.0;
      }
    }
  }

//   std::cout << "Update EM_states: " << std::endl;
//   for (auto &x : EM_states) {
//     std::cout << x << std::endl << std::endl;
//   }

  std::vector<std::map<std::string, double>> origCounts(buffer->nChildren());

  int total = 0;
  i = 0;
  for (auto &b : buffers) {
    int j = 0;
    for (auto &x : permutations) {
      origCounts[i][x] = (double)b->getMeasurementCounts()[x];
      int count = floor(shots * EM_states[i](j) + 0.5);
      j++;
      total += count;
      b->appendMeasurement(x, count);
      b->addExtraInfo("unmitigated-counts", origCounts[i]);
    }
    i++;
  }

  return;
} // execute (vectorized)

} // namespace quantum
} // namespace xacc


// Copyright (c) 2015-2020 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <bench/bench.h>

#include <crypto/sha256.h>
#include <key.h>
#include <random.h>
#include <stacktraces.h>
#include <util/strencodings.h>
#include <util/system.h>
#include <validation.h>

#include <boost/lexical_cast.hpp>

#include <memory>

#include <bls/bls.h>

const std::function<std::string(const char*)> G_TRANSLATION_FUN = nullptr;

static const char* DEFAULT_BENCH_FILTER = ".*";

void InitBLSTests();
void CleanupBLSTests();

static fs::path SetDataDir()
{
    fs::path ret = fs::temp_directory_path() / "bench_dash" / fs::unique_path();
    fs::create_directories(ret);
    gArgs.ForceSetArg("-datadir", ret.string());
    return ret;
}
static void SetupBenchArgs()
{
    gArgs.AddArg("-?", "Print this help message and exit", false, OptionsCategory::OPTIONS);
    gArgs.AddArg("-list", "List benchmarks without executing them", false, OptionsCategory::OPTIONS);
    gArgs.AddArg("-filter=<regex>", strprintf("Regular expression filter to select benchmark by name (default: %s)", DEFAULT_BENCH_FILTER), false, OptionsCategory::OPTIONS);
    gArgs.AddArg("-asymptote=n1,n2,n3,...", strprintf("Test asymptotic growth of the runtime of an algorithm, if supported by the benchmark"), false, OptionsCategory::OPTIONS);
    gArgs.AddArg("-output_csv=<output.csv>", "Generate CSV file with the most important benchmark results.", false, OptionsCategory::OPTIONS);
    gArgs.AddArg("-output_json=<output.json>", "Generate JSON file with all benchmark results.", false, OptionsCategory::OPTIONS);

    // Hidden
    gArgs.AddArg("-h", "", false, OptionsCategory::HIDDEN);
    gArgs.AddArg("-help", "", false, OptionsCategory::HIDDEN);
}

// parses a comma separated list like "10,20,30,50"
static std::vector<double> parseAsymptote(const std::string& str) {
    std::stringstream ss(str);
    std::vector<double> numbers;
    double d;
    char c;
    while (ss >> d) {
        numbers.push_back(d);
        ss >> c;
    }
    return numbers;
}

int main(int argc, char** argv)
{
    SetupBenchArgs();
    std::string error;
    if (!gArgs.ParseParameters(argc, argv, error)) {
        fprintf(stderr, "Error parsing command line arguments: %s\n", error.c_str());
        return EXIT_FAILURE;
    }

    if (HelpRequested(gArgs)) {
        std::cout << gArgs.GetHelpMessage();
        return EXIT_SUCCESS;
    }

    // Set the datadir after parsing the bench options
    const fs::path bench_datadir{SetDataDir()};

    SHA256AutoDetect();

    RegisterPrettySignalHandlers();
    RegisterPrettyTerminateHander();

    RandomInit();
    ECC_Start();
    ECCVerifyHandle verifyHandle;

    BLSInit();
    InitBLSTests();
    SetupEnvironment();

    benchmark::Args args;
    args.regex_filter = gArgs.GetArg("-filter", DEFAULT_BENCH_FILTER);
    args.is_list_only = gArgs.GetBoolArg("-list", false);
    args.asymptote = parseAsymptote(gArgs.GetArg("-asymptote", ""));
    args.output_csv = gArgs.GetArg("-output_csv", "");
    args.output_json = gArgs.GetArg("-output_json", "");

    benchmark::BenchRunner::RunAll(args);

    fs::remove_all(bench_datadir);

    // need to be called before global destructors kick in (PoolAllocator is needed due to many BLSSecretKeys)
    CleanupBLSTests();

    ECC_Stop();

    return EXIT_SUCCESS;
}


#include "AST.h"
#include "InterpreterContext.h"
#include "StringPool.h"
#include "VariablesScope.h"
#include <iostream>
#include <boost/range/adaptor/reversed.hpp>
#include <boost/range/algorithm.hpp>
#include <cmath>

namespace
{

template <class TFunction>
CValue ExecuteSafely(TFunction && fn)
{
    try
    {
        return fn();
    }
    catch (std::exception const&)
    {
        return CValue::FromError(std::current_exception());
    }
}

class CSinFunction : public IFunctionAst
{
public:
    CValue Call(CInterpreterContext &context, const std::vector<CValue> &arguments) const override
    {
        (void)context;
        return ExecuteSafely([&] {
            double radians = arguments.at(0).AsDouble();
            return CValue::FromDouble(std::sin(radians));
        });
    }

    unsigned GetNameId() const override
    {
        return 0;
    }
};

class CRandFunction : public IFunctionAst
{
public:
    CValue Call(CInterpreterContext &context, const std::vector<CValue> &arguments) const override
    {
        (void)context;
        return ExecuteSafely([&] {
            double minimum = arguments.at(0).AsDouble();
            double maximum = arguments.at(1).AsDouble();
            if (minimum > maximum)
            {
                return CValue::FromErrorMessage("invalid arguments for rand - range maximum is lesser than minimum.");
            }
            double rand0to1 = double(std::rand()) / std::numeric_limits<unsigned>::max();

            return CValue::FromDouble((maximum - minimum) * rand0to1 + minimum);
        });
    }

    unsigned GetNameId() const override
    {
        return 0;
    }
};

}

CInterpreterContext::CInterpreterContext(std::ostream &output, std::ostream &errors, CStringPool &pool)
    : m_pool(pool)
    , m_output(output)
    , m_errors(errors)
{
    AddBuiltin("sin", std::unique_ptr<IFunctionAst>(new CSinFunction));
    AddBuiltin("rand", std::unique_ptr<IFunctionAst>(new CRandFunction));
}

CInterpreterContext::~CInterpreterContext()
{
}

void CInterpreterContext::DefineVariable(unsigned nameId, const CValue &value)
{
    if (ValidateValue(value))
    {
        m_scopes.back()->AssignVariable(nameId, value);
    }
}

void CInterpreterContext::AssignVariable(unsigned nameId, const CValue &value)
{
    if (ValidateValue(value))
    {
        if (CVariablesScope *pScope = FindScopeWithVariable(nameId))
        {
            pScope->AssignVariable(nameId, value);
        }
        else
        {
            DefineVariable(nameId, value);
        }
    }
}

CValue CInterpreterContext::GetVariableValue(unsigned nameId) const
{
    if (CVariablesScope *pScope = FindScopeWithVariable(nameId))
    {
        return *pScope->GetVariableValue(nameId);
    }
    return CValue::FromErrorMessage("unknown variable " + m_pool.GetString(nameId));
}

void CInterpreterContext::PushScope(std::unique_ptr<CVariablesScope> &&scope)
{
    m_scopes.emplace_back(std::move(scope));
}

std::unique_ptr<CVariablesScope> CInterpreterContext::PopScope()
{
    std::unique_ptr<CVariablesScope> ret(m_scopes.back().release());
    m_scopes.pop_back();
    return ret;
}

size_t CInterpreterContext::GetScopesCount() const
{
    return m_scopes.size();
}

IFunctionAst *CInterpreterContext::GetFunction(unsigned nameId) const
{
    try
    {
        return m_functions.at(nameId);
    }
    catch (std::exception const&)
    {
        return nullptr;
    }
}

void CInterpreterContext::AddFunction(unsigned nameId, IFunctionAst *function)
{
    if (function)
    {
        m_functions[nameId] = function;
    }
}

std::string CInterpreterContext::GetStringLiteral(unsigned stringId) const
{
    return m_pool.GetString(stringId);
}

void CInterpreterContext::PrintResults(const std::vector<CValue> &values)
{
    m_output << "  ";
    for (const auto &value : values)
    {
        if (ValidateValue(value))
        {
            m_output << value.ToString();
        }
    }
    m_output << std::endl;
}

void CInterpreterContext::PrintError(const std::string &message)
{
    m_errors << "  Error: " << message << std::endl;
}

bool CInterpreterContext::ValidateValue(const CValue &value)
{
    try
    {
        value.RethrowIfException();
        return true;
    }
    catch (std::exception const& ex)
    {
        PrintError(ex.what());
        return false;
    }
}

void CInterpreterContext::SetReturnValue(boost::optional<CValue> const& valueOpt)
{
    m_returnValueOpt = valueOpt;
}

boost::optional<CValue> CInterpreterContext::GetReturnValue() const
{
    return m_returnValueOpt;
}

CVariablesScope *CInterpreterContext::FindScopeWithVariable(unsigned nameId) const
{
    auto range = boost::adaptors::reverse(m_scopes);
    auto it = boost::find_if(range, [=](const auto &pScope) {
        return pScope->HasVariable(nameId);
    });
    if (it != range.end())
    {
        return it->get();
    }
    return nullptr;
}

void CInterpreterContext::AddBuiltin(const std::string &name, std::unique_ptr<IFunctionAst> &&function)
{
    m_builtins.emplace_back(std::move(function));
    unsigned nameRand = m_pool.Insert(name);
    m_functions[nameRand] = m_builtins.back().get();
}


/*
 * Implementation of EDM methods, including S-map and cross-mapping
 *
 * - Patrick Laub, Department of Management and Marketing,
 *   The University of Melbourne, patrick.laub@unimelb.edu.au
 * - Edoardo Tescari, Melbourne Data Analytics Platform,
 *  The University of Melbourne, e.tescari@unimelb.edu.au
 *
 */

#pragma warning(disable : 4018)

#include "edm.h"
#include "cpu.h"
#include "distances.h"
#include "stats.h" // for correlation and mean_absolute_error
#include "thread_pool.h"
#include "train_predict_split.h"

#ifndef FMT_HEADER_ONLY
#define FMT_HEADER_ONLY
#endif

#define EIGEN_NO_DEBUG
#define EIGEN_DONT_PARALLELIZE
#include <Eigen/SVD>
#include <algorithm> // std::partial_sort
#include <cmath>
#include <fstream> // just to create low-level input dumps

std::atomic<int> numTasksStarted = 0;
std::atomic<int> numTasksFinished = 0;
ThreadPool workerPool(0), taskRunnerPool(0);

std::vector<std::future<Prediction>> launch_task_group(const ManifoldGenerator& generator, Options opts,
                                                       const std::vector<int>& Es, const std::vector<int>& libraries,
                                                       int k, int numReps, int crossfold, bool explore, bool full,
                                                       bool saveFinalPredictions, bool saveFinalCoPredictions,
                                                       bool saveSMAPCoeffs, bool copredictMode,
                                                       const std::vector<bool>& usable, const std::string& rngState,
                                                       IO* io, bool keep_going(), void all_tasks_finished())
{

  workerPool.set_num_workers(opts.nthreads);
  taskRunnerPool.set_num_workers(num_physical_cores());

  // Construct the instance which will (repeatedly) split the data
  // into either the training manifold or the prediction manifold.
  TrainPredictSplitter splitter = TrainPredictSplitter(explore, full, crossfold, usable, rngState);

  int numLibraries = (explore ? 1 : libraries.size());

  opts.numTasks = numReps * Es.size() * numLibraries;
  opts.configNum = 0;
  opts.taskNum = 0;

  int maxE = Es[Es.size() - 1];

  std::vector<bool> cousable;

  if (copredictMode) {
    opts.numTasks *= 2;
    cousable = generator.generate_usable(maxE, true);
  }

  int E, kAdj, library, trainSize;

  std::vector<std::future<Prediction>> futures;

  bool newTrainPredictSplit = true;

  // Note: the 'numReps' either refers to the 'replicate' option
  // used for bootstrap resampling, or the 'crossfold' number of
  // cross-validation folds. Both options can't be used together,
  // so numReps = max(replicate, crossfold).
  for (int iter = 1; iter <= numReps; iter++) {
    if (explore) {
      newTrainPredictSplit = true;
      trainSize = splitter.next_training_size(iter);
    }

    for (int i = 0; i < Es.size(); i++) {
      E = Es[i];

      // 'libraries' is implicitly set to one value in explore mode
      // though in xmap mode it is a user-supplied list which we loop over.
      for (int l = 0; l == 0 || l < libraries.size(); l++) {
        if (!explore) {
          newTrainPredictSplit = true;
        }

        if (explore) {
          library = trainSize;
        } else {
          library = libraries[l];
        }

        // Set the number of neighbours to use
        if (k > 0) {
          kAdj = k;
        } else if (k < 0) {
          kAdj = -1; // Leave a sentinel value so we know to skip the nearest neighbours calculation
        } else if (k == 0) {
          bool isSMap = opts.algorithm == Algorithm::SMap;
          int defaultK = generator.E_actual(E) + 1 + isSMap;
          kAdj = defaultK < library ? defaultK : library;
        }

        bool lastConfig = (E == maxE) && (l + 1 == numLibraries);

        if (explore) {
          opts.savePrediction = saveFinalPredictions && ((iter == numReps) || (crossfold > 0)) && lastConfig;
        } else {
          opts.savePrediction = saveFinalPredictions && (iter == numReps) && lastConfig;
        }
        opts.saveSMAPCoeffs = saveSMAPCoeffs;

        if (newTrainPredictSplit) {
          splitter.update_train_predict_split(library, iter);
          newTrainPredictSplit = false;
        }

        opts.copredict = false;
        opts.k = kAdj;

        futures.emplace_back(launch_edm_task(generator, opts, E, splitter.trainingRows(), splitter.predictionRows(), io,
                                             keep_going, all_tasks_finished));

        opts.taskNum += 1;

        if (copredictMode) {
          opts.copredict = true;
          if (explore) {
            opts.savePrediction = saveFinalCoPredictions && ((iter == numReps) || (crossfold > 0)) && lastConfig;
          } else {
            opts.savePrediction = saveFinalCoPredictions && ((iter == numReps)) && lastConfig;
          }
          opts.saveSMAPCoeffs = false;
          futures.emplace_back(
            launch_edm_task(generator, opts, E, splitter.trainingRows(), cousable, io, keep_going, all_tasks_finished));

          opts.taskNum += 1;
        }

        opts.configNum += opts.thetas.size();
      }
    }
  }

  return futures;
}

std::future<Prediction> launch_edm_task(const ManifoldGenerator& generator, Options opts, int E,
                                        const std::vector<bool>& trainingRows, const std::vector<bool>& predictionRows,
                                        IO* io, bool keep_going(), void all_tasks_finished())
{
  // Expand the 'metrics' vector now that we know the value of E.
  std::vector<Metric> metrics;

  // For the Wasserstein distance, it's more convenient to have one 'metric' for each variable (before taking lags).
  // However, for the L^1 / L^2 distances, it's more convenient to have one 'metric' for each individual
  // point of each observations, so metrics.size() == M.E_actual().
  if (opts.distance == Distance::Wasserstein) {
    // Add a metric for the main variable and for the dt variable.
    // These are always treated as a continuous values (though perhaps in the future this will change).
    metrics.push_back(Metric::Diff);
    if (generator.E_dt(E) > 0) {
      metrics.push_back(Metric::Diff);
    }

    // Add in the metrics for the 'extra' variables as they were supplied to us.
    for (int k = 0; k < generator.numExtras(); k++) {
      metrics.push_back(opts.metrics[k]);
    }
  } else {
    // Add metrics for the main variable and the dt variable and their lags.
    // These are always treated as a continuous values (though perhaps in the future this will change).
    for (int lagNum = 0; lagNum < E + generator.E_dt(E); lagNum++) {
      metrics.push_back(Metric::Diff);
    }

    // The user specified how to treat the extra variables.
    for (int k = 0; k < generator.numExtras(); k++) {
      int numLags = (k < generator.numExtrasLagged()) ? E : 1;
      for (int lagNum = 0; lagNum < numLags; lagNum++) {
        metrics.push_back(opts.metrics[k]);
      }
    }
  }

  opts.metrics = metrics;

  if (opts.taskNum == 0) {
    numTasksStarted = 0;
    numTasksFinished = 0;
  }

  numTasksStarted += 1;

  // This hack is simply to dump some really low level data structures
  // purely for the purpose of generating microbenchmarks.
  if (io != nullptr && io->verbosity > 4) {
    json lowLevelInputDump;
    lowLevelInputDump["generator"] = generator;
    lowLevelInputDump["opts"] = opts;
    lowLevelInputDump["E"] = E;
    lowLevelInputDump["trainingRows"] = trainingRows;
    lowLevelInputDump["predictionRows"] = predictionRows;

    std::ofstream o("lowLevelInputDump.json");
    o << lowLevelInputDump << std::endl;
  }

  // Note, we can't have missing data inside the training manifold when using the S-Map algorithm
  bool skipMissing = (opts.algorithm == Algorithm::SMap);

  Manifold M = generator.create_manifold(E, trainingRows, opts.copredict, false, opts.dtWeight, skipMissing);
  Manifold Mp = generator.create_manifold(E, predictionRows, opts.copredict, true, opts.dtWeight);

  return taskRunnerPool.enqueue([opts, M, Mp, predictionRows, io, keep_going, all_tasks_finished] {
    return edm_task(opts, M, Mp, predictionRows, io, keep_going, all_tasks_finished);
  });
}

Prediction edm_task(const Options opts, const Manifold M, const Manifold Mp, const std::vector<bool> predictionRows,
                    IO* io, bool keep_going(), void all_tasks_finished())
{
  bool multiThreaded = opts.nthreads > 1;
  int numThetas = (int)opts.thetas.size();
  int numPredictions = Mp.nobs();
  int numCoeffCols = M.E_actual() + 1;

  auto ystar = std::make_unique<double[]>(numThetas * numPredictions);
  std::fill_n(ystar.get(), numThetas * numPredictions, MISSING_SENTINEL);
  Eigen::Map<MatrixXd> ystarView(ystar.get(), numThetas, numPredictions);

  // If we're saving the coefficients (i.e. in xmap mode), then we're not running with multiple 'theta' values.
  auto coeffs = std::make_unique<double[]>(numPredictions * numCoeffCols);
  std::fill_n(coeffs.get(), numPredictions * numCoeffCols, MISSING_SENTINEL);
  Eigen::Map<MatrixXd> coeffsView(coeffs.get(), numPredictions, numCoeffCols);

  auto rc = std::make_unique<retcode[]>(numThetas * numPredictions);
  std::fill_n(rc.get(), numThetas * numPredictions, UNKNOWN_ERROR);
  Eigen::Map<MatrixXi> rcView(rc.get(), numThetas, numPredictions);

  std::vector<int> kUsed;
  for (int i = 0; i < numPredictions; i++) {
    kUsed.push_back(-1);
  }

  if (opts.numTasks > 1 && opts.taskNum == 0) {
    io->progress_bar(0.0);
  }

  if (multiThreaded) {
    std::vector<std::future<void>> results(numPredictions);
    for (int i = 0; i < numPredictions; i++) {
      results[i] = workerPool.enqueue(
        [&, i] { make_prediction(i, opts, M, Mp, ystarView, rcView, coeffsView, &(kUsed[i]), keep_going); });
    }

    if (opts.numTasks == 1) {
      io->progress_bar(0.0);
    }
    for (int i = 0; i < numPredictions; i++) {
      results[i].get();
      if (opts.numTasks == 1) {
        io->progress_bar((i + 1) / ((double)numPredictions));
      }
    }
  } else {
    if (opts.numTasks == 1) {
      io->progress_bar(0.0);
    }
    for (int i = 0; i < numPredictions; i++) {
      if (keep_going != nullptr && keep_going() == false) {
        break;
      }
      make_prediction(i, opts, M, Mp, ystarView, rcView, coeffsView, &(kUsed[i]), keep_going);
      if (opts.numTasks == 1) {
        io->progress_bar((i + 1) / ((double)numPredictions));
      }
    }
  }

  Prediction pred;

  // Store the results, so long as we weren't interrupted by a 'break'.
  if (keep_going == nullptr || keep_going() == true) {
    // Start by calculating the MAE & rho of prediction, if requested
    for (int t = 0; t < numThetas * opts.calcRhoMAE; t++) {
      PredictionStats stats;

      std::vector<double> y1, y2;

      for (int i = 0; i < Mp.ySize(); i++) {
        if (Mp.y(i) != MISSING_SENTINEL && ystarView(t, i) != MISSING_SENTINEL) {
          y1.push_back(Mp.y(i));
          y2.push_back(ystarView(t, i));
        }
      }

      if (!(y1.empty() || y2.empty())) {
        stats.mae = mean_absolute_error(y1, y2);
        stats.rho = correlation(y1, y2);
      } else {
        stats.mae = MISSING_SENTINEL;
        stats.rho = MISSING_SENTINEL;
      }

      pred.stats.push_back(stats);
    }

    pred.configNum = opts.configNum;

    // Check if any make_prediction call failed, and if so find the most serious error
    pred.rc = *std::max_element(rc.get(), rc.get() + numThetas * numPredictions);

    // If we're storing the prediction and/or the SMAP coefficients, put them
    // into the resulting Prediction struct. Otherwise, let them be deleted.
    if (opts.savePrediction) {
      // Take only the predictions for the largest theta value.
      if (numThetas == 1) {
        pred.ystar = std::move(ystar);
      } else {
        pred.ystar = std::make_unique<double[]>(numPredictions);
        for (int i = 0; i < numPredictions; i++) {
          pred.ystar[i] = ystarView(numThetas - 1, i);
        }
      }
    } else {
      pred.ystar = nullptr;
    }

    if (opts.saveSMAPCoeffs) {
      pred.coeffs = std::move(coeffs);
    } else {
      pred.coeffs = nullptr;
    }

    if (opts.savePrediction || opts.saveSMAPCoeffs) {
      pred.predictionRows = std::move(predictionRows);
    }

    if (opts.saveKUsed) {
      pred.kUsed = kUsed;
    }

    pred.cmdLine = opts.cmdLine;
    pred.copredict = opts.copredict;

    pred.numThetas = numThetas;
    pred.numPredictions = numPredictions;
    pred.numCoeffCols = numCoeffCols;

    if (opts.numTasks > 1) {
      io->progress_bar((numTasksFinished + 1) / ((double)opts.numTasks));
    }
  }

  numTasksFinished += 1;

  if (numTasksFinished == opts.numTasks) {
    if (all_tasks_finished != nullptr) {
      all_tasks_finished();
    }
  }

  return pred;
}

// Use a training manifold 'M' to make a prediction about the prediction manifold 'Mp'.
// Specifically, predict the 'Mp_i'-th value of the prediction manifold 'Mp'.
//
// The predicted value is stored in 'ystar', along with any return codes in 'rc'.
// Optionally, the user may ask to store some S-map intermediate values in 'coeffs'.
//
// The 'opts' value specifies the kind of prediction to make (e.g. S-map, or simplex method).
// This function is usually run in a worker thread, and the 'keep_going' callback is frequently called to
// see whether the user still wants this result, or if they have given up & simply want the execution
// to terminate.
//
// We sometimes let 'M' and 'Mp' be the same manifold, so we train and predict using the same values.
// In this case, the algorithm may cheat by pulling out the identical trajectory from the training manifold
// and using this as the prediction. As such, we throw away any neighbours which have a distance of 0 from
// the target point.
void make_prediction(int Mp_i, const Options& opts, const Manifold& M, const Manifold& Mp, Eigen::Map<MatrixXd> ystar,
                     Eigen::Map<MatrixXi> rc, Eigen::Map<MatrixXd> coeffs, int* kUsed, bool keep_going())
{
  // An impatient user may want to cancel a long-running EDM command, so we occasionally check using this
  // callback to see whether we ought to keep going with this EDM command. Of course, this adds a tiny inefficiency,
  // but there doesn't seem to be a simple way to easily kill running worker threads across all OSs.
  if (keep_going != nullptr && keep_going() == false) {
    rc(0, Mp_i) = BREAK_HIT;
    return;
  }

  // Create a list of indices which may potentially be the neighbours of Mp(Mp_i,.)
  std::vector<int> tryInds = potential_neighbour_indices(Mp_i, opts, M, Mp);

  DistanceIndexPairs potentialNN;
  if (opts.distance == Distance::Wasserstein) {
    potentialNN = wasserstein_distances(Mp_i, opts, M, Mp, tryInds);
  } else {
    potentialNN = lp_distances(Mp_i, opts, M, Mp, tryInds);
  }

  if (keep_going != nullptr && keep_going() == false) {
    rc(0, Mp_i) = BREAK_HIT;
    return;
  }

  // Do we have enough distances to find k neighbours?
  int numValidDistances = potentialNN.inds.size();
  int k = opts.k;
  *kUsed = numValidDistances;
  if (k > numValidDistances) {
    if (opts.forceCompute) {
      k = numValidDistances;
    } else {
      rc(0, Mp_i) = INSUFFICIENT_UNIQUE;
      return;
    }
  }

  if (k == 0) {
    // Whether we throw an error or just silently ignore this prediction
    // depends on whether we are in 'strict' mode or not.
    rc(0, Mp_i) = opts.forceCompute ? SUCCESS : INSUFFICIENT_UNIQUE;
    return;
  }

  // If we asked for all of the neighbours to be considered (e.g. with k = -1), return this index vector directly.
  DistanceIndexPairs kNNs;
  if (k < 0 || k == potentialNN.inds.size()) {
    kNNs = potentialNN;
  } else {
    kNNs = kNearestNeighbours(potentialNN, k);
  }

  if (keep_going != nullptr && keep_going() == false) {
    rc(0, Mp_i) = BREAK_HIT;
    return;
  }

  if (opts.algorithm == Algorithm::Simplex) {
    for (int t = 0; t < opts.thetas.size(); t++) {
      simplex_prediction(Mp_i, t, opts, M, kNNs.dists, kNNs.inds, ystar, rc, kUsed);
    }
  } else if (opts.algorithm == Algorithm::SMap) {
    for (int t = 0; t < opts.thetas.size(); t++) {
      smap_prediction(Mp_i, t, opts, M, Mp, kNNs.dists, kNNs.inds, ystar, coeffs, rc, kUsed);
    }
  } else {
    rc(0, Mp_i) = INVALID_ALGORITHM;
  }
}

std::vector<int> potential_neighbour_indices(int Mp_i, const Options& opts, const Manifold& M, const Manifold& Mp)
{
  bool skipOtherPanels = opts.panelMode && (opts.idw < 0);

  std::vector<int> inds;

  for (int i = 0; i < M.nobs(); i++) {
    if (skipOtherPanels && (M.panel(i) != Mp.panel(Mp_i))) {
      continue;
    }

    inds.push_back(i);
  }

  return inds;
}

// For a given point, find the k nearest neighbours of this point.
//
// If there are many potential neighbours with the exact same distances, we
// prefer the neighbours with the smallest index value. This corresponds
// to a stable sort in C++ STL terminology.
//
// In typical use-cases of 'edm explore' the value of 'k' is small, like 5-20.
// However for a typical 'edm xmap' the value of 'k' is set as large as possible.
// If 'k' is small, the partial_sort is efficient as it only finds the 'k' smallest
// distances. If 'k' is larger, then it is faster to simply sort the entire distance
// vector.
DistanceIndexPairs kNearestNeighbours(const DistanceIndexPairs& potentialNeighbours, int k)
{
  std::vector<int> idx(potentialNeighbours.inds.size());
  std::iota(idx.begin(), idx.end(), 0);

  if (k >= (int)(idx.size() / 2)) {
    auto comparator = [&potentialNeighbours](int i1, int i2) {
      return potentialNeighbours.dists[i1] < potentialNeighbours.dists[i2];
    };
    std::stable_sort(idx.begin(), idx.end(), comparator);
  } else {
    auto stableComparator = [&potentialNeighbours](int i1, int i2) {
      if (potentialNeighbours.dists[i1] != potentialNeighbours.dists[i2])
        return potentialNeighbours.dists[i1] < potentialNeighbours.dists[i2];
      else
        return i1 < i2;
    };
    std::partial_sort(idx.begin(), idx.begin() + k, idx.end(), stableComparator);
  }

  std::vector<int> kNNInds(k);
  std::vector<double> kNNDists(k);

  for (int i = 0; i < k; i++) {
    kNNInds[i] = potentialNeighbours.inds[idx[i]];
    kNNDists[i] = potentialNeighbours.dists[idx[i]];
  }

  return { kNNInds, kNNDists };
}

// An alternative version of 'kNearestNeighbours' which doesn't sort the neighbours.
// This version splits ties differently on different OS's, so it can't be used directly,
// though perhaps a platform-independent implementation of std::nth_element would solve this problem.
DistanceIndexPairs kNearestNeighboursUnstable(const DistanceIndexPairs& potentialNeighbours, int k)
{
  std::vector<int> indsToPartition(potentialNeighbours.inds.size());
  std::iota(indsToPartition.begin(), indsToPartition.end(), 0);

  auto comparator = [&potentialNeighbours](int i1, int i2) {
    return potentialNeighbours.dists[i1] < potentialNeighbours.dists[i2];
  };
  std::nth_element(indsToPartition.begin(), indsToPartition.begin() + k, indsToPartition.end(), comparator);

  std::vector<int> kNNInds(k);
  std::vector<double> kNNDists(k);

  for (int i = 0; i < k; i++) {
    kNNInds[i] = potentialNeighbours.inds[indsToPartition[i]];
    kNNDists[i] = potentialNeighbours.dists[indsToPartition[i]];
  }

  return { kNNInds, kNNDists };
}

void simplex_prediction(int Mp_i, int t, const Options& opts, const Manifold& M, const std::vector<double>& dists,
                        const std::vector<int>& kNNInds, Eigen::Map<MatrixXd> ystar, Eigen::Map<MatrixXi> rc,
                        int* kUsed)
{
  int k = kNNInds.size();

  // Find the smallest distance (closest neighbour) among the supplied neighbours.
  double minDist = *std::min_element(dists.begin(), dists.end());

  // Calculate our weighting of each neighbour, and the total sum of these weights.
  std::vector<double> w(k);
  double sumw = 0.0;
  const double theta = opts.thetas[t];

  int numNonZeroWeights = 0;
  for (int j = 0; j < k; j++) {
    w[j] = exp(-theta * (dists[j] / minDist));
    sumw = sumw + w[j];
    numNonZeroWeights += (w[j] > 0);
  }

  // For the sake of debugging, count how many neighbours we end up with.
  if (opts.saveKUsed) {
    *kUsed = numNonZeroWeights;
  }

  // Make the simplex projection/prediction.
  double r = 0.0;
  for (int j = 0; j < k; j++) {
    r = r + M.y(kNNInds[j]) * (w[j] / sumw);
  }

  // Store the results & return value.
  ystar(t, Mp_i) = r;
  rc(t, Mp_i) = SUCCESS;
}

void smap_prediction(int Mp_i, int t, const Options& opts, const Manifold& M, const Manifold& Mp,
                     const std::vector<double>& dists, const std::vector<int>& kNNInds, Eigen::Map<MatrixXd> ystar,
                     Eigen::Map<MatrixXd> coeffs, Eigen::Map<MatrixXi> rc, int* kUsed)
{
  int k = kNNInds.size();

  // Pull out the nearest neighbours from the manifold, and
  // simultaneously prepend a column of ones in front of the manifold data.
  MatrixXd X_ls_cj(k, M.E_actual() + 1);
  X_ls_cj << Eigen::VectorXd::Ones(k), M.map()(kNNInds, Eigen::all);

  // Calculate the weight for each neighbour
  Eigen::Map<const Eigen::VectorXd> distsMap(&(dists[0]), dists.size());
  Eigen::VectorXd w = Eigen::exp(-opts.thetas[t] * (distsMap.array() / distsMap.mean()));

  // For the sake of debugging, count how many neighbours we end up with.
  if (opts.saveKUsed) {
    int numNonZeroWeights = 0;
    for (double& w_i : w) {
      if (w_i > 0) {
        numNonZeroWeights += 1;
      }
    }
    *kUsed = numNonZeroWeights;
  }

  // Scale everything by our weights vector
  X_ls_cj.array().colwise() *= w.array();
  Eigen::VectorXd y_ls = M.yMap()(kNNInds).array() * w.array();

  // The old way to solve this system:
  // Eigen::BDCSVD<MatrixXd> svd(X_ls_cj, Eigen::ComputeThinU | Eigen::ComputeThinV);
  //  Eigen::VectorXd ics = svd.solve(y_ls);

  // The pseudo-inverse of X can be calculated as (X^T * X)^(-1) * X^T
  // see https://scicomp.stackexchange.com/a/33375
  const int svdOpts = Eigen::ComputeThinU | Eigen::ComputeThinV; // 'ComputeFull*' would probably work identically here.
  Eigen::BDCSVD<MatrixXd> svd(X_ls_cj.transpose() * X_ls_cj, svdOpts);
  Eigen::VectorXd ics = svd.solve(X_ls_cj.transpose() * y_ls);

  double r = ics(0);
  for (int j = 0; j < M.E_actual(); j++) {
    if (Mp(Mp_i, j) != MISSING_SENTINEL) {
      r += Mp(Mp_i, j) * ics(j + 1);
    }
  }

  // If the 'savesmap' option is given, save the 'ics' coefficients
  // for the largest value of theta.
  if (opts.saveSMAPCoeffs && t == opts.thetas.size() - 1) {
    for (int j = 0; j < M.E_actual() + 1; j++) {
      if (ics(j) == 0.) {
        coeffs(Mp_i, j) = MISSING_SENTINEL;
      } else {
        coeffs(Mp_i, j) = ics(j);
      }
    }
  }

  ystar(t, Mp_i) = r;
  rc(t, Mp_i) = SUCCESS;
}


//
// Copyright (c) 2016 - 2017 Mesh Consultants Inc.
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

#include "Geomlib_TriMeshThicken.h"

#include <iostream>

#include <Eigen/Core>

#pragma warning(push, 0)
#include <igl/is_boundary_edge.h>
#include <igl/edge_flaps.h>
#include <igl/edges.h>
#pragma warning(pop)

#include <Urho3D/Core/Variant.h>

#include "ConversionUtilities.h"
#include "TriMesh.h"

#pragma warning(disable : 4244)

using Urho3D::Variant;
using Urho3D::VariantMap;
using Urho3D::VariantVector;
using Urho3D::Vector3;

bool Geomlib::TriMeshThicken(
	const Urho3D::Variant& meshIn,
	float d,
	Urho3D::Variant& meshOut
)
{
	// meshIn: verify and parse
	if (!TriMesh_Verify(meshIn)) {
		std::cerr << "ERROR: Geomlib::Thicken --- meshIn unverified\n";
		meshOut = Variant();
		return false;
	}
	VariantMap meshMap = meshIn.GetVariantMap();
	const VariantVector vertexList = meshMap["vertices"].GetVariantVector();
	const VariantVector faceList = meshMap["faces"].GetVariantVector();

	// Compute vertices for outer part of solid, using unit normals
	VariantVector vertNormals = TriMesh_ComputeVertexNormals(meshIn, true);
	assert(vertNormals.Size() == vertexList.Size());
	VariantVector outerVertexList;
	for (unsigned i = 0; i < vertexList.Size(); ++i) {
		Vector3 vert = vertexList[i].GetVector3();
		Vector3 unorm = vertNormals[i].GetVector3();

		Vector3 newVert = vert + d * unorm;
		outerVertexList.Push(Variant(newVert));
	}
	// Compute faces for outer part of solid
	// (using convention that outer vertices are listed in the 2nd half of final vertex list
	VariantVector outerFaceList;
	unsigned numVertices = vertexList.Size();
	for (unsigned i = 0; i < faceList.Size(); ++i) {
		unsigned v = faceList[i].GetUInt();
		outerFaceList.Push(numVertices + faceList[i].GetUInt());
	}

	// Compute faces to form panels connecting exposed edges
	// on inner and outer surfaces.
	VariantVector sideFaceList;

	// Determine which edges are boundary edges, using libigl
	Eigen::MatrixXf V;
	Eigen::MatrixXi F;
	bool matrixSuccess = IglMeshToMatrices(meshIn, V, F);
	assert(matrixSuccess);
	Eigen::MatrixXi EE; // EE is used to precompute # of edges, then is discarded
	igl::edges(F, EE);
	Eigen::MatrixXi E, EI;
	Eigen::MatrixXi EF(EE.rows(), 2);
	EF = Eigen::MatrixXi::Constant(EE.rows(), 2, -1); // we use the -1 default values to help identify boundary edges
	Eigen::VectorXi EMAP;
	igl::edge_flaps(F, E, EMAP, EF, EI);
	Eigen::VectorXi B;
	igl::is_boundary_edge(E, F, B);
	assert(E.rows() == B.rows());

	// Loop over boundary edges, constructing "thickened" outer edges
	unsigned numFaces = F.rows();
	for (unsigned e = 0; e < B.rows(); ++e) {
		if (B(e) == 1) {
			// Boundary found! Edge e is a boundary edge.
			// E
			unsigned f0 = EF(e, 0);
			unsigned f1 = EF(e, 1);

			assert(
				(f0 == -1 && (f1 >= 0 && f1 < numFaces)) ||
				((f0 >= 0 && f0 < numFaces) && f1 == -1)
			);

			unsigned v0 = -1;
			unsigned v1 = -1;

			if (f0 == -1) {
				// f1 is the face with boundary edge e opposite vertex EI(e, 1)
				unsigned v = EI(e, 1);

				v0 = F(f1, (v + 1) % 3);
				v1 = F(f1, (v + 2) % 3);
			}
			else if (f1 == -1) {
				// f0 is the face with boundary edge e opposite vertex EI(e, 0)
				unsigned v = EI(e, 0);

				v0 = F(f0, (v + 1) % 3);
				v1 = F(f0, (v + 2) % 3);
			}
			else {
				meshOut = Variant();
				return false;
			}

			// v0, v1: indices into vertexList of vertices (in order) of edge on the boundary
			unsigned v2 = v0 + vertexList.Size();
			unsigned v3 = v1 + vertexList.Size();

			sideFaceList.Push(Variant(v0));
			sideFaceList.Push(Variant(v1));
			sideFaceList.Push(Variant(v2));
			//
			sideFaceList.Push(Variant(v2));
			sideFaceList.Push(Variant(v1));
			sideFaceList.Push(Variant(v3));
		}
	}

	VariantVector innerFaceList;
	for (unsigned i = 0; i < faceList.Size(); i += 3) {
		innerFaceList.Push(faceList[i]);
		innerFaceList.Push(faceList[i + 2]);
		innerFaceList.Push(faceList[i + 1]);
	}

	// Collect all the inner, outer, and side panel data
	/*
	Speculative: if d < 0, reverse all orientations....
	*/

	VariantVector newVertexList = vertexList;
	newVertexList.Push(outerVertexList);
	Variant newVertices(newVertexList);

	VariantVector newFaceList = innerFaceList;
	newFaceList.Push(outerFaceList);
	newFaceList.Push(sideFaceList);
	if (d < 0.0f) {
		for (unsigned i = 0; i < newFaceList.Size(); i += 3) {
			int i1 = newFaceList[i + 1].GetInt();
			int i2 = newFaceList[i + 2].GetInt();

			newFaceList[i + 1] = Variant(i2);
			newFaceList[i + 2] = Variant(i1);
		}
	}
	Variant newFaces(newFaceList);

	meshOut = TriMesh_Make(newVertices, newFaces);
	return true;
}


/*
 * Primitive.cpp
 *
 *  Created on: 2012-02-19
 *      Author: leo
 */

#include <fruitpunch/Graphics/Vertex.h>
#include <fruitpunch/Graphics/Matrices.h>
#include <fruitpunch/Graphics/Renderer.h>
#include <fruitpunch/Graphics/Window.h>
#include <fruitpunch/Graphics/Primitive.h>
#include <fruitpunch/Graphics/Image.h>
#include <fruitpunch/Common/Runtime.h>
#include <fruitpunch/Elements/Scene.h>

#include <boost/shared_ptr.hpp>
#include <list>
#include <stdio.h>
#include <armadillo>
#include <vector>
#include <algorithm>
#include <functional>

using namespace std;
using namespace arma;
using namespace fp_core::common;

namespace fp_core {

Primitive::Primitive() : Element(), m_trans_vertices(NULL) {

  point zero;
  zero.x = 0;
  zero.y = 0;

  point one;
  one.x = 1;
  one.y = 1;
  m_center = zero;
  m_rotation = 0;
  m_scale = one;
  m_position = zero;

  point center;
  center.x = 0.5f;
  center.y = 0.5f;
  mTexScale = one;
  mTexPosition = zero;
  mTexRotation = 0;
  mTexCenter = center;

  m_alpha = 1;

  m_window = Runtime::getCurrent()->getWindow().get();
}

//______________________________________________________________________________

Primitive::Primitive(const Primitive & p) : m_trans_vertices(NULL) {
  copy_from(p);
}

//______________________________________________________________________________

Primitive::~Primitive() {

}

//______________________________________________________________________________

Primitive & Primitive::operator =(const Primitive & p) {
  copy_from(p);
  return *this;
}

//______________________________________________________________________________

void Primitive::addVertex(const Vertex& vertex) {
  m_vertices.push_back(vertex);

  // if we have at least three vertices, generate a center
  if (m_vertices.size() >= 3) {
    m_center = verticesCenter();
  }
}

//______________________________________________________________________________

point * Primitive::pointArray() const {

  // Allocates memory
  point * result = new point[m_vertices.size()];

  // Gets the transformation matrix
  fmat33 main = transformationMatrix();

  // create point array
  for (unsigned int i = 0; i < m_vertices.size(); i++) {

    // gets the point
    point p = m_vertices[i].getPosition();

    // transforms the point
    p = getWindow()->getRenderer()->transform_point(p, main);
    result[i] = p;
  }

  return result;
}

//______________________________________________________________________________

void Primitive::move(float x, float y) {
  m_position.x = m_position.x + x;
  m_position.y = m_position.y + y;
}

//______________________________________________________________________________

void Primitive::rotate(float amount) {
  m_rotation = m_rotation + amount;
}

//______________________________________________________________________________

void Primitive::scale(float x, float y) {
  m_scale.x = m_scale.x * x;
  m_scale.y = m_scale.y * y;
}

//______________________________________________________________________________

void Primitive::copy_from(const Primitive & p) {
  m_center = p.getCenter();
  m_rotation = p.getRotation();
  m_position = p.getPosition();
  m_scale = p.getScale();
  m_stick = p.isStick();
  m_vertices = p.getVertices();
  m_window = p.getWindow();
  mImage = p.getImage();
  mTexCenter = p.getTexCenter();
  mTexPosition = p.getTexPosition();
  mTexRotation = p.getTexRotation();
  mTexScale = p.getTexScale();
  m_alpha = p.getAlpha();
}

//______________________________________________________________________________

std::vector<point> Primitive::boundaries() const {
  vector<point> positions = Vertex::positionPointVector(m_vertices);
  return point::boundaries(positions);
}

//______________________________________________________________________________

Box Primitive::boundingBox() const {
  vector<point> bounds = boundaries();
  Box b;
  point bottom_left;
  point bottom_right;
  point top_left;
  point top_right;

  bottom_left.x = bounds[0].x;
  bottom_left.y = bounds[0].y;
  bottom_right.x = bounds[1].x;
  bottom_right.y = bounds[0].y;
  top_left.x = bounds[0].x;
  top_left.y = bounds[1].y;
  top_right.x = bounds[1].x;
  top_right.y = bounds[1].y;

  b.setCoordinates(top_left, top_right, bottom_left, bottom_right);
  return b;
}

//______________________________________________________________________________

point Primitive::verticesCenter() const {
  vector<point> bounds = boundaries();
  point center;
  point min = bounds[0];
  point max = bounds[1];
  center.x = (min.x + max.x) / 2;
  center.y = (min.y + max.y) / 2;
  return center;
}

//______________________________________________________________________________

point * Primitive::uvPointArray() const {
  vector<Vertex>::const_iterator it;

  fmat33 main;
  main.eye(); // assigns identity

  // ---------------------------------------------------------------------------
  // Uv Transformations
  // ---------------------------------------------------------------------------

  // goes to local space
   main = (arma::fmat33) (to_local_space_matrix(mTexCenter) * main);

   // concatenates the transformations
   main = (arma::fmat33) (rotation_matrix(mTexRotation) * main);
   main = (arma::fmat33) (scale_matrix(mTexScale) * main);

   // goes to world space
   main = (arma::fmat33) (to_global_space_matrix(mTexCenter) * main);

   // moves to position
   main = (arma::fmat33) (translation_matrix(mTexPosition) * main);

   point * result = new point[m_vertices.size()];
   int i = 0;
  for (it = m_vertices.begin(); it != m_vertices.end(); it++) {
    point uv = it->getUvmap();

    // transforms the UV coordinates before pushing it
    uv = getWindow()->getRenderer()->transform_point(uv, main);

    result[i] = uv;
    i++;
  }
  return result;
}

//______________________________________________________________________________

struct Primitive::comparator {
  boost::shared_ptr<Renderer> renderer;
  point center;

  bool operator ()(const point& i, const point& j) const {
    Renderer& r = *renderer;
    return point::comesBeforeCC(i, j, center, r);
  }
};

bool Primitive::contains(point p) {

  vector<point> V = this->pointVector();
  if (V.size() < 3) return false;

  // checks if the point intersects any of the vertices
  vector<point>::iterator it = V.begin();
  for (; it != V.end(); it++) {
      if (it->x == p.x && it->y == p.y) {
          return true;
      }
  }

  /*
   * Checks if the point is inside the primitive (polygon) by using the
   * first algorithm described in
   * http://local.wasp.uwa.edu.au/~pbourke/geometry/insidepoly/
   */

  // Order points counter clockwise
  comparator comp;
  comp.renderer = getWindow()->getRenderer();
  comp.center = getCenter();
  std::sort(V.begin(), V.end(), comp);

  // close the polygon
  V.push_back(V[0]);

  int crossings = 0;

  // loops through all the edges
  for (unsigned int i=0; i<V.size()-1; i++) {
    point vA = V[i];
    point vB = V[i+1];

    // orders vA and vB according to which one is at the bottom - biggest
    // Y coordinate. vB will always be the highest Y.
    if (vA.y > vB.y) {
      point temp = vA;
      vA = vB;
      vB = temp;
    }

    // gets the minimum and maximum x,y values and checks if the point
    // falls right on the edge
    vector<point> edge;
    edge.push_back(vA);
    edge.push_back(vB);
    vector<point> bounds = point::boundaries(edge);
    if (p.x <= bounds[1].x && p.x >= bounds[0].x && p.y <= bounds[1].y && p.y >= bounds[0].y) {
        return true;
    }

    //printf("vA: (%f,%f), vB: (%f,%f) \n", vA.x, vA.y, vB.x, vB.y);
    // Checks if the point is between the Ys of the edges
    if (vA.y == vB.y) continue; // no point doing anything if its a horizontal line
    if (p.y > vB.y || p.y < vA.y) continue;

    // Checks if its a vertical edge and if the point comes to the left of the
    // edge
    if (vA.x == vB.x && p.x <= vA.x) {
      crossings++;
      continue;
    }

    // If it gets here, its because the edge is not a horizontal or vertical
    // line. In this case, we have to get a line equation in the form
    // f(x) = ax + b in order to get the X value of the edge that intercepts
    // the Y value of our point.
    float a = (vA.y - vB.y) / (vA.x - vB.x);
    float b = vB.y - a * vB.x;
    float ex = (p.y - b) / a;
    if (p.x <= ex) crossings++;
  }

  return (crossings % 2 != 0); // true if its odd
}

//______________________________________________________________________________

std::vector<point> Primitive::transformedBounds() {
  vector<point> points = pointVector();
  return point::boundaries(points);
}

//______________________________________________________________________________

vector<point> Primitive::pointVector() {
	vector<point> result;
	point * points = pointArray();
	for (int i =0; i<m_vertices.size(); i++) {
		result.push_back(points[i]);
	}
	delete points;
	return result;
}

//______________________________________________________________________________

vector<point> Primitive::uvPointVector() {
	vector<point> result;
	point * points = uvPointArray();
	for (int i =0; i<m_vertices.size(); i++) {
		result.push_back(points[i]);
	}
	delete points;
	return result;
}

//______________________________________________________________________________

arma::fmat33 Primitive::transformationMatrix() const {
  fmat33 main;
  main.eye(); // assigns identity

  // goes to local space
  main = (arma::fmat33) (to_local_space_matrix(m_center) * main);

  // concatenates the transformations
  main = (arma::fmat33) (rotation_matrix(m_rotation) * main);
  main = (arma::fmat33) (scale_matrix(m_scale) * main);

  // goes to world space
  main = (arma::fmat33) (to_global_space_matrix(m_center) * main);

  // moves to position
  main = (arma::fmat33) (translation_matrix(m_position) * main);

  return main;
}

// ---------------------------------------------------------------------------
// Getters and setters
// ---------------------------------------------------------------------------

float Primitive::getWidth() const {
  vector<point> bounds = boundaries();

  // original width is the maximum X - minimum X
  float origWidth = bounds[1].x - bounds[0].x;

  return origWidth * m_scale.x;
}

float Primitive::getHeight() const {
  vector<point> bounds = boundaries();

  // original width is the maximum X - minimum X
  float origHeight = bounds[1].y - bounds[0].y;

  return origHeight * m_scale.y;
}

void Primitive::setWidth(float width) {
  vector<point> bounds = boundaries();

  // original width is the maximum X - minimum X
  float origWidth = bounds[1].x - bounds[0].x;

  m_scale.x = width / origWidth;
}

void Primitive::setHeight(float height) {
  vector<point> bounds = boundaries();

  // original width is the maximum X - minimum X
  float origHeight = bounds[1].y - bounds[0].y;

  m_scale.y = height / origHeight;
}

float Primitive::getTop() {
  return transformedBounds()[0].y;
}

float Primitive::getLeft() {
  return transformedBounds()[0].x;
}

void Primitive::setTop(float top) {
  // we need minimum Y
  vector<point> bounds = transformedBounds();
  float minY = bounds[0].y;

  // calculates the difference
  float delta = top - minY;

  // corrects it
  move(0, delta);

}

void Primitive::setLeft(float left) {
  // we need minimum X
  vector<point> bounds = transformedBounds();
  float minX = bounds[0].x;

  // calculates the difference
  float delta = left - minX;

  // corrects it
  move(delta, 0);
}

float Primitive::getX() const {
  return m_position.x;
}

float Primitive::getY() const {
  return m_position.y;
}

void Primitive::setX(float x) {
  m_position.x = x;
}

void Primitive::setY(float y) {
  m_position.y = y;
}

point Primitive::getCenter() const {
  return m_center;
}

point Primitive::getPosition() const {
  return m_position;
}

float Primitive::getRotation() const {
  return m_rotation;
}

point Primitive::getScale() const {
  return m_scale;
}

vector<Vertex> Primitive::getVertices() const {
  return m_vertices;
}

Window *Primitive::getWindow() const {
  return m_window;
}

void Primitive::setCenter(point center) {
  m_center = center;
}

void Primitive::setPosition(point position) {
  m_position = position;
}

void Primitive::setRotation(float rotation) {
  m_rotation = rotation;
}

void Primitive::setScale(point scale) {
  m_scale = scale;
}

void Primitive::setWindow(Window *window) {
  m_window = window;
}

bool Primitive::isStick() const {
  return m_stick;
}

void Primitive::setStick(bool stick) {
  m_stick = stick;
}

point * Primitive::getTransVertices() const
{
    return m_trans_vertices;
}

void Primitive::setTransVertices(point * transVertices)
{
	if (m_trans_vertices != NULL) {
		delete[] m_trans_vertices;
	}

	m_trans_vertices = new point[m_vertices.size()];
	for (int i=0; i<m_vertices.size(); i++) {
		m_trans_vertices[i] = transVertices[i];
	}
}

Vertex& Primitive::getVertex(int index) {
	return m_vertices[index];
}

boost::shared_ptr<Image> Primitive::getImage() const {
  return mImage;
}

void Primitive::setImage(boost::shared_ptr<Image> image) {
  mImage = image;
}

point Primitive::getTexCenter() const {
  return mTexCenter;
}

float Primitive::getAlpha() const {
	return m_alpha;
}
point Primitive::getTexPosition() const {
  return mTexPosition;
}

void Primitive::setAlpha(float alpha) {
	m_alpha = alpha;
}
point Primitive::getTexScale() const {
  return mTexScale;
}

void Primitive::setTexCenter(point texCenter) {
  mTexCenter = texCenter;
}

void Primitive::setTexPosition(point texPosition) {
  mTexPosition = texPosition;
}

void Primitive::setTexScale(point texScale) {
  mTexScale = texScale;
}

float Primitive::getTexRotation() const {
  return mTexRotation;
}

void Primitive::setTexRotation(float texRotation) {
  mTexRotation = texRotation;
}

void fp_core::Primitive::onAdd() {
  Runtime::getCurrent()->getWindow()->addPrimitive(getThis<Primitive>());
}

void fp_core::Primitive::onRemove() {
	getScene()->getWindow()->removePrimitive(getThis<Primitive>());
}

float Primitive::getTexX() {
	return mTexPosition.x;
}

float Primitive::getTexY() {
	return mTexPosition.y;
}

void Primitive::setTexX(float val) {
	mTexPosition.x = val;
}

void Primitive::setTexY(float val) {
	mTexPosition.y = val;
}

/* namespace fp_core */
}