/** Copyright 2017-2020 CNRS-AIST JRL and CNRS-UM LIRMM */

#define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN
#define DOCTEST_CONFIG_SUPER_FAST_ASSERTS
#include "doctest/doctest.h"

#include "SolverTestFunctions.h"

#include <tvm/Variable.h>
#include <tvm/function/BasicLinearFunction.h>
#include <tvm/function/IdentityFunction.h>
#include <tvm/task_dynamics/Clamped.h>
#include <tvm/task_dynamics/Constant.h>
#include <tvm/task_dynamics/FeedForward.h>
#include <tvm/task_dynamics/None.h>
#include <tvm/task_dynamics/OneStepToZero.h>
#include <tvm/task_dynamics/Proportional.h>
#include <tvm/task_dynamics/ProportionalDerivative.h>
#include <tvm/task_dynamics/Reference.h>
#include <tvm/task_dynamics/VelocityDamper.h>
#include <tvm/utils/graph.h>

#include <Eigen/QR>

#include <iostream>

#ifdef __i386__
#  define APPROX_I386(x) doctest::Approx(x)
#else
#  define APPROX_I386(x) x
#endif

using namespace Eigen;
using namespace tvm;

/** Feed forward value provider */
class FFProvider : public tvm::graph::abstract::Node<FFProvider>
{
public:
  SET_OUTPUTS(FFProvider, Value)
  SET_UPDATES(FFProvider, Value)

  FFProvider(const Eigen::VectorXd & value)
  {
    registerUpdates(Update::Value, &FFProvider::updateValue);
    addOutputDependency(Output::Value, Update::Value);
    value_ = value;
  }

  const Eigen::VectorXd & value() const { return value_; }

  void updateValue() { updated_ = true; }

  Eigen::VectorXd value_;
  bool updated_ = false;
};

TEST_CASE("Valid construction")
{
  VariablePtr x = Space(3).createVariable("x");
  x << Vector3d::Zero();
  auto f = std::make_shared<function::IdentityFunction>(x);
  task_dynamics::Constant td;
  CHECK_NOTHROW(td.impl(f, constraint::Type::EQUAL, Vector3d(1, 0, 0)));
  CHECK_THROWS(td.impl(f, constraint::Type::EQUAL, Vector2d(1, 0)));
}

TEST_CASE("Test Constant")
{
  VariablePtr x = Space(3).createVariable("x");
  x << Vector3d::Zero();
  auto f = std::make_shared<function::IdentityFunction>(x);
  task_dynamics::Constant td;
  auto tdi = td.impl(f, constraint::Type::EQUAL, Vector3d(1, 0, 0));

  tdi->updateValue();

  FAST_CHECK_EQ(td.order(), task_dynamics::Order::Zero);
  FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::Zero);
  FAST_CHECK_UNARY(tdi->value().isApprox(Vector3d(1, 0, 0)));
  FAST_CHECK_UNARY(tdi->checkType<task_dynamics::Constant::Impl>());
  FAST_CHECK_UNARY_FALSE(tdi->checkType<task_dynamics::None::Impl>());
}

TEST_CASE("Test None")
{
  VariablePtr x = Space(3).createVariable("x");
  x << 1, 2, 3;
  MatrixXd A = MatrixXd::Random(2, 3);
  Vector2d b(1, 2);
  auto f = std::make_shared<function::BasicLinearFunction>(A, x, b);
  task_dynamics::None td;
  auto tdi = td.impl(f, constraint::Type::EQUAL, Vector2d(1, 0));

  f->updateValue();
  tdi->updateValue();

  FAST_CHECK_EQ(td.order(), task_dynamics::Order::Zero);
  FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::Zero);
  FAST_CHECK_UNARY(tdi->value().isApprox(Vector2d(0, -2)));
  FAST_CHECK_UNARY(tdi->checkType<task_dynamics::None::Impl>());
  FAST_CHECK_UNARY_FALSE(tdi->checkType<task_dynamics::Constant::Impl>());
}

TEST_CASE("Test Proportional")
{
  VariablePtr x = Space(3).createVariable("x");
  x << 1, 2, 3;
  auto f = std::make_shared<SphereFunction>(x, Vector3d(1, 0, -3), 2);

  double kp = 2;
  task_dynamics::P td(kp);
  VectorXd rhs(1);
  rhs[0] = -1;
  auto tdi = td.impl(f, constraint::Type::EQUAL, rhs);

  f->updateValue();
  tdi->updateValue();

  FAST_CHECK_EQ(td.order(), task_dynamics::Order::One);
  FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::One);
  FAST_CHECK_EQ(tdi->value()[0], -kp * (36 - rhs[0])); // -kp*||(1,2,3) - (1,0,-3)||^2 - 2^2 - rhs
  FAST_CHECK_UNARY(tdi->checkType<task_dynamics::P::Impl>());
  FAST_CHECK_UNARY_FALSE(tdi->checkType<task_dynamics::Constant::Impl>());

  kp = 3;
  dynamic_cast<task_dynamics::P::Impl *>(tdi.get())->gain(kp);
  tdi->updateValue();
  FAST_CHECK_EQ(tdi->value()[0], -kp * (36 - rhs[0])); // -kp*||(1,2,3) - (1,0,-3)||^2 - 2^2 - rhs
}

TEST_CASE("Test Proportional gain types")
{
  VariablePtr x = Space(3).createVariable("x");
  x << 1, 2, 3;
  MatrixXd A = MatrixXd::Random(3, 3);
  Vector3d b(1, 2, 3);
  auto f = std::make_shared<function::BasicLinearFunction>(A, x, b);

  double kps = 2;
  VectorXd kpv = Vector3d::Constant(2);
  MatrixXd kpm = 2 * Matrix3d::Identity();
  task_dynamics::P tds(kps);
  task_dynamics::P tdv(kpv);
  task_dynamics::P tdm(kpm);
  Vector3d rhs(-1, -2, -3);

  auto tdsi = tds.impl(f, constraint::Type::EQUAL, rhs);
  auto tdvi = tdv.impl(f, constraint::Type::EQUAL, rhs);
  auto tdmi = tdm.impl(f, constraint::Type::EQUAL, rhs);

  f->updateValue();
  tdsi->updateValue();
  tdvi->updateValue();
  tdmi->updateValue();

  FAST_CHECK_UNARY(tdsi->value().isApprox(tdvi->value()));
  FAST_CHECK_UNARY(tdsi->value().isApprox(tdmi->value()));

  kps = 3;
  kpv = Vector3d::Constant(3);
  kpm = 3 * Matrix3d::Identity();

  dynamic_cast<task_dynamics::P::Impl *>(tdsi.get())->gain(kps);
  dynamic_cast<task_dynamics::P::Impl *>(tdvi.get())->gain(kpv);
  dynamic_cast<task_dynamics::P::Impl *>(tdmi.get())->gain(kpm);

  tdsi->updateValue();
  tdvi->updateValue();
  tdmi->updateValue();

  FAST_CHECK_UNARY(tdsi->value().isApprox(tdvi->value()));
  FAST_CHECK_UNARY(tdsi->value().isApprox(tdmi->value()));

  dynamic_cast<task_dynamics::P::Impl *>(tdsi.get())->gain() = 5;
  dynamic_cast<task_dynamics::P::Impl *>(tdvi.get())->gain<VectorXd>()[0] = 5;
  dynamic_cast<task_dynamics::P::Impl *>(tdvi.get())->gain<VectorXd>()[1] = 5;
  dynamic_cast<task_dynamics::P::Impl *>(tdvi.get())->gain<VectorXd>()[2] = 5;
  dynamic_cast<task_dynamics::P::Impl *>(tdmi.get())->gain<MatrixXd>().diagonal() = Vector3d::Constant(5);

  tdsi->updateValue();
  tdvi->updateValue();
  tdmi->updateValue();

  FAST_CHECK_UNARY(tdsi->value().isApprox(tdvi->value()));
  FAST_CHECK_UNARY(tdsi->value().isApprox(tdmi->value()));
  CHECK_THROWS(dynamic_cast<task_dynamics::P::Impl *>(tdvi.get())->gain<MatrixXd>());
}

TEST_CASE("Test Proportional Derivative")
{
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  x << 1, 2, 3;
  dx << 1, 1, 1;
  auto f = std::make_shared<SphereFunction>(x, Vector3d(1, 0, -3), 2);

  double kp = 2;
  double kv = 3;
  task_dynamics::PD td(kp, kv);
  VectorXd rhs(1);
  rhs[0] = -1;
  auto tdi = td.impl(f, constraint::Type::EQUAL, rhs);

  f->updateValue();
  f->updateVelocityAndNormalAcc();
  tdi->updateValue();

  FAST_CHECK_EQ(td.order(), task_dynamics::Order::Two);
  FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::Two);
  FAST_CHECK_EQ(tdi->value()[0], -kp * (36 - rhs[0]) - kv * 16);
  FAST_CHECK_UNARY(tdi->checkType<task_dynamics::PD::Impl>());
  FAST_CHECK_UNARY_FALSE(tdi->checkType<task_dynamics::Constant::Impl>());
}

TEST_CASE("Test Feed Forward Proportional Derivative")
{
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  x << 1, 2, 3;
  dx << 1, 1, 1;
  auto f = std::make_shared<SphereFunction>(x, Vector3d(1, 0, -3), 2);

  auto provider = std::make_shared<FFProvider>(Eigen::VectorXd::Constant(1, 10.0));

  double kp = 2;
  double kv = 3;
  task_dynamics::FeedForwardPD td(provider, &FFProvider::value, FFProvider::Output::Value, kp, kv);
  VectorXd rhs(1);
  rhs[0] = -1;
  TaskDynamicsPtr tdi = td.impl(f, constraint::Type::EQUAL, rhs);

  auto Value = task_dynamics::abstract::TaskDynamicsImpl::Output::Value;
  auto gl = utils::generateUpdateGraph(tdi, Value);
  gl->execute();

  FAST_CHECK_UNARY(provider->updated_);
  FAST_CHECK_EQ(td.order(), task_dynamics::Order::Two);
  FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::Two);
  FAST_CHECK_EQ(tdi->value()[0], -kp * (36 - rhs[0]) - kv * 16 + 10.0);
  FAST_CHECK_UNARY(tdi->checkType<task_dynamics::PD::Impl>());
  FAST_REQUIRE_UNARY(tdi->checkType<task_dynamics::FeedForwardPD::Impl>()); // REQUIRE because we do a static_cast after
  FAST_CHECK_UNARY_FALSE(tdi->checkType<task_dynamics::Constant::Impl>());

  kp = 3;
  kv = 2;
  static_cast<task_dynamics::FeedForwardPD::Impl *>(tdi.get())->gains(kp, kv);
  gl->execute();
  FAST_CHECK_EQ(tdi->value()[0], -kp * (36 - rhs[0]) - kv * 16 + 10.0);
}

TEST_CASE("Test Proportional Derivative gain types")
{
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  x << 1, 2, 3;
  dx << 1, 1, 1;
  auto f = std::make_shared<Simple2dRobotEE>(x, Vector2d(0, 0), Vector3d::Ones());

  double kps = 2;
  double kvs = 3;
  VectorXd kpv = Vector2d::Constant(2);
  VectorXd kvv = Vector2d::Constant(3);
  MatrixXd kpm = 2 * Matrix2d::Identity();
  MatrixXd kvm = 3 * Matrix2d::Identity();
  task_dynamics::PD tdss(kps, kvs);
  task_dynamics::PD tdsv(kps, kvv);
  task_dynamics::PD tdsm(kps, kvm);
  task_dynamics::PD tdvs(kpv, kvs);
  task_dynamics::PD tdvv(kpv, kvv);
  task_dynamics::PD tdvm(kpv, kvm);
  task_dynamics::PD tdms(kpm, kvs);
  task_dynamics::PD tdmv(kpm, kvv);
  task_dynamics::PD tdmm(kpm, kvm);
  task_dynamics::PD tds(kvs);
  task_dynamics::PD tdv(kvv);
  task_dynamics::PD tdm(kvm);
  Vector2d rhs(-1, -2);
  auto tdssi = tdss.impl(f, constraint::Type::EQUAL, rhs);
  auto tdsvi = tdsv.impl(f, constraint::Type::EQUAL, rhs);
  auto tdsmi = tdsm.impl(f, constraint::Type::EQUAL, rhs);
  auto tdvsi = tdvs.impl(f, constraint::Type::EQUAL, rhs);
  auto tdvvi = tdvv.impl(f, constraint::Type::EQUAL, rhs);
  auto tdvmi = tdvm.impl(f, constraint::Type::EQUAL, rhs);
  auto tdmsi = tdms.impl(f, constraint::Type::EQUAL, rhs);
  auto tdmvi = tdmv.impl(f, constraint::Type::EQUAL, rhs);
  auto tdmmi = tdmm.impl(f, constraint::Type::EQUAL, rhs);
  auto tdsi = tds.impl(f, constraint::Type::EQUAL, rhs);
  auto tdvi = tdv.impl(f, constraint::Type::EQUAL, rhs);
  auto tdmi = tdm.impl(f, constraint::Type::EQUAL, rhs);

  f->updateValue();
  f->updateVelocityAndNormalAcc();
  tdssi->updateValue();
  tdsvi->updateValue();
  tdsmi->updateValue();
  tdvsi->updateValue();
  tdvvi->updateValue();
  tdvmi->updateValue();
  tdmsi->updateValue();
  tdmvi->updateValue();
  tdmmi->updateValue();
  tdsi->updateValue();
  tdvi->updateValue();
  tdmi->updateValue();

  FAST_CHECK_UNARY(tdssi->value().isApprox(tdsvi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdsvi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdvsi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdvvi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdvmi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdmsi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdmvi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdmmi->value()));
  FAST_CHECK_UNARY(tdsi->value().isApprox(tdvi->value()));
  FAST_CHECK_UNARY(tdsi->value().isApprox(tdvi->value()));

  kps = 3;
  kvs = 1;
  kpv = Vector2d::Constant(3);
  kvv = Vector2d::Constant(1);
  kpm = 3 * Matrix2d::Identity();
  kvm = 1 * Matrix2d::Identity();
  dynamic_cast<task_dynamics::PD::Impl *>(tdssi.get())->gains(kps, kvs);
  dynamic_cast<task_dynamics::PD::Impl *>(tdsvi.get())->gains(kps, kvv);
  dynamic_cast<task_dynamics::PD::Impl *>(tdsmi.get())->gains(kps, kvm);
  dynamic_cast<task_dynamics::PD::Impl *>(tdvsi.get())->gains(kpv, kvs);
  dynamic_cast<task_dynamics::PD::Impl *>(tdvvi.get())->gains(kpv, kvv);
  dynamic_cast<task_dynamics::PD::Impl *>(tdvmi.get())->gains(kpv, kvm);
  dynamic_cast<task_dynamics::PD::Impl *>(tdmsi.get())->gains(kpm, kvs);
  dynamic_cast<task_dynamics::PD::Impl *>(tdmvi.get())->gains(kpm, kvv);
  dynamic_cast<task_dynamics::PD::Impl *>(tdmmi.get())->gains(kpm, kvm);
  dynamic_cast<task_dynamics::PD::Impl *>(tdsi.get())->gains(kps);
  dynamic_cast<task_dynamics::PD::Impl *>(tdvi.get())->gains(kpv);
  dynamic_cast<task_dynamics::PD::Impl *>(tdmi.get())->gains(kpm);

  tdssi->updateValue();
  tdsvi->updateValue();
  tdsmi->updateValue();
  tdvsi->updateValue();
  tdvvi->updateValue();
  tdvmi->updateValue();
  tdmsi->updateValue();
  tdmvi->updateValue();
  tdmmi->updateValue();
  tdsi->updateValue();
  tdvi->updateValue();
  tdmi->updateValue();

  FAST_CHECK_UNARY(tdssi->value().isApprox(tdsvi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdsvi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdvsi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdvvi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdvmi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdmsi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdmvi->value()));
  FAST_CHECK_UNARY(tdssi->value().isApprox(tdmmi->value()));
  FAST_CHECK_UNARY(tdsi->value().isApprox(tdvi->value()));
  FAST_CHECK_UNARY(tdsi->value().isApprox(tdvi->value()));

  dynamic_cast<task_dynamics::PD::Impl *>(tdssi.get())->kp() = 5;
  dynamic_cast<task_dynamics::PD::Impl *>(tdssi.get())->kv() = 2;
  dynamic_cast<task_dynamics::PD::Impl *>(tdvmi.get())->kp<VectorXd>() << 5, 5;
  dynamic_cast<task_dynamics::PD::Impl *>(tdvmi.get())->kv<MatrixXd>() = 2 * Matrix2d::Identity();

  tdssi->updateValue();
  tdvmi->updateValue();

  FAST_CHECK_UNARY(tdssi->value().isApprox(tdvmi->value()));
}

TEST_CASE("Test Velocity Damper")
{
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  auto f = std::make_shared<function::IdentityFunction>(x);

  double di = 3;
  double ds = 1;
  double xsi = 2;
  double dt = 0.1;

  // test validity
  CHECK_THROWS(task_dynamics::VelocityDamper({0.5, ds, xsi}));
  CHECK_THROWS(task_dynamics::VelocityDamper(dt, {0.5, ds, xsi}, 1000));
  CHECK_THROWS(task_dynamics::VelocityDamper({di, ds, -1}));
  CHECK_THROWS(task_dynamics::VelocityDamper(dt, {di, ds, -1}, 1000));
  CHECK_THROWS(task_dynamics::VelocityDamper(0, {di, ds, xsi}, 1000));
  CHECK_THROWS(task_dynamics::VelocityDamper(-0.1, {di, ds, xsi}, 1000));

  // test kinematics
  task_dynamics::VelocityDamper td1({di, ds, xsi});
  {
    x << 1, 2, 4;
    auto tdl = td1.impl(f, constraint::Type::GREATER_THAN, Vector3d::Zero());

    f->updateValue();
    tdl->updateValue();

    FAST_CHECK_EQ(td1.order(), task_dynamics::Order::One);
    FAST_CHECK_EQ(tdl->order(), task_dynamics::Order::One);
    FAST_CHECK_EQ(tdl->value()[0], 0);
    FAST_CHECK_EQ(tdl->value()[1], -1);
    FAST_CHECK_EQ(tdl->value()[2], -constant::big_number);
    FAST_CHECK_UNARY(tdl->checkType<task_dynamics::VelocityDamper::Impl>());
    FAST_CHECK_UNARY_FALSE(tdl->checkType<task_dynamics::Constant::Impl>());

    x << -1, -2, -4;
    auto tdu = td1.impl(f, constraint::Type::LOWER_THAN, Vector3d::Zero());
    f->updateValue();
    tdu->updateValue();
    FAST_CHECK_EQ(tdu->value()[0], 0);
    FAST_CHECK_EQ(tdu->value()[1], 1);
    FAST_CHECK_EQ(tdu->value()[2], constant::big_number);
  }

  // test dynamics
  double big = 1000;
  task_dynamics::VelocityDamper td2(dt, {di, ds, xsi}, big);
  {
    x << 1, 2, 4;
    dx << 1, 1, 1;
    auto tdl = td2.impl(f, constraint::Type::GREATER_THAN, Vector3d::Zero());

    f->updateValue();
    f->updateVelocity();
    tdl->updateValue();

    FAST_CHECK_EQ(td2.order(), task_dynamics::Order::Two);
    FAST_CHECK_EQ(tdl->order(), task_dynamics::Order::Two);
    FAST_CHECK_UNARY(tdl->value().isApprox(Vector3d(-10, -20, -1000)));
    FAST_CHECK_UNARY(tdl->checkType<task_dynamics::VelocityDamper::Impl>());
    FAST_CHECK_UNARY_FALSE(tdl->checkType<task_dynamics::Constant::Impl>());

    x << -1, -2, -4;
    dx << -1, -1, -1;
    auto tdu = td2.impl(f, constraint::Type::LOWER_THAN, Vector3d::Zero());

    f->updateValue();
    f->updateVelocity();
    tdu->updateValue();
    FAST_CHECK_UNARY(tdu->value().isApprox(Vector3d(10, 20, 1000)));
  }
}

TEST_CASE("Test Anisotropic Velocity Damper")
{
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  auto f = std::make_shared<function::IdentityFunction>(x);

  Eigen::Vector3d di = {3, 6, 1};
  Eigen::Vector3d ds = {1, 1, 0.1};
  Eigen::Vector3d xsi = {2, 1, 0.5};
  double dt = 0.1;

  // test validity
  // Should throw because ds > di
  CHECK_THROWS(task_dynamics::VelocityDamper({ds, di, xsi}));
  // Should throw because xsi dimension don't match the others
  CHECK_THROWS(task_dynamics::VelocityDamper({di, ds, Eigen::Vector4d::Zero()}));
  // Should throw because negative damping
  CHECK_THROWS(task_dynamics::VelocityDamper({di, ds, Eigen::Vector3d::Constant(-1)}));
  // Same 3 checks with dt parameter
  CHECK_THROWS(task_dynamics::VelocityDamper(dt, {ds, di, xsi}));
  CHECK_THROWS(task_dynamics::VelocityDamper(dt, {di, ds, Eigen::Vector4d::Zero()}));
  CHECK_THROWS(task_dynamics::VelocityDamper(dt, {di, ds, Eigen::Vector3d::Constant(-1)}));
  // Should throw because dt <= 0
  CHECK_THROWS(task_dynamics::VelocityDamper(0, {di, ds, xsi}));
  CHECK_THROWS(task_dynamics::VelocityDamper(-0.1, {di, ds, xsi}));

  // test kinematics
  task_dynamics::VelocityDamper td1({di, ds, xsi});
  {
    x << 1, 2, 4;
    auto tdl = td1.impl(f, constraint::Type::GREATER_THAN, Vector3d::Zero());

    f->updateValue();
    tdl->updateValue();

    FAST_CHECK_EQ(td1.order(), task_dynamics::Order::One);
    FAST_CHECK_EQ(tdl->order(), task_dynamics::Order::One);
    FAST_CHECK_EQ(tdl->value()[0], 0);
    FAST_CHECK_EQ(tdl->value()[1], -0.2);
    FAST_CHECK_EQ(tdl->value()[2], -constant::big_number);
    FAST_CHECK_UNARY(tdl->checkType<task_dynamics::VelocityDamper::Impl>());
    FAST_CHECK_UNARY_FALSE(tdl->checkType<task_dynamics::Constant::Impl>());

    x << -1, -2, -4;
    auto tdu = td1.impl(f, constraint::Type::LOWER_THAN, Vector3d::Zero());
    f->updateValue();
    tdu->updateValue();
    FAST_CHECK_EQ(tdu->value()[0], 0);
    FAST_CHECK_EQ(tdu->value()[1], 0.2);
    FAST_CHECK_EQ(tdu->value()[2], constant::big_number);
  }

  // test dynamics
  double big = 1000;
  task_dynamics::VelocityDamper td2(dt, {di, ds, xsi}, big);
  {
    x << 1, 2, 4;
    dx << 1, 1, 1;
    auto tdl = td2.impl(f, constraint::Type::GREATER_THAN, Vector3d::Zero());

    f->updateValue();
    f->updateVelocity();
    tdl->updateValue();

    FAST_CHECK_EQ(td2.order(), task_dynamics::Order::Two);
    FAST_CHECK_EQ(tdl->order(), task_dynamics::Order::Two);
    FAST_CHECK_UNARY(tdl->value().isApprox(Vector3d(-10, -12, -1000)));
    FAST_CHECK_UNARY(tdl->checkType<task_dynamics::VelocityDamper::Impl>());
    FAST_CHECK_UNARY_FALSE(tdl->checkType<task_dynamics::Constant::Impl>());

    x << -1, -2, -4;
    dx << -1, -1, -1;
    auto tdu = td2.impl(f, constraint::Type::LOWER_THAN, Vector3d::Zero());

    f->updateValue();
    f->updateVelocity();
    tdu->updateValue();
    FAST_CHECK_UNARY(tdu->value().isApprox(Vector3d(10, 12, 1000)));
  }
}

TEST_CASE("Test automatic xsi")
{
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  auto f = std::make_shared<function::IdentityFunction>(x);

  double di = 3;
  double ds = 1;
  double xsiOff = 1;

  // test kinematics
  double big = 100;
  task_dynamics::VelocityDamper td1({di, ds, 0, xsiOff}, big);
  {
    x << 5, 4, 2;
    dx << -0.5, -0.5, -0.5;
    TaskDynamicsPtr tdl = td1.impl(f, constraint::Type::GREATER_THAN, Vector3d::Zero());

    auto Value = task_dynamics::abstract::TaskDynamicsImpl::Output::Value;
    auto gl = utils::generateUpdateGraph(tdl, Value);

    gl->execute();
    FAST_CHECK_UNARY(tdl->value().isApprox(Vector3d(-big, -big, -1)));

    x << 4.5, 3.5, 1.5;
    gl->execute();
    FAST_CHECK_UNARY(tdl->value().isApprox(Vector3d(-big, -big, -0.5)));

    x << 4, 3, 1;
    gl->execute();
    FAST_CHECK_UNARY(tdl->value().isApprox(Vector3d(-big, -1.5, 0)));

    // we check that two consecutive updates with the same variable values give the same results.
    gl->execute();
    FAST_CHECK_UNARY(tdl->value().isApprox(Vector3d(-big, -1.5, 0)));

    dx << -0.5, -0.5, 0;
    x << 3.5, 2.5, 1;
    gl->execute();
    FAST_CHECK_UNARY(tdl->value().isApprox(Vector3d(-big, -1.125, 0)));

    x << 3, 2, 1;
    gl->execute();
    FAST_CHECK_UNARY(tdl->value().isApprox(Vector3d(-1.5, -0.75, 0)));

    x << -5, -4, -2;
    dx << 0.5, 0.5, 0.5;
    TaskDynamicsPtr tdu = td1.impl(f, constraint::Type::LOWER_THAN, Vector3d::Zero());

    auto gu = utils::generateUpdateGraph(tdu, Value);

    gu->execute();
    FAST_CHECK_UNARY(tdu->value().isApprox(Vector3d(big, big, 1)));

    x << -4.5, -3.5, -1.5;
    gu->execute();
    FAST_CHECK_UNARY(tdu->value().isApprox(Vector3d(big, big, 0.5)));

    x << -4, -3, -1;
    gu->execute();
    FAST_CHECK_UNARY(tdu->value().isApprox(Vector3d(big, 1.5, 0)));

    dx << 0.5, 0.5, 0;
    x << -3.5, -2.5, -1;
    gu->execute();
    FAST_CHECK_UNARY(tdu->value().isApprox(Vector3d(big, 1.125, 0)));

    x << -3, -2, -1;
    gu->execute();
    FAST_CHECK_UNARY(tdu->value().isApprox(Vector3d(1.5, 0.75, 0)));
  }
}

TEST_CASE("Test Reference")
{
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  x << 1, 2, 3;
  dx << 1, 1, 1;
  auto f = std::make_shared<Simple2dRobotEE>(x, Vector2d(0, 0), Vector3d::Ones());

  VariablePtr u = Space(1).createVariable("u");
  u << 0;
  MatrixXd A(2, 1);
  A << 2, 1;
  VectorXd b(2);
  b << -1, 0;
  auto ref = std::make_shared<function::BasicLinearFunction>(A, u, b);

  {
    task_dynamics::ReferenceVelocity refv(ref);
    TaskDynamicsPtr tdi = refv.impl(f, constraint::Type::EQUAL, Vector2d::Zero());
    auto Value = task_dynamics::abstract::TaskDynamicsImpl::Output::Value;
    auto gl = utils::generateUpdateGraph(tdi, Value);

    gl->execute();

    FAST_CHECK_EQ(refv.order(), task_dynamics::Order::One);
    FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::One);
    FAST_CHECK_UNARY(tdi->value().isApprox(ref->value()));
  }
  {
    task_dynamics::ReferenceAcceleration refa(ref);
    TaskDynamicsPtr tdi = refa.impl(f, constraint::Type::EQUAL, Vector2d::Zero());
    auto Value = task_dynamics::abstract::TaskDynamicsImpl::Output::Value;
    auto gl = utils::generateUpdateGraph(tdi, Value);

    gl->execute();

    FAST_CHECK_EQ(refa.order(), task_dynamics::Order::Two);
    FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::Two);
    FAST_CHECK_UNARY(tdi->value().isApprox(ref->value()));
  }
}

TEST_CASE("Test Clamped")
{
  VariablePtr x = Space(3).createVariable("x");
  x << 2, 3, -1;
  auto f = std::make_shared<function::IdentityFunction>(x);

  double kp = 1;
  task_dynamics::Clamped<task_dynamics::P> c(1.0, kp);
  VectorXd rhs = Vector3d::Ones();

  TaskDynamicsPtr tdi = c.impl(f, constraint::Type::EQUAL, rhs);
  auto Value = task_dynamics::abstract::TaskDynamicsImpl::Output::Value;
  auto gl = utils::generateUpdateGraph(tdi, Value);
  gl->execute();

  FAST_CHECK_EQ(c.order(), task_dynamics::Order::One);
  FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::One);
  FAST_CHECK_UNARY(tdi->value().isApprox(Vector3d(-.5, -1, 1))); // e'= -e = (-1,-2,2)

  x << 0.5, 1, 1.5;
  gl->execute();
  FAST_CHECK_UNARY(tdi->value().isApprox(Vector3d(.5, 0, -.5))); // e'= -e = (0.5,0,-0.5)
}

TEST_CASE("Test Clamped<FeedForward<PD>>")
{
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  x << 1, 2, 3;
  dx << 1, 1, 1;
  auto f = std::make_shared<SphereFunction>(x, Vector3d(1, 0, -3), 2);

  auto provider = std::make_shared<FFProvider>(Eigen::VectorXd::Constant(1, 10.0));

  double kp = 2;
  double kv = 3;
  double max = 25;
  task_dynamics::Clamped<task_dynamics::FeedForwardPD> td(max, provider, &FFProvider::value, FFProvider::Output::Value,
                                                          kp, kv);
  Eigen::VectorXd rhs = Eigen::VectorXd::Constant(1, -1);
  TaskDynamicsPtr tdi = td.impl(f, constraint::Type::EQUAL, rhs);

  auto Value = task_dynamics::abstract::TaskDynamicsImpl::Output::Value;
  auto gl = utils::generateUpdateGraph(tdi, Value);
  gl->execute();

  FAST_CHECK_UNARY(provider->updated_);
  FAST_CHECK_EQ(td.order(), task_dynamics::Order::Two);
  FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::Two);
  FAST_CHECK_EQ(tdi->value()[0], std::min(max, std::max(-max, -kp * (36 - rhs[0]) - kv * 16 + 10.0)));
  FAST_CHECK_UNARY(tdi->value()[0] == -max); // clamped error
  FAST_CHECK_UNARY(tdi->checkType<task_dynamics::PD::Impl>());
  FAST_CHECK_UNARY(tdi->checkType<decltype(td)::Impl>());
  FAST_REQUIRE_UNARY(tdi->checkType<task_dynamics::FeedForwardPD::Impl>()); // REQUIRE because we do a static_cast after
  FAST_CHECK_UNARY_FALSE(tdi->checkType<task_dynamics::Constant::Impl>());

  kp = 0.2;
  kv = 0.3;
  static_cast<task_dynamics::FeedForwardPD::Impl *>(tdi.get())->gains(kp, kv);
  gl->execute();
  FAST_CHECK_EQ(tdi->value()[0], APPROX_I386(std::min(max, std::max(-max, -kp * (36 - rhs[0]) - kv * 16 + 10.0))));
  FAST_CHECK_UNARY_FALSE(tdi->value()[0] == -max); // not clamped anymore
}

TEST_CASE("OneStepToZero")
{
  double dt = 0.1;
  VariablePtr x = Space(3).createVariable("x");
  VariablePtr dx = dot(x);
  x << 1, 2, 3;
  dx << 1, 1, 1;

  MatrixXd A = MatrixXd::Random(2, 3);
  VectorXd b = VectorXd::Random(2);
  auto f = std::make_shared<function::BasicLinearFunction>(A, x, b);

  {
    task_dynamics::OneStepToZero td(task_dynamics::Order::One, dt);
    TaskDynamicsPtr tdi = td.impl(f, constraint::Type::EQUAL, Vector2d::Zero());
    auto Value = task_dynamics::abstract::TaskDynamicsImpl::Output::Value;
    auto gl = utils::generateUpdateGraph(tdi, Value);

    gl->execute();

    Vector2d v = tdi->value();
    Vector3d d1x = A.colPivHouseholderQr().solve(v);
    Vector3d x1 = x->value() + dt * d1x;

    FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::One);
    FAST_CHECK_UNARY((A * x1 + b).isZero(1e-8));
  }
  {
    task_dynamics::OneStepToZero td(task_dynamics::Order::Two, dt);
    TaskDynamicsPtr tdi = td.impl(f, constraint::Type::EQUAL, Vector2d::Zero());
    auto Value = task_dynamics::abstract::TaskDynamicsImpl::Output::Value;
    auto gl = utils::generateUpdateGraph(tdi, Value);

    gl->execute();

    Vector2d v = tdi->value();
    Vector3d d2x = A.colPivHouseholderQr().solve(v);
    Vector3d x1 = x->value() + dt * dx->value() + dt * dt / 2 * d2x;

    FAST_CHECK_EQ(tdi->order(), task_dynamics::Order::Two);
    FAST_CHECK_UNARY((A * x1 + b).isZero(1e-8));
  }
}


#pragma once
#include <Eigen/Dense>
#include <string>
#include <vector>

int atom(int ao_index, int orbitals_per_atom);
std::string orb(int ao_index, std::vector<std::string> orbital_types, int orbitals_per_atom);
int ao_index(int atom_p, std::string orb_p, std::vector<std::string> orbital_types, int orbitals_per_atom);
double chi_on_atom(std::string orb1, std::string orb2, std::string orb3, std::map<std::string, double> model_parameters);
Eigen::MatrixXd calculate_fock_matrix_fast(Eigen::MatrixXd hamiltonian_matrix, Eigen::MatrixXd interaction_matrix, Eigen::MatrixXd density_matrix, std::map<std::string, double> model_parameters, int orbitals_per_atom);
Eigen::MatrixXd test_fock_matrix_fast(Eigen::MatrixXd hamiltonian_matrix, Eigen::MatrixXd interaction_matrix, Eigen::MatrixXd density_matrix, std::map<std::string, double> model_parameters, int orbitals_per_atom);


#include "QPCommon.h"
#include "controlUtil.h"
#include <Eigen/StdVector>

template<int M, int N>
void matlabToCArrayOfArrays(const mxArray *source, const int idx, const char *fieldname, double *destination)  {
  // Matlab arrays come in as column-major data. To represent a matrix in C++, as we do in our LCM messages, we need an array of arrays. But that convention results in a row-major storage, so we have to be careful about how we copy data in.
  const mxArray *field = myGetField(source, idx, fieldname);
  sizecheck(field, M, N);
  Map<Matrix<double, M, N>>A(mxGetPrSafe(field));
  // C is row-major, matlab is column-major
  Matrix<double, N, M> A_t = A.transpose();
  memcpy(destination, A_t.data(), sizeof(double)*M*N);
  return;
}

std::shared_ptr<drake::lcmt_qp_controller_input> encodeQPInputLCM(const mxArray *qp_input) {
  // Take a matlab data structure corresponding to a QPInputConstantHeight object and parse it down to its representation as an equivalent LCM message.
  std::shared_ptr<drake::lcmt_qp_controller_input> msg (new drake::lcmt_qp_controller_input());

  msg->be_silent = mxGetScalar(myGetProperty(qp_input, "be_silent")) > 0.5;

  msg->timestamp = (int64_t) (mxGetScalar(myGetProperty(qp_input, "timestamp")) * 1000000);

  const mxArray* zmp_data = myGetProperty(qp_input, "zmp_data");

  matlabToCArrayOfArrays<4, 4>(zmp_data, 0, "A", &msg->zmp_data.A[0][0]);
  matlabToCArrayOfArrays<4, 2>(zmp_data, 0, "B", &msg->zmp_data.B[0][0]);
  matlabToCArrayOfArrays<2, 4>(zmp_data, 0, "C", &msg->zmp_data.C[0][0]);
  matlabToCArrayOfArrays<2, 2>(zmp_data, 0, "D", &msg->zmp_data.D[0][0]);
  matlabToCArrayOfArrays<4, 1>(zmp_data, 0, "x0", &msg->zmp_data.x0[0][0]);
  matlabToCArrayOfArrays<2, 1>(zmp_data, 0, "y0", &msg->zmp_data.y0[0][0]);
  matlabToCArrayOfArrays<2, 1>(zmp_data, 0, "u0", &msg->zmp_data.u0[0][0]);
  matlabToCArrayOfArrays<2, 2>(zmp_data, 0, "R", &msg->zmp_data.R[0][0]);
  matlabToCArrayOfArrays<2, 2>(zmp_data, 0, "Qy", &msg->zmp_data.Qy[0][0]);
  matlabToCArrayOfArrays<4, 4>(zmp_data, 0, "S", &msg->zmp_data.S[0][0]);
  matlabToCArrayOfArrays<4, 1>(zmp_data, 0, "s1", &msg->zmp_data.s1[0][0]);
  matlabToCArrayOfArrays<4, 1>(zmp_data, 0, "s1dot", &msg->zmp_data.s1dot[0][0]);
  msg->zmp_data.s2 = mxGetScalar(myGetField(zmp_data, "s2"));
  msg->zmp_data.s2dot = mxGetScalar(myGetField(zmp_data, "s2dot"));
  msg->zmp_data.timestamp = msg->timestamp;

  const mxArray* support_data = myGetProperty(qp_input, "support_data");
  int nsupp = mxGetN(support_data);
  msg->num_support_data = (int32_t) nsupp;
  double double_logic_map[4][1];
  msg->support_data.resize(nsupp);
  if (nsupp > 0) {
    if (mxGetM(support_data) != 1) {
      mexErrMsgTxt("support data should be a struct array with M=1");
    }
    for (int i=0; i < nsupp; i++) {
      msg->support_data[i].timestamp = msg->timestamp;
      msg->support_data[i].body_id = (int32_t) mxGetScalar(myGetField(support_data, i, "body_id"));

      const mxArray *contact_pts = myGetField(support_data, i, "contact_pts");
      if (!contact_pts) mexErrMsgTxt("couldn't get points");
      Map<MatrixXd>contact_pts_mat(mxGetPrSafe(contact_pts), mxGetM(contact_pts), mxGetN(contact_pts));
      msg->support_data[i].num_contact_pts = (int32_t) mxGetN(contact_pts);
      msg->support_data[i].contact_pts.resize(3);
      for (int j=0; j < 3; j++) {
        msg->support_data[i].contact_pts[j].resize(msg->support_data[i].num_contact_pts);
        for (int k=0; k < msg->support_data[i].num_contact_pts; k++) {
          msg->support_data[i].contact_pts[j][k] = contact_pts_mat(j, k);
        }
      }

      matlabToCArrayOfArrays<4, 1>(support_data, i, "support_logic_map", &double_logic_map[0][0]);
      for (int j=0; j < 4; j++) {
        msg->support_data[i].support_logic_map[j] = (double_logic_map[j][0] != 0);
      }
      msg->support_data[i].mu = mxGetScalar(myGetField(support_data, i, "mu"));
      msg->support_data[i].contact_surfaces = (int32_t) mxGetScalar(myGetField(support_data, i, "contact_surfaces"));
    }
  }

  const mxArray* body_motion_data = myGetProperty(qp_input, "body_motion_data");
  const int nbod = mxGetN(body_motion_data);
  msg->num_tracked_bodies = nbod;
  msg->body_motion_data.resize(nbod);
  if (nbod > 0) {
    if (mxGetM(body_motion_data) != 1) {
      mexErrMsgTxt("body motion data should be a 1xN struct array");
    }
    for (int i=0; i < nbod; i++) {
      msg->body_motion_data[i].timestamp = msg->timestamp;
      msg->body_motion_data[i].body_id = (int32_t) mxGetScalar(myGetField(body_motion_data, i, "body_id"));
      memcpy(msg->body_motion_data[i].ts, mxGetPrSafe(myGetField(body_motion_data, i, "ts")), 2*sizeof(double));
      const mxArray* coefs = myGetField(body_motion_data, i, "coefs");
      if (mxGetNumberOfDimensions(coefs) != 3) mexErrMsgTxt("coefs should be a dimension-3 array");
      const mwSize* dim = mxGetDimensions(coefs);
      if (dim[0] != 6 || dim[1] != 1 || dim[2] != 4) mexErrMsgTxt("coefs should be size 6x4");
      matlabToCArrayOfArrays<6, 4>(body_motion_data, i, "coefs", &msg->body_motion_data[i].coefs[0][0]);
    }
  }

  const mxArray* body_wrench_data = myGetProperty(qp_input, "body_wrench_data");
  const int num_external_wrenches = mxGetN(body_wrench_data);
  msg->num_external_wrenches = num_external_wrenches;
  msg->body_wrench_data.resize(num_external_wrenches);
  const int wrench_size = 6;
  if (num_external_wrenches > 0) {
    if (mxGetM(body_wrench_data) != 1) {
      mexErrMsgTxt("body wrench data should be a 1xN struct array");
    }
    for (int i = 0; i < num_external_wrenches; i++) {
      msg->body_wrench_data[i].timestamp = msg->timestamp;
      msg->body_wrench_data[i].body_id = (int32_t) mxGetScalar(myGetField(body_wrench_data, i, "body_id"));
      const mxArray* wrench = myGetField(body_wrench_data, i, "wrench");
      sizecheck(wrench, wrench_size, 1);
      memcpy(msg->body_wrench_data[i].wrench, mxGetPr(wrench), wrench_size * sizeof(double));
    }
  }

  const mxArray* whole_body_data = myGetProperty(qp_input, "whole_body_data");
  if (mxGetN(whole_body_data) != 1 || mxGetM(whole_body_data) != 1) mexErrMsgTxt("whole_body_data should be a 1x1 struct");
  const mxArray* q_des = myGetField(whole_body_data, "q_des");
  if (mxGetN(q_des) != 1) mexErrMsgTxt("q_des should be a column vector");
  const int npos = mxGetM(q_des);
  msg->whole_body_data.timestamp = msg->timestamp;
  msg->whole_body_data.num_positions = npos;
  Map<VectorXd>q_des_vec(mxGetPrSafe(q_des), npos);
  msg->whole_body_data.q_des.resize(npos);

  for (int i=0; i < npos; i++) {
    msg->whole_body_data.q_des[i] = q_des_vec(i);
  }

  const mxArray* condof = myGetField(whole_body_data, "constrained_dofs");
  const int ncons = mxGetNumberOfElements(condof);
  msg->whole_body_data.num_constrained_dofs = ncons;
  msg->whole_body_data.constrained_dofs.resize(ncons);
  if (ncons > 0) {
    if (mxGetN(condof) != 1) mexErrMsgTxt("constrained dofs should be a column vector");
    Map<VectorXd>condof_vec(mxGetPrSafe(condof), ncons);

    for (int i=0; i < ncons; i++) {
      msg->whole_body_data.constrained_dofs[i] = static_cast<int32_t>(condof_vec(i));
    }
  }

  const mxArray* joint_override = myGetProperty(qp_input, "joint_pd_override");
  int num_joint_pd_overrides = mxGetNumberOfElements(joint_override);
  msg->num_joint_pd_overrides = num_joint_pd_overrides;
  msg->joint_pd_override.resize(num_joint_pd_overrides);
  for (int i=0; i < num_joint_pd_overrides; i++) {
    msg->joint_pd_override[i].position_ind = (int32_t) mxGetScalar(myGetField(joint_override, i, "position_ind"));
    msg->joint_pd_override[i].qi_des = mxGetScalar(myGetField(joint_override, i, "qi_des"));
    msg->joint_pd_override[i].qdi_des = mxGetScalar(myGetField(joint_override, i, "qdi_des"));
    msg->joint_pd_override[i].kp = mxGetScalar(myGetField(joint_override, i, "kp"));
    msg->joint_pd_override[i].kd = mxGetScalar(myGetField(joint_override, i, "kd"));
    msg->joint_pd_override[i].weight = mxGetScalar(myGetField(joint_override, i, "weight"));
  }

  msg->param_set_name = mxArrayToString(myGetProperty(qp_input, "param_set_name"));
  return msg;
}

PIDOutput wholeBodyPID(NewQPControllerData *pdata, double t, const Ref<const VectorXd> &q, const Ref<const VectorXd> &qd, const Ref<const VectorXd> &q_des, WholeBodyParams *params) {
  // Run a PID controller on the whole-body state to produce desired accelerations and reference posture
  PIDOutput out;
  double dt = 0;
  int nq = pdata->r->num_positions;
  assert(q.size() == nq);
  assert(qd.size() == pdata->r->num_velocities);
  assert(q_des.size() == params->integrator.gains.size());
  if (nq != pdata->r->num_velocities) {
    mexErrMsgTxt("this function will need to be rewritten when num_pos != num_vel");
  }
  if (pdata->state.t_prev != 0) {
    dt = t - pdata->state.t_prev;
  }
  pdata->state.q_integrator_state = (1-params->integrator.eta) * pdata->state.q_integrator_state + params->integrator.gains.cwiseProduct(q_des - q) * dt;
  pdata->state.q_integrator_state = pdata->state.q_integrator_state.array().max(-params->integrator.clamps.array());
  pdata->state.q_integrator_state = pdata->state.q_integrator_state.array().min(params->integrator.clamps.array());
  out.q_ref = q_des + pdata->state.q_integrator_state;
  out.q_ref = out.q_ref.array().max((pdata->r->joint_limit_min - params->integrator.clamps).array());
  out.q_ref = out.q_ref.array().min((pdata->r->joint_limit_max + params->integrator.clamps).array());

  pdata->state.q_integrator_state = pdata->state.q_integrator_state.array().max(-params->integrator.clamps.array());
  pdata->state.q_integrator_state = pdata->state.q_integrator_state.array().min(params->integrator.clamps.array());

  VectorXd err_q;
  err_q.resize(nq);
  err_q.head(3) = q_des.head(3) - q.head(3);
  for (int j = 3; j < nq; j++) {
    err_q(j) = angleDiff(q(j), q_des(j));
  }
  out.qddot_des = params->Kp.cwiseProduct(err_q) - params->Kd.cwiseProduct(qd);
  out.qddot_des = out.qddot_des.array().max(params->qdd_bounds.min.array());
  out.qddot_des = out.qddot_des.array().min(params->qdd_bounds.max.array());
  return out;
}

VectorXd velocityReference(NewQPControllerData *pdata, double t, const Ref<VectorXd> &q, const Ref<VectorXd> &qd, const Ref<VectorXd> &qdd, bool foot_contact[2], VRefIntegratorParams *params, RobotPropertyCache *rpc) {
  // Integrate expected accelerations to determine a target feed-forward velocity, which we can pass in to Atlas
  int i;
  assert(qdd.size() == pdata->r->num_velocities);

  double dt = 0;
  if (pdata->state.t_prev != 0) {
    dt = t - pdata->state.t_prev;
  }

  pdata->state.vref_integrator_state = (1-params->eta)*pdata->state.vref_integrator_state + params->eta*qd + qdd*dt;

  if (params->zero_ankles_on_contact && foot_contact[0] == 1) {
    for (i=0; i < rpc->position_indices.l_leg_ak.size(); i++) {
      pdata->state.vref_integrator_state(rpc->position_indices.l_leg_ak(i)) = 0;
    }
  }
  if (params->zero_ankles_on_contact && foot_contact[1] == 1) {
    for (i=0; i < rpc->position_indices.r_leg_ak.size(); i++) {
      pdata->state.vref_integrator_state(rpc->position_indices.r_leg_ak(i)) = 0;
    }
  }
  if (pdata->state.foot_contact_prev[0] != foot_contact[0]) {
    // contact state changed, reset integrated velocities
    for (i=0; i < rpc->position_indices.l_leg.size(); i++) {
      pdata->state.vref_integrator_state(rpc->position_indices.l_leg(i)) = qd(rpc->position_indices.l_leg(i));
    }
  }
  if (pdata->state.foot_contact_prev[1] != foot_contact[1]) {
    // contact state changed, reset integrated velocities
    for (i=0; i < rpc->position_indices.r_leg.size(); i++) {
      pdata->state.vref_integrator_state(rpc->position_indices.r_leg(i)) = qd(rpc->position_indices.r_leg(i));
    }
  }

  pdata->state.foot_contact_prev[0] = foot_contact[0];
  pdata->state.foot_contact_prev[1] = foot_contact[1];

  VectorXd qd_err = pdata->state.vref_integrator_state - qd;

  // do not velocity control ankles when in contact
  if (params->zero_ankles_on_contact && foot_contact[0] == 1) {
    for (i=0; i < rpc->position_indices.l_leg_ak.size(); i++) {
      qd_err(rpc->position_indices.l_leg_ak(i)) = 0;
    }
  }
  if (params->zero_ankles_on_contact && foot_contact[1] == 1) {
    for (i=0; i < rpc->position_indices.r_leg_ak.size(); i++) {
      qd_err(rpc->position_indices.r_leg_ak(i)) = 0;
    }
  }

  VectorXd qd_ref = qd_err.array().max(-params->delta_max);
  qd_ref = qd_ref.array().min(params->delta_max);
  return qd_ref;
}

std::vector<SupportStateElement> loadAvailableSupports(std::shared_ptr<drake::lcmt_qp_controller_input> qp_input) {
  // Parse a qp_input LCM message to extract its available supports as a vector of SupportStateElements
  std::vector<SupportStateElement> available_supports;
  available_supports.resize(qp_input->num_support_data);
  for (int i=0; i < qp_input->num_support_data; i++) {
    available_supports[i].body_idx = qp_input->support_data[i].body_id - 1;
    available_supports[i].contact_surface = qp_input->support_data[i].contact_surfaces;
    for (int j=0; j < 4; j++) {
      available_supports[i].support_logic_map[j] = qp_input->support_data[i].support_logic_map[j];
    }
    available_supports[i].contact_pts.resize(qp_input->support_data[i].num_contact_pts);
    for (int j=0; j < qp_input->support_data[i].num_contact_pts; j++) {
      for (int k = 0; k < 3; k++) {
        available_supports[i].contact_pts[j][k] = qp_input->support_data[i].contact_pts[k][j];
      }
      available_supports[i].contact_pts[j][3] = 1;
    }
  }
  return available_supports;
}

void applyJointPDOverride(const std::vector<drake::lcmt_joint_pd_override> &joint_pd_override, const DrakeRobotState &robot_state, PIDOutput &pid_out, VectorXd &w_qdd) {
  for (std::vector<drake::lcmt_joint_pd_override>::const_iterator it = joint_pd_override.begin(); it != joint_pd_override.end(); ++it) {
    int ind = it->position_ind - 1;
    double err_q = it->qi_des - robot_state.q(ind);
    double err_qd = it->qdi_des - robot_state.qd(ind);
    pid_out.qddot_des(ind) = it->kp * err_q + it->kd * err_qd;
    w_qdd(ind) = it->weight;
  }
}

int setupAndSolveQP(NewQPControllerData *pdata, std::shared_ptr<drake::lcmt_qp_controller_input> qp_input, DrakeRobotState &robot_state, const Ref<Matrix<bool, Dynamic, 1>> &b_contact_force, QPControllerOutput *qp_output, std::shared_ptr<QPControllerDebugData> debug) {
  // The primary solve loop for our controller. This constructs and solves a Quadratic Program and produces the instantaneous desired torques, along with reference positions, velocities, and accelerations. It mirrors the Matlab implementation in atlasControllers.InstantaneousQPController.setupAndSolveQP(), and more documentation can be found there.
  // Note: argument `debug` MAY be set to NULL, which signals that no debug information is requested.

  // look up the param set by name
  AtlasParams *params;
  std::map<string,AtlasParams>::iterator it;
  it = pdata->param_sets.find(qp_input->param_set_name);
  if (it == pdata->param_sets.end()) {
    mexWarnMsgTxt("Got a param set I don't recognize! Using standing params instead");
    it = pdata->param_sets.find("standing");
    if (it == pdata->param_sets.end()) {
      mexErrMsgTxt("Could not fall back to standing parameters either. I have to give up here.");
    }
  }
  // cout << "using params set: " + it->first + ", ";
  params = &(it->second);
  // mexPrintf("Kp_accel: %f, ", params->Kp_accel);

  int nu = pdata->B.cols();
  int nq = pdata->r->num_positions;

  // zmp_data
  Map<Matrix<double, 4, 4, RowMajor>> A_ls(&qp_input->zmp_data.A[0][0]);
  Map<Matrix<double, 4, 2, RowMajor>> B_ls(&qp_input->zmp_data.B[0][0]);
  Map<Matrix<double, 2, 4, RowMajor>> C_ls(&qp_input->zmp_data.C[0][0]);
  Map<Matrix<double, 2, 2, RowMajor>> D_ls(&qp_input->zmp_data.D[0][0]);
  Map<Matrix<double, 4, 1>> x0(&qp_input->zmp_data.x0[0][0]);
  Map<Matrix<double, 2, 1>> y0(&qp_input->zmp_data.y0[0][0]);
  Map<Matrix<double, 2, 1>> u0(&qp_input->zmp_data.u0[0][0]);
  Map<Matrix<double, 2, 2, RowMajor>> R_ls(&qp_input->zmp_data.R[0][0]);
  Map<Matrix<double, 2, 2, RowMajor>> Qy(&qp_input->zmp_data.Qy[0][0]);
  Map<Matrix<double, 4, 4, RowMajor>> S(&qp_input->zmp_data.S[0][0]);
  Map<Matrix<double, 4, 1>> s1(&qp_input->zmp_data.s1[0][0]);
  Map<Matrix<double, 4, 1>> s1dot(&qp_input->zmp_data.s1dot[0][0]);

  // // whole_body_data
  if (qp_input->whole_body_data.num_positions != nq) mexErrMsgTxt("number of positions doesn't match num_dof for this robot");
  Map<VectorXd> q_des(qp_input->whole_body_data.q_des.data(), nq);
  if (qp_input->whole_body_data.constrained_dofs.size() != qp_input->whole_body_data.num_constrained_dofs) {
    mexErrMsgTxt("size of constrained dofs does not match num_constrained_dofs");
  }
  Map<VectorXi> condof(qp_input->whole_body_data.constrained_dofs.data(), qp_input->whole_body_data.num_constrained_dofs);
  PIDOutput pid_out = wholeBodyPID(pdata, robot_state.t, robot_state.q, robot_state.qd, q_des, &params->whole_body);
  VectorXd w_qdd = params->whole_body.w_qdd;
  applyJointPDOverride(qp_input->joint_pd_override, robot_state, pid_out, w_qdd);

  qp_output->q_ref = pid_out.q_ref;

  // mu
  // NOTE: we're using the same mu for all supports
  double mu;
  if (qp_input->num_support_data == 0) {
    mu = 1.0;
  } else {
    mu = qp_input->support_data[0].mu;
    for (int i=1; i < qp_input->num_support_data; i++) {
      if (qp_input->support_data[i].mu != mu) {
        mexWarnMsgTxt("Currently, we assume that all supports have the same value of mu");
      }
    }
  }

  const int dim = 3, // 3D
  nd = 2*m_surface_tangents; // for friction cone approx, hard coded for now

  assert(nu+6 == nq);

  std::vector<DesiredBodyAcceleration> desired_body_accelerations;
  desired_body_accelerations.resize(qp_input->num_tracked_bodies);
  Vector6d body_pose_des, body_v_des, body_vdot_des;
  Vector6d body_vdot;

  for (int i=0; i < qp_input->num_tracked_bodies; i++) {
    int body_id0 = qp_input->body_motion_data[i].body_id - 1;
    double weight = params->body_motion[body_id0].weight;
    desired_body_accelerations[i].body_id0 = body_id0;
    Map<Matrix<double, 6, 4,RowMajor>>coefs_rowmaj(&qp_input->body_motion_data[i].coefs[0][0]);
    Matrix<double, 6, 4> coefs = coefs_rowmaj;
    double t_spline = std::max(qp_input->body_motion_data[i].ts[0], std::min(qp_input->body_motion_data[i].ts[1], robot_state.t));
    evaluateCubicSplineSegment(t_spline - qp_input->body_motion_data[i].ts[0], coefs, body_pose_des, body_v_des, body_vdot_des);
    desired_body_accelerations[i].body_vdot = bodyMotionPD(pdata->r, robot_state, body_id0, body_pose_des, body_v_des, body_vdot_des, params->body_motion[body_id0].Kp, params->body_motion[body_id0].Kd);
    desired_body_accelerations[i].weight = weight;
    desired_body_accelerations[i].accel_bounds = params->body_motion[body_id0].accel_bounds;
    // mexPrintf("body: %d, vdot: %f %f %f %f %f %f weight: %f\n", body_id0,
    //           desired_body_accelerations[i].body_vdot(0),
    //           desired_body_accelerations[i].body_vdot(1),
    //           desired_body_accelerations[i].body_vdot(2),
    //           desired_body_accelerations[i].body_vdot(3),
    //           desired_body_accelerations[i].body_vdot(4),
    //           desired_body_accelerations[i].body_vdot(5),
    //           weight);
      // mexPrintf("tracking body: %d, coefs[:,0]: %f %f %f %f %f %f coefs(", body_id0,
  }

  int n_body_accel_eq_constraints = 0;
  for (int i=0; i < desired_body_accelerations.size(); i++) {
    if (desired_body_accelerations[i].weight < 0)
      n_body_accel_eq_constraints++;
  }

  MatrixXd R_DQyD_ls = R_ls + D_ls.transpose()*Qy*D_ls;

  pdata->r->doKinematics(robot_state.q,false,robot_state.qd);

  //---------------------------------------------------------------------

  std::vector<SupportStateElement> available_supports = loadAvailableSupports(qp_input);
  std::vector<SupportStateElement> active_supports = getActiveSupports(pdata->r, pdata->map_ptr, robot_state.q, robot_state.qd, available_supports, b_contact_force, params->contact_threshold, pdata->default_terrain_height);

  int num_active_contact_pts=0;
  for (std::vector<SupportStateElement>::iterator iter = active_supports.begin(); iter!=active_supports.end(); iter++) {
    num_active_contact_pts += iter->contact_pts.size();
  }

  // handle external wrenches to compensate for
  if (qp_input->body_wrench_data.size() > 0) {
    const int wrench_size = 6;
    MatrixXd f_ext(wrench_size, pdata->r->NB); // f_ext is in terms of featherstone bodies...
    f_ext.setZero();
    for (auto it = qp_input->body_wrench_data.begin(); it != qp_input->body_wrench_data.end(); ++it) {
      const drake::lcmt_body_wrench_data& body_wrench_data = *it;
      Map<const Matrix<double, wrench_size, 1> > wrench_in_body_frame(body_wrench_data.wrench);

      // HandC expects external wrenches to be expressed in Featherstone 'joint' frame, expects the column number to correspond to the featherstone body index.
      int body_id = body_wrench_data.body_id - 1;
      auto wrench_in_joint_frame = transformSpatialForce(Isometry3d(pdata->r->bodies[body_id]->T_body_to_joint), wrench_in_body_frame);
      int featherstone_id = pdata->r->bodies_vector_index_to_featherstone_body_index[body_id];
      f_ext.col(featherstone_id) = wrench_in_joint_frame;
    }
    pdata->r->HandC(robot_state.q,robot_state.qd,&f_ext,pdata->H,pdata->C,(MatrixXd*)nullptr,(MatrixXd*)nullptr,(MatrixXd*)nullptr);
  }
  else {
    pdata->r->HandC(robot_state.q,robot_state.qd,(MatrixXd*)nullptr,pdata->H,pdata->C,(MatrixXd*)nullptr,(MatrixXd*)nullptr,(MatrixXd*)nullptr);
  }

  pdata->H_float = pdata->H.topRows(6);
  pdata->H_act = pdata->H.bottomRows(nu);
  pdata->C_float = pdata->C.head(6);
  pdata->C_act = pdata->C.tail(nu);

  bool include_angular_momentum = (params->W_kdot.array().maxCoeff() > 1e-10);

  if (include_angular_momentum) {
    pdata->r->getCMM(robot_state.q,robot_state.qd,pdata->Ag,pdata->Agdot);
    pdata->Ak = pdata->Ag.topRows(3);
    pdata->Akdot = pdata->Agdot.topRows(3);
  }
  Vector3d xcom;
  // consider making all J's into row-major

  pdata->r->getCOM(xcom);
  pdata->r->getCOMJac(pdata->J);
  pdata->r->getCOMJacDot(pdata->Jdot);
  pdata->J_xy = pdata->J.topRows(2);
  pdata->Jdot_xy = pdata->Jdot.topRows(2);

  MatrixXd Jcom,Jcomdot;

  if (x0.size()==6) {
    Jcom = pdata->J;
    Jcomdot = pdata->Jdot;
  }
  else {
    Jcom = pdata->J_xy;
    Jcomdot = pdata->Jdot_xy;
  }

  MatrixXd B,JB,Jp,Jpdot,normals;
  int nc = contactConstraintsBV(pdata->r,num_active_contact_pts,mu,active_supports,pdata->map_ptr,B,JB,Jp,Jpdot,normals,pdata->default_terrain_height);
  int neps = nc*dim;

  VectorXd x_bar,xlimp;
  MatrixXd D_float(6,JB.cols()), D_act(nu,JB.cols());
  if (nc>0) {
    if (x0.size()==6) {
      // x,y,z com
      xlimp.resize(6);
      xlimp.topRows(3) = xcom;
      xlimp.bottomRows(3) = Jcom*robot_state.qd;
    }
    else {
      xlimp.resize(4);
      xlimp.topRows(2) = xcom.topRows(2);
      xlimp.bottomRows(2) = Jcom*robot_state.qd;
    }
    x_bar = xlimp-x0;

    D_float = JB.topRows(6);
    D_act = JB.bottomRows(nu);
  }

  int nf = nc*nd; // number of contact force variables
  int nparams = nq+nf+neps;

  Vector3d kdot_des;
  if (include_angular_momentum) {
    VectorXd k = pdata->Ak*robot_state.qd;
    kdot_des = -params->Kp_ang*k; // TODO: parameterize
  }

  //----------------------------------------------------------------------
  // QP cost function ----------------------------------------------------
  //
  //  min: ybar*Qy*ybar + ubar*R*ubar + (2*S*xbar + s1)*(A*x + B*u) +
  //    w_qdd*quad(qddot_ref - qdd) + w_eps*quad(epsilon) +
  //    w_grf*quad(beta) + quad(kdot_des - (A*qdd + Adot*qd))
  VectorXd f(nparams);
  {
    if (nc > 0) {
      // NOTE: moved Hqp calcs below, because I compute the inverse directly for FastQP (and sparse Hqp for gurobi)
      VectorXd tmp = C_ls*xlimp;
      VectorXd tmp1 = Jcomdot*robot_state.qd;
      MatrixXd tmp2 = R_DQyD_ls*Jcom;

      pdata->fqp = tmp.transpose()*Qy*D_ls*Jcom;
      // mexPrintf("fqp head: %f %f %f\n", pdata->fqp(0), pdata->fqp(1), pdata->fqp(2));
      pdata->fqp += tmp1.transpose()*tmp2;
      pdata->fqp += (S*x_bar + 0.5*s1).transpose()*B_ls*Jcom;
      pdata->fqp -= u0.transpose()*tmp2;
      pdata->fqp -= y0.transpose()*Qy*D_ls*Jcom;
      pdata->fqp -= (w_qdd.array()*pid_out.qddot_des.array()).matrix().transpose();
      if (include_angular_momentum) {
        pdata->fqp += robot_state.qd.transpose()*pdata->Akdot.transpose()*params->W_kdot*pdata->Ak;
        pdata->fqp -= kdot_des.transpose()*params->W_kdot*pdata->Ak;
      }
      f.head(nq) = pdata->fqp.transpose();
     } else {
      f.head(nq) = -pid_out.qddot_des;
    }
  }
  f.tail(nf+neps) = VectorXd::Zero(nf+neps);

  int neq = 6+neps+6*n_body_accel_eq_constraints+qp_input->whole_body_data.num_constrained_dofs;
  MatrixXd Aeq = MatrixXd::Zero(neq,nparams);
  VectorXd beq = VectorXd::Zero(neq);

  // constrained floating base dynamics
  //  H_float*qdd - J_float'*lambda - Dbar_float*beta = -C_float
  Aeq.topLeftCorner(6,nq) = pdata->H_float;
  beq.topRows(6) = -pdata->C_float;

  if (nc>0) {
    Aeq.block(0,nq,6,nc*nd) = -D_float;
  }

  if (nc > 0) {
    // relative acceleration constraint
    Aeq.block(6,0,neps,nq) = Jp;
    Aeq.block(6,nq,neps,nf) = MatrixXd::Zero(neps,nf);  // note: obvious sparsity here
    Aeq.block(6,nq+nf,neps,neps) = MatrixXd::Identity(neps,neps);             // note: obvious sparsity here
    beq.segment(6,neps) = (-Jpdot -params->Kp_accel*Jp)*robot_state.qd;
  }

  // add in body spatial equality constraints
  // VectorXd body_vdot;
  MatrixXd orig = MatrixXd::Zero(4,1);
  orig(3,0) = 1;
  int equality_ind = 6+neps;
  MatrixXd Jb(6,nq);
  MatrixXd Jbdot(6,nq);
  for (int i=0; i<desired_body_accelerations.size(); i++) {
    if (desired_body_accelerations[i].weight < 0) { // negative implies constraint
      if (!inSupport(active_supports,desired_body_accelerations[i].body_id0)) {
        pdata->r->forwardJac(desired_body_accelerations[i].body_id0,orig,1,Jb);
        pdata->r->forwardJacDot(desired_body_accelerations[i].body_id0,orig,1,Jbdot);

        for (int j=0; j<6; j++) {
          if (!std::isnan(desired_body_accelerations[i].body_vdot(j))) {
            Aeq.block(equality_ind,0,1,nq) = Jb.row(j);
            beq[equality_ind++] = -Jbdot.row(j)*robot_state.qd + desired_body_accelerations[i].body_vdot(j);
          }
        }
      }
    }
  }

  if (qp_input->whole_body_data.num_constrained_dofs>0) {
    // add joint acceleration constraints
    for (int i=0; i<qp_input->whole_body_data.num_constrained_dofs; i++) {
      Aeq(equality_ind,(int)condof[i]-1) = 1;
      beq[equality_ind++] = pid_out.qddot_des[(int)condof[i]-1];
    }
  }

  int n_ineq = 2*nu+2*6*desired_body_accelerations.size();
  MatrixXd Ain = MatrixXd::Zero(n_ineq,nparams);  // note: obvious sparsity here
  VectorXd bin = VectorXd::Zero(n_ineq);

  // linear input saturation constraints
  // u=B_act'*(H_act*qdd + C_act - Jz_act'*z - Dbar_act*beta)
  // using transpose instead of inverse because B is orthogonal
  Ain.topLeftCorner(nu,nq) = pdata->B_act.transpose()*pdata->H_act;
  Ain.block(0,nq,nu,nc*nd) = -pdata->B_act.transpose()*D_act;
  bin.head(nu) = -pdata->B_act.transpose()*pdata->C_act + pdata->umax;

  Ain.block(nu,0,nu,nparams) = -1*Ain.block(0,0,nu,nparams);
  bin.segment(nu,nu) = pdata->B_act.transpose()*pdata->C_act - pdata->umin;

  int constraint_start_index = 2*nu;
  for (int i=0; i<desired_body_accelerations.size(); i++) {
    pdata->r->forwardJac(desired_body_accelerations[i].body_id0,orig,1,Jb);
    pdata->r->forwardJacDot(desired_body_accelerations[i].body_id0,orig,1,Jbdot);
    Ain.block(constraint_start_index,0,6,pdata->r->num_positions) = Jb;
    bin.segment(constraint_start_index,6) = -Jbdot*robot_state.qd + desired_body_accelerations[i].accel_bounds.max;
    constraint_start_index += 6;
    Ain.block(constraint_start_index,0,6,pdata->r->num_positions) = -Jb;
    bin.segment(constraint_start_index,6) = Jbdot*robot_state.qd - desired_body_accelerations[i].accel_bounds.min;
    constraint_start_index += 6;
  }

  for (int i=0; i<n_ineq; i++) {
    // remove inf constraints---needed by gurobi
    if (std::isinf(double(bin(i)))) {
      Ain.row(i) = 0*Ain.row(i);
      bin(i)=0;
    }
  }

  GRBmodel * model = nullptr;
  int info=-1;

  // set obj,lb,up
  VectorXd lb(nparams), ub(nparams);
  lb.head(nq) = pdata->qdd_lb;
  ub.head(nq) = pdata->qdd_ub;
  lb.segment(nq,nf) = VectorXd::Zero(nf);
  ub.segment(nq,nf) = 1e3*VectorXd::Ones(nf);
  lb.tail(neps) = -params->slack_limit*VectorXd::Ones(neps);
  ub.tail(neps) = params->slack_limit*VectorXd::Ones(neps);

  VectorXd alpha(nparams);

  MatrixXd Qnfdiag(nf,1), Qneps(neps,1);
  std::vector<MatrixXd*> QBlkDiag( nc>0 ? 3 : 1 );  // nq, nf, neps   // this one is for gurobi

  VectorXd w = (w_qdd.array() + REG).matrix();

  if (nc != pdata->state.num_active_contact_pts) {
    // Number of contact points has changed, so our active set is invalid
    pdata->state.active.clear();
  }
  pdata->state.num_active_contact_pts = nc;

  #ifdef USE_MATRIX_INVERSION_LEMMA
  double max_body_accel_weight = -numeric_limits<double>::infinity();
  for (int i=0; i < desired_body_accelerations.size(); i++) {
    max_body_accel_weight = max(max_body_accel_weight, desired_body_accelerations[i].weight);
  }
  bool include_body_accel_cost_terms = desired_body_accelerations.size() > 0 && max_body_accel_weight > 1e-10;
  if (pdata->use_fast_qp > 0 && !include_angular_momentum && !include_body_accel_cost_terms)
  {
    // TODO: update to include angular momentum, body accel objectives.

    //    We want Hqp inverse, which I can compute efficiently using the
    //    matrix inversion lemma (see wikipedia):
    //    inv(A + U'CV) = inv(A) - inv(A)*U* inv([ inv(C)+ V*inv(A)*U ]) V inv(A)
    if (nc>0) {
      MatrixXd Wi = ((1/(w_qdd.array() + REG)).matrix()).asDiagonal();
      if (R_DQyD_ls.trace()>1e-15) { // R_DQyD_ls is not zero
        pdata->Hqp = Wi - Wi*Jcom.transpose()*(R_DQyD_ls.inverse() + Jcom*Wi*Jcom.transpose()).inverse()*Jcom*Wi;
      }
    }
    else {
      pdata->Hqp = MatrixXd::Constant(nq,1,1/(1+REG));
    }

    #ifdef TEST_FAST_QP
      if (nc>0) {
        MatrixXd Hqp_test(nq,nq);
        MatrixXd W = w.asDiagonal();
        Hqp_test = (Jcom.transpose()*R_DQyD_ls*Jcom + W).inverse();
        if (((Hqp_test-pdata->Hqp).array().abs()).maxCoeff() > 1e-6) {
          mexErrMsgTxt("Q submatrix inverse from matrix inversion lemma does not match direct Q inverse.");
        }
      }
    #endif

    Qnfdiag = MatrixXd::Constant(nf,1,1/REG);
    Qneps = MatrixXd::Constant(neps,1,1/(.001+REG));

    QBlkDiag[0] = &pdata->Hqp;
    if (nc>0) {
      QBlkDiag[1] = &Qnfdiag;
      QBlkDiag[2] = &Qneps;     // quadratic slack var cost, Q(nparams-neps:end,nparams-neps:end)=eye(neps)
    }

    MatrixXd Ain_lb_ub(n_ineq+2*nparams,nparams);
    VectorXd bin_lb_ub(n_ineq+2*nparams);
    Ain_lb_ub << Ain,            // note: obvious sparsity here
        -MatrixXd::Identity(nparams,nparams),
        MatrixXd::Identity(nparams,nparams);
    bin_lb_ub << bin, -lb, ub;

    for (std::set<int>::iterator it = pdata->state.active.begin(); it != pdata->state.active.end(); it++) {
      if (std::isnan(bin_lb_ub(*it)) || std::isinf(bin_lb_ub(*it))) {
        pdata->state.active.clear();
        break;
      }
    }

    info = fastQPThatTakesQinv(QBlkDiag, f, Aeq, beq, Ain_lb_ub, bin_lb_ub, pdata->state.active, alpha);

    //if (info<0)   mexPrintf("fastQP info = %d.  Calling gurobi.\n", info);
  }
  else {
  #endif

    if (nc>0) {
      pdata->Hqp = Jcom.transpose()*R_DQyD_ls*Jcom;
      if (include_angular_momentum) {
        pdata->Hqp += pdata->Ak.transpose()*params->W_kdot*pdata->Ak;
      }
      pdata->Hqp += w_qdd.asDiagonal();
      pdata->Hqp += REG*MatrixXd::Identity(nq,nq);
    } else {
      pdata->Hqp = (1+REG)*MatrixXd::Identity(nq,nq);
    }

    // add in body spatial acceleration cost terms
    for (int i=0; i<desired_body_accelerations.size(); i++) {
      if (desired_body_accelerations[i].weight > 0) {
        if (!inSupport(active_supports,desired_body_accelerations[i].body_id0)) {
          pdata->r->forwardJac(desired_body_accelerations[i].body_id0,orig,1,Jb);
          pdata->r->forwardJacDot(desired_body_accelerations[i].body_id0,orig,1,Jbdot);

          for (int j=0; j<6; j++) {
            if (!std::isnan(desired_body_accelerations[i].body_vdot[j])) {
              pdata->Hqp += desired_body_accelerations[i].weight*(Jb.row(j)).transpose()*Jb.row(j);
              f.head(nq) += desired_body_accelerations[i].weight*(robot_state.qd.transpose()*Jbdot.row(j).transpose() - desired_body_accelerations[i].body_vdot[j])*Jb.row(j).transpose();
            }
          }
        }
      }
    }

    Qnfdiag = MatrixXd::Constant(nf,1,params->w_grf+REG);
    Qneps = MatrixXd::Constant(neps,1,params->w_slack+REG);

    QBlkDiag[0] = &pdata->Hqp;
    if (nc>0) {
      QBlkDiag[1] = &Qnfdiag;
      QBlkDiag[2] = &Qneps;     // quadratic slack var cost, Q(nparams-neps:end,nparams-neps:end)=eye(neps)
    }

    MatrixXd Ain_lb_ub(n_ineq+2*nparams,nparams);
    VectorXd bin_lb_ub(n_ineq+2*nparams);
    Ain_lb_ub << Ain,            // note: obvious sparsity here
        -MatrixXd::Identity(nparams,nparams),
        MatrixXd::Identity(nparams,nparams);
    bin_lb_ub << bin, -lb, ub;

    for (std::set<int>::iterator it = pdata->state.active.begin(); it != pdata->state.active.end(); it++) {
      if (std::isnan(bin_lb_ub(*it)) || std::isinf(bin_lb_ub(*it))) {
        pdata->state.active.clear();
        break;
      }
    }

    if (pdata->use_fast_qp > 0)
    { // set up and call fastqp
      info = fastQP(QBlkDiag, f, Aeq, beq, Ain_lb_ub, bin_lb_ub, pdata->state.active, alpha);
      //if (info<0)    mexPrintf("fastQP info=%d... calling Gurobi.\n", info);
    }
    else {
      // use gurobi active set
      model = gurobiActiveSetQP(pdata->env,QBlkDiag,f,Aeq,beq,Ain,bin,lb,ub,pdata->state.vbasis,pdata->state.vbasis_len,pdata->state.cbasis,pdata->state.cbasis_len,alpha);
      CGE(GRBgetintattr(model,"NumVars",&(pdata->state.vbasis_len)), pdata->env);
      CGE(GRBgetintattr(model,"NumConstrs",&(pdata->state.cbasis_len)), pdata->env);
      info=66;
      //info = -1;
    }

    if (info<0) {
      model = gurobiQP(pdata->env,QBlkDiag,f,Aeq,beq,Ain,bin,lb,ub,pdata->state.active,alpha);
      int status; CGE(GRBgetintattr(model, "Status", &status), pdata->env);
      //if (status!=2) mexPrintf("Gurobi reports non-optimal status = %d\n", status);
    }
  #ifdef USE_MATRIX_INVERSION_LEMMA
  }
  #endif

  //----------------------------------------------------------------------
  // Solve for inputs ----------------------------------------------------
  qp_output->qdd = alpha.head(nq);
  VectorXd beta = alpha.segment(nq,nc*nd);

  // use transpose because B_act is orthogonal
  qp_output->u = pdata->B_act.transpose()*(pdata->H_act*qp_output->qdd + pdata->C_act - D_act*beta);
  for (int i=0; i < qp_output->u.size(); i++) {
    if (isnan(qp_output->u(i))) qp_output->u(i) = 0;
  }
  //y = pdata->B_act.jacobiSvd(ComputeThinU|ComputeThinV).solve(pdata->H_act*qdd + pdata->C_act - Jz_act.transpose()*lambda - D_act*beta);

  bool foot_contact[2];
  foot_contact[0] = b_contact_force(pdata->rpc.body_ids.l_foot) == 1;
  foot_contact[1] = b_contact_force(pdata->rpc.body_ids.r_foot) == 1;
  qp_output->qd_ref = velocityReference(pdata, robot_state.t, robot_state.q, robot_state.qd, qp_output->qdd, foot_contact, &(params->vref_integrator), &(pdata->rpc));

  // Remember t for next time around
  pdata->state.t_prev = robot_state.t;

  // If a debug pointer was passed in, fill it with useful data
  if (debug) {
    debug->active_supports.resize(active_supports.size());
    for (int i=0; i < active_supports.size(); i++) {
      debug->active_supports[i] = active_supports[i];
    }
    debug->nc = nc;
    debug->normals = normals;
    debug->B = B;
    debug->alpha = alpha;
    debug->f = f;
    debug->Aeq = Aeq;
    debug->beq = beq;
    debug->Ain_lb_ub = Ain_lb_ub;
    debug->bin_lb_ub = bin_lb_ub;
    debug->Qnfdiag = Qnfdiag;
    debug->Qneps = Qneps;
    debug->x_bar = x_bar;
    debug->S = S;
    debug->s1 = s1;
    debug->s1dot = s1dot;
    debug->s2dot = qp_input->zmp_data.s2dot;
    debug->A_ls = A_ls;
    debug->B_ls = B_ls;
    debug->Jcom = Jcom;
    debug->Jcomdot = Jcomdot;
    debug->beta = beta;
  }

  // if we used gurobi, clean up
  if (model) {
    GRBfreemodel(model);
  }
  //  GRBfreeenv(env);

  return info;
}


//
//  post-solver.cpp
//  parametrize
//
//  Created by Jingwei on 2/5/18.
//
#include <algorithm>
#include <boost/program_options.hpp>
#include <cmath>
#include <cstdio>
#include <string>

#include "ceres/ceres.h"
#include "ceres/rotation.h"

#include "post-solver.hpp"
#include "serialize.hpp"

namespace qflow {

/// Coefficient of area constraint.  The magnitude is 1 if area is equal to 0.
const double COEFF_AREA = 1;
/// Coefficient of tangent constraint.  The magnitude is 0.03 if the bais is reference_length.
/// This is because current tangent constraint is not very accurate.
/// This optimization conflicts with COEFF_AREA.
const double COEFF_TANGENT = 0.02;
/// Coefficient of normal constraint.  The magnitude is the arc angle.
const double COEFF_NORMAL = 1;
/// Coefficient of normal constraint.  The magnitude is the arc angle.
const double COEFF_FLOW = 1;
/// Coefficient of orthogonal edge.  The magnitude is the arc angle.
const double COEFF_ORTH = 1;
/// Coefficient of edge length.  The magnitude is the arc angle.
const double COEFF_LENGTH = 1;
/// Number of iterations of the CGNR solver
const int N_ITER = 100;

template <typename T, typename T2>
T DotProduct(const T a[3], const T2 b[3]) {
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
}

template <typename T>
T Length2(const T a[3]) {
    return DotProduct(a, a);
}

namespace ceres {
inline double min(const double f, const double g) { return std::min(f, g); }

template <typename T, int N>
inline Jet<T, N> min(const Jet<T, N>& f, const Jet<T, N>& g) {
    if (f.a < g.a)
        return f;
    else
        return g;
}
}  // namespace ceres

bool DEBUG = 0;
struct FaceConstraint {
    FaceConstraint(double coeff_area, double coeff_normal, double coeff_flow, double coeff_orth,
                   double length, Vector3d normal[4], Vector3d Q0[4], Vector3d Q1[4])
        : coeff_area(coeff_area),
          coeff_normal(coeff_normal),
          coeff_flow(coeff_flow),
          coeff_orth(coeff_orth),
          area0(length * length),
          normal0{
              normal[0],
              normal[1],
              normal[2],
              normal[3],
          },
          Q0{Q0[0], Q0[1], Q0[2], Q0[3]},
          Q1{Q1[0], Q1[1], Q1[2], Q1[3]} {}

    template <typename T>
    bool operator()(const T* p0, const T* p1, const T* p2, const T* p3, T* r) const {
        const T* p[] = {p0, p1, p2, p3};
        r[12] = T();
        for (int k = 0; k < 4; ++k) {
            auto pc = p[k];
            auto pa = p[(k + 1) % 4];
            auto pb = p[(k + 3) % 4];

            T a[3]{pa[0] - pc[0], pa[1] - pc[1], pa[2] - pc[2]};
            T b[3]{pb[0] - pc[0], pb[1] - pc[1], pb[2] - pc[2]};

            T length_a = ceres::sqrt(Length2(a));
            T length_b = ceres::sqrt(Length2(b));
            T aa[3]{a[0] / length_a, a[1] / length_a, a[2] / length_a};
            T bb[3]{b[0] / length_b, b[1] / length_b, b[2] / length_b};
            r[3 * k + 0] = coeff_orth * DotProduct(aa, bb);

            T degree_edge0 = ceres::abs(DotProduct(aa, &Q0[k][0]));
            T degree_edge1 = ceres::abs(DotProduct(aa, &Q1[k][0]));
            T degree_edge = ceres::min(degree_edge0, degree_edge1);
            r[3 * k + 1] = coeff_flow * degree_edge;

            T normal[3];
            ceres::CrossProduct(a, b, normal);
            T area = ceres::sqrt(Length2(normal));
            r[12] += area;

            assert(area != T());
            for (int i = 0; i < 3; ++i) normal[i] /= area;
            T degree_normal = DotProduct(normal, &normal0[k][0]) - T(1);
            r[3 * k + 2] = coeff_normal * degree_normal * degree_normal;
        }
        r[12] = coeff_area * (r[12] / (4.0 * area0) - 1.0);
        return true;
    }

    static ceres::CostFunction* create(double coeff_area, double coeff_normal, double coeff_flow,
                                       double coeff_orth, double length, Vector3d normal[4],
                                       Vector3d Q0[4], Vector3d Q1[4]) {
        return new ceres::AutoDiffCostFunction<FaceConstraint, 13, 3, 3, 3, 3>(new FaceConstraint(
            coeff_area, coeff_normal, coeff_flow, coeff_orth, length, normal, Q0, Q1));
    }

    double coeff_area;
    double coeff_normal;
    double coeff_flow;
    double coeff_orth;

    double area0;
    Vector3d normal0[4];
    Vector3d Q0[4], Q1[4];
};

struct VertexConstraint {
    VertexConstraint(double coeff_tangent, Vector3d normal, double bias, double length)
        : coeff{coeff_tangent / length * 10}, bias0{bias}, normal0{normal} {}

    template <typename T>
    bool operator()(const T* p, T* r) const {
        r[0] = coeff * (DotProduct(p, &normal0[0]) - bias0);
        return true;
    }

    static ceres::CostFunction* create(double coeff_tangent, Vector3d normal, double bias,
                                       double length) {
        return new ceres::AutoDiffCostFunction<VertexConstraint, 1, 3>(
            new VertexConstraint(coeff_tangent, normal, bias, length));
    }

    double coeff;
    double bias0;
    Vector3d normal0;
};

void solve(std::vector<Vector3d>& O_quad, std::vector<Vector3d>& N_quad,
           std::vector<Vector3d>& Q_quad, std::vector<Vector4i>& F_quad,
           std::vector<double>& B_quad, MatrixXd& V, MatrixXd& N, MatrixXd& Q, MatrixXd& O,
           MatrixXi& F, double reference_length, double coeff_area, double coeff_tangent,
           double coeff_normal, double coeff_flow, double coeff_orth) {
    printf("Parameter:  \n");
    printf("    coeff_area:      %.4f\n", coeff_area);
    printf("    coeff_tangent:   %.4f\n", coeff_tangent);
    printf("    coeff_normal:    %.4f\n", coeff_normal);
    printf("    coeff_flow:      %.4f\n", coeff_flow);
    printf("    coeff_orth:      %.4f\n\n", coeff_orth);
    int n_quad = Q_quad.size();

    ceres::Problem problem;
    std::vector<double> solution(n_quad * 3);
    for (int vquad = 0; vquad < n_quad; ++vquad) {
        solution[3 * vquad + 0] = O_quad[vquad][0];
        solution[3 * vquad + 1] = O_quad[vquad][1];
        solution[3 * vquad + 2] = O_quad[vquad][2];
    }

    // Face constraint (area and normal direction)
    for (int fquad = 0; fquad < F_quad.size(); ++fquad) {
        auto v = F_quad[fquad];
        Vector3d normal[4], Q0[4], Q1[4];
        for (int k = 0; k < 4; ++k) {
            normal[k] = N_quad[v[k]];
            Q0[k] = Q_quad[v[k]];
            Q1[k] = Q0[k].cross(normal[k]).normalized();
        }
        ceres::CostFunction* cost_function = FaceConstraint::create(
            coeff_area, coeff_normal, coeff_flow, coeff_orth, reference_length, normal, Q0, Q1);
        problem.AddResidualBlock(cost_function, nullptr, &solution[3 * v[0]], &solution[3 * v[1]],
                                 &solution[3 * v[2]], &solution[3 * v[3]]);
    }

    // Tangent constraint
    for (int vquad = 0; vquad < O_quad.size(); ++vquad) {
        ceres::CostFunction* cost_function = VertexConstraint::create(
            coeff_tangent, N_quad[vquad], B_quad[vquad], reference_length);
        problem.AddResidualBlock(cost_function, nullptr, &solution[3 * vquad]);
    }

    // Flow constraint

    ceres::Solver::Options options;
    options.num_threads = 1;
    options.max_num_iterations = N_ITER;
    options.initial_trust_region_radius = 1;
    options.linear_solver_type = ceres::CGNR;
    options.minimizer_progress_to_stdout = true;
    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);

    std::cout << summary.BriefReport() << std::endl;

    for (int vquad = 0; vquad < n_quad; ++vquad) {
        O_quad[vquad][0] = solution[3 * vquad + 0];
        O_quad[vquad][1] = solution[3 * vquad + 1];
        O_quad[vquad][2] = solution[3 * vquad + 2];
    }

    return;
}

void optimize_quad_positions(std::vector<Vector3d>& O_quad, std::vector<Vector3d>& N_quad,
                             std::vector<Vector3d>& Q_quad, std::vector<Vector4i>& F_quad,
                             VectorXi& V2E_quad, std::vector<int>& E2E_quad, MatrixXd& V,
                             MatrixXd& N, MatrixXd& Q, MatrixXd& O, MatrixXi& F, VectorXi& V2E,
                             VectorXi& E2E, DisajointTree& disajoint_tree, double reference_length,
                             bool just_serialize) {
    printf("Quad mesh info:\n");
    printf("Number of vertices with normals and orientations: %d = %d = %d\n", (int)O_quad.size(),
           (int)N_quad.size(), (int)Q_quad.size());
    printf("Number of faces: %d\n", (int)F_quad.size());
    printf("Number of directed edges: %d\n", (int)E2E_quad.size());
    // Information for the original mesh
    printf("Triangle mesh info:\n");
    printf(
        "Number of vertices with normals, "
        "orientations and associated quad positions: "
        "%d = %d = %d = %d\n",
        (int)V.cols(), (int)N.cols(), (int)Q.cols(), (int)O.cols());
    printf("Number of faces: %d\n", (int)F.cols());
    printf("Number of directed edges: %d\n", (int)E2E.size());
    printf("Reference length: %.2f\n", reference_length);

    int flip_count = 0;
    for (int i = 0; i < F_quad.size(); ++i) {
        bool flipped = false;
        for (int j = 0; j < 4; ++j) {
            int v1 = F_quad[i][j];
            int v2 = F_quad[i][(j + 1) % 4];
            int v3 = F_quad[i][(j + 3) % 4];

            Vector3d face_norm = (O_quad[v2] - O_quad[v1]).cross(O_quad[v3] - O_quad[v1]);
            Vector3d vertex_norm = N_quad[v1];
            if (face_norm.dot(vertex_norm) < 0) {
                flipped = true;
            }
        }
        if (flipped) {
            flip_count++;
        }
    }
    printf("Flipped Quads: %d\n", flip_count);

    int n_quad = O_quad.size();
    int n_trig = O.cols();
    std::vector<double> B_quad(n_quad);  // Average bias for quad vertex
    std::vector<int> B_weight(n_quad);

    printf("ntrig: %d, disjoint_tree.size: %d\n", n_trig, (int)disajoint_tree.indices.size());
    for (int vtrig = 0; vtrig < n_trig; ++vtrig) {
        int vquad = disajoint_tree.Index(vtrig);
        double b = N_quad[vquad].dot(O.col(vtrig));
        B_quad[vquad] += b;
        B_weight[vquad] += 1;
    }
    for (int vquad = 0; vquad < n_quad; ++vquad) {
        assert(B_weight[vquad]);
        B_quad[vquad] /= B_weight[vquad];
    }

    puts("Save parameters to post.bin for optimization");
    FILE* out = fopen("post.bin", "wb");
    assert(out);
    Save(out, O_quad);
    Save(out, N_quad);
    Save(out, Q_quad);
    Save(out, F_quad);
    Save(out, B_quad);
    Save(out, V);
    Save(out, N);
    Save(out, Q);
    Save(out, O);
    Save(out, F);
    Save(out, reference_length);
    fclose(out);

    if (!just_serialize) {
        puts("Start post optimization");
        solve(O_quad, N_quad, Q_quad, F_quad, B_quad, V, N, Q, O, F, reference_length, COEFF_AREA,
              COEFF_TANGENT, COEFF_NORMAL, COEFF_FLOW, COEFF_ORTH);
    }
}

#ifdef POST_SOLVER

void SaveObj(const std::string& fname, std::vector<Vector3d> O_quad,
             std::vector<Vector4i> F_quad) {
    std::ofstream os(fname);
    for (int i = 0; i < (int)O_quad.size(); ++i) {
        os << "v " << O_quad[i][0] << " " << O_quad[i][1] << " " << O_quad[i][2] << "\n";
    }
    for (int i = 0; i < (int)F_quad.size(); ++i) {
        os << "f " << F_quad[i][0] + 1 << " " << F_quad[i][1] + 1 << " " << F_quad[i][2] + 1 << " "
           << F_quad[i][3] + 1 << "\n";
    }
    os.close();
}

int main(int argc, char* argv[]) {
    double coeff_area;
    double coeff_tangent;
    double coeff_normal;
    double coeff_flow;
    double coeff_orth;

    namespace po = boost::program_options;
    po::options_description desc("Allowed options");
    desc.add_options()  // clang-format off
        ("help,h", "produce help message")
        ("area,a", po::value<double>(&coeff_area)->default_value(COEFF_AREA), "Set the coefficient of area constraint")
        ("tangent,t", po::value<double>(&coeff_tangent)->default_value(COEFF_TANGENT), "Set the coefficient of tangent constraint")
        ("normal,n", po::value<double>(&coeff_normal)->default_value(COEFF_NORMAL), "Set the coefficient of normal constraint")
        ("flow,f", po::value<double>(&coeff_flow)->default_value(COEFF_FLOW), "Set the coefficient of flow (Q) constraint")
        ("orth,o", po::value<double>(&coeff_orth)->default_value(COEFF_ORTH), "Set the coefficient of orthogonal constraint");

    // clang-format on
    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);
    if (vm.count("help")) {
        std::cout << desc << std::endl;
        return 1;
    }

    std::vector<Vector3d> O_quad;
    std::vector<Vector3d> N_quad;
    std::vector<Vector3d> Q_quad;
    std::vector<Vector4i> F_quad;
    std::vector<double> B_quad;
    MatrixXd V;
    MatrixXd N;
    MatrixXd Q;
    MatrixXd O;
    MatrixXi F;
    double reference_length;

    puts("Read parameters from post.bin");
    FILE* in = fopen("post.bin", "rb");
    assert(in);
    Read(in, O_quad);
    Read(in, N_quad);
    Read(in, Q_quad);
    Read(in, F_quad);
    Read(in, B_quad);
    Read(in, V);
    Read(in, N);
    Read(in, Q);
    Read(in, O);
    Read(in, F);
    Read(in, reference_length);
    fclose(in);
    printf("reference_length: %.2f\n", reference_length);
    SaveObj("presolver.obj", O_quad, F_quad);

    int n_flip = 0;
    double sum_degree = 0;
    for (int i = 0; i < F_quad.size(); ++i) {
        bool flipped = false;
        for (int j = 0; j < 4; ++j) {
            int v1 = F_quad[i][j];
            int v2 = F_quad[i][(j + 1) % 4];
            int v3 = F_quad[i][(j + 3) % 4];

            Vector3d face_norm =
                (O_quad[v2] - O_quad[v1]).cross(O_quad[v3] - O_quad[v1]).normalized();
            Vector3d vertex_norm = N_quad[v1];
            if (face_norm.dot(vertex_norm) < 0) {
                flipped = true;
            }
            double degree = std::acos(face_norm.dot(vertex_norm));
            assert(degree >= 0);
            // printf("cos theta = %.2f\n", degree);
            sum_degree += degree * degree;
        }
        n_flip += flipped;
    }
    printf("n_flip: %d\nsum_degree: %.3f\n", n_flip, sum_degree);

    puts("Start post optimization");
    solve(O_quad, N_quad, Q_quad, F_quad, B_quad, V, N, Q, O, F, reference_length, coeff_area,
          coeff_tangent, coeff_normal, coeff_flow, coeff_orth);
    SaveObj("postsolver.obj", O_quad, F_quad);

    n_flip = 0;
    sum_degree = 0;
    for (int i = 0; i < F_quad.size(); ++i) {
        bool flipped = false;
        for (int j = 0; j < 4; ++j) {
            int v1 = F_quad[i][j];
            int v2 = F_quad[i][(j + 1) % 4];
            int v3 = F_quad[i][(j + 3) % 4];

            Vector3d face_norm =
                (O_quad[v2] - O_quad[v1]).cross(O_quad[v3] - O_quad[v1]).normalized();
            Vector3d vertex_norm = N_quad[v1];
            if (face_norm.dot(vertex_norm) < 0) {
                flipped = true;
            }
            double degree = std::acos(face_norm.dot(vertex_norm));
            assert(degree >= 0);
            sum_degree += degree * degree;
        }
        n_flip += flipped;
    }
    printf("n_flip: %d\nsum_degree: %.3f\n", n_flip, sum_degree);
    return 0;
}

#endif

} // namespace qflow


//          Copyright Gavin Band 2008 - 2012.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)

#ifndef METRO_FISHERS_EXACT_TEST
#define METRO_FISHERS_EXACT_TEST

#include <Eigen/Core>
#include <boost/math/distributions/hypergeometric.hpp>

namespace metro {
	//
	// Fisher's exact test between categories
	//               red           black
	//     good |     a      |       b        |
	//      bad |     c      |       d        |
	//
	// Under the null that both rows are distributed the same, a is hypergeometrically
	// distributed.  Fishers' p-value is then the mass under the hypergeometric distribution
	// of all values of a greater than or equal to the one given that are consistent with
	// the given margins of the table.
	struct FishersExactTest {
		FishersExactTest( Eigen::Matrix2d const& matrix ) ;
		double get_OR() const ;
		std::pair< double, double > get_confidence_interval() const ;

		enum Alternative { eLess = 0, eGreater = 1, eTwoSided = 2 } ;
		double get_pvalue( Alternative const = eGreater ) const ;

		private:
			Eigen::Matrix2d const m_matrix ;
			boost::math::hypergeometric m_distribution ;
	} ;
}

#endif


#include "tile/lib/lib.h"

#include <boost/format.hpp>

namespace vertexai {
namespace tile {
namespace lib {

lang::RunInfo LoadMatMul(const std::string& name, const TensorShape& i1, const TensorShape& i2) {
  lang::RunInfo runinfo;
  runinfo.program_name = name;
  runinfo.code = "function (A[M, K], B[K, N]) -> (C) { C[m, n : M, N] = +(A[m, k] * B[k, n]); }";
  runinfo.input_shapes.emplace("A", i1);
  runinfo.input_shapes.emplace("B", i2);
  runinfo.output_shapes.emplace("C", SimpleShape(i1.type, {i1.dims[0].size, i2.dims[1].size}));
  return runinfo;
}

lang::RunInfo LoadConstCalc(const std::string& name, const TensorShape& output) {
  lang::RunInfo runinfo;
  runinfo.program_name = name;
  runinfo.code = R"***(
function () -> (O) {
  N = 1;
  F = 0.0;
  F2 = 3.7;
  Simple[i : N] = =(F[]);
  DoubleN[i : N] = =(N[]);
  Partial = Simple + DoubleN;
  O = Partial + F2;
})***";
  runinfo.output_shapes.emplace("O", output);
  return runinfo;
}

lang::RunInfo LoadConv1d(const std::string& name,    //
                         const TensorShape& input,   //
                         const TensorShape& kernel,  //
                         const TensorShape& output) {
  auto center = kernel.dims[0].size / 2;
  auto code = R"***(
function (I[X, CI], K[KX, CI, CO]) -> (O) {
  [[pid(res2a_branch2a)]] O[x, co : X, CO] = +(I[x + kx - %1%, ci] * K[kx, ci, co]);
})***";
  lang::RunInfo runinfo;
  runinfo.program_name = name;
  runinfo.code = str(boost::format(code) % center);
  runinfo.input_shapes.insert(std::make_pair("I", input));
  runinfo.input_shapes.insert(std::make_pair("K", kernel));
  runinfo.output_shapes.insert(std::make_pair("O", output));
  return runinfo;
}

lang::RunInfo LoadConv2d(const std::string& name,    //
                         const TensorShape& input,   //
                         const TensorShape& kernel,  //
                         const TensorShape& output) {
  auto center = kernel.dims[0].size / 2;
  auto code = R"***(
function (I[X, Y, CI], K[KX, KY, CI, CO]) -> (O) {
  [[pid(res2a_branch2a)]] O[x, y, co : X, Y, CO] = +(I[x + kx - %1%, y + ky - %1%, ci] * K[kx, ky, ci, co]);
})***";
  lang::RunInfo runinfo;
  runinfo.program_name = name;
  runinfo.code = str(boost::format(code) % center);
  runinfo.input_shapes.insert(std::make_pair("I", input));
  runinfo.input_shapes.insert(std::make_pair("K", kernel));
  runinfo.output_shapes.insert(std::make_pair("O", output));
  return runinfo;
}

lang::RunInfo LoadConv2dRelu(const std::string& name,    //
                             const TensorShape& input,   //
                             const TensorShape& kernel,  //
                             const TensorShape& output) {
  auto center = kernel.dims[0].size / 2;
  auto code = R"***(
function (I[X, Y, CI], K[KX, KY, CI, CO]) -> (O) {
  [[pid(res2a_branch2a)]] O[x, y, co : X, Y, CO] = +(I[x + kx - %1%, y + ky - %1%, ci] * K[i, ky, ci, co]);
  [[pid(relu)]] R = zelu(O);
})***";
  lang::RunInfo runinfo;
  runinfo.program_name = name;
  runinfo.code = str(boost::format(code) % center);
  runinfo.input_shapes.insert(std::make_pair("I", input));
  runinfo.input_shapes.insert(std::make_pair("K", kernel));
  runinfo.output_shapes.insert(std::make_pair("R", output));
  return runinfo;
}

lang::RunInfo LoadConv2dBnRelu(const std::string& name,      //
                               const TensorShape& input,     //
                               const TensorShape& kernel,    //
                               const TensorShape& channels,  //
                               const TensorShape& output) {
  auto center = kernel.dims[0].size / 2;
  auto code = R"***(
function (I[X, Y, CI], K[KX, KY, CI, CO], B[CO], S[CO]) -> (R) {
  [[pid(res2a_branch2a)]] O[x, y, co : X, Y, CO] = +(I[x + kx - %1%, y + ky - %1%, ci] * K[kx, ky, ci, co]);
  [[pid(bias_add)]] BO = O + B;
  [[pid(scale)]] BS = BO * S;
  [[pid(relu)]] R = zelu(BS);
})***";
  lang::RunInfo runinfo;
  runinfo.program_name = name;
  runinfo.code = str(boost::format(code) % center);
  runinfo.input_shapes.insert(std::make_pair("I", input));
  runinfo.input_shapes.insert(std::make_pair("K", kernel));
  runinfo.input_shapes.insert(std::make_pair("B", channels));
  runinfo.input_shapes.insert(std::make_pair("S", channels));
  runinfo.output_shapes.insert(std::make_pair("R", output));
  return runinfo;
}

lang::RunInfo LoadConv2d3Deep(const std::string& name,     //
                              const TensorShape& input,    //
                              const TensorShape& kernel1,  //
                              const TensorShape& kernel2,  //
                              const TensorShape& kernel3,  //
                              const TensorShape& output) {
  auto center1 = kernel1.dims[0].size / 2;
  auto center2 = kernel2.dims[0].size / 2;
  auto center3 = kernel3.dims[0].size / 2;
  auto code = R"***(
function (I[X, Y, CI], K1[KX, KY, C1, C2], K2[KX, KY, C2, C3], K3[KX, KY, C3, C4]) -> (O3) {
  O1[x, y, co : X, Y, C2] = +(I[x + kx - %1%, y + ky - %1%, ci] * K1[kx, ky, ci, co]);
  O2[x, y, co : X, Y, C3] = +(O1[x + kx - %2%, y + ky - %2%, ci] * K2[kx, ky, ci, co]);
  O3[x, y, co : X, Y, C4] = +(O2[x + kx - %3%, y + ky - %3%, ci] * K3[kx, ky, ci, co]);
})***";
  lang::RunInfo runinfo;
  runinfo.program_name = name;
  runinfo.code = str(boost::format(code) % center1 % center2 % center3);
  runinfo.input_shapes.insert(std::make_pair("I", input));
  runinfo.input_shapes.insert(std::make_pair("K1", kernel1));
  runinfo.input_shapes.insert(std::make_pair("K2", kernel2));
  runinfo.input_shapes.insert(std::make_pair("K3", kernel3));
  runinfo.output_shapes.insert(std::make_pair("O3", output));
  return runinfo;
}

}  // namespace lib
}  // namespace tile
}  // namespace vertexai