/*
 * Copyright 2012-2019 CNRS-UM LIRMM, CNRS-AIST JRL
 */

// associated header
#include "Tasks/QPTasks.h"

// includes
// std
#include <cmath>
#include <iterator>
#include <set>

// Eigen
#include <Eigen/Geometry>

// RBDyn
#include <RBDyn/MultiBody.h>
#include <RBDyn/MultiBodyConfig.h>

// Tasks
#include <Tasks/Bounds.h>

namespace tasks
{

namespace qp
{

/**
 *														SetPointTaskCommon
 */

SetPointTaskCommon::SetPointTaskCommon(const std::vector<rbd::MultiBody> & mbs,
                                       int rI,
                                       HighLevelTask * hlTask,
                                       double weight)
: Task(weight), hlTask_(hlTask), error_(hlTask->dim()), dimWeight_(Eigen::VectorXd::Ones(hlTask->dim())),
  robotIndex_(rI), alphaDBegin_(0), Q_(mbs[rI].nrDof(), mbs[rI].nrDof()), C_(mbs[rI].nrDof()),
  preQ_(hlTask->dim(), mbs[rI].nrDof()), preC_(hlTask->dim())
{
}

SetPointTaskCommon::SetPointTaskCommon(const std::vector<rbd::MultiBody> & mbs,
                                       int rI,
                                       HighLevelTask * hlTask,
                                       const Eigen::VectorXd & dimWeight,
                                       double weight)
: Task(weight), hlTask_(hlTask), error_(hlTask->dim()), dimWeight_(dimWeight), robotIndex_(rI), alphaDBegin_(0),
  Q_(mbs[rI].nrDof(), mbs[rI].nrDof()), C_(mbs[rI].nrDof()), preQ_(hlTask->dim(), mbs[rI].nrDof()), preC_(hlTask->dim())
{
}

void SetPointTaskCommon::dimWeight(const Eigen::VectorXd & dim)
{
  dimWeight_ = dim;
}

void SetPointTaskCommon::updateNrVars(const std::vector<rbd::MultiBody> & /* mbs */, const SolverData & data)
{
  alphaDBegin_ = data.alphaDBegin(robotIndex_);
}

void SetPointTaskCommon::computeQC(Eigen::VectorXd & error)
{
  const Eigen::MatrixXd & J = hlTask_->jac();
  const Eigen::VectorXd & normalAcc = hlTask_->normalAcc();

  error.noalias() -= normalAcc;
  preC_.noalias() = dimWeight_.asDiagonal() * error;
  C_.noalias() = -J.transpose() * preC_;

  preQ_.noalias() = dimWeight_.asDiagonal() * J;
  Q_.noalias() = J.transpose() * preQ_;
}

const Eigen::MatrixXd & SetPointTaskCommon::Q() const
{
  return Q_;
}

const Eigen::VectorXd & SetPointTaskCommon::C() const
{
  return C_;
}

/**
 *														SetPointTask
 */

SetPointTask::SetPointTask(const std::vector<rbd::MultiBody> & mbs,
                           int rI,
                           HighLevelTask * hlTask,
                           double stiffness,
                           double weight)
: SetPointTaskCommon(mbs, rI, hlTask, weight), stiffness_(stiffness), stiffnessSqrt_(2. * std::sqrt(stiffness))
{
}

SetPointTask::SetPointTask(const std::vector<rbd::MultiBody> & mbs,
                           int rI,
                           HighLevelTask * hlTask,
                           double stiffness,
                           const Eigen::VectorXd & dimWeight,
                           double weight)
: SetPointTaskCommon(mbs, rI, hlTask, dimWeight, weight), stiffness_(stiffness),
  stiffnessSqrt_(2. * std::sqrt(stiffness))
{
}

void SetPointTask::stiffness(double stiffness)
{
  stiffness_ = stiffness;
  stiffnessSqrt_ = 2. * std::sqrt(stiffness);
}

void SetPointTask::update(const std::vector<rbd::MultiBody> & mbs,
                          const std::vector<rbd::MultiBodyConfig> & mbcs,
                          const SolverData & data)
{
  hlTask_->update(mbs, mbcs, data);

  const Eigen::VectorXd & err = hlTask_->eval();
  const Eigen::VectorXd & speed = hlTask_->speed();

  error_.noalias() = stiffness_ * err;
  error_.noalias() -= stiffnessSqrt_ * speed;
  computeQC(error_);
}

/**
 *														TrackingTask
 */

TrackingTask::TrackingTask(const std::vector<rbd::MultiBody> & mbs,
                           int rI,
                           HighLevelTask * hlTask,
                           double gainPos,
                           double gainVel,
                           double weight)
: SetPointTaskCommon(mbs, rI, hlTask, weight), gainPos_(gainPos), gainVel_(gainVel),
  errorPos_(Eigen::VectorXd::Zero(hlTask->dim())), errorVel_(Eigen::VectorXd::Zero(hlTask->dim())),
  refAccel_(Eigen::VectorXd::Zero(hlTask->dim()))
{
}

TrackingTask::TrackingTask(const std::vector<rbd::MultiBody> & mbs,
                           int rI,
                           HighLevelTask * hlTask,
                           double gainPos,
                           double gainVel,
                           const Eigen::VectorXd & dimWeight,
                           double weight)
: SetPointTaskCommon(mbs, rI, hlTask, dimWeight, weight), gainPos_(gainPos), gainVel_(gainVel),
  errorPos_(Eigen::VectorXd::Zero(hlTask->dim())), errorVel_(Eigen::VectorXd::Zero(hlTask->dim())),
  refAccel_(Eigen::VectorXd::Zero(hlTask->dim()))
{
}

void TrackingTask::setGains(double gainPos, double gainVel)
{
  gainPos_ = gainPos;
  gainVel_ = gainVel;
}

void TrackingTask::errorPos(const Eigen::VectorXd & errorPos)
{
  errorPos_ = errorPos;
}

void TrackingTask::errorVel(const Eigen::VectorXd & errorVel)
{
  errorVel_ = errorVel;
}

void TrackingTask::refAccel(const Eigen::VectorXd & refAccel)
{
  refAccel_ = refAccel;
}

void TrackingTask::update(const std::vector<rbd::MultiBody> & mbs,
                          const std::vector<rbd::MultiBodyConfig> & mbcs,
                          const SolverData & data)
{
  hlTask_->update(mbs, mbcs, data);

  error_.noalias() = gainPos_ * errorPos_;
  error_.noalias() += gainVel_ * errorVel_;
  error_.noalias() += refAccel_;
  computeQC(error_);
}

/**
 *														TrajectoryTask
 */

TrajectoryTask::TrajectoryTask(const std::vector<rbd::MultiBody> & mbs,
                               int rI,
                               HighLevelTask * hlTask,
                               double gainPos,
                               double gainVel,
                               double weight)
: SetPointTaskCommon(mbs, rI, hlTask, weight), stiffness_(gainPos * Eigen::VectorXd::Ones(hlTask->dim())),
  damping_(gainVel * Eigen::VectorXd::Ones(hlTask->dim())), refVel_(Eigen::VectorXd::Zero(hlTask->dim())),
  refAccel_(Eigen::VectorXd::Zero(hlTask->dim()))
{
}

TrajectoryTask::TrajectoryTask(const std::vector<rbd::MultiBody> & mbs,
                               int rI,
                               HighLevelTask * hlTask,
                               double gainPos,
                               double gainVel,
                               const Eigen::VectorXd & dimWeight,
                               double weight)
: SetPointTaskCommon(mbs, rI, hlTask, dimWeight, weight), stiffness_(gainPos * Eigen::VectorXd::Ones(hlTask->dim())),
  damping_(gainVel * Eigen::VectorXd::Ones(hlTask->dim())), refVel_(Eigen::VectorXd::Zero(hlTask->dim())),
  refAccel_(Eigen::VectorXd::Zero(hlTask->dim()))
{
}

void TrajectoryTask::setGains(double gainPos, double gainVel)
{
  stiffness_ = gainPos * Eigen::VectorXd::Ones(hlTask_->dim());
  damping_ = gainVel * Eigen::VectorXd::Ones(hlTask_->dim());
}

void TrajectoryTask::setGains(const Eigen::VectorXd & stiffness, const Eigen::VectorXd & damping)
{
  assert(stiffness.size() == stiffness_.size());
  assert(damping.size() == damping_.size());
  stiffness_ = stiffness;
  damping_ = damping;
}

void TrajectoryTask::stiffness(double gainPos)
{
  stiffness_ = gainPos * Eigen::VectorXd::Ones(hlTask_->dim());
}

void TrajectoryTask::stiffness(const Eigen::VectorXd & stiffness)
{
  assert(stiffness.size() == stiffness_.size());
  stiffness_ = stiffness;
}

const Eigen::VectorXd & TrajectoryTask::stiffness() const
{
  return stiffness_;
}

void TrajectoryTask::damping(double gainVel)
{
  damping_ = gainVel * Eigen::VectorXd::Ones(hlTask_->dim());
}

void TrajectoryTask::damping(const Eigen::VectorXd & damping)
{
  assert(damping.size() == damping_.size());
  damping_ = damping;
}

const Eigen::VectorXd & TrajectoryTask::damping() const
{
  return damping_;
}

void TrajectoryTask::refVel(const Eigen::VectorXd & refVel)
{
  refVel_ = refVel;
}

const Eigen::VectorXd & TrajectoryTask::refVel() const
{
  return refVel_;
}

void TrajectoryTask::refAccel(const Eigen::VectorXd & refAccel)
{
  refAccel_ = refAccel;
}

const Eigen::VectorXd & TrajectoryTask::refAccel() const
{
  return refAccel_;
}

void TrajectoryTask::update(const std::vector<rbd::MultiBody> & mbs,
                            const std::vector<rbd::MultiBodyConfig> & mbcs,
                            const SolverData & data)
{
  hlTask_->update(mbs, mbcs, data);

  const Eigen::VectorXd & err = hlTask_->eval();
  const Eigen::VectorXd & speed = hlTask_->speed();

  error_.noalias() = stiffness_.asDiagonal() * err;
  error_.noalias() += damping_.asDiagonal() * (refVel_ - speed);
  error_.noalias() += refAccel_;
  computeQC(error_);
}

/**
 *														PIDTask
 */

PIDTask::PIDTask(const std::vector<rbd::MultiBody> & mbs,
                 int rI,
                 HighLevelTask * hlTask,
                 double P,
                 double I,
                 double D,
                 double weight)
: SetPointTaskCommon(mbs, rI, hlTask, weight), P_(P), I_(I), D_(D), error_(Eigen::VectorXd::Zero(hlTask->dim())),
  errorD_(Eigen::VectorXd::Zero(hlTask->dim())), errorI_(Eigen::VectorXd::Zero(hlTask->dim()))
{
}

PIDTask::PIDTask(const std::vector<rbd::MultiBody> & mbs,
                 int rI,
                 HighLevelTask * hlTask,
                 double P,
                 double I,
                 double D,
                 const Eigen::VectorXd & dimWeight,
                 double weight)
: SetPointTaskCommon(mbs, rI, hlTask, dimWeight, weight), P_(P), I_(I), D_(D),
  error_(Eigen::VectorXd::Zero(hlTask->dim())), errorD_(Eigen::VectorXd::Zero(hlTask->dim())),
  errorI_(Eigen::VectorXd::Zero(hlTask->dim()))
{
}

double PIDTask::P() const
{
  return P_;
}

void PIDTask::P(double p)
{
  P_ = p;
}

double PIDTask::I() const
{
  return I_;
}

void PIDTask::I(double i)
{
  I_ = i;
}

double PIDTask::D() const
{
  return D_;
}

void PIDTask::D(double d)
{
  D_ = d;
}

void PIDTask::error(const Eigen::VectorXd & err)
{
  error_ = err;
}

void PIDTask::errorD(const Eigen::VectorXd & errD)
{
  errorD_ = errD;
}

void PIDTask::errorI(const Eigen::VectorXd & errI)
{
  errorI_ = errI;
}

void PIDTask::update(const std::vector<rbd::MultiBody> & mbs,
                     const std::vector<rbd::MultiBodyConfig> & mbcs,
                     const SolverData & data)
{
  hlTask_->update(mbs, mbcs, data);

  error_.noalias() = P_ * error_;
  error_.noalias() -= D_ * errorD_;
  error_.noalias() -= I_ * errorI_;
  computeQC(error_);
}

/**
 *														TargetObjectiveTask
 */

TargetObjectiveTask::TargetObjectiveTask(const std::vector<rbd::MultiBody> & mbs,
                                         int rI,
                                         HighLevelTask * hlTask,
                                         double timeStep,
                                         double dur,
                                         const Eigen::VectorXd & objDot,
                                         double weight)
: Task(weight), hlTask_(hlTask), dt_(timeStep), objDot_(objDot), dimWeight_(Eigen::VectorXd::Ones(hlTask->dim())),
  robotIndex_(rI), alphaDBegin_(0), phi_(hlTask->dim()), psi_(hlTask->dim()), Q_(mbs[rI].nrDof(), mbs[rI].nrDof()),
  C_(mbs[rI].nrDof()), preQ_(hlTask->dim(), mbs[rI].nrDof()), CVecSum_(hlTask->dim()), preC_(hlTask->dim())
{
  duration(dur);
}

TargetObjectiveTask::TargetObjectiveTask(const std::vector<rbd::MultiBody> & mbs,
                                         int rI,
                                         HighLevelTask * hlTask,
                                         double timeStep,
                                         double dur,
                                         const Eigen::VectorXd & objDot,
                                         const Eigen::VectorXd & dimWeight,
                                         double weight)
: Task(weight), hlTask_(hlTask), dt_(timeStep), objDot_(objDot), dimWeight_(dimWeight), robotIndex_(rI),
  alphaDBegin_(0), phi_(hlTask->dim()), psi_(hlTask->dim()), Q_(mbs[rI].nrDof(), mbs[rI].nrDof()), C_(mbs[rI].nrDof()),
  preQ_(hlTask->dim(), mbs[rI].nrDof()), CVecSum_(hlTask->dim()), preC_(hlTask->dim())
{
  duration(dur);
}

double TargetObjectiveTask::duration() const
{
  return (nrIter_ - iter_) * dt_;
}

void TargetObjectiveTask::duration(double d)
{
  nrIter_ = static_cast<int>(std::round(d / dt_));
  iter_ = 0;
}

void TargetObjectiveTask::updateNrVars(const std::vector<rbd::MultiBody> & /* mbs */, const SolverData & data)
{
  alphaDBegin_ = data.alphaDBegin(robotIndex_);
}

void TargetObjectiveTask::update(const std::vector<rbd::MultiBody> & mbs,
                                 const std::vector<rbd::MultiBodyConfig> & mbcs,
                                 const SolverData & data)
{
  using namespace Eigen;

  hlTask_->update(mbs, mbcs, data);

  const MatrixXd & J = hlTask_->jac();
  const VectorXd & err = hlTask_->eval();
  const VectorXd & speed = hlTask_->speed();
  const VectorXd & normalAcc = hlTask_->normalAcc();

  // M·[phi, psi]^T = Obj

  // M =
  //  ⎡          2            2⎤
  //  ⎢(-t₀ + tf)   (-t₀ + tf) ⎥
  //  ⎢───────────  ───────────⎥
  //  ⎢     3            6     ⎥
  //  ⎢                        ⎥
  //  ⎢   t₀   tf      t₀   tf ⎥
  //  ⎢ - ── + ──    - ── + ── ⎥
  //  ⎣   2    2       2    2  ⎦

  // M^I =
  //  ⎡          6               2    ⎤
  //  ⎢ ───────────────────   ─────── ⎥
  //  ⎢   2               2   t₀ - tf ⎥
  //  ⎢ t₀  - 2⋅t₀⋅tf + tf            ⎥
  //  ⎢                               ⎥
  //  ⎢          6               4    ⎥
  //  ⎢─────────────────────  ────────⎥
  //  ⎢    2               2  -t₀ + tf⎥
  //  ⎣- t₀  + 2⋅t₀⋅tf - tf           ⎦

  // Obj = [ err - (tf - t₀)·J α, objDot - J·α ]

  double d = (nrIter_ - iter_) * dt_;
  double ds = std::pow(d, 2);

  Matrix2d MI;
  Vector2d Obj;

  MI << 6. / ds, 2. / (-d), 6. / (-ds), 4. / d;

  for(int i = 0; i < hlTask_->dim(); ++i)
  {
    Obj << err(i) - d * speed(i), objDot_(i) - speed(i);
    Vector2d pp(MI * Obj);
    phi_(i) = pp(0);
    psi_(i) = pp(1);
  }

  preQ_.noalias() = dimWeight_.asDiagonal() * J;
  Q_.noalias() = J.transpose() * preQ_;

  CVecSum_.noalias() = phi_ - normalAcc;
  preC_.noalias() = dimWeight_.asDiagonal() * CVecSum_;
  C_.noalias() = -J.transpose() * preC_;

  ++iter_;
}

const Eigen::MatrixXd & TargetObjectiveTask::Q() const
{
  return Q_;
}

const Eigen::VectorXd & TargetObjectiveTask::C() const
{
  return C_;
}

/**
 *												JointsSelector
 */

JointsSelector JointsSelector::ActiveJoints(const std::vector<rbd::MultiBody> & mbs,
                                            int robotIndex,
                                            HighLevelTask * hl,
                                            const std::vector<std::string> & activeJointsName,
                                            const std::map<std::string, std::vector<std::array<int, 2>>> & activeDofs)
{
  return JointsSelector(mbs, robotIndex, hl, activeJointsName, activeDofs);
}

JointsSelector JointsSelector::UnactiveJoints(
    const std::vector<rbd::MultiBody> & mbs,
    int robotIndex,
    HighLevelTask * hl,
    const std::vector<std::string> & unactiveJointsName,
    const std::map<std::string, std::vector<std::array<int, 2>>> & unactiveDofs)
{
  using namespace std::placeholders;
  const rbd::MultiBody & mb = mbs[robotIndex];

  std::vector<std::string> activeJointsName;
  std::map<std::string, std::vector<std::array<int, 2>>> activeDofs;
  // sort unactiveJointsName by puting them into a set
  std::set<std::string> unactiveJointsNameSet(unactiveJointsName.begin(), unactiveJointsName.end());
  // create a set with all joints id
  std::set<std::string> jointsNameSet;
  std::transform(mb.joints().begin(), mb.joints().end(), std::inserter(jointsNameSet, jointsNameSet.begin()),
                 std::bind(&rbd::Joint::name, _1));

  // remove unactive joints from the set
  std::set_difference(jointsNameSet.begin(), jointsNameSet.end(), unactiveJointsNameSet.begin(),
                      unactiveJointsNameSet.end(), std::inserter(activeJointsName, activeJointsName.begin()));

  // Transform unactive dofs into active dofs
  for(const auto & uad : unactiveDofs)
  {
    const auto & jN = uad.first;
    const auto & jDofs = uad.second;
    const auto & j = mb.joint(mb.jointIndexByName(jN));
    int ji = 0;
    for(size_t i = 0; i < jDofs.size(); ++i)
    {
      auto dofStart = jDofs[i][0];
      if(dofStart > ji)
      {
        activeDofs[jN].push_back({{ji, dofStart - ji}});
      }
      ji = dofStart + jDofs[i][1];
    }
    auto finalDof = jDofs.back()[0] + jDofs.back()[1];
    if(finalDof < j.dof())
    {
      activeDofs[jN].push_back({{finalDof, j.dof() - finalDof}});
    }
  }

  return JointsSelector(mbs, robotIndex, hl, activeJointsName);
}

JointsSelector::JointsSelector(const std::vector<rbd::MultiBody> & mbs,
                               int robotIndex,
                               HighLevelTask * hl,
                               const std::vector<std::string> & selectedJointsName,
                               const std::map<std::string, std::vector<std::array<int, 2>>> & activeDofs)
: jac_(Eigen::MatrixXd::Zero(hl->dim(), mbs[robotIndex].nrDof())), selectedJoints_(), hl_(hl)
{
  const rbd::MultiBody & mb = mbs[robotIndex];
  selectedJoints_.reserve(selectedJointsName.size());
  for(const std::string & jName : selectedJointsName)
  {
    int index = mb.jointIndexByName(jName);
    if(activeDofs.count(jName))
    {
      auto pInDof = mb.jointPosInDof(index);
      for(const auto & jdof : activeDofs.at(jName))
      {
        assert(jdof[0] + jdof[1] < mb.joint(index).dof());
        selectedJoints_.push_back({pInDof + jdof[0], jdof[1]});
      }
    }
    else
    {
      selectedJoints_.push_back({mb.jointPosInDof(index), mb.joint(index).dof()});
    }
  }
  // sort data in posInDof order
  std::sort(selectedJoints_.begin(), selectedJoints_.end(),
            [](const SelectedData & s1, const SelectedData & s2) { return s1.posInDof < s2.posInDof; });
}

int JointsSelector::dim()
{
  return hl_->dim();
}

void JointsSelector::update(const std::vector<rbd::MultiBody> & mbs,
                            const std::vector<rbd::MultiBodyConfig> & mbcs,
                            const SolverData & data)
{
  hl_->update(mbs, mbcs, data);
  const Eigen::MatrixXd & jac = hl_->jac();
  for(SelectedData sd : selectedJoints_)
  {
    jac_.block(0, sd.posInDof, jac_.rows(), sd.dof) = jac.block(0, sd.posInDof, jac_.rows(), sd.dof);
  }
}

const Eigen::MatrixXd & JointsSelector::jac()
{
  return jac_;
}

const Eigen::VectorXd & JointsSelector::eval()
{
  return hl_->eval();
}

const Eigen::VectorXd & JointsSelector::speed()
{
  return hl_->speed();
}

const Eigen::VectorXd & JointsSelector::normalAcc()
{
  return hl_->normalAcc();
}

/** Torque Task **/
TorqueTask::TorqueTask(const std::vector<rbd::MultiBody> & mbs, int robotIndex, const TorqueBound & tb, double weight)
: TorqueTask(mbs, robotIndex, tb, TorqueDBound{}, 0, weight)
{
}

TorqueTask::TorqueTask(const std::vector<rbd::MultiBody> & mbs,
                       int robotIndex,
                       const TorqueBound & tb,
                       const Eigen::VectorXd & jointSelect,
                       double weight)
: TorqueTask(mbs, robotIndex, tb, TorqueDBound{}, 0, jointSelect, weight)
{
}

TorqueTask::TorqueTask(const std::vector<rbd::MultiBody> & mbs,
                       int robotIndex,
                       const TorqueBound & tb,
                       const std::string & efName,
                       double weight)
: TorqueTask(mbs, robotIndex, tb, TorqueDBound{}, 0, efName, weight)
{
}

TorqueTask::TorqueTask(const std::vector<rbd::MultiBody> & mbs,
                       int robotIndex,
                       const TorqueBound & tb,
                       const TorqueDBound & tdb,
                       double dt,
                       double weight)
: Task(weight), robotIndex_(robotIndex), alphaDBegin_(-1), lambdaBegin_(-1), motionConstr(mbs, robotIndex, tb, tdb, dt),
  jointSelector_(mbs[robotIndex].nrDof()), Q_(mbs[robotIndex].nrDof(), mbs[robotIndex].nrDof()),
  C_(mbs[robotIndex].nrDof())
{
  jointSelector_.setOnes();
}

TorqueTask::TorqueTask(const std::vector<rbd::MultiBody> & mbs,
                       int robotIndex,
                       const TorqueBound & tb,
                       const TorqueDBound & tdb,
                       double dt,
                       const Eigen::VectorXd & jointSelect,
                       double weight)
: Task(weight), robotIndex_(robotIndex), alphaDBegin_(-1), lambdaBegin_(-1), motionConstr(mbs, robotIndex, tb, tdb, dt),
  jointSelector_(jointSelect), Q_(mbs[robotIndex].nrDof(), mbs[robotIndex].nrDof()), C_(mbs[robotIndex].nrDof())
{
}

TorqueTask::TorqueTask(const std::vector<rbd::MultiBody> & mbs,
                       int robotIndex,
                       const TorqueBound & tb,
                       const TorqueDBound & tdb,
                       double dt,
                       const std::string & efName,
                       double weight)
: Task(weight), robotIndex_(robotIndex), alphaDBegin_(-1), lambdaBegin_(-1), motionConstr(mbs, robotIndex, tb, tdb, dt),
  jointSelector_(mbs[robotIndex].nrDof()), Q_(mbs[robotIndex].nrDof(), mbs[robotIndex].nrDof()),
  C_(mbs[robotIndex].nrDof())
{
  rbd::Jacobian jac(mbs[robotIndex], efName);
  jointSelector_.setZero();
  for(auto i : jac.jointsPath())
  {
    // Do not add root joint !
    if(i != 0)
    {
      jointSelector_.segment(mbs[robotIndex].jointPosInDof(i), mbs[robotIndex].joint(i).dof()).setOnes();
    }
  }
}

void TorqueTask::updateNrVars(const std::vector<rbd::MultiBody> & mbs, const SolverData & data)
{
  motionConstr.updateNrVars(mbs, data);
  alphaDBegin_ = data.alphaDBegin(robotIndex_);
  lambdaBegin_ = data.lambdaBegin();
  Q_.resize(data.nrVars(), data.nrVars());
  C_.resize(data.nrVars());
}

void TorqueTask::update(const std::vector<rbd::MultiBody> & mbs,
                        const std::vector<rbd::MultiBodyConfig> & mbcs,
                        const SolverData & data)
{
  motionConstr.update(mbs, mbcs, data);
  Q_.noalias() = motionConstr.matrix().transpose() * jointSelector_.asDiagonal() * motionConstr.matrix();
  C_.noalias() = motionConstr.fd().C().transpose() * jointSelector_.asDiagonal() * motionConstr.matrix();
  // C_.setZero();
}

/**
 *												PostureTask
 */

PostureTask::PostureTask(const std::vector<rbd::MultiBody> & mbs,
                         int rI,
                         std::vector<std::vector<double>> q,
                         double stiffness,
                         double weight)
: Task(weight), pt_(mbs[rI], q), stiffness_(stiffness), damping_(2. * std::sqrt(stiffness)), robotIndex_(rI),
  alphaDBegin_(0), jointDatas_(), Q_(mbs[rI].nrDof(), mbs[rI].nrDof()), C_(mbs[rI].nrDof()), alphaVec_(mbs[rI].nrDof()),
  refVel_(mbs[rI].nrDof()), refAccel_(mbs[rI].nrDof())
{
  refVel_.setZero();
  refAccel_.setZero();
}

void PostureTask::stiffness(double stiffness)
{
  stiffness_ = stiffness;
  damping_ = 2. * std::sqrt(stiffness);
}

void PostureTask::gains(double stiffness)
{
  stiffness_ = stiffness;
  damping_ = 2. * std::sqrt(stiffness);
}

void PostureTask::gains(double stiffness, double damping)
{
  stiffness_ = stiffness;
  damping_ = damping;
}

void PostureTask::jointsStiffness(const std::vector<rbd::MultiBody> & mbs, const std::vector<JointStiffness> & jsv)
{
  jointDatas_.clear();
  jointDatas_.reserve(jsv.size());

  const rbd::MultiBody & mb = mbs[robotIndex_];
  for(const JointStiffness & js : jsv)
  {
    int jointIndex = mb.jointIndexByName(js.jointName);
    jointDatas_.push_back(
        {js.stiffness, 2. * std::sqrt(js.stiffness), mb.jointPosInDof(jointIndex), mb.joint(jointIndex).dof()});
  }
}

void PostureTask::jointsGains(const std::vector<rbd::MultiBody> & mbs, const std::vector<JointGains> & jgv)
{
  jointDatas_.clear();
  jointDatas_.reserve(jgv.size());

  const rbd::MultiBody & mb = mbs[robotIndex_];
  for(const JointGains & jg : jgv)
  {
    int jointIndex = mb.jointIndexByName(jg.jointName);
    jointDatas_.push_back({jg.stiffness, jg.damping, mb.jointPosInDof(jointIndex), mb.joint(jointIndex).dof()});
  }
}

void PostureTask::updateNrVars(const std::vector<rbd::MultiBody> & /* mbs */, const SolverData & data)
{
  alphaDBegin_ = data.alphaDBegin(robotIndex_);
}

void PostureTask::update(const std::vector<rbd::MultiBody> & mbs,
                         const std::vector<rbd::MultiBodyConfig> & mbcs,
                         const SolverData & /* data */)
{
  const rbd::MultiBody & mb = mbs[robotIndex_];
  const rbd::MultiBodyConfig & mbc = mbcs[robotIndex_];

  pt_.update(mb, mbc);
  rbd::paramToVector(mbc.alpha, alphaVec_);

  Q_ = pt_.jac();
  C_.setZero();

  int deb = mb.jointPosInDof(1);
  int end = mb.nrDof() - deb;
  // joint
  C_.segment(deb, end) = -stiffness_ * pt_.eval().segment(deb, end)
                         + damping_ * (alphaVec_.segment(deb, end) - refVel_.segment(deb, end))
                         - refAccel_.segment(deb, end);

  for(const JointData & pjd : jointDatas_)
  {
    C_.segment(pjd.start, pjd.size) =
        -pjd.stiffness * pt_.eval().segment(pjd.start, pjd.size)
        + pjd.damping * (alphaVec_.segment(pjd.start, pjd.size) - refVel_.segment(pjd.start, pjd.size))
        - refAccel_.segment(pjd.start, pjd.size);
  }
}

const Eigen::MatrixXd & PostureTask::Q() const
{
  return Q_;
}

const Eigen::VectorXd & PostureTask::C() const
{
  return C_;
}

const Eigen::VectorXd & PostureTask::eval() const
{
  return pt_.eval();
}

/**
 *											PositionTask
 */

PositionTask::PositionTask(const std::vector<rbd::MultiBody> & mbs,
                           int rI,
                           const std::string & bodyName,
                           const Eigen::Vector3d & pos,
                           const Eigen::Vector3d & bodyPoint)
: pt_(mbs[rI], bodyName, pos, bodyPoint), robotIndex_(rI)
{
}

int PositionTask::dim()
{
  return 3;
}

void PositionTask::update(const std::vector<rbd::MultiBody> & mbs,
                          const std::vector<rbd::MultiBodyConfig> & mbcs,
                          const SolverData & data)
{
  pt_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & PositionTask::jac()
{
  return pt_.jac();
}

const Eigen::VectorXd & PositionTask::eval()
{
  return pt_.eval();
}

const Eigen::VectorXd & PositionTask::speed()
{
  return pt_.speed();
}

const Eigen::VectorXd & PositionTask::normalAcc()
{
  return pt_.normalAcc();
}

/**
 *																OrientationTask
 */

OrientationTask::OrientationTask(const std::vector<rbd::MultiBody> & mbs,
                                 int rI,
                                 const std::string & bodyName,
                                 const Eigen::Quaterniond & ori)
: ot_(mbs[rI], bodyName, ori), robotIndex_(rI)
{
}

OrientationTask::OrientationTask(const std::vector<rbd::MultiBody> & mbs,
                                 int rI,
                                 const std::string & bodyName,
                                 const Eigen::Matrix3d & ori)
: ot_(mbs[rI], bodyName, ori), robotIndex_(rI)
{
}

int OrientationTask::dim()
{
  return 3;
}

void OrientationTask::update(const std::vector<rbd::MultiBody> & mbs,
                             const std::vector<rbd::MultiBodyConfig> & mbcs,
                             const SolverData & data)
{
  ot_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & OrientationTask::jac()
{
  return ot_.jac();
}

const Eigen::VectorXd & OrientationTask::eval()
{
  return ot_.eval();
}

const Eigen::VectorXd & OrientationTask::speed()
{
  return ot_.speed();
}

const Eigen::VectorXd & OrientationTask::normalAcc()
{
  return ot_.normalAcc();
}

/**
 *											SurfaceTransformTask
 */

SurfaceTransformTask::SurfaceTransformTask(const std::vector<rbd::MultiBody> & mbs,
                                           int robotIndex,
                                           const std::string & bodyName,
                                           const sva::PTransformd & X_0_t,
                                           const sva::PTransformd & X_b_p)
: TransformTaskCommon(mbs, robotIndex, bodyName, X_0_t, X_b_p)
{
}

void SurfaceTransformTask::update(const std::vector<rbd::MultiBody> & mbs,
                                  const std::vector<rbd::MultiBodyConfig> & mbcs,
                                  const SolverData & data)
{
  tt_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

/**
 *											TransformTask
 */

TransformTask::TransformTask(const std::vector<rbd::MultiBody> & mbs,
                             int robotIndex,
                             const std::string & bodyName,
                             const sva::PTransformd & X_0_t,
                             const sva::PTransformd & X_b_p,
                             const Eigen::Matrix3d & E_0_c)
: TransformTaskCommon(mbs, robotIndex, bodyName, X_0_t, X_b_p)
{
  tt_.E_0_c(E_0_c);
}

void TransformTask::E_0_c(const Eigen::Matrix3d & E_0_c)
{
  tt_.E_0_c(E_0_c);
}

const Eigen::Matrix3d & TransformTask::E_0_c() const
{
  return tt_.E_0_c();
}

void TransformTask::update(const std::vector<rbd::MultiBody> & mbs,
                           const std::vector<rbd::MultiBodyConfig> & mbcs,
                           const SolverData & data)
{
  tt_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

/**
 *																SurfaceOrientationTask
 */

SurfaceOrientationTask::SurfaceOrientationTask(const std::vector<rbd::MultiBody> & mbs,
                                               int rI,
                                               const std::string & bodyName,
                                               const Eigen::Quaterniond & ori,
                                               const sva::PTransformd & X_b_s)
: ot_(mbs[rI], bodyName, ori, X_b_s), robotIndex_(rI)
{
}

SurfaceOrientationTask::SurfaceOrientationTask(const std::vector<rbd::MultiBody> & mbs,
                                               int rI,
                                               const std::string & bodyName,
                                               const Eigen::Matrix3d & ori,
                                               const sva::PTransformd & X_b_s)
: ot_(mbs[rI], bodyName, ori, X_b_s), robotIndex_(rI)
{
}

int SurfaceOrientationTask::dim()
{
  return 3;
}

void SurfaceOrientationTask::update(const std::vector<rbd::MultiBody> & mbs,
                                    const std::vector<rbd::MultiBodyConfig> & mbcs,
                                    const SolverData & data)
{
  ot_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & SurfaceOrientationTask::jac()
{
  return ot_.jac();
}

const Eigen::VectorXd & SurfaceOrientationTask::eval()
{
  return ot_.eval();
}

const Eigen::VectorXd & SurfaceOrientationTask::speed()
{
  return ot_.speed();
}

const Eigen::VectorXd & SurfaceOrientationTask::normalAcc()
{
  return ot_.normalAcc();
}

/**
 *																GazeTask
 */

GazeTask::GazeTask(const std::vector<rbd::MultiBody> & mbs,
                   int robotIndex,
                   const std::string & bodyName,
                   const Eigen::Vector2d & point2d,
                   double depthEstimate,
                   const sva::PTransformd & X_b_gaze,
                   const Eigen::Vector2d & point2d_ref)
: gazet_(mbs[robotIndex], bodyName, point2d, depthEstimate, X_b_gaze, point2d_ref), robotIndex_(robotIndex)
{
}

GazeTask::GazeTask(const std::vector<rbd::MultiBody> & mbs,
                   int robotIndex,
                   const std::string & bodyName,
                   const Eigen::Vector3d & point3d,
                   const sva::PTransformd & X_b_gaze,
                   const Eigen::Vector2d & point2d_ref)
: gazet_(mbs[robotIndex], bodyName, point3d, X_b_gaze, point2d_ref), robotIndex_(robotIndex)
{
}

int GazeTask::dim()
{
  return 2;
}

void GazeTask::update(const std::vector<rbd::MultiBody> & mbs,
                      const std::vector<rbd::MultiBodyConfig> & mbcs,
                      const SolverData & data)
{
  gazet_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & GazeTask::jac()
{
  return gazet_.jac();
}

const Eigen::VectorXd & GazeTask::eval()
{
  return gazet_.eval();
}

const Eigen::VectorXd & GazeTask::speed()
{
  return gazet_.speed();
}

const Eigen::VectorXd & GazeTask::normalAcc()
{
  return gazet_.normalAcc();
}

/**
 *																PositionBasedVisServoTask
 */

PositionBasedVisServoTask::PositionBasedVisServoTask(const std::vector<rbd::MultiBody> & mbs,
                                                     int robotIndex,
                                                     const std::string & bodyName,
                                                     const sva::PTransformd & X_t_s,
                                                     const sva::PTransformd & X_b_s)
: pbvst_(mbs[robotIndex], bodyName, X_t_s, X_b_s), robotIndex_(robotIndex)
{
}

int PositionBasedVisServoTask::dim()
{
  return 6;
}

void PositionBasedVisServoTask::update(const std::vector<rbd::MultiBody> & mbs,
                                       const std::vector<rbd::MultiBodyConfig> & mbcs,
                                       const SolverData & data)
{
  pbvst_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & PositionBasedVisServoTask::jac()
{
  return pbvst_.jac();
}

const Eigen::VectorXd & PositionBasedVisServoTask::eval()
{
  return pbvst_.eval();
}

const Eigen::VectorXd & PositionBasedVisServoTask::speed()
{
  return pbvst_.speed();
}

const Eigen::VectorXd & PositionBasedVisServoTask::normalAcc()
{
  return pbvst_.normalAcc();
}

/**
 *													CoMTask
 */

CoMTask::CoMTask(const std::vector<rbd::MultiBody> & mbs, int rI, const Eigen::Vector3d & com)
: ct_(mbs[rI], com), robotIndex_(rI)
{
}

CoMTask::CoMTask(const std::vector<rbd::MultiBody> & mbs,
                 int rI,
                 const Eigen::Vector3d & com,
                 std::vector<double> weight)
: ct_(mbs[rI], com, std::move(weight)), robotIndex_(rI)
{
}

void CoMTask::updateInertialParameters(const std::vector<rbd::MultiBody> & mbs)
{
  ct_.updateInertialParameters(mbs[robotIndex_]);
}

int CoMTask::dim()
{
  return 3;
}

void CoMTask::update(const std::vector<rbd::MultiBody> & mbs,
                     const std::vector<rbd::MultiBodyConfig> & mbcs,
                     const SolverData & data)
{
  ct_.update(mbs[robotIndex_], mbcs[robotIndex_], rbd::computeCoM(mbs[robotIndex_], mbcs[robotIndex_]),
             data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & CoMTask::jac()
{
  return ct_.jac();
}

const Eigen::VectorXd & CoMTask::eval()
{
  return ct_.eval();
}

const Eigen::VectorXd & CoMTask::speed()
{
  return ct_.speed();
}

const Eigen::VectorXd & CoMTask::normalAcc()
{
  return ct_.normalAcc();
}

/**
 *													MultiCoMTask
 */

MultiCoMTask::MultiCoMTask(const std::vector<rbd::MultiBody> & mbs,
                           std::vector<int> rI,
                           const Eigen::Vector3d & com,
                           double stiffness,
                           double weight)
: Task(weight), alphaDBegin_(-1), stiffness_(stiffness), stiffnessSqrt_(2. * std::sqrt(stiffness)),
  dimWeight_(Eigen::Vector3d::Ones()), posInQ_(rI.size()), mct_(mbs, std::move(rI), com), Q_(), C_(), CSum_(), preQ_()
{
  init(mbs);
}

MultiCoMTask::MultiCoMTask(const std::vector<rbd::MultiBody> & mbs,
                           std::vector<int> rI,
                           const Eigen::Vector3d & com,
                           double stiffness,
                           const Eigen::Vector3d & dimWeight,
                           double weight)
: Task(weight), alphaDBegin_(-1), stiffness_(stiffness), stiffnessSqrt_(2. * std::sqrt(stiffness)),
  dimWeight_(dimWeight), posInQ_(rI.size()), mct_(mbs, std::move(rI), com), Q_(), C_(), CSum_(), preQ_()
{
  init(mbs);
}

void MultiCoMTask::updateInertialParameters(const std::vector<rbd::MultiBody> & mbs)
{
  mct_.updateInertialParameters(mbs);
}

void MultiCoMTask::stiffness(double stiffness)
{
  stiffness_ = stiffness;
  stiffnessSqrt_ = 2. * std::sqrt(stiffness);
}

void MultiCoMTask::dimWeight(const Eigen::Vector3d & dim)
{
  dimWeight_ = dim;
}

void MultiCoMTask::updateNrVars(const std::vector<rbd::MultiBody> & /* mbs */, const SolverData & data)
{
  auto minMaxIndex = std::minmax_element(mct_.robotIndexes().begin(), mct_.robotIndexes().end());
  alphaDBegin_ = data.alphaDBegin(*(minMaxIndex.first));
  int lastBegin = data.alphaDBegin(*(minMaxIndex.second));
  int lastAlphaD = data.alphaD(*(minMaxIndex.second));
  int size = lastBegin + lastAlphaD - alphaDBegin_;

  Q_.setZero(size, size);
  C_.setZero(size);

  posInQ_.clear();
  for(int r : mct_.robotIndexes())
  {
    posInQ_.push_back(data.alphaDBegin(r) - alphaDBegin_);
  }
}

void MultiCoMTask::update(const std::vector<rbd::MultiBody> & mbs,
                          const std::vector<rbd::MultiBodyConfig> & mbcs,
                          const SolverData & data)
{
  mct_.update(mbs, mbcs, data.normalAccB());
  CSum_ = stiffness_ * mct_.eval();
  CSum_ -= stiffnessSqrt_ * mct_.speed();
  CSum_ -= mct_.normalAcc();
  for(int i = 0; i < int(posInQ_.size()); ++i)
  {
    int r = mct_.robotIndexes()[i];
    int begin = posInQ_[i];
    int dof = data.alphaD(r);

    const Eigen::MatrixXd & J = mct_.jac(i);
    preQ_.block(0, 0, 3, dof).noalias() = dimWeight_.asDiagonal() * J;

    Q_.block(begin, begin, dof, dof).noalias() = J.transpose() * preQ_.block(0, 0, 3, dof);
    C_.segment(begin, dof).noalias() = -J.transpose() * dimWeight_.asDiagonal() * CSum_;
  }
}

const Eigen::MatrixXd & MultiCoMTask::Q() const
{
  return Q_;
}

const Eigen::VectorXd & MultiCoMTask::C() const
{
  return C_;
}

const Eigen::VectorXd & MultiCoMTask::eval() const
{
  return mct_.eval();
}

const Eigen::VectorXd & MultiCoMTask::speed() const
{
  return mct_.speed();
}

void MultiCoMTask::init(const std::vector<rbd::MultiBody> & mbs)
{
  int maxDof = 0;
  for(int r : mct_.robotIndexes())
  {
    maxDof = std::max(maxDof, mbs[r].nrDof());
  }
  preQ_.resize(3, maxDof);
}

/**
 *													MultiRobotTransformTask
 */

MultiRobotTransformTask::MultiRobotTransformTask(const std::vector<rbd::MultiBody> & mbs,
                                                 int r1Index,
                                                 int r2Index,
                                                 const std::string & r1BodyName,
                                                 const std::string & r2BodyName,
                                                 const sva::PTransformd & X_r1b_r1s,
                                                 const sva::PTransformd & X_r2b_r2s,
                                                 double stiffness,
                                                 double weight)
: Task(weight), alphaDBegin_(-1), stiffness_(stiffness), stiffnessSqrt_(2. * std::sqrt(stiffness)),
  dimWeight_(Eigen::Vector6d::Ones()), posInQ_(2, -1), robotIndexes_{{r1Index, r2Index}},
  mrtt_(mbs, r1Index, r2Index, r1BodyName, r2BodyName, X_r1b_r1s, X_r2b_r2s), Q_(), C_(),
  CSum_(Eigen::Vector6d::Zero()), preQ_()
{
  int maxDof = 0;
  for(int r : robotIndexes_)
  {
    maxDof = std::max(maxDof, mbs[r].nrDof());
  }
  preQ_.resize(6, maxDof);
}

void MultiRobotTransformTask::X_r1b_r1s(const sva::PTransformd & X_r1b_r1s)
{
  mrtt_.X_r1b_r1s(X_r1b_r1s);
}

const sva::PTransformd & MultiRobotTransformTask::X_r1b_r1s() const
{
  return mrtt_.X_r1b_r1s();
}

void MultiRobotTransformTask::X_r2b_r2s(const sva::PTransformd & X_r2b_r2s)
{
  mrtt_.X_r2b_r2s(X_r2b_r2s);
}

const sva::PTransformd & MultiRobotTransformTask::X_r2b_r2s() const
{
  return mrtt_.X_r2b_r2s();
}

void MultiRobotTransformTask::stiffness(double stiffness)
{
  stiffness_ = stiffness;
  stiffnessSqrt_ = 2. * std::sqrt(stiffness);
}

void MultiRobotTransformTask::dimWeight(const Eigen::Vector6d & dim)
{
  dimWeight_ = dim;
}

void MultiRobotTransformTask::updateNrVars(const std::vector<rbd::MultiBody> & /* mbs */, const SolverData & data)
{
  auto minMaxIndex = std::minmax_element(robotIndexes_.begin(), robotIndexes_.end());
  alphaDBegin_ = data.alphaDBegin(*(minMaxIndex.first));
  int lastBegin = data.alphaDBegin(*(minMaxIndex.second));
  int lastAlphaD = data.alphaD(*(minMaxIndex.second));
  int size = lastBegin + lastAlphaD - alphaDBegin_;

  Q_.setZero(size, size);
  C_.setZero(size);

  posInQ_.clear();
  for(int r : robotIndexes_)
  {
    posInQ_.push_back(data.alphaDBegin(r) - alphaDBegin_);
  }
}

void MultiRobotTransformTask::update(const std::vector<rbd::MultiBody> & mbs,
                                     const std::vector<rbd::MultiBodyConfig> & mbcs,
                                     const SolverData & data)
{
  mrtt_.update(mbs, mbcs, data.normalAccB());
  CSum_.noalias() = stiffness_ * mrtt_.eval();
  CSum_.noalias() -= stiffnessSqrt_ * mrtt_.speed();
  CSum_.noalias() -= mrtt_.normalAcc();

  // first we set to zero used part of Q and C
  for(int i = 0; i < int(posInQ_.size()); ++i)
  {
    int r = robotIndexes_[i];
    int begin = posInQ_[i];
    int dof = data.alphaD(r);
    Q_.block(begin, begin, dof, dof).setZero();
    C_.segment(begin, dof).setZero();
  }

  for(int i = 0; i < int(posInQ_.size()); ++i)
  {
    int r = robotIndexes_[i];
    int begin = posInQ_[i];
    int dof = data.alphaD(r);

    const Eigen::MatrixXd & J = mrtt_.jac(i);
    preQ_.block(0, 0, 6, dof).noalias() = dimWeight_.asDiagonal() * J;

    // scince the two robot index could be the same
    // we had to increment the Q and C matrix
    Q_.block(begin, begin, dof, dof).noalias() += J.transpose() * preQ_.block(0, 0, 6, dof);
    C_.segment(begin, dof).noalias() -= J.transpose() * dimWeight_.asDiagonal() * CSum_;
  }
}

const Eigen::MatrixXd & MultiRobotTransformTask::Q() const
{
  return Q_;
}

const Eigen::VectorXd & MultiRobotTransformTask::C() const
{
  return C_;
}

const Eigen::VectorXd & MultiRobotTransformTask::eval() const
{
  return mrtt_.eval();
}

const Eigen::VectorXd & MultiRobotTransformTask::speed() const
{
  return mrtt_.speed();
}

/**
 *													MomentumTask
 */

MomentumTask::MomentumTask(const std::vector<rbd::MultiBody> & mbs, int rI, const sva::ForceVecd & mom)
: momt_(mbs[rI], mom), robotIndex_(rI)
{
}

int MomentumTask::dim()
{
  return 6;
}

void MomentumTask::update(const std::vector<rbd::MultiBody> & mbs,
                          const std::vector<rbd::MultiBodyConfig> & mbcs,
                          const SolverData & data)
{
  momt_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & MomentumTask::jac()
{
  return momt_.jac();
}

const Eigen::VectorXd & MomentumTask::eval()
{
  return momt_.eval();
}

const Eigen::VectorXd & MomentumTask::speed()
{
  return momt_.speed();
}

const Eigen::VectorXd & MomentumTask::normalAcc()
{
  return momt_.normalAcc();
}

/**
 *														ContactTask
 */

void ContactTask::error(const Eigen::Vector3d & error)
{
  error_ = error;
}

void ContactTask::errorD(const Eigen::Vector3d & errorD)
{
  errorD_ = errorD;
}

void ContactTask::updateNrVars(const std::vector<rbd::MultiBody> & /* mbs */, const SolverData & data)
{
  int nrLambda = 0;
  begin_ = data.lambdaBegin();
  std::vector<FrictionCone> cones;
  int curLambda = 0;

  if(nrLambda == 0)
  {
    for(const BilateralContact & uc : data.allContacts())
    {
      curLambda = uc.nrLambda();
      if(uc.contactId == contactId_)
      {
        nrLambda = curLambda;
        cones = uc.r1Cones;
        break;
      }

      begin_ += curLambda;
    }
  }

  conesJac_.resize(3, nrLambda);
  int index = 0;
  for(const FrictionCone & fc : cones)
  {
    for(const Eigen::Vector3d & gen : fc.generators)
    {
      conesJac_.col(index) = gen;
      ++index;
    }
  }

  Q_.resize(nrLambda, nrLambda);
  Q_.noalias() = conesJac_.transpose() * conesJac_;
  C_.setZero(nrLambda);
}

void ContactTask::update(const std::vector<rbd::MultiBody> & /* mbs */,
                         const std::vector<rbd::MultiBodyConfig> & /* mbcs */,
                         const SolverData & /* data */)
{
  /*C_.noalias() = -conesJac_.transpose()*
      (stiffness_*error_ - stiffnessSqrt_*errorD_);*/
  C_.noalias() = -conesJac_.transpose() * error_;
}

const Eigen::MatrixXd & ContactTask::Q() const
{
  return Q_;
}

const Eigen::VectorXd & ContactTask::C() const
{
  return C_;
}

/**
 *														GripperTorqueTask
 */

void GripperTorqueTask::updateNrVars(const std::vector<rbd::MultiBody> & /* mbs */, const SolverData & data)
{
  using namespace Eigen;
  bool found = false;

  begin_ = data.bilateralBegin();
  for(const BilateralContact & bc : data.bilateralContacts())
  {
    int curLambda = 0;
    // compute the number of lambda needed by the current bilateral
    for(std::size_t i = 0; i < bc.r1Points.size(); ++i)
    {
      curLambda += bc.nrLambda(static_cast<int>(i));
    }

    if(bc.contactId == contactId_)
    {
      found = true;
      Q_.setZero(curLambda, curLambda);
      C_.resize(curLambda);

      int pos = 0;
      // minimize Torque applied on the gripper motor
      // min Sum_i^nrF  T_i·( p_i^T_o x f_i)
      for(std::size_t i = 0; i < bc.r1Cones.size(); ++i)
      {
        Vector3d T_o_p = bc.r1Points[i] - origin_;
        for(std::size_t j = 0; j < bc.r1Cones[i].generators.size(); ++j)
        {
          // we use abs because the contact force cannot apply
          // negative torque on the gripper
          C_(pos) = std::abs(axis_.transpose() * (T_o_p.cross(bc.r1Cones[i].generators[j])));
          ++pos;
        }
      }
      break;
    }

    begin_ += curLambda;
  }

  // if no contact was found we don't activate the task
  // (safe position and empty matrix)
  if(!found)
  {
    begin_ = 0;
    Q_.resize(0, 0);
    C_.resize(0);
  }
}

void GripperTorqueTask::update(const std::vector<rbd::MultiBody> & /* mbs */,
                               const std::vector<rbd::MultiBodyConfig> & /* mbcs */,
                               const SolverData & /* data */)
{
}

const Eigen::MatrixXd & GripperTorqueTask::Q() const
{
  return Q_;
}

const Eigen::VectorXd & GripperTorqueTask::C() const
{
  return C_;
}

/**
 *											LinVelocityTask
 */

LinVelocityTask::LinVelocityTask(const std::vector<rbd::MultiBody> & mbs,
                                 int rI,
                                 const std::string & bodyName,
                                 const Eigen::Vector3d & speed,
                                 const Eigen::Vector3d & bodyPoint)
: pt_(mbs[rI], bodyName, speed, bodyPoint), robotIndex_(rI)
{
}

int LinVelocityTask::dim()
{
  return 3;
}

void LinVelocityTask::update(const std::vector<rbd::MultiBody> & mbs,
                             const std::vector<rbd::MultiBodyConfig> & mbcs,
                             const SolverData & data)
{
  pt_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & LinVelocityTask::jac()
{
  return pt_.jac();
}

const Eigen::VectorXd & LinVelocityTask::eval()
{
  return pt_.eval();
}

const Eigen::VectorXd & LinVelocityTask::speed()
{
  return pt_.speed();
}

const Eigen::VectorXd & LinVelocityTask::normalAcc()
{
  return pt_.normalAcc();
}

/**
 *											OrientationTrackingTask
 */

OrientationTrackingTask::OrientationTrackingTask(const std::vector<rbd::MultiBody> & mbs,
                                                 int rI,
                                                 const std::string & bodyName,
                                                 const Eigen::Vector3d & bodyPoint,
                                                 const Eigen::Vector3d & bodyAxis,
                                                 const std::vector<std::string> & trackingJointsName,
                                                 const Eigen::Vector3d & trackedPoint)
: robotIndex_(rI), ott_(mbs[rI], bodyName, bodyPoint, bodyAxis, trackingJointsName, trackedPoint),
  alphaVec_(mbs[rI].nrDof()), speed_(3), normalAcc_(3)
{
}

int OrientationTrackingTask::dim()
{
  return 3;
}

void OrientationTrackingTask::update(const std::vector<rbd::MultiBody> & mbs,
                                     const std::vector<rbd::MultiBodyConfig> & mbcs,
                                     const SolverData & /* data */)
{
  ott_.update(mbs[robotIndex_], mbcs[robotIndex_]);
  rbd::paramToVector(mbcs[robotIndex_].alpha, alphaVec_);

  speed_.noalias() = ott_.jac() * alphaVec_;
  normalAcc_.noalias() = ott_.jacDot() * alphaVec_;
}

const Eigen::MatrixXd & OrientationTrackingTask::jac()
{
  return ott_.jac();
}

const Eigen::VectorXd & OrientationTrackingTask::eval()
{
  return ott_.eval();
}

const Eigen::VectorXd & OrientationTrackingTask::speed()
{
  return speed_;
}

const Eigen::VectorXd & OrientationTrackingTask::normalAcc()
{
  return normalAcc_;
}

/**
 *											RelativeDistTask
 */

RelativeDistTask::RelativeDistTask(const std::vector<rbd::MultiBody> & mbs,
                                   const int rIndex,
                                   const double timestep,
                                   tasks::RelativeDistTask::rbInfo & rbi1,
                                   tasks::RelativeDistTask::rbInfo & rbi2,
                                   const Eigen::Vector3d & u1,
                                   const Eigen::Vector3d & u2)
: rIndex_(rIndex), rdt_(mbs[rIndex], timestep, rbi1, rbi2, u1, u2)
{
}

int RelativeDistTask::dim()
{
  return 1;
}

void RelativeDistTask::update(const std::vector<rbd::MultiBody> & mbs,
                              const std::vector<rbd::MultiBodyConfig> & mbcs,
                              const SolverData & data)
{
  rdt_.update(mbs[rIndex_], mbcs[rIndex_], data.normalAccB(rIndex_));
}

const Eigen::MatrixXd & RelativeDistTask::jac()
{
  return rdt_.jac();
}

const Eigen::VectorXd & RelativeDistTask::eval()
{
  return rdt_.eval();
}

const Eigen::VectorXd & RelativeDistTask::speed()
{
  return rdt_.speed();
}

const Eigen::VectorXd & RelativeDistTask::normalAcc()
{
  return rdt_.normalAcc();
}

/**
 *											VectorOrientationTask
 */

VectorOrientationTask::VectorOrientationTask(const std::vector<rbd::MultiBody> & mbs,
                                             int rI,
                                             const std::string & bodyName,
                                             const Eigen::Vector3d & bodyVector,
                                             const Eigen::Vector3d & targetVector)
: vot_(mbs[rI], bodyName, bodyVector, targetVector), robotIndex_(rI)
{
}

int VectorOrientationTask::dim()
{
  return 3;
}

void VectorOrientationTask::update(const std::vector<rbd::MultiBody> & mbs,
                                   const std::vector<rbd::MultiBodyConfig> & mbcs,
                                   const SolverData & data)
{
  vot_.update(mbs[robotIndex_], mbcs[robotIndex_], data.normalAccB(robotIndex_));
}

const Eigen::MatrixXd & VectorOrientationTask::jac()
{
  return vot_.jac();
}

const Eigen::VectorXd & VectorOrientationTask::eval()
{
  return vot_.eval();
}

const Eigen::VectorXd & VectorOrientationTask::speed()
{
  return vot_.speed();
}

const Eigen::VectorXd & VectorOrientationTask::normalAcc()
{
  return vot_.normalAcc();
}

} // namespace qp

} // namespace tasks


/*
@copyright Louis Dionne 2015
Distributed under the Boost Software License, Version 1.0.
(See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
 */

#include <boost/hana.hpp>

#include <boost/hana/assert.hpp>
#include <boost/hana/concept/ring.hpp>
#include <boost/hana/tuple.hpp>

#include <laws/ring.hpp>
using namespace boost::hana;

int main() {
    test::TestRing<int>{make<tuple_tag>(0,1,2,3,4,5)};
    test::TestRing<long>{make<tuple_tag>(0l,1l,2l,3l,4l,5l)};

    // one
    static_assert(one<int>() == 1, "");

    // mult
    static_assert(mult(6, 4) == 6 * 4, "");
}


#include "pair_flare.h"
#include "atom.h"
#include "comm.h"
#include "error.h"
#include "force.h"
#include "memory.h"
#include "neigh_list.h"
#include "neigh_request.h"
#include "neighbor.h"
#include <Eigen/Dense>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <vector>
#include <sys/time.h>

// flare++ modules
#include "cutoffs.h"
#include "lammps_descriptor.h"
#include "radial.h"
#include "y_grad.h"

using namespace LAMMPS_NS;

#define MAXLINE 1024

typedef unsigned long long timestamp_t;

static timestamp_t
get_timestamp ()
{
  struct timeval now;
  gettimeofday (&now, NULL);
  return  now.tv_usec + (timestamp_t)now.tv_sec * 1000000;
}

/* ---------------------------------------------------------------------- */

PairFLARE::PairFLARE(LAMMPS *lmp) : Pair(lmp) {
  restartinfo = 0;
  manybody_flag = 1;

  beta = NULL;
}

/* ----------------------------------------------------------------------
   check if allocated, since class can be destructed when incomplete
------------------------------------------------------------------------- */

PairFLARE::~PairFLARE() {
  if (copymode)
    return;

  memory->destroy(beta);

  if (allocated) {
    memory->destroy(setflag);
    memory->destroy(cutsq);
  }
}

/* ---------------------------------------------------------------------- */

void PairFLARE::compute(int eflag, int vflag) {
  int i, j, ii, jj, inum, jnum, itype, jtype, n_inner, n_count;
  double evdwl, delx, dely, delz, xtmp, ytmp, ztmp, rsq;
  double *coeff;
  int *ilist, *jlist, *numneigh, **firstneigh;

  evdwl = 0.0;
  ev_init(eflag, vflag);

  double **x = atom->x;
  double **f = atom->f;
  int *type = atom->type;
  int nlocal = atom->nlocal;
  int nall = nlocal + atom->nghost;
  int newton_pair = force->newton_pair;

  inum = list->inum;
  ilist = list->ilist;
  numneigh = list->numneigh;
  firstneigh = list->firstneigh;

  int beta_init, beta_counter;
  double B2_norm_squared, B2_val_1, B2_val_2;
  Eigen::VectorXd single_bond_vals, B2_vals, B2_env_dot, u;
  Eigen::MatrixXd single_bond_env_dervs, B2_env_dervs;
  double empty_thresh = 1e-8;

  for (ii = 0; ii < inum; ii++) {
    i = list->ilist[ii];
    itype = type[i];
    jnum = numneigh[i];
    xtmp = x[i][0];
    ytmp = x[i][1];
    ztmp = x[i][2];
    jlist = firstneigh[i];

    // Count the atoms inside the cutoff.
    n_inner = 0;
    for (int jj = 0; jj < jnum; jj++) {
      j = jlist[jj];
      int s = type[j] - 1;
      double cutoff_val = cutoff_matrix(itype-1, s);

      delx = x[j][0] - xtmp;
      dely = x[j][1] - ytmp;
      delz = x[j][2] - ztmp;
      rsq = delx * delx + dely * dely + delz * delz;
      if (rsq < (cutoff_val * cutoff_val))
        n_inner++;
    }

    // Compute covariant descriptors.
    double secs;
    single_bond_multiple_cutoffs(x, type, jnum, n_inner, i, xtmp, ytmp, ztmp,
                                 jlist, basis_function, cutoff_function,
                                 n_species, n_max, l_max, radial_hyps,
                                 cutoff_hyps, single_bond_vals,
                                 single_bond_env_dervs, cutoff_matrix);

    // Compute invariant descriptors.
    B2_descriptor(B2_vals, B2_norm_squared,
                  single_bond_vals, n_species, n_max, l_max);

    compute_energy_and_u(B2_vals, B2_norm_squared, single_bond_vals, power,
           n_species, n_max, l_max, beta_matrices[itype - 1], u, &evdwl);

    // Continue if the environment is empty.
    if (B2_norm_squared < empty_thresh)
      continue;

    // Update energy, force and stress arrays.
    n_count = 0;
    double fxsum = 0, fysum = 0, fzsum = 0;
    for (int jj = 0; jj < jnum; jj++) {
      j = jlist[jj];
      int s = type[j] - 1;
      double cutoff_val = cutoff_matrix(itype-1, s);
      delx = xtmp - x[j][0];
      dely = ytmp - x[j][1];
      delz = ztmp - x[j][2];
      rsq = delx * delx + dely * dely + delz * delz;

      if (rsq < (cutoff_val * cutoff_val)) {
        // Compute partial force f_ij = u * dA/dr_ij
        double fx = single_bond_env_dervs.row(n_count * 3).dot(u);
        double fy = single_bond_env_dervs.row(n_count * 3 + 1).dot(u);
        double fz = single_bond_env_dervs.row(n_count * 3 + 2).dot(u);
        // Compute local energy and partial forces.

        f[i][0] += fx;
        f[i][1] += fy;
        f[i][2] += fz;
        f[j][0] -= fx;
        f[j][1] -= fy;
        f[j][2] -= fz;
        fxsum += fx;
        fysum += fy;
        fzsum += fz;
        //printf("i = %d, j = %d, f = %g %g %g\n", i, j, fx, fy, fz);

        if (vflag) {
          ev_tally_xyz(i, j, nlocal, newton_pair, 0.0, 0.0, fx, fy, fz, delx,
                       dely, delz);
        }
        n_count++;
      }
    }
      //printf("i = %d, Fsum = %g %g %g\n", i, fxsum, fysum, fzsum);

    // Compute local energy.
    if (eflag)
      ev_tally_full(i, 2.0 * evdwl, 0.0, 0.0, 0.0, 0.0, 0.0);
  }

  if (vflag_fdotr)
    virial_fdotr_compute();
}

/* ----------------------------------------------------------------------
   allocate all arrays
------------------------------------------------------------------------- */

void PairFLARE::allocate() {
  allocated = 1;
  int n = atom->ntypes;

  memory->create(setflag, n + 1, n + 1, "pair:setflag");

  // Set the diagonal of setflag to 1 (otherwise pair.cpp will throw an error)
  for (int i = 1; i <= n; i++)
    setflag[i][i] = 1;

  // Create cutsq array (used in pair.cpp)
  memory->create(cutsq, n + 1, n + 1, "pair:cutsq");
}

/* ----------------------------------------------------------------------
   global settings
------------------------------------------------------------------------- */

void PairFLARE::settings(int narg, char ** /*arg*/) {
  // "flare" should be the only word after "pair_style" in the input file.
  if (narg > 0)
    error->all(FLERR, "Illegal pair_style command");
}

/* ----------------------------------------------------------------------
   set coeffs for one or more type pairs
   read DYNAMO funcfl file
------------------------------------------------------------------------- */

void PairFLARE::coeff(int narg, char **arg) {
  if (!allocated)
    allocate();

  // Should be exactly 3 arguments following "pair_coeff" in the input file.
  if (narg != 3)
    error->all(FLERR, "Incorrect args for pair coefficients");

  // Ensure I,J args are "* *".
  if (strcmp(arg[0], "*") != 0 || strcmp(arg[1], "*") != 0)
    error->all(FLERR, "Incorrect args for pair coefficients");

  read_file(arg[2]);
}

/* ----------------------------------------------------------------------
   init specific to this pair style
------------------------------------------------------------------------- */

void PairFLARE::init_style() {
  // Require newton on.
  if (force->newton_pair == 0)
    error->all(FLERR, "Pair style requires newton pair on");

  // Request a full neighbor list.
  int irequest = neighbor->request(this, instance_me);
  neighbor->requests[irequest]->half = 0;
  neighbor->requests[irequest]->full = 1;
}

/* ----------------------------------------------------------------------
   init for one type pair i,j and corresponding j,i
------------------------------------------------------------------------- */

double PairFLARE::init_one(int i, int j) {
  // init_one is called for each i, j pair in pair.cpp after calling init_style.

  return cutoff;
}

/* ----------------------------------------------------------------------
   read potential values from a DYNAMO single element funcfl file
------------------------------------------------------------------------- */

void PairFLARE::read_file(char *filename) {
  int me = comm->me;
  char line[MAXLINE], radial_string[MAXLINE], cutoff_string[MAXLINE];
  int radial_string_length, cutoff_string_length;
  FILE *fptr;

  // Check that the potential file can be opened.
  if (me == 0) {
    fptr = utils::open_potential(filename,lmp,nullptr);
    if (fptr == NULL) {
      char str[128];
      snprintf(str, 128, "Cannot open potential file %s", filename);
      error->one(FLERR, str);
    }
  }

  int tmp, nwords;
  if (me == 0) {
    fgets(line, MAXLINE, fptr); // Date and contributor

    fgets(line, MAXLINE, fptr); // Power, use integer instead of double for simplicity
    sscanf(line, "%i", &power);

    fgets(line, MAXLINE, fptr);
    sscanf(line, "%s", radial_string); // Radial basis set
    radial_string_length = strlen(radial_string);

    fgets(line, MAXLINE, fptr);
    sscanf(line, "%i %i %i %i", &n_species, &n_max, &l_max, &beta_size);

    fgets(line, MAXLINE, fptr);
    sscanf(line, "%s", cutoff_string); // Cutoff function
    cutoff_string_length = strlen(cutoff_string);
  }

  MPI_Bcast(&power, 1, MPI_INT, 0, world);
  MPI_Bcast(&n_species, 1, MPI_INT, 0, world);
  MPI_Bcast(&n_max, 1, MPI_INT, 0, world);
  MPI_Bcast(&l_max, 1, MPI_INT, 0, world);
  MPI_Bcast(&beta_size, 1, MPI_INT, 0, world);
  MPI_Bcast(&cutoff, 1, MPI_DOUBLE, 0, world);
  MPI_Bcast(&radial_string_length, 1, MPI_INT, 0, world);
  MPI_Bcast(&cutoff_string_length, 1, MPI_INT, 0, world);
  MPI_Bcast(radial_string, radial_string_length + 1, MPI_CHAR, 0, world);
  MPI_Bcast(cutoff_string, cutoff_string_length + 1, MPI_CHAR, 0, world);

  // Parse the cutoffs.
  int n_cutoffs = n_species * n_species;
  memory->create(cutoffs, n_cutoffs, "pair:cutoffs");
  if (me == 0)
    grab(fptr, n_cutoffs, cutoffs);
  MPI_Bcast(cutoffs, n_cutoffs, MPI_DOUBLE, 0, world);

  // Fill in the cutoff matrix.
  cutoff = -1;
  cutoff_matrix = Eigen::MatrixXd::Zero(n_species, n_species);
  int cutoff_count = 0;
  for (int i = 0; i < n_species; i++){
    for (int j = 0; j < n_species; j++){
      double cutoff_val = cutoffs[cutoff_count];
      cutoff_matrix(i, j) = cutoff_val;
      if (cutoff_val > cutoff) cutoff = cutoff_val;
      cutoff_count ++;
    }
  }

  // Set number of descriptors.
  int n_radial = n_max * n_species;
  n_descriptors = (n_radial * (n_radial + 1) / 2) * (l_max + 1);

  // Check the relationship between the power spectrum and beta.
  int beta_check;
  if (power == 1) {
    beta_check = n_descriptors;
  } else if (power == 2) {
    beta_check = n_descriptors * (n_descriptors + 1) / 2;
  } else {
    error->all(FLERR, "Power should be 1 or 2.");
  }
  if (beta_check != beta_size)
    error->all(FLERR, "Beta size doesn't match the number of descriptors.");

  // Set the radial basis.
  if (!strcmp(radial_string, "chebyshev")) {
    basis_function = chebyshev;
    radial_hyps = std::vector<double>{0, cutoff};
  }

  // Set the cutoff function.
  if (!strcmp(cutoff_string, "quadratic"))
    cutoff_function = quadratic_cutoff;
  else if (!strcmp(cutoff_string, "cosine"))
    cutoff_function = cos_cutoff;

  // Parse the beta vectors.
  memory->create(beta, beta_size * n_species, "pair:beta");
  if (me == 0)
    grab(fptr, beta_size * n_species, beta);
  MPI_Bcast(beta, beta_size * n_species, MPI_DOUBLE, 0, world);

  // Fill in the beta matrix.
  // TODO: Remove factor of 2 from beta.
  Eigen::MatrixXd beta_matrix;
  int beta_count = 0;
  double beta_val;

  if (power == 1) {
    for (int k = 0; k < n_species; k++) {
      beta_matrix = Eigen::MatrixXd::Zero(n_descriptors, 1);
      for (int i = 0; i < n_descriptors; i++) {
        beta_matrix(i, 0) = beta[beta_count];
        beta_count++;
      }
      beta_matrices.push_back(beta_matrix);
    }
  } else if (power == 2) {
    for (int k = 0; k < n_species; k++) {
      beta_matrix = Eigen::MatrixXd::Zero(n_descriptors, n_descriptors);
      for (int i = 0; i < n_descriptors; i++) {
        for (int j = i; j < n_descriptors; j++) {
          if (i == j)
            beta_matrix(i, j) = beta[beta_count];
          else if (i != j) {
            beta_val = beta[beta_count] / 2;
            beta_matrix(i, j) = beta_val;
            beta_matrix(j, i) = beta_val;
          }
          beta_count++;
        }
      }
      beta_matrices.push_back(beta_matrix);
    }
  }
}

/* ----------------------------------------------------------------------
   grab n values from file fp and put them in list
   values can be several to a line
   only called by proc 0
------------------------------------------------------------------------- */

void PairFLARE::grab(FILE *fptr, int n, double *list) {
  char *ptr;
  char line[MAXLINE];

  int i = 0;
  while (i < n) {
    fgets(line, MAXLINE, fptr);
    ptr = strtok(line, " \t\n\r\f");
    list[i++] = atof(ptr);
    while ((ptr = strtok(NULL, " \t\n\r\f")))
      list[i++] = atof(ptr);
  }
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_ARG_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_ARG_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!
  @ingroup group-arithmetic
    Function object implementing arg capabilities

    Computes the angular orientation of its parameter.

    @par semantic:
    For any given value @c x of floating type @c T:

    @code
    T r = arg(x);
    @endcode

    is equivalent to:

    @code
    T r = (is_nan(x)) ? x : (is_negative(x)) ? Pi<T>() : Zero<T>();
    @endcode

    @par Note:

    - Returns \f$+0\f$ or \f$\pi\f$  (or @ref Nan if and only if the input is @ref Nan)
    according to the input @ref sign bit,

    - arg(x, use_signbit_) always returns \f$+0\f$ or \f$\pi\f$ taking into account the bit of sign
    even in the Nan case and is a bit speedier.

    - This function is the restriction to real numbers of the complex
    arg(ument) function.

    - This function is not defined for integral types.

    @par Alias:

    angle

    @see is_negative,  is_ltz,  Pi

  **/
  const boost::dispatch::functor<tag::arg_> arg = {};
} }
#endif

#include <boost/simd/function/scalar/arg.hpp>
#include <boost/simd/function/simd/arg.hpp>

#endif


// Copyright Stephan T. Lavavej, http://nuwen.net .
// Distributed under the Boost Software License, Version 1.0.
// See accompanying file LICENSE_1_0.txt or copy at
// http://boost.org/LICENSE_1_0.txt .

#include "fibonacci.hh"
#include "gluon.hh"
#include "test.hh"
#include "typedef.hh"
#include "vector.hh"

#include "external_begin.hh"
    #include <boost/iterator/counting_iterator.hpp>
#include "external_end.hh"

using namespace boost;
using namespace nuwen;

int main() {
    {
        const vuc_t v(make_counting_iterator(1), make_counting_iterator(256));
        const vuc_t f = fibonacci_encode(v);

        NUWEN_TEST("fibonacci1", fibonacci_decode<uc_t>(f) == v)
        NUWEN_TEST("fibonacci2", fibonacci_encode(sequence_cast<duc_t>(v)) == f)
        NUWEN_TEST("fibonacci3", fibonacci_encode(sequence_cast<luc_t>(v)) == f)
    }

    {
        const vus_t v(make_counting_iterator(1), make_counting_iterator(65536));

        NUWEN_TEST("fibonacci4", fibonacci_decode<us_t>(fibonacci_encode(v)) == v)
    }

    {
        const vul_t v = vec(glu<ul_t>(0x01234567UL)(0xFFFFFFFFUL)(0x17762161UL));

        NUWEN_TEST("fibonacci5", fibonacci_decode<ul_t>(fibonacci_encode(v)) == v)
    }

    {
        const vull_t v = vec(glu<ull_t>(0x0123456789ABCDEFULL)(0xFFFFFFFFFFFFFFFFULL)(0x2161CAFE1776BEEFULL));

        NUWEN_TEST("fibonacci6", fibonacci_decode<ull_t>(fibonacci_encode(v)) == v)
    }

    {
        const vus_t v = vec(glu<us_t>(137)(2161)(1776)(47));

        const vuc_t f = vuc_from_hex("8571253812182C");

        NUWEN_TEST("fibonacci7", fibonacci_encode(v) == f)
        NUWEN_TEST("fibonacci8", fibonacci_decode<us_t>(f) == v)
    }

    NUWEN_TEST("fibonacci9", fibonacci_decode<uc_t>(vuc_from_hex("FD80")) == vec(glu<uc_t>(1)(1)(1)(2)))
}


#include "gtest/gtest.h"
#include <Eigen/Core>
#include "MeshFactory.h"
#include "Paths.h"

namespace {
  using namespace Geotree;

  class PathsTest : public ::testing::Test {
  };

  TEST_F(PathsTest, CubeSide)
  {
    MeshFactory factory;
    Mesh cube0 = factory.cube(10,10,10);
    Mesh cube1 = factory.cube(10,1,1);
    Mesh cube2 = factory.cube(12,1,1);
    cube1.translate(Vector3d(1, 4.5, 1));
    cube2.translate(Vector3d(-1, 4.5, 1));
    IntersectionPoints points0(cube0, cube1);
    IntersectionPoints points1(cube0, cube2);

    Paths paths0(points0);
    Paths paths1(points1);

    EXPECT_EQ(1, paths0.count());
    EXPECT_EQ(2, paths1.count());
  }
}


#ifndef __DISTANCE_HPP__
#define __DISTANCE_HPP__

#include <armadillo>
#include "pauth_types.hpp"

namespace pauth {

/*! \brief Get the relative position between molecules i and j
 *
 * \param     r_i           Position of first molecule
 * \param     r_j           Position of second molecule
 * \param                   Placeholder for simulation box edge lengths
 * \return                  Relative position between the molecules
 */
inline arma::vec euclidean_rij(const arma::vec &r_i, const arma::vec &r_j,
                               const arma::vec &) {
  return r_i - r_j;
}

/*! \brief Get the square of the distance between molecules i and j
 *
 * Calculates the square of the distance between molecules i and j
 * Guarantees: >= 0
 *
 * \param     r_i           Position of first molecule
 * \param     r_j           Position of second molecule
 * \param                   Placeholder for simulation box edge lengths
 * \return                  Distance between the molecules squared
 */
inline double euclidean_m(const arma::vec &r_i, const arma::vec &r_j,
                          const arma::vec &edge_lengths) {
  const auto r_ij = euclidean_rij(r_i, r_j, edge_lengths);
  return arma::dot(r_ij, r_ij);
}

static metric euclidean(euclidean_rij, euclidean_m);

/*! \brief Get the relative position between molecules i and j
 *
 * \param     r_i           Position of first molecule
 * \param     r_j           Position of second molecule
 * \param                   Placeholder for simulation box edge lengths
 * \return                  Relative position between the molecules
 */
arma::vec periodic_euclidean_rij(const arma::vec &r_i, const arma::vec &r_j,
                                 const arma::vec &edge_lengths);

/*! \brief Get the square of the distance between molecules i and j
 *
 * Calculates the square of the distance between molecules i and j
 * Guarantees: >= 0
 *
 * \param     r_i           Position of first molecule
 * \param     r_j           Position of second molecule
 * \param     edge_lengths  Simulation box edge lengths
 * \return                  Distance between the molecules squared
 */
double periodic_euclidean_m(const arma::vec &r_i, const arma::vec &r_j,
                            const arma::vec &edge_lengths);

static metric periodic_euclidean(periodic_euclidean_rij, periodic_euclidean_m);

} // namespace pauth

#endif


/********************************************************************************
*  This file is part of CinoLib                                                 *
*  Copyright(C) 2016: Marco Livesu                                              *
*                                                                               *
*  The MIT License                                                              *
*                                                                               *
*  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 NON INFRINGEMENT. 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.                                                             *
*                                                                               *
*  Author(s):                                                                   *
*                                                                               *
*     Marco Livesu (marco.livesu@gmail.com)                                     *
*     http://pers.ge.imati.cnr.it/livesu/                                       *
*                                                                               *
*     Italian National Research Council (CNR)                                   *
*     Institute for Applied Mathematics and Information Technologies (IMATI)    *
*     Via de Marini, 6                                                          *
*     16149 Genoa,                                                              *
*     Italy                                                                     *
*********************************************************************************/
#include <cinolib/polygon_kernel.h>
#include <cinolib/min_max_inf.h>
#include <cinolib/geometry/polygon_utils.h>

#ifdef CINOLIB_USES_BOOST
#include <boost/geometry/geometries/polygon.hpp>
#include <boost/geometry/geometries/point_xy.hpp>
#include <boost/geometry.hpp>
//
typedef boost::geometry::model::d2::point_xy<double> BoostPoint;
typedef boost::geometry::model::polygon<BoostPoint>  BoostPolygon;
#endif

namespace cinolib
{

CINO_INLINE
double polygon_kernel(const std::vector<vec3d> & poly,   // will discard z component
                            std::vector<vec3d> & kernel) // z component will be zero
{
    std::vector<vec2d> poly_2d, kernel_2d;
    for(auto p : poly) poly_2d.push_back(vec2d(p));

    double area = polygon_kernel(poly_2d, kernel_2d);

    if(area > 0)
    {
        kernel.clear();
        for(auto p : kernel_2d) kernel.push_back(vec3d(p.x(), p.y(), 0));
    }

    return area;
}

//::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

#ifdef CINOLIB_USES_BOOST

CINO_INLINE
double polygon_kernel(std::vector<vec2d>   poly,
                      std::vector<vec2d> & kernel)
{
    kernel.clear();
    if(poly.empty()) return 0;

    if(!polygon_is_CCW(poly)) std::reverse(poly.begin(), poly.end());

    // define 2d axis aligned bbox
    vec2d min( inf_double,  inf_double);
    vec2d max(-inf_double, -inf_double);
    for(const vec2d & p : poly)
    {
        min = min.min(p);
        max = max.max(p);
    }
    double delta = min.dist(max);

    // define half spaces
    std::vector<BoostPolygon> half_spaces;
    for(uint i=0; i<poly.size(); ++i)
    {
        vec2d A   = poly.at(i);
        vec2d B   = poly.at((i+1)%poly.size());
        vec2d u = B - A; u.normalize(); // edge direction
        vec2d v = vec2d(-u.y(), u.x()); // direction orthogonal to u (rotated CCW)
        A -= u * delta;
        B += u * delta;
        vec2d C = B + v * delta;
        vec2d D = A + v * delta;

        BoostPolygon p;
        boost::geometry::append(p, BoostPoint(A.x(), A.y()));
        boost::geometry::append(p, BoostPoint(B.x(), B.y()));
        boost::geometry::append(p, BoostPoint(C.x(), C.y()));
        boost::geometry::append(p, BoostPoint(D.x(), D.y()));
        boost::geometry::correct(p);
        half_spaces.push_back(p);
    }

    // define kernel as intersection of half-spaces
    BoostPolygon kernel_boost = half_spaces.front();
    for(uint i=1; i<half_spaces.size(); ++i)
    {
        std::vector<BoostPolygon> res;
        boost::geometry::intersection(kernel_boost, half_spaces.at(i), res);
        if (res.empty()) return 0; // no kernel
        assert(res.size()==1);     // this should always hold. I am intersecting convex polygons...
        kernel_boost = res.front();
    }

    // convert from BoostPolygon to vec2d array
    for(uint i=0; i<kernel_boost.outer().size()-1; ++i)
    {
        kernel.push_back(vec2d(boost::geometry::get<0>(kernel_boost.outer()[i]),
                               boost::geometry::get<1>(kernel_boost.outer()[i])));
    }

    // Boost uses CW ordered vertices, so here I make the kernel CCW
    std::reverse(kernel.begin(), kernel.end());

    return boost::geometry::area(kernel_boost);
}

#else

CINO_INLINE
double polygon_kernel(const std::vector<vec2d> &,
                            std::vector<vec2d> &)
{
    std::cerr << "ERROR : Boost Polygon disabled. Recompile defining symbol CINOLIB_USES_BOOST" << std::endl;
    exit(-1);
}

#endif

}


//=======================================================================
// Copyright 2001 Jeremy G. Siek, Andrew Lumsdaine, Lie-Quan Lee,
//
// 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)
//=======================================================================
#include <boost/config.hpp>
#include <boost/concept/assert.hpp>
#include <iostream>
#include <fstream>
#include <stack>
#include <map>
#include <boost/lexical_cast.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/depth_first_search.hpp>
#include <boost/graph/graphviz.hpp>
#include <boost/graph/copy.hpp>
#include <boost/graph/reverse_graph.hpp>

using namespace boost;

template < typename OutputIterator >
class back_edge_recorder : public default_dfs_visitor
{
public:
  back_edge_recorder(OutputIterator out):m_out(out) { }

  template < typename Edge, typename Graph >
  void back_edge(Edge e, const Graph &)
  {
    *m_out++ = e;
  }
private:
  OutputIterator m_out;
};

// object generator function
template < typename OutputIterator >
back_edge_recorder < OutputIterator >
make_back_edge_recorder(OutputIterator out)
{
  return back_edge_recorder < OutputIterator > (out);
}

template < typename Graph, typename Loops > void
find_loops(typename graph_traits < Graph >::vertex_descriptor entry,
           const Graph & g,
           Loops & loops)    // A container of sets of vertices
{
  BOOST_CONCEPT_ASSERT(( BidirectionalGraphConcept<Graph> ));
  typedef typename graph_traits < Graph >::edge_descriptor Edge;
  std::vector < Edge > back_edges;
  std::vector < default_color_type > color_map(num_vertices(g));
  depth_first_visit(g, entry,
                    make_back_edge_recorder(std::back_inserter(back_edges)),
                    make_iterator_property_map(color_map.begin(),
                                               get(vertex_index, g), color_map[0]));

  for (typename std::vector < Edge >::size_type i = 0; i < back_edges.size(); ++i) {
    typename Loops::value_type x;
    loops.push_back(x);
    compute_loop_extent(back_edges[i], g, loops.back());
  }
}

template < typename Graph, typename Set > void
compute_loop_extent(typename graph_traits <
                    Graph >::edge_descriptor back_edge, const Graph & g,
                    Set & loop_set)
{
  BOOST_CONCEPT_ASSERT(( BidirectionalGraphConcept<Graph> ));
  typedef typename graph_traits < Graph >::vertex_descriptor Vertex;
  typedef color_traits < default_color_type > Color;

  Vertex loop_head, loop_tail;
  loop_tail = source(back_edge, g);
  loop_head = target(back_edge, g);

  std::vector < default_color_type >
    reachable_from_head(num_vertices(g), Color::white());
  default_color_type c;
  depth_first_visit(g, loop_head, default_dfs_visitor(),
                    make_iterator_property_map(reachable_from_head.begin(),
                                               get(vertex_index, g), c));

  std::vector < default_color_type > reachable_to_tail(num_vertices(g));
  reverse_graph < Graph > reverse_g(g);
  depth_first_visit(reverse_g, loop_tail, default_dfs_visitor(),
                    make_iterator_property_map(reachable_to_tail.begin(),
                                               get(vertex_index, g), c));

  typename graph_traits < Graph >::vertex_iterator vi, vi_end;
  for (boost::tie(vi, vi_end) = vertices(g); vi != vi_end; ++vi)
    if (reachable_from_head[*vi] != Color::white()
        && reachable_to_tail[*vi] != Color::white())
      loop_set.insert(*vi);
}

int
main(int argc, char *argv[])
{
  if (argc < 3) {
    std::cerr << "usage: loops_dfs <in-file> <out-file>" << std::endl;
    return -1;
  }
  GraphvizDigraph g_in;
  read_graphviz(argv[1], g_in);

  typedef adjacency_list < vecS, vecS, bidirectionalS,
    GraphvizVertexProperty,
    GraphvizEdgeProperty, GraphvizGraphProperty > Graph;
  typedef graph_traits < Graph >::vertex_descriptor Vertex;

  Graph g;
#if defined(BOOST_MSVC) && BOOST_MSVC <= 1300
  // VC++ has trouble with the get_property() function
  get_property(g, graph_name) = "loops";
#endif

  copy_graph(g_in, g);

  typedef std::set < Vertex > set_t;
  typedef std::list < set_t > list_of_sets_t;
  list_of_sets_t loops;
  Vertex entry = *vertices(g).first;

  find_loops(entry, g, loops);

  property_map<Graph, vertex_attribute_t>::type vattr_map = get(vertex_attribute, g);
  property_map<Graph, edge_attribute_t>::type eattr_map = get(edge_attribute, g);
  graph_traits < Graph >::edge_iterator ei, ei_end;

  for (list_of_sets_t::iterator i = loops.begin(); i != loops.end(); ++i) {
    std::vector < bool > in_loop(num_vertices(g), false);
    for (set_t::iterator j = (*i).begin(); j != (*i).end(); ++j) {
      vattr_map[*j]["color"] = "gray";
      in_loop[*j] = true;
    }
    for (boost::tie(ei, ei_end) = edges(g); ei != ei_end; ++ei)
      if (in_loop[source(*ei, g)] && in_loop[target(*ei, g)])
        eattr_map[*ei]["color"] = "gray";
  }

  std::ofstream loops_out(argv[2]);
#if defined(BOOST_MSVC) && BOOST_MSVC <= 1300
  // VC++ has trouble with the get_property() functions
  loops_out << "digraph loops {\n"
            << "size=\"3,3\"\n"
            << "ratio=\"fill\"\n"
            << "shape=\"box\"\n";
  graph_traits<Graph>::vertex_iterator vi, vi_end;
  for (boost::tie(vi, vi_end) = vertices(g); vi != vi_end; ++vi) {
    loops_out << *vi << "[";
    for (std::map<std::string,std::string>::iterator ai = vattr_map[*vi].begin();
         ai != vattr_map[*vi].end(); ++ai) {
      loops_out << ai->first << "=" << ai->second;
      if (next(ai) != vattr_map[*vi].end())
        loops_out << ", ";
    }
    loops_out<< "]";
  }

  for (boost::tie(ei, ei_end) = edges(g); ei != ei_end; ++ei) {
    loops_out << source(*ei, g) << " -> " << target(*ei, g) << "[";
    std::map<std::string,std::string>& attr_map = eattr_map[*ei];
    for (std::map<std::string,std::string>::iterator eai = attr_map.begin();
         eai != attr_map.end(); ++eai) {
      loops_out << eai->first << "=" << eai->second;
      if (next(eai) != attr_map.end())
        loops_out << ", ";
    }
    loops_out<< "]";
  }
  loops_out << "}\n";
#else
  get_property(g, graph_graph_attribute)["size"] = "3,3";
  get_property(g, graph_graph_attribute)["ratio"] = "fill";
  get_property(g, graph_vertex_attribute)["shape"] = "box";

  write_graphviz(loops_out, g,
                 make_vertex_attributes_writer(g),
                 make_edge_attributes_writer(g),
                 make_graph_attributes_writer(g));
#endif
  return 0;
}


#include <iostream>
#include <vector>
#include <queue>

#include "plan.hpp"

#include "boost/algorithm/string.hpp"
#include "boost/functional/hash.hpp"
#include "boost/graph/adjacency_list.hpp"
#include "boost/graph/connected_components.hpp"
#include "boost/graph/copy.hpp"
#include <boost/graph/mcgregor_common_subgraphs.hpp>
#include "boost/graph/vf2_sub_graph_iso.hpp"
#include "boost/property_map/property_map.hpp"

using EdgeProperty = boost::property<boost::edge_name_t, unsigned int>;
using VertexProperty = boost::property<boost::vertex_name_t, unsigned int, boost::property<boost::vertex_index_t, int> >;
using Graph = boost::adjacency_list<boost::vecS, boost::vecS, boost::bidirectionalS, VertexProperty, EdgeProperty>;

std::ostream& operator<<(std::ostream& out, const Graph& g) {
    //out<<"---Print Graph Start ----"<< std::endl;
    auto labelling_vertex = boost::get(boost::vertex_name, g);
    out << boost::num_vertices(g)<< " " << boost::num_edges(g) << std::endl;
    for (unsigned i = 0; i < boost::num_vertices(g); i++) {
      out << labelling_vertex[i] << " ";
    }
    out << std::endl;
    for (auto ep = boost::edges(g); ep.first != ep.second; ++ep.first) {
      unsigned int source_ = boost::source(*ep.first, g);
      unsigned int target_ = boost::target(*ep.first, g);
      out << source_ << " " << target_ << std::endl;
    }
    //out<<"---Print Graph End----"<< std::endl;
    return out;
}

struct GraphHash {
  std::size_t operator()(Graph const &Graph) const;
};

std::size_t GraphHash::operator()(Graph const &Graph) const {
  std::size_t res = 0;
  auto labelling_vertex = boost::get(boost::vertex_name, Graph);
  auto labelling_edge = boost::get(boost::edge_name, Graph);
  unsigned int edge_xor = 1;
  unsigned int vertex_xor = 1;
  for (auto ep = boost::edges(Graph); ep.first != ep.second; ++ep.first) {
    unsigned int source = boost::source(*ep.first, Graph);
    unsigned int target = boost::target(*ep.first, Graph);
    edge_xor = edge_xor ^ labelling_edge[*ep.first];
    vertex_xor = vertex_xor ^ labelling_vertex[source] ^ labelling_vertex[target];
  }
  unsigned int multi = edge_xor + vertex_xor;
  boost::hash_combine(res, multi);
  boost::hash_combine(res, boost::num_vertices(Graph));
  boost::hash_combine(res, boost::num_edges(Graph));
  return res;
}

bool check_iso(const Graph& small_graph, const Graph& large_graph) {
  // fast check at beginning
  if (boost::num_vertices(small_graph) != boost::num_vertices(large_graph) || boost::num_edges(small_graph) != boost::num_edges(large_graph)) {
    return false;
  }

  auto vertex_name_map1 = boost::get(boost::vertex_name, small_graph);
  auto vertex_name_map2 = boost::get(boost::vertex_name, large_graph);
  auto edge_name_map1 = boost::get(boost::edge_name, small_graph);
  auto edge_name_map2 = boost::get(boost::edge_name, large_graph);

  auto vertex_comp = boost::make_property_map_equivalent(vertex_name_map1, vertex_name_map2);
  auto edge_comp = boost::make_property_map_equivalent(edge_name_map1, edge_name_map2);
  // callback that do nothing
  auto cb = [&](auto &&f, auto &&) {
    // do nothing
    return true;
  };
  // boost::vf2_print_callback <Graph, Graph> callback(small_graph, large_graph);
  return boost::vf2_subgraph_iso(small_graph, large_graph, cb, boost::vertex_order_by_mult(small_graph),
                                 boost::edges_equivalent(edge_comp).vertices_equivalent(vertex_comp));
}

struct CmpGraph {
  inline bool operator()(const Graph &a, const Graph &b) const {
    return check_iso(a, b);
  }
};

struct CountKeyHash{
  std::size_t operator()(std::pair<unsigned, std::vector<unsigned>> const &pair) const {
      auto node_id = pair.first;
      auto& labels = pair.second;
      std::size_t res = 0;
      boost::hash_combine(res, boost::hash_range(labels.begin(), labels.end()));
      boost::hash_combine(res, node_id);
      return res;
  }
};

std::vector<Graph> get_id_graph_map_from_plan(Plan plan);

int main(int argc, char* argv[]) {
  Plan plan(argv[1]);
  std::ifstream count_file(argv[2]);
  std::vector<Graph> id_graph_map = get_id_graph_map_from_plan(plan);

  std::vector<unsigned> id_vertex_num_map = plan.get_id_vertex_num();
  std::unordered_map<Graph, unsigned, GraphHash, CmpGraph> labeled_query_count;
  std::unordered_map<std::pair<unsigned, std::vector<unsigned>>, unsigned, CountKeyHash> raw_count;

//deduplicate count record
  unsigned node_id;
  while (count_file >> node_id) {
    unsigned node_num = id_vertex_num_map[node_id];
    std::vector<unsigned> labels;
    labels.resize(node_num);
    for (unsigned i = 0; i < node_num; i++){
      unsigned label;
      count_file >> labels[i];
    }
    unsigned count;
    count_file >> count;
    raw_count[std::make_pair(node_id,labels)] = count;
  }
  count_file.close();
  //combine isomorphic labeled queries

  for (auto iter = raw_count.begin(); iter != raw_count.end(); iter++) {
    auto& key = iter->first;
    auto node_id = key.first;
    auto& labels = key.second;
    auto count = iter->second;
    Graph g =  id_graph_map[node_id];
    auto labelling_vertex = boost::get(boost::vertex_name, g);
    for (unsigned i = 0; i < labels.size(); i++){
      labelling_vertex[i] = labels[i];
    }
    labeled_query_count[g] = labeled_query_count[g] + count;
  }
  for (auto iter = labeled_query_count.begin(); iter != labeled_query_count.end(); iter++) {
    std::cout << "Count:" << iter->second << std::endl;
    std::cout << iter->first << std::endl;
  }

  return 0;
}

std::vector<Graph> get_id_graph_map_from_plan(Plan plan){
  std::vector<Graph> ret;
  ret.resize(plan.nodes.size());
  unsigned root_node_id = plan.root_node_id;
  std::queue<unsigned> q; //proccessed node_id
  q.push(root_node_id);
  boost::add_vertex(ret[root_node_id]);
  boost::add_vertex(ret[root_node_id]);
  boost::add_edge(0, 1, ret[root_node_id]);
  while (!q.empty()) {
    auto cur = q.front();
    auto& cur_node = plan.nodes[cur];
    q.pop();
    unsigned start_edge = cur_node.edge_start_idx;
    unsigned end_edge = cur_node.edge_start_idx + cur_node.num_edges;
    for (unsigned i = start_edge; i < end_edge; i++) {
      auto& edge = plan.edges[i];
      auto child = edge.dst;
      boost::copy_graph(ret[cur], ret[child->idx]);
      if (cur_node.subgraph_num_vertices < child->subgraph_num_vertices) {
        boost::add_vertex(ret[child->idx]);
      }
      for (auto& opt: edge.operations) {
        if (opt.is_forward){
          boost::add_edge(opt.src_key, opt.dst_key, ret[child->idx]);
        }else{
          boost::add_edge(opt.dst_key, opt.src_key, ret[child->idx]);
        }
      }
      q.push(child->idx);
    }
  }
  return ret;
}


#pragma once
#include "sco/modeling.hpp"
#include "sco/modeling_utils.hpp"
#include "sco/sco_fwd.hpp"
#include <Eigen/Core>
#include "trajopt/common.hpp"
#include <openrave/openrave.h>
namespace trajopt {

using namespace sco;
typedef BasicArray<Var> VarArray;

void makeTrajVariablesAndBounds(int n_steps, const RobotAndDOF& manip, OptProb& prob_out, VarArray& vars_out);

class CartPoseCost : public CostFromErrFunc, public Plotter {
public:
  CartPoseCost(const VarVector& vars, const OR::Transform& pose, const Vector3d& rot_coeffs, const Vector3d& pos_coeffs, RobotAndDOFPtr manip, KinBody::LinkPtr link);
  void Plot(const DblVec& x, OR::EnvironmentBase& env, std::vector<OR::GraphHandlePtr>& handles);
};

class CartPoseConstraint : public ConstraintFromFunc, public Plotter {
public:
  CartPoseConstraint(const VarVector& vars, const OR::Transform& pose, RobotAndDOFPtr manip, KinBody::LinkPtr link, const VectorXd& coeffs);
  void Plot(const DblVec& x, OR::EnvironmentBase& env, std::vector<OR::GraphHandlePtr>& handles);
};

class CartVelConstraint : public ConstraintFromFunc {
public:
  CartVelConstraint(const VarVector& step0vars, const VarVector& step1vars, RobotAndDOFPtr manip, KinBody::LinkPtr link, double distlimit);
};

class JointPosCost : public Cost {
public:
  JointPosCost(const VarVector& vars, const VectorXd& vals, const VectorXd& coeffs);
  virtual ConvexObjectivePtr convex(const vector<double>& x, Model* model);
  virtual double value(const vector<double>&);
private:
  VarVector vars_;
  VectorXd vals_, coeffs_;
  QuadExpr expr_;
};

class JointVelCost : public Cost {
public:
  JointVelCost(const VarArray& traj, const VectorXd& coeffs);
  virtual ConvexObjectivePtr convex(const vector<double>& x, Model* model);
  virtual double value(const vector<double>&);
private:
  VarArray vars_;
  VectorXd coeffs_;
  QuadExpr expr_;
};

class JointAccCost : public Cost {
public:
  JointAccCost(const VarArray& traj, const VectorXd& coeffs);
  virtual ConvexObjectivePtr convex(const vector<double>& x, Model* model);
  virtual double value(const vector<double>&);
private:
  VarArray vars_;
  VectorXd coeffs_;
  QuadExpr expr_;
};

}