#include <openrave-core.h>
#include <openrave/openrave.h>
#include <boost/assign.hpp>

#include "trajopt/problem_description.hpp"
#include "trajopt/kinematic_terms.hpp"
#include "trajopt/rave_utils.hpp"
#include "trajopt/utils.hpp"
#include "trajopt/trajectory_costs.hpp"
#include "utils/eigen_conversions.hpp"
#include "sco/expr_op_overloads.hpp"
#include "sco/optimizers.hpp"
#include "quat_ops.hpp"
#include "dynamics_utils.hpp"
#include "sco/modeling_utils.hpp"
#include "osgviewer/osgviewer.hpp"
#include "trajopt/collision_terms.hpp"
#include "utils/stl_to_string.hpp"
#include "utils/config.hpp"
#include "o3.hpp"
#include "static_object.hpp"
#include "incremental_rb.hpp"
#include "utils/logging.hpp"
#include "composite_config.hpp"
#include "sco/expr_ops.hpp"
using namespace boost::assign;
using namespace OpenRAVE;
using namespace trajopt;
using namespace sco;
using namespace Eigen;
using namespace std;
using namespace util;

extern void PolygonToEquations(const MatrixX2d& pts, MatrixX2d& ab,VectorXd& c);

double mu = 1; // TODO make it a parameter
bool idle = false;
bool constrain_all_poses = true;
int nSteps=10;

enum ContactType {
  SLIDING, STICKING
};

struct Face {
  MatrixX2d poly;
  OpenRAVE::Transform Tlf;
  KinBody::LinkPtr link;
  Face() {}
  Face(const MatrixX2d& _poly, const OpenRAVE::Transform& _Tlf, KinBody::LinkPtr _link) : poly(_poly), Tlf(_Tlf), link(_link) {}
};

struct Contact {
  ContactType m_type;
  // contact at a single timestep
  VarVector m_dofvars;
  VarVector m_ptA, m_ptB; // local contacts inside faces
  Var m_fn; // normal force B to A
  VarVector m_fr; // friction force B to A
  ConfigurationPtr m_config; //composite config
  Face m_faceA, m_faceB;
  Contact(ContactType type, ConfigurationPtr config, const Face& faceA, const Face& faceB, const VarVector& dofvars,
      const VarVector& ptA, const VarVector& ptB, const Var& fn, const VarVector& fr)
  : m_type(type), m_dofvars(dofvars),m_ptA(ptA), m_ptB(ptB), m_fn(fn), m_fr(fr), m_config(config),  m_faceA(faceA), m_faceB(faceB) {}
  virtual ~Contact() {}
};
typedef boost::shared_ptr<Contact> ContactPtr;

struct FaceContactErrorCalculator: public VectorOfVector {
  ConfigurationPtr m_config;
  Face m_faceA, m_faceB;
  int m_nDof;
  FaceContactErrorCalculator(ConfigurationPtr config, const Face& faceA, const Face& faceB) :
    m_config(config), m_faceA(faceA), m_faceB(faceB), m_nDof(m_config->GetDOF()) {}

  void CalcWorldPoints(const VectorXd& x, OR::Vector& ptWorldA, OR::Vector& ptWorldB) const {
    OR::Vector nA, nB;
    CalcWorldPointsAndNormal(x, ptWorldA, ptWorldB, nA, nB);
  }
  void CalcWorldPointsAndNormal(const VectorXd& x, OR::Vector& ptWorldA, OR::Vector& ptWorldB, OR::Vector& nWorldA, OR::Vector& nWorldB)  const {
    DblVec dofs(x.data(), x.data() + m_nDof);
    OpenRAVE::Vector ptLocalA, ptLocalB;
    ptLocalA.x = x(m_nDof);
    ptLocalA.y = x(m_nDof + 1);
    ptLocalB.x = x(m_nDof + 2);
    ptLocalB.y = x(m_nDof + 3);

    m_config->SetDOFValues(dofs);
    ptWorldA = m_faceA.link->GetTransform() * (m_faceA.Tlf * ptLocalA);
    ptWorldB = m_faceB.link->GetTransform() * (m_faceB.Tlf * ptLocalB);

    nWorldA = geometry::quatRotate(geometry::quatMultiply(m_faceA.link->GetTransform().rot,m_faceA.Tlf.rot), OR::Vector(0,0,1));
    nWorldB = geometry::quatRotate(geometry::quatMultiply(m_faceB.link->GetTransform().rot,m_faceB.Tlf.rot), OR::Vector(0,0,1));

  }
  VectorXd operator()(const VectorXd& x) const {
    m_config->SetDOFValues(toDblVec(x.topRows(m_nDof)));
    OR::Vector ptWorldA, ptWorldB, nWorldA, nWorldB;
    CalcWorldPoints(x, ptWorldA, ptWorldB);
     return toVector3d(ptWorldA - ptWorldB);

  }

};

struct FaceContactConstraint: public ConstraintFromFunc, public Plotter {
  FaceContactConstraint(ConfigurationPtr config, const VarVector& dofvars, const VarVector& ptAvars, const VarVector& ptBvars,
      const Face& faceA, const Face& faceB) :
      ConstraintFromFunc(VectorOfVectorPtr(new FaceContactErrorCalculator(config, faceA, faceB)),
          concat(dofvars, ptAvars, ptBvars), VectorXd::Zero(0), EQ, faceA.link->GetName()+"_"+faceB.link->GetName()+"contact") {}

  void Plot(const DblVec& xvec, OR::EnvironmentBase& env, std::vector<OR::GraphHandlePtr>& handles);
};

void FaceContactConstraint::Plot(const DblVec& xvec, OR::EnvironmentBase& env, std::vector<OR::GraphHandlePtr>& handles) {
  FaceContactErrorCalculator* calc = static_cast<FaceContactErrorCalculator*>(f_.get());
  OR::Vector ptA, ptB;
  VectorXd x = getVec(xvec, vars_);
  calc->m_config->SetDOFValues(toDblVec(x.topRows(calc->m_nDof)));
  calc->CalcWorldPoints(x, ptA, ptB);
  handles.push_back(env.drawarrow(ptA, ptB, .001, OR::Vector(1, 0, 1, 1)));
}

struct ForceBalanceConstraint : public EqConstraint {

  ConfigurationPtr m_config;
  KinBody::LinkPtr m_link;
  vector<ContactPtr> m_contacts;
  VarVector m_dofs;

  ForceBalanceConstraint(ConfigurationPtr config, KinBody::LinkPtr link, const vector<ContactPtr>& contacts, const VarVector& dofs) : m_config(config), m_link(link), m_contacts(contacts), m_dofs(dofs) {
    name_ = "forcebalance";
  }

  AffExprVector CalcWrenchExpr(const DblVec& x, Contact& c);

  DblVec CalcWrenchValue(const DblVec& x, Contact& c);

  vector<double> value(const vector<double>& x);

  ConvexConstraintsPtr convex(const vector<double>& x, Model* model);
};

AffExprVector ForceBalanceConstraint::CalcWrenchExpr(const DblVec& x, Contact& c) {
  DblVec dofs =getDblVec(x, c.m_dofvars);
  c.m_config->SetDOFValues(dofs);

  OpenRAVE::Transform Tlf = c.m_faceA.Tlf;
  AffExprVector fExprFace(3), rExprFace(3);
  if (c.m_faceA.link == m_link) {
    fExprFace[0] = -AffExpr(c.m_fr[0]);
    fExprFace[1] = -AffExpr(c.m_fr[1]);
    fExprFace[2] = -AffExpr(c.m_fn);
    rExprFace[0] = AffExpr(c.m_ptA[0]);
    rExprFace[1] = AffExpr(c.m_ptA[1]);
  }
  else if (c.m_faceB.link == m_link) {
    fExprFace[0] =  AffExpr(c.m_fr[0]);
    fExprFace[1] =  AffExpr(c.m_fr[1]);
    fExprFace[2] =  AffExpr(c.m_fn);
    rExprFace[0] = AffExpr(c.m_ptB[0]);
    rExprFace[1] = AffExpr(c.m_ptB[1]);
  }
  else assert(0);

  AffExprVector fExprLink = rotateExpr(Tlf.rot, fExprFace);
  AffExprVector rExprLink = transformExpr(Tlf, rExprFace);

  AffExprVector tExprLink = linearizedCrossProduct(x, rExprLink, fExprLink);
  // cout << c.m_faceA.link->GetName() << c.m_faceB.link->GetName() << endl;
  // cout << "rexprlink" << rExprLink[0].value(x) << " " << rExprLink[1].value(x)<< " "<<rExprLink[2].value(x)<< endl;
  // cout << "rexprface" << rExprFace[0].value(x) << " " << rExprFace[1].value(x)<< " "<<rExprFace[2].value(x)<< endl;
  // cout << "fexprlink" << fExprLink[0].value(x) << " " << fExprLink[1].value(x)<< " "<<fExprLink[2].value(x)<< endl;
  // cout << "fexprface" << fExprFace[0].value(x) << " " << fExprFace[1].value(x)<< " "<<fExprFace[2].value(x)<< endl;
  // cout << "Tlf" << c.m_faceA.Tlf << endl;

  return concat(fExprLink, tExprLink);
}

DblVec ForceBalanceConstraint::CalcWrenchValue(const DblVec& x, Contact& c) {
  DblVec dofs =getDblVec(x, c.m_dofvars);

  c.m_config->SetDOFValues(dofs);

  OpenRAVE::Transform Tlf = c.m_faceA.Tlf;
  OR::Vector fFace, rFace;
  if (c.m_faceA.link == m_link) {
    fFace[0] = -c.m_fr[0].value(x);
    fFace[1] = -c.m_fr[1].value(x);
    fFace[2] = -c.m_fn.value(x);
    rFace[0] = c.m_ptA[0].value(x);
    rFace[1] = c.m_ptA[1].value(x);
  }
  else if (c.m_faceB.link == m_link) {
    fFace[0] =  c.m_fr[0].value(x);
    fFace[1] =  c.m_fr[1].value(x);
    fFace[2] = c.m_fn.value(x);
    rFace[0] = c.m_ptB[0].value(x);
    rFace[1] = c.m_ptB[1].value(x);
  }
  else assert(0);

  OR::Vector fLink = OR::geometry::quatRotate(Tlf.rot, fFace);
  OR::Vector rLink = Tlf * rFace;

  OR::Vector tLink = rLink.cross(fLink);

  DblVec out(6);
  out[0] = fLink[0];
  out[1] = fLink[1];
  out[2] = fLink[2];
  out[3] = tLink[0];
  out[4] = tLink[1];
  out[5] = tLink[2];
  return out;
}

vector<double> ForceBalanceConstraint::value(const vector<double>& x) {
  DblVec ftotal(6, 0);
  ftotal[2] += -9.8;
  BOOST_FOREACH(ContactPtr& contact, m_contacts) {
    DblVec fcontact = CalcWrenchValue(x, *contact);
    for (int j=0; j < 6; ++j) ftotal[j] += fcontact[j];
  }
  return ftotal;
}

ConvexConstraintsPtr ForceBalanceConstraint::convex(const vector<double>& x, Model* model) {
  AffExprVector ftotal(6);
  ftotal[2].constant += -9.8;
  BOOST_FOREACH(ContactPtr& contact, m_contacts) {
    AffExprVector fcontact = CalcWrenchExpr(x, *contact);
    exprInc(ftotal, fcontact);
  }
  ConvexConstraintsPtr out(new ConvexConstraints(model));
  for (int i=0; i < 6; ++i) {
    out->addEqCnt(ftotal[i]);
  }
  return out;
}

struct SlidingFrictionErrorCalc : public VectorOfVector {
  // dofs cp ffr fn

  int m_nDof;
  double m_mu;
  Face m_face;
  KinBody::LinkPtr m_link;
  ConfigurationPtr m_config;

  SlidingFrictionErrorCalc(ConfigurationPtr config, double mu, const Face& face)
    : m_nDof(config->GetDOF()), m_mu(mu), m_face(face), m_link(face.link),m_config(config) {
    }
  VectorXd operator()(const VectorXd& x) const;

};

VectorXd SlidingFrictionErrorCalc::operator()(const VectorXd& x) const {
  VectorXd dofs0 = x.topRows(m_nDof);
  VectorXd dofs1 = x.middleRows(m_nDof, m_nDof);

  OR::Vector localp;
  localp[0] = x(2*m_nDof + 0);
  localp[1] = x(2*m_nDof + 1);

  double ffrx = x(2*m_nDof + 2 + 0);
  double ffry = x(2*m_nDof + 2 + 1);
  double fn = x(2*m_nDof + 2 + 2);

  m_config->SetDOFValues(toDblVec(x.middleRows(m_nDof, m_nDof)));
  OR::Vector p1world = m_link->GetTransform().trans;
  m_config->SetDOFValues(toDblVec(x.topRows(m_nDof)));
  OR::Vector p0world = m_link->GetTransform().trans;
  // pppp << p1world << " | " <<  p0world << endl;
  OR::Vector localvel0 = geometry::quatRotate(
    geometry::quatMultiply(geometry::quatInverse(m_face.Tlf.rot),
                           geometry::quatInverse(m_link->GetTransform().rot)),
              (p1world - p0world));
  double normvel = sqrtf(localvel0.lengthsqr2());
  // double normf = sqrtf(ffrx*ffrx + ffry*ffry);
  Vector2d err;
  err(0) = normvel * ffrx - m_mu * fn * localvel0[0] ;
  err(1) = normvel * ffry - m_mu * fn * localvel0[1] ;
  return err;
}

struct QuadraticCost : public Cost {
  QuadExpr m_expr;
  QuadraticCost(const QuadExpr& expr) : m_expr(expr) {}
  double value(const DblVec& x) {return m_expr.value(x);}
  ConvexObjectivePtr convex(const DblVec& x, Model* model)  {
    ConvexObjectivePtr out(new ConvexObjective(model));
    exprInc(out->quad_, m_expr);
    return out;
  }

};

struct MechanicsProblem: public OptProb {
  int m_nSteps;
  typedef vector<ContactPtr> ContactVec;
  vector<ContactVec> m_timestep2contacts;

  RobotAndDOFPtr m_robotConfig;
  BasicArray<IncrementalRBPtr> m_dynObjConfigs;
  vector<StaticObjectPtr> m_staticObjConfigs;
  vector<KinBodyPtr> m_dynBodies, m_staticBodies;

  VarArray m_th; // robot
  vector<VarArray> m_obj2posvars; // position
  vector<MatrixXd> m_obj2q; // quaternion
  vector<VarArray> m_obj2r; // incremental rotation

  MechanicsProblem(int nSteps, const vector<KinBodyPtr>& dynamicBodies, const vector<KinBodyPtr>& staticBodies, RobotAndDOFPtr robotConfig);

  void AddContacts(ContactType ctype, int iFirst, int iLast, const Face& faceA, const Face& faceB);
  void AddMotionCosts(float jointvel_coeff, float obj_linvel_coeff, float obj_angvel_coeff);
  void AddDynamicsConstraints();
  void AddCollisionCosts(double coeff);
  ConfigurationPtr GetConfig(KinBody::LinkPtr, int timestep);
  VarVector GetDOFVars(KinBody::LinkPtr, int timestep);
  void Callback(DblVec&);
  void AnimateSolution(const DblVec& x);
  vector<ContactPtr> GetContacts(KinBody::LinkPtr link, int timestep);

private:
  void AddContact(int iStep, const Face& faceA, const Face& faceB, const VarVector& ptAvars, const VarVector& ptBvars, const Var& fn, const VarVector& ffr);
  void AddContactVariables(ContactType ctype, int iFirst, int iLast, const Face& faceA, const Face& faceB, VarArray& ptAVars, VarArray& ptBVars, VarArray& fnVars, VarArray& ffrVars);
  vector<ContactPtr> GetContacts(KinBody::LinkPtr);
};

vector<ContactPtr> MechanicsProblem::GetContacts(KinBody::LinkPtr link, int timestep) {
  vector<ContactPtr> out;
  BOOST_FOREACH(const ContactPtr& contact, m_timestep2contacts[timestep]) {
    if (contact->m_faceA.link == link || contact->m_faceB.link == link) out.push_back(contact);
  }
  return out;
}

ConfigurationPtr MechanicsProblem::GetConfig(KinBody::LinkPtr link, int timestep) {
  KinBodyPtr body = link->GetParent();
  if (GetLinkMaybeAttached(m_robotConfig->GetRobot(), link->GetName())) return m_robotConfig;
  for (int i=0; i < m_dynBodies.size(); ++i) if (m_dynBodies[i] == body) return m_dynObjConfigs(timestep, i);
  for (int i=0; i < m_staticBodies.size(); ++i) if (m_staticBodies[i] == body) return m_staticObjConfigs[i];
  return ConfigurationPtr();
}

VarVector MechanicsProblem::GetDOFVars(KinBody::LinkPtr link, int timestep) {
  KinBodyPtr body = link->GetParent();
  if (GetLinkMaybeAttached(m_robotConfig->GetRobot(), link->GetName())) return m_th.row(timestep);
  for (int i=0; i < m_dynBodies.size(); ++i) if (m_dynBodies[i] == body) return m_obj2posvars[i].row(timestep);
  for (int i=0; i < m_staticBodies.size(); ++i) if (m_staticBodies[i] == body) return VarVector();
  assert(0);
  return VarVector();
}

MechanicsProblem::MechanicsProblem(int nSteps, const vector<KinBodyPtr>& dynBodies, const vector<KinBodyPtr>& staticBodies, RobotAndDOFPtr robotConfig) :
    m_nSteps(nSteps),
    m_robotConfig(robotConfig),
    m_dynBodies(dynBodies),
    m_staticBodies(staticBodies)
{
  m_dynObjConfigs.resize(m_nSteps, dynBodies.size());
  for (int iStep = 0; iStep  < m_nSteps; ++iStep) {
    for (int iBody = 0; iBody < m_dynBodies.size(); ++iBody) {
      m_dynObjConfigs.at(iStep, iBody).reset(new IncrementalRB(m_dynBodies[iBody]));
    }
  }

  m_staticObjConfigs.resize(staticBodies.size());
  for (int iBody = 0; iBody < m_staticBodies.size(); ++iBody) {
    m_staticObjConfigs[iBody].reset(new StaticObject(m_staticBodies[iBody]));
  }

  m_timestep2contacts.resize(m_nSteps);

  m_obj2posvars.resize(dynBodies.size());
  m_obj2q.resize(dynBodies.size());
  m_obj2r.resize(dynBodies.size());

  vector<VarArray*> arrs;
  vector<string> prefixes;
  vector<int> cols;
  arrs += &m_th;
  prefixes += "th";
  cols += m_robotConfig->GetDOF();
  for (int i=0; i < dynBodies.size(); ++i) {
    arrs += &m_obj2posvars[i];
    prefixes += (boost::format("obj%i")%i).str();
    cols += 6;
  }
  AddVarArrays(*this, m_nSteps, cols, prefixes, arrs);

  for (int iBody=0; iBody < dynBodies.size(); ++iBody) {
    m_obj2r[iBody] = m_obj2posvars[iBody].block(0, 3, m_nSteps, 3);
    MatrixXd& q = m_obj2q[iBody];
    q.resize(m_nSteps,4);
    for (int iStep = 0; iStep < m_nSteps; ++iStep) {
      q.row(iStep) = Vector4d(1, 0, 0, 0).transpose();
    }
  }

}

void MechanicsProblem::AddContacts(ContactType ctype, int iFirst, int iLast, const Face& faceA, const Face& faceB) {
  VarArray ptAVars, ptBVars, fnVars, ffrVars;
  AddContactVariables(ctype, iFirst, iLast, faceA, faceB, ptAVars, ptBVars, fnVars, ffrVars);
  for (int iStep=iFirst; iStep <= iLast; ++iStep) {
    int iCP = (ctype == SLIDING) ? iStep : 0;
    AddContact(iStep, faceA, faceB, ptAVars.row(iCP), ptBVars.row(iCP), fnVars(iStep,0), ffrVars.row(iStep));
    if (ctype==SLIDING) {
      // addLinearConstraint(ffrVars(iStep,0));
      // addLinearConstraint(ffrVars(iStep,1));
    }

  }

  if (ctype == SLIDING) {
    for (int iStep=iFirst; iStep< iLast; ++iStep) {
      ConfigurationPtr configA = GetConfig(faceA.link, iStep);
      ConfigurationPtr configB = GetConfig(faceB.link, iStep);
      ConfigurationPtr compositeConfig(new CompositeConfig(configA, configB));
      VarVector compositeDOFVars0 = concat(GetDOFVars(faceA.link, iStep), GetDOFVars(faceB.link, iStep));
      VarVector compositeDOFVars1 = concat(GetDOFVars(faceA.link, iStep+1), GetDOFVars(faceB.link, iStep+1));
      VarVector vars = concat(compositeDOFVars0, compositeDOFVars1, ptAVars.row(iStep), ffrVars.row(iStep), fnVars.row(iStep));
      VectorOfVectorPtr f(new SlidingFrictionErrorCalc(GetConfig(faceA.link, iStep), mu, faceA));
      addConstraint(ConstraintPtr(new ConstraintFromFunc(f, vars, VectorXd::Ones(2), INEQ, "slipfric")));
    }
  }

  QuadExpr forcesquared;
  for (int i=0; i < fnVars.size(); ++i) exprInc(forcesquared, exprSquare(fnVars.m_data[i]));
  addCost(CostPtr(new QuadraticCost(exprMult(forcesquared,.1))));
}

// after you've created contact point variables and constrained them to lie in the face
void MechanicsProblem::AddContact(int iStep, const Face& faceA, const Face& faceB, const VarVector& ptAVars, const VarVector& ptBVars, const Var& fn, const VarVector& ffr) {
  ConfigurationPtr configA = GetConfig(faceA.link, iStep);
  ConfigurationPtr configB = GetConfig(faceB.link, iStep);
  ConfigurationPtr compositeConfig(new CompositeConfig(configA, configB));
  VarVector compositeDOFVars = concat(GetDOFVars(faceA.link, iStep), GetDOFVars(faceB.link, iStep));
  ContactPtr contact(new Contact(SLIDING, compositeConfig, faceA, faceB, compositeDOFVars, ptAVars, ptBVars, fn, ffr));
  addConstraint(ConstraintPtr(new FaceContactConstraint(compositeConfig, compositeDOFVars, ptAVars, ptBVars, faceA, faceB)));
  m_timestep2contacts[iStep].push_back(contact);
}

void MechanicsProblem::AddContactVariables(ContactType ctype, int iFirst, int iLast, const Face& faceA, const Face& faceB, VarArray& ptAVars, VarArray& ptBVars, VarArray& fnVars, VarArray& ffrVars) {

  if (ctype == SLIDING) {
    vector<int> cols; cols += 1,2,2,2;
    vector<string> names; names += "fn", "ffr", "ptA", "ptB";
    vector<VarArray*> arrayPtrs; arrayPtrs += &fnVars, &ffrVars, &ptAVars, &ptBVars;
    AddVarArrays(*this, iLast - iFirst + 1, cols, names, arrayPtrs);
  }
  else {
    {
      vector<int> cols; cols += 1,2;
      vector<string> names; names += "fn", "ffr";
      vector<VarArray*> arrayPtrs; arrayPtrs += &fnVars, &ffrVars;
      AddVarArrays(*this, iLast - iFirst + 1, cols, names, arrayPtrs);
    }
    {
      vector<int> cols; cols += 2,2;
      vector<string> names; names += "ptA", "ptB";
      vector<VarArray*> arrayPtrs; arrayPtrs += &ptAVars, &ptBVars;
      AddVarArrays(*this, 1, cols, names, arrayPtrs);
    }
  }

  setLowerBounds(vector<double>(fnVars.size(), 0), fnVars.flatten());

  for (int i = iFirst; i <= iLast; ++i) {
    // TODO add quadratic friction cone constraint
    // norm(f) < mu n
    // addLinearConstr(- mu * fnVars(i,0) + ffrVars(i,0), INEQ);
    // addLinearConstr(- mu * fnVars(i,0) - ffrVars(i,0), INEQ);
    // addLinearConstr(- mu * fnVars(i,0) + ffrVars(i,1), INEQ);
    // addLinearConstr(- mu * fnVars(i,0) - ffrVars(i,1), INEQ);
    getModel()->addIneqCnt(exprSquare(ffrVars(i,0)) + exprSquare(ffrVars(i,1)) - exprMult(exprSquare(fnVars(i,0)),mu), "fc");
  }

  VectorXd Ac, Bc;
  MatrixX2d Aab, Bab;
  PolygonToEquations(faceA.poly, Aab, Ac);
  PolygonToEquations(faceB.poly, Bab, Bc);
  int nContacts = ptAVars.rows();
  for (int i=0; i < nContacts; ++i) {
    for (int iSide = 0; iSide < Aab.rows(); ++iSide) {
      addLinearConstraint(Aab(iSide, 0) * ptAVars(i,0) + Aab(iSide, 1) * ptAVars(i,1) + Ac(iSide),INEQ);
    }
    for (int iSide = 0; iSide < Bab.rows(); ++iSide) {
      addLinearConstraint(Bab(iSide, 0) * ptBVars(i,0) + Bab(iSide, 1) * ptBVars(i,1) + Bc(iSide),INEQ);
    }
  }

}

void MechanicsProblem::AddMotionCosts(float jointvel_coeff, float obj_linvel_coeff, float obj_angvel_coeff) {
  addCost(CostPtr(new JointVelCost(m_th, jointvel_coeff * VectorXd::Ones(m_th.cols()))));
  for (int i=0; i < m_dynBodies.size(); ++i) {
    // xxx if this is the only usage of m_obj2r, then don't make it
    if (obj_linvel_coeff > 0) addCost(CostPtr(new JointVelCost(m_obj2posvars[i].block(0,0,m_nSteps,3), obj_linvel_coeff*VectorXd::Ones(3) )));
    if (obj_angvel_coeff > 0) addCost(CostPtr(new AngVelCost(m_dynObjConfigs.col(i), m_obj2r[i], obj_angvel_coeff )));
  }

}

void MechanicsProblem::AddCollisionCosts(double coeff) {
  // XXX assumes there's two of them
  for (int iStep=0; iStep < m_nSteps; ++iStep) {
    ConfigurationPtr compositeConfig(new CompositeConfig(m_robotConfig, m_dynObjConfigs(iStep,0)));
    VarVector compositeDOFVars = concat(m_th.row(iStep), m_obj2posvars[0].row(iStep));
    addCost(CostPtr(new CollisionCost(0, coeff, compositeConfig, compositeDOFVars)));
  }
}

void MechanicsProblem::AddDynamicsConstraints() {
  // for each non-robot object
  for (int iStep = 0; iStep < m_nSteps; ++iStep) {
    for (int iObj = 0; iObj < m_dynBodies.size(); ++iObj) {
      ConfigurationPtr config = m_dynObjConfigs(iStep, iObj);
      KinBody::LinkPtr link = m_dynBodies[iObj]->GetLinks()[0];
      VarVector dofvars = GetDOFVars(link, iStep);
      addConstraint(ConstraintPtr(new ForceBalanceConstraint(config, link, GetContacts(link, iStep), dofvars)));
    }
  }
}

void PlotFace(EnvironmentBase& env, const Face& face, vector<OR::GraphHandlePtr>& handles) {
  int npts = face.poly.rows();
  MatrixX3f ptdata(npts, 3);
  OR::Transform Twf = face.link->GetTransform()*face.Tlf;
  for (int i=0; i < npts; ++i) {
    OR::Vector ptlocal(face.poly(i,0), face.poly(i,1), 0);
    OR::Vector ptworld = Twf*ptlocal;
    for (int j=0; j < 3; ++j) ptdata(i,j) = ptworld[j];
  }
  handles.push_back(env.drawlinestrip(ptdata.data(), ptdata.rows(), sizeof(float)*3, 4, OR::RaveVector<float>(1,0,0,1)));

  Vector3f ecw = ptdata.colwise().mean();
  OR::Vector centerworld(ecw(0), ecw(1), ecw(2));
  OR::Vector normalface(0,0,1);
  OR::Vector normalworld = OR::geometry::quatRotate(Twf.rot, normalface);

  handles.push_back(env.drawarrow(centerworld, centerworld + .05*normalworld, .001, OR::RaveVector<float>(1,1,0,1)));
}

void MechanicsProblem::AnimateSolution(const DblVec& x) {
  EnvironmentBasePtr env = m_robotConfig->GetEnv();
  OSGViewerPtr viewer = OSGViewer::GetOrCreate(env);

  for (int iStep=0; iStep < m_nSteps; ++iStep) {
    m_robotConfig->SetDOFValues(getDblVec(x, m_th.row(iStep)));
    for (int iObj=0; iObj < m_dynBodies.size(); ++iObj) {
      DblVec objdofs = getDblVec(x, m_obj2posvars[iObj].row(iStep));
      m_dynObjConfigs(iStep, iObj)->SetDOFValues(objdofs);
    }
    printf("step %i\n",iStep);
    viewer->Idle();
  }

}

void MechanicsProblem::Callback(DblVec& x) {

  // Update quats
  for (int iObj=0; iObj < m_dynBodies.size(); ++iObj) {
    MatrixXd& q = m_obj2q[iObj];
    MatrixXd rvals = getTraj(x, m_obj2r[iObj]);
    for (int iStep = 0; iStep < q.rows(); ++iStep) {
      q.row(iStep) = quatMult(quatExp(rvals.row(iStep)), q.row(iStep));
      // cout << quatMult(quatExp(rvals.row(iStep)), q.row(iStep)).transpose() << " =?= " << geometry::quatMultiply(geometry::quatFromAxisAngle(toRave(rvals.row(iStep))), toRaveQuat(q.row(iStep))) << endl;
      // cout << quatExp(rvals.row(iStep)).transpose() << " =?= " << OR::geometry::quatFromAxisAngle(toRave(rvals.row(iStep))) << endl;
      m_dynObjConfigs(iStep, iObj)->m_q = toRaveQuat(q.row(iStep));
      m_dynObjConfigs(iStep, iObj)->m_r = OR::Vector(0,0,0);
    }
    setVec(x, m_obj2r[iObj].flatten(), DblVec(m_obj2r[iObj].size(), 0));
  }

  // Plotting
  RobotBasePtr robot = boost::dynamic_pointer_cast<RobotAndDOF>(m_robotConfig)->GetRobot();
  EnvironmentBasePtr env = robot->GetEnv();
  OSGViewerPtr viewer = OSGViewer::GetOrCreate(env);

  assert(robot);

  vector<GraphHandlePtr> handles;
  BOOST_FOREACH(const CostPtr& cost, getCosts()) {
    if (Plotter* plotter = dynamic_cast<Plotter*>(cost.get())) {
      plotter->Plot(x, *env, handles);
      plotter->Plot(x, *env, handles);
    }
  }
  BOOST_FOREACH(const ConstraintPtr& cnt, getConstraints()) {
    if (Plotter* plotter = dynamic_cast<Plotter*>(cnt.get())) {
      plotter->Plot(x, *env, handles);
    }
  }

  TrajArray robotTraj = getTraj(x, m_th);
//  cout << "robotTraj\n" << robotTraj << endl;

  for (int iStep=0; iStep < m_nSteps; ++iStep) {
    m_robotConfig->SetDOFValues(toDblVec(robotTraj.row(iStep)));

    vector<ContactPtr> contacts = m_timestep2contacts[iStep];
    BOOST_FOREACH(ContactPtr contact, contacts) {
      if (contact->m_faceA.link->GetParent() == robot) PlotFace(*env, contact->m_faceA, handles);
      else if (contact->m_faceB.link->GetParent() == robot) PlotFace(*env, contact->m_faceB, handles);
    }

    handles.push_back(viewer->PlotKinBody(robot));
    SetTransparency(handles.back(), .35);
  }

  for (int iObj=0; iObj < m_dynBodies.size(); ++iObj) {
    TrajArray objTraj = getTraj(x, m_obj2posvars[iObj]);
//    cout << "objTraj\n" << objTraj << endl;
    KinBodyPtr body = m_dynBodies[iObj];
    KinBody::LinkPtr link = body->GetLinks()[0];
    for (int iStep = 0; iStep < m_nSteps; ++iStep) {

      m_dynObjConfigs(iStep, iObj)->SetDOFValues(toDblVec(objTraj.row(iStep)));

      vector<ContactPtr> contacts = GetContacts(link, iStep);
      BOOST_FOREACH(const ContactPtr& contact, contacts) {
        PlotFace(*env, contact->m_faceA, handles);
        PlotFace(*env, contact->m_faceB, handles);
      }

      handles.push_back(viewer->PlotKinBody(body));

      SetTransparency(handles.back(), .35);
    }
  }

  BOOST_FOREACH(const ConstraintPtr& cnt, getConstraints()) {
    if (cnt->name()=="slipfric") {
      // cout << "slipfric val: " << CSTR(cnt->value(x)) << endl;
    }
  }

  if (idle) viewer->Idle();

}

void Callback(OptProb* prob, DblVec& x) {
  dynamic_cast<MechanicsProblem*>(prob)->Callback(x);
}

void Setup3DOFPush(EnvironmentBasePtr env, boost::shared_ptr<MechanicsProblem>& prob, DblVec& xinit) {

  Face armtip, boxface, boxbottom, tabletop;//xxx
  env->Load(string(DATA_DIR) + "/3link_pushing.env.xml");

  RobotBasePtr robot = GetRobotByName(*env, "3DOFRobot");
  assert(robot);
  KinBodyPtr obj = GetBodyByName(*env, "box");
  assert(obj);
  KinBodyPtr table = GetBodyByName(*env, "table");
  assert(table);

  vector<int> robotDofInds;
  robotDofInds += 0, 1, 2;
  RobotAndDOFPtr robotConfig(new RobotAndDOF(robot, robotDofInds));

  {
    boxface.poly.resize(4,2);
    boxface.poly << -.1, -.1, -.1, .1, .1, .1, .1, -.1;
    boxface.Tlf.trans.x = .1;
    boxface.Tlf.rot = OpenRAVE::geometry::quatFromAxisAngle(OpenRAVE::Vector(0, 1, 0), M_PI / 2);
    boxface.link = obj->GetLinks()[0];
  }
  {
    armtip.poly.resize(4,2);
    armtip.poly << -.1, -.1, -.1, .1, .1, .1, .1, -.1;
    armtip.poly *= .05;
    armtip.Tlf.trans.x = .08;
    armtip.Tlf.rot =  OpenRAVE::geometry::quatFromAxisAngle(OpenRAVE::Vector(0, 1, 0), M_PI / 2);
    armtip.link = robot->GetLinks()[3];
  }
  {
    boxbottom = boxface;
    boxbottom.Tlf.rot = OpenRAVE::geometry::quatFromAxisAngle(OpenRAVE::Vector(0, 1, 0), M_PI);
    boxbottom.Tlf.trans = OR::Vector(0,0,-.1);
  }
  {
    tabletop.poly.resize(4,2);
    tabletop.poly << -1.0, -1.0,  -1.0, 1.0,  1.0, 1.0,  1.0, -1.0;
    tabletop.Tlf.trans.z = .1;
    tabletop.link = table->GetLinks()[0];
  }

  vector<KinBodyPtr> staticBodies, dynamicBodies;
  dynamicBodies.push_back(obj);
  staticBodies.push_back(table);

  prob.reset(new MechanicsProblem(nSteps, dynamicBodies, staticBodies, robotConfig));

//
  prob->AddContacts(STICKING, 0, nSteps-1, boxface, armtip);
  prob->AddContacts(SLIDING, 0, nSteps-1, boxbottom, tabletop);
  prob->AddDynamicsConstraints();
  prob->AddMotionCosts(10,10,10);
   prob->AddCollisionCosts(1);
  BasicTrustRegionSQP opt(prob);

  xinit = DblVec(prob->getNumVars(), 0);

  VectorXd startPt(6);
  startPt.topRows(3) = toVector3d(obj->GetTransform().trans);
  VectorXd endPt = startPt;
  endPt[0] += -.1;

  VarArray& posvars = prob->m_obj2posvars[0];
  for (int j = 0; j < 6; ++j) {
    setVec(xinit, posvars.col(j),VectorXd::LinSpaced(nSteps, startPt(j), endPt(j)));
    VectorXd ps = VectorXd::LinSpaced(nSteps, startPt(j), endPt(j));

    // for (int i=0; i < nSteps; ++i) {
    //   prob->addLinearConstr(exprSub(AffExpr(posvars(i, j)), ps(i)), EQ);
    // }
    prob->addLinearConstraint(exprSub(AffExpr(posvars(0, j)), startPt[j]), EQ);
    if (j < 3)
      prob->addLinearConstraint(exprSub(AffExpr(posvars(nSteps - 1, j)), endPt[j]),EQ);
    else
      prob->addLinearConstraint(AffExpr(posvars(nSteps - 1, j)),EQ);
  }

  // // TODO disabling rotation since something screwy is going on
  // for (int iStep=0; iStep < nSteps; ++iStep) {
  //   for (int j=3; j < 6; ++j) {
  //     prob->addLinearConstr(AffExpr(prob->m_obj2posvars[0](iStep, j)), EQ);
  //   }
  // }

}

Face GetFace(KinBody::LinkPtr link, const OR::Vector& dir) {
  Face f;
  OR::AABB bb = link->ComputeLocalAABB();
  OR::Vector extents = bb.extents;
  double lens[2];
  int i=-1;
  int sgn=0;
  for (int j=0; j < 3; ++j) if (fabs(dir[j]) > 1e-5) {
    i = j;
    sgn = dir[j];
  }
  assert(i != -1);
  f.Tlf.rot = geometry::quatRotateDirection(OR::Vector(0,0,1), dir);
  cout << "dir: " << dir << "tlf.rot: " << f.Tlf.rot << endl;
  OR::Vector fextents = f.Tlf.inverse()*extents;

  f.Tlf.trans = bb.pos;
  f.Tlf.trans[i] += sgn*extents[i];

  lens[0] = fextents[0];
  lens[1] = fextents[1];

  f.poly.resize(4,2);
  f.poly << lens[0], lens[1], -lens[0], lens[1], -lens[0], -lens[1], lens[0], -lens[1];
  f.link = link;
  return f;
}

void SetupPR2Push(EnvironmentBasePtr env, boost::shared_ptr<MechanicsProblem>& prob, DblVec& xinit) {

  env->Load(string(DATA_DIR) + "/pr2_pushing.env.xml");

  RobotBasePtr robot = GetRobotByName(*env, "pr2");

  KinBodyPtr box = GetBodyByName(*env, "box");
  KinBodyPtr table = GetBodyByName(*env, "table");
  assert(robot && box && table);

  RobotAndDOFPtr robotConfig(new RobotAndDOF(robot, GetManipulatorByName(*robot,"rightarm")->GetArmIndices()));

  KinBody::LinkPtr fingerlink0 = robot->GetLink("r_gripper_l_finger_tip_link");
  KinBody::LinkPtr fingerlink1 = robot->GetLink("r_gripper_r_finger_tip_link");

  KinBodyPtr fingerbox0 = GetBodyByName(*env, "fingerbox0");
  KinBodyPtr fingerbox1 = GetBodyByName(*env, "fingerbox1");
  fingerbox0->SetTransform(fingerlink0->GetTransform());
  fingerbox1->SetTransform(fingerlink1->GetTransform());
  robot->Grab(fingerbox0, fingerlink0);
  robot->Grab(fingerbox1, fingerlink1);

  Face finger0 = GetFace(fingerbox0->GetLinks()[0], OR::Vector(1,0,0));
  Face finger1 = GetFace(fingerbox1->GetLinks()[0], OR::Vector(1,0,0));
  // finger0.poly /= 2;
  // finger1.poly /= 2;

  cout << "finger0 poly" << finger0.poly << endl;

  Face boxfront = GetFace(box->GetLinks()[0], OR::Vector(-1,0,0));
  Face boxbottom = GetFace(box->GetLinks()[0], OR::Vector(0,0,-1));
  Face tabletop = GetFace(table->GetLinks()[0], OR::Vector(0,0,1));

  vector<KinBodyPtr> staticBodies = list_of(table), dynamicBodies = list_of(box);

  prob.reset(new MechanicsProblem(nSteps, dynamicBodies, staticBodies, robotConfig));

  prob->AddContacts(STICKING, 0, nSteps-1, boxfront, finger0);
  prob->AddContacts(STICKING, 0, nSteps-1, boxfront, finger1);
  prob->AddContacts(SLIDING, 0, nSteps-1, boxbottom, tabletop);
  prob->AddDynamicsConstraints();
  prob->AddMotionCosts(10,10,10);
   prob->AddCollisionCosts(200);
  BasicTrustRegionSQP opt(prob);

  xinit = DblVec(prob->getNumVars(), 0);
  VectorXd startPt = VectorXd::Zero(6);
  startPt.topRows(3) = toVector3d(box->GetTransform().trans);
  VectorXd endPt = startPt;
  endPt[0] += .12;
  // endPt[1] = .05;

  cout << "box start: " << startPt.transpose() << endl;
  cout << "table transform" << table->GetLinks()[0]->GetTransform() << endl;

  DblVec dofvals = robotConfig->GetDOFValues();
  for (int j=0; j < robotConfig->GetDOF(); ++j) {
    setVec(xinit, prob->m_th.col(j), VectorXd::Ones(nSteps)*dofvals[j]);
  }
  prob->Callback(xinit);

  VarArray& posvars = prob->m_obj2posvars[0];
  for (int i=0; i < nSteps; ++i) {
    for (int j=0; j < 6; ++j) {
      double ptarg = (startPt(j) * (nSteps-1-i) + endPt(j) * i) / (nSteps-1);
      xinit[posvars(i,j).var_rep->index] = ptarg;
      if (constrain_all_poses || (i==0) || (i==nSteps-1)) prob->addLinearConstraint(exprSub(AffExpr(posvars(i, j)), ptarg), EQ);
    }
  }

}

void SetupPR2Lift(EnvironmentBasePtr env, boost::shared_ptr<MechanicsProblem>& prob, DblVec& xinit) {

  env->Load(string(DATA_DIR) + "/pr2_lifting.env.xml");

  RobotBasePtr robot = GetRobotByName(*env, "pr2");

  KinBodyPtr box = GetBodyByName(*env, "box");
  KinBodyPtr table = GetBodyByName(*env, "table");
  assert(robot && box && table);

  vector<int> arminds = concat( GetManipulatorByName(*robot,"rightarm")->GetArmIndices(),
                                GetManipulatorByName(*robot,"leftarm")->GetArmIndices());

  RobotAndDOFPtr robotConfig(new RobotAndDOF(robot, arminds));

  KinBody::LinkPtr rfingerlink0 = robot->GetLink("r_gripper_l_finger_tip_link");
  KinBody::LinkPtr rfingerlink1 = robot->GetLink("r_gripper_r_finger_tip_link");
  KinBody::LinkPtr lfingerlink0 = robot->GetLink("l_gripper_l_finger_tip_link");
  KinBody::LinkPtr lfingerlink1 = robot->GetLink("l_gripper_r_finger_tip_link");

  KinBodyPtr rfingerbox0 = GetBodyByName(*env, "rfingerbox0");
  KinBodyPtr rfingerbox1 = GetBodyByName(*env, "rfingerbox1");
  KinBodyPtr lfingerbox0 = GetBodyByName(*env, "lfingerbox0");
  KinBodyPtr lfingerbox1 = GetBodyByName(*env, "lfingerbox1");

  rfingerbox0->SetTransform(rfingerlink0->GetTransform());
  rfingerbox1->SetTransform(rfingerlink1->GetTransform());
  robot->Grab(rfingerbox0, rfingerlink0);
  robot->Grab(rfingerbox1, rfingerlink1);

  lfingerbox0->SetTransform(lfingerlink0->GetTransform());
  lfingerbox1->SetTransform(lfingerlink1->GetTransform());
  robot->Grab(lfingerbox0, lfingerlink0);
  robot->Grab(lfingerbox1, lfingerlink1);

  Face rfinger0 = GetFace(rfingerbox0->GetLinks()[0], OR::Vector(1,0,0));
  Face rfinger1 = GetFace(rfingerbox1->GetLinks()[0], OR::Vector(1,0,0));
  Face lfinger0 = GetFace(lfingerbox0->GetLinks()[0], OR::Vector(1,0,0));
  Face lfinger1 = GetFace(lfingerbox1->GetLinks()[0], OR::Vector(1,0,0));
  // finger0.poly /= 2;
  // finger1.poly /= 2;

  Face boxleft = GetFace(box->GetLinks()[0], OR::Vector(0,1,0));
  Face boxright = GetFace(box->GetLinks()[0], OR::Vector(0,-1,0));

  vector<KinBodyPtr> staticBodies = list_of(table), dynamicBodies = list_of(box);

  prob.reset(new MechanicsProblem(nSteps, dynamicBodies, staticBodies, robotConfig));

  prob->AddContacts(STICKING, 0, nSteps-1, boxleft, lfinger0);
  prob->AddContacts(STICKING, 0, nSteps-1, boxleft, lfinger1);
  prob->AddContacts(STICKING, 0, nSteps-1, boxright, rfinger0);
  prob->AddContacts(STICKING, 0, nSteps-1, boxright, rfinger1);
  prob->AddDynamicsConstraints();
  prob->AddMotionCosts(10,10,10);
   prob->AddCollisionCosts(200);
  BasicTrustRegionSQP opt(prob);

  xinit = DblVec(prob->getNumVars(), 0);
  VectorXd startPt = VectorXd::Zero(6);
  startPt.topRows(3) = toVector3d(box->GetTransform().trans);
  VectorXd endPt = startPt;
  endPt[2] += .1;

  cout << "box start: " << startPt.transpose() << endl;
  cout << "table transform" << table->GetLinks()[0]->GetTransform() << endl;

  DblVec dofvals = robotConfig->GetDOFValues();
  for (int j=0; j < robotConfig->GetDOF(); ++j) {
    setVec(xinit, prob->m_th.col(j), VectorXd::Ones(nSteps)*dofvals[j]);
  }

  VarArray& posvars = prob->m_obj2posvars[0];
  for (int i=0; i < nSteps; ++i) {
    for (int j=0; j < 6; ++j) {
      double ptarg = (startPt(j) * (nSteps-1-i) + endPt(j) * i) / (nSteps-1);
      xinit[posvars(i,j).var_rep->index] = ptarg;
      if (constrain_all_poses || (i==0) || (i==nSteps-1)) prob->addLinearConstraint(exprSub(AffExpr(posvars(i, j)), ptarg), EQ);
    }
  }

}

#if 0
void SetupPR2Pickup(EnvironmentBasePtr env, boost::shared_ptr<MechanicsProblem>& prob, DblVec& xinit) {

  env->Load(string(DATA_DIR) + "/pr2_pushing.env.xml");

  RobotBasePtr robot = GetRobotByName(*env, "pr2");
  KinBodyPtr box = GetBodyByName(*env, "box");
  KinBodyPtr table = GetBodyByName(*env, "table");
  assert(robot && box && table);

  int nSteps = 10;

  RobotAndDOFPtr robotConfig(new RobotAndDOF(robot, GetManipulatorByName(*robot,"rightarm")->GetArmIndices()));

  Face finger0 = GetFace(robot->GetLink("r_gripper_l_finger_tip_link"), OR::Vector(1,0,0));
  Face finger1 = GetFace(robot->GetLink("r_gripper_r_finger_tip_link"), OR::Vector(1,0,0));

  Face boxfront = GetFace(box->GetLinks()[0], OR::Vector(-1,0,0));
  Face boxbottom = GetFace(box->GetLinks()[0], OR::Vector(0,0,-1));
  Face tabletop = GetFace(table->GetLinks()[0], OR::Vector(0,0,1));

  vector<KinBodyPtr> staticBodies = list_of(table), dynamicBodies = list_of(box);

  prob.reset(new MechanicsProblem(nSteps, dynamicBodies, staticBodies, robotConfig));

  prob->AddContacts(STICKING, 0, nSteps-1, finger0, boxfront);
  prob->AddContacts(STICKING, 0, nSteps-1, finger1, boxfront);
  prob->AddContacts(SLIDING, 0, nSteps-1, boxbottom, tabletop);
  prob->AddDynamicsConstraints();
  prob->AddMotionCosts(10,10,10);
   prob->AddCollisionCosts(1);
  BasicTrustRegionSQP opt(prob);

  xinit = DblVec(prob->getNumVars(), 0);
  VectorXd startPt(6);
  startPt.topRows(3) = toVector3d(box->GetTransform().trans);
  VectorXd endPt = startPt;
  endPt[0] += .1;

  VarArray& posvars = prob->m_obj2posvars[0];
  for (int i=0; i < nSteps; ++i) {
    for (int j=0; j < 6; ++j) {
      double ptarg = (startPt(j) * (nSteps-i) + endPt(j) * i) / (nSteps-1);
      xinit[posvars(i,j).var_rep->index] = ptarg;
      if (constrain_all_poses || (i==0) || (i==nSteps-1)) prob->addLinearConstraint(exprSub(AffExpr(posvars(i, j)), ptarg), EQ);
    }
  }

}
#endif

/**
Problems:
1. slide plate towards robot
2. lift a box
3. push a rolling chair
4. tip a book out of a bookshelf
5. grasp an object with a three-fingered hand

4. is interesting. Just initialize it with a contact between its finger and the top. and tell the optimizer it needs to slide the thing out.
so it'll first just try to push down and slide it. Let's first have a maybe top contact, and then we have two maybe side contacts.

Assuming one of the contacting objects is a face/surface, we can use that

So each contact has a surface face. Contact knows how to calculate the world wrench, and the expression for it.
But I decided that it was better to do the force balance in the link frame so everything isn't dependent on the origin.
Or we could compromise by doing force balance in a world-aligned frame centered at the link COM.

So next generation: fix force balance, contact knows about which face to use
Then a later generation will accommodate non-faces

*/

int main(int argc, char** argv) {

  vector<string> validProblems; validProblems += "3linkpush", "pr2push", "pr2lift";
  string problem = validProblems[0];

  Config config;
  config.add(new Parameter<double>("mu", &mu, "mu (friction coefficient)"));
  config.add(new Parameter<string>("problem", &problem, "problem: [3linkpush, pr2push]"));
  config.add(new Parameter<bool>("idle", &idle, "idle after stuff"));
  config.add(new Parameter<bool>("constrain_all_poses", &constrain_all_poses, "constrain pose trajectory, as opposed to just start and end"));
  config.add(new Parameter<int>("nSteps", &nSteps, "number of timesteps"));
  CommandParser parser(config);
  parser.read(argc, argv);

  if (std::find(validProblems.begin(), validProblems.end(), problem) == validProblems.end()) {
    PRINT_AND_THROW("invalid problem name");
  }

  RaveInitialize(false, Level_Debug);
  OpenRAVE::EnvironmentBasePtr env = RaveCreateEnvironment();
  env->StopSimulation();

  DblVec xinit;
  boost::shared_ptr<MechanicsProblem> prob;

  if (problem == "3linkpush") {
    Setup3DOFPush(env, prob, xinit);
  }
  else if (problem == "pr2push") {
    SetupPR2Push(env, prob, xinit);
  }
  else if (problem == "pr2lift") {
    SetupPR2Lift(env, prob, xinit);
  }

  OSGViewerPtr viewer = OSGViewer::GetOrCreate(env);

  BasicTrustRegionSQP opt(prob);
  opt.initialize(xinit);
  opt.max_iter_ = 1000;
  opt.addCallback(&Callback);

  opt.optimize();
  prob->AnimateSolution(opt.x());

  RaveDestroy();
}


// TopicModelLocalStepCPPX.cpp
// Define this symbol to enable runtime tests for allocations
#define EIGEN_RUNTIME_NO_MALLOC

#include <math.h>
#include "Eigen/Dense"
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/gamma.hpp>

using namespace Eigen;
using namespace std;

// ======================================================== Declare funcs
// ======================================================== visible externally

extern "C" {
    void sparseLocalStepSingleDoc(
        double* ElogLik_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int d,
        int D,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT
        );

    void sparseLocalStepSingleDoc_ActiveOnly(
        double* ElogLik_d_IN,
        double* wc_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int d,
        int D,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT,
        int doTrackELBO,
        double* elboVec_OUT,
        int numRestarts,
        int REVISE_FIRST,
        int REVISE_EVERY,
        int *rAcceptVec_IN, int* rTrialVec_IN,
        int verbose
        );

    void sparseLocalStepSingleDocWithWordCounts(
        double* wordcounts_d_IN,
        double* ElogLik_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT
        );
}

// ======================================================== Custom Type Defs
// ========================================================
// Simple names for array types
typedef Matrix<double, Dynamic, Dynamic, RowMajor> Mat2D_d;
typedef Matrix<double, 1, Dynamic, RowMajor> Mat1D_d;
typedef Array<double, Dynamic, Dynamic, RowMajor> Arr2D_d;
typedef Array<double, 1, Dynamic, RowMajor> Arr1D_d;
typedef Array<int, 1, Dynamic, RowMajor> Arr1D_i;

// Simple names for array types with externally allocated memory
typedef Map<Mat2D_d> ExtMat2D_d;
typedef Map<Mat1D_d> ExtMat1D_d;
typedef Map<Arr2D_d> ExtArr2D_d;
typedef Map<Arr1D_d> ExtArr1D_d;
typedef Map<Arr1D_i> ExtArr1D_i;

double calcELBOForSingleDoc_V2(
    Arr1D_d alphaEbeta,
    Arr1D_d topicCount_d,
    Arr1D_d spResp_data,
    Arr1D_i spResp_colids,
    Arr2D_d ElogLik_d,
    int K,
    int N,
    int nnzPerRow
    )
{
    double ELBO = 0.0;
    // Lalloc
    for (int k = 0; k < K; k++) {
        ELBO += boost::math::lgamma(topicCount_d(k) + alphaEbeta(k));
    }

    // Ldata and Lentropy
    for (int n = 0; n < N; n++) {
        for (int nzk = n * nnzPerRow; nzk < (n+1) * nnzPerRow; nzk++) {
            int k = spResp_colids(nzk);
            ELBO += spResp_data(nzk) * ElogLik_d(n,k);
            if (spResp_data(nzk) > 1e-9) {
                ELBO -= spResp_data(nzk) * log(spResp_data(nzk));
            }
        }
    }
    return ELBO;
}

double calcELBOForSingleDoc_V1(
    ExtArr1D_d alphaEbeta,
    ExtArr1D_d topicCount_d,
    Arr1D_d ElogProb_d,
    Arr1D_i activeTopics_d,
    double totalLogSumResp,
    double sum_gammalnalphaEbeta,
    int Kactive
    )
{
    // Compute Lalloc = \sum_k gammaln(\theta_dk)
    double ELBO = sum_gammalnalphaEbeta;
    for (int ka = 0; ka < Kactive; ka++) {
        int k = activeTopics_d(ka);
        ELBO += (
            boost::math::lgamma(topicCount_d(k) + alphaEbeta(k))
            - boost::math::lgamma(alphaEbeta(k))
            - topicCount_d(k) * ElogProb_d(k)
            );
    }
    ELBO += totalLogSumResp;
    return ELBO;
}

double updateAssignments_FixPerTokenActiveSet(
    ExtArr2D_d ElogLik_d,
    ExtArr1D_d wc_d,
    ExtArr1D_d alphaEbeta,
    Arr1D_i activeTopics_d,
    ExtArr1D_d& topicCount_d, // & makes sure we can edit passed array in place
    Arr1D_d& ElogProb_d,
    Arr1D_d& logScores_n,
    ExtArr1D_d& spResp_data,
    ExtArr1D_i& spResp_colids,
    int N,
    int K,
    int Kactive,
    int nnzPerRow,
    int doTrackELBO
    )
{
    assert(nnzPerRow > 1);
    double totalLogSumResp = 0.0;
    // UPDATE ElogProb_d using input doc-topic counts
    for (int ka = 0; ka < Kactive; ka++) {
        int k = activeTopics_d(ka);
        ElogProb_d(k) = boost::math::digamma(
            topicCount_d(k) + alphaEbeta(k));
    }
    // RESET topicCounts to all zeros
    topicCount_d.fill(0);
    // UPDATE assignments, obeying sparsity constraint
    for (int n = 0; n < N; n++) {
        double sumResp_n = 0.0;
        int m = n * nnzPerRow;
        double maxScore_n;
        for (int ka = 0; ka < nnzPerRow; ka++) {
            int k = spResp_colids(m + ka);
            logScores_n(ka) = ElogProb_d(k) + ElogLik_d(n,k);
            if (ka == 0 || logScores_n(ka) > maxScore_n) {
                maxScore_n = logScores_n(ka);
            }
        }
        for (int nzk = 0; nzk < nnzPerRow; nzk++) {
            spResp_data(m + nzk) = \
                exp(logScores_n(nzk) - maxScore_n);
            sumResp_n += spResp_data(m + nzk);
        }

        for (int nzk = m; nzk < m + nnzPerRow; nzk++) {
            spResp_data(nzk) /= sumResp_n;
            topicCount_d(spResp_colids(nzk)) += \
                wc_d(n) * spResp_data(nzk);
        }
        if (doTrackELBO) {
            totalLogSumResp += wc_d(n) * (maxScore_n + log(sumResp_n));
        }
    } // end for loop over tokens n
    assert(abs(wc_d.sum() - topicCount_d.sum()) < .000001);
    return totalLogSumResp;
}

double updateAssignments_ActiveOnly(
    ExtArr2D_d ElogLik_d,
    ExtArr1D_d wc_d,
    ExtArr1D_d alphaEbeta,
    Arr1D_i activeTopics_d,
    ExtArr1D_d& topicCount_d, // & makes sure we can edit passed array in place
    Arr1D_d& ElogProb_d,
    Arr1D_d& logScores_n,
    Arr1D_d& tempScores_n,
    ExtArr1D_d& spResp_data,
    ExtArr1D_i& spResp_colids,
    int N,
    int K,
    int Kactive,
    int nnzPerRow,
    int iter,
    int initProbsToEbeta,
    int doTrackELBO
    )
{
    double totalLogSumResp = 0.0;

    // UPDATE ElogProb_d using input doc-topic counts
    if (iter == 0 and initProbsToEbeta == 1) {
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            ElogProb_d(k) = log(alphaEbeta(k));
        }
    } else {
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            ElogProb_d(k) = boost::math::digamma(
                topicCount_d(k) + alphaEbeta(k));
        }
    }
    // RESET topicCounts to all zeros
    topicCount_d.fill(0);
    // UPDATE assignments, obeying sparsity constraint
    for (int n = 0; n < N; n++) {
        int m = n * nnzPerRow;
        int argmax_n = 0;
        double maxScore_n;
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            logScores_n(ka) = ElogProb_d(k) + ElogLik_d(n,k);
            if (ka == 0 || logScores_n(ka) > maxScore_n) {
                maxScore_n = logScores_n(ka);
                if (nnzPerRow == 1) {
                    argmax_n = k;
                }
            }
        }
        if (nnzPerRow == 1) {
            spResp_data(m) = 1.0;
            spResp_colids(m) = argmax_n;
            topicCount_d(argmax_n) += wc_d(n);
            if (doTrackELBO) {
                totalLogSumResp += wc_d(n) * maxScore_n;
            }
        } else {
            // Find the top L entries in logScores_n
            // Copy current row over into a temp buffer
            std::copy(
                logScores_n.data(),
                logScores_n.data() + Kactive,
                tempScores_n.data());
            // Sort the data in the temp buffer (in place)
            std::nth_element(
                tempScores_n.data(),
                tempScores_n.data() + Kactive - nnzPerRow,
                tempScores_n.data() + Kactive);
            // Walk thru this row and find the top "nnzPerRow" positions
            double pivotScore = tempScores_n(Kactive - nnzPerRow);

            int nzk = m;
            double sumResp_n = 0.0;
            for (int ka = 0; ka < Kactive; ka++) {
                if (logScores_n(ka) >= pivotScore) {
                    spResp_data(nzk) = \
                        exp(logScores_n(ka) - maxScore_n);
                    spResp_colids(nzk) = activeTopics_d(ka);
                    sumResp_n += spResp_data(nzk);
                    nzk += 1;
                }
            }
            //assert(nzk - m == nnzPerRow);

            for (nzk = m; nzk < m + nnzPerRow; nzk++) {
                spResp_data(nzk) /= sumResp_n;
                topicCount_d(spResp_colids(nzk)) += \
                    wc_d(n) * spResp_data(nzk);
            }
            if (doTrackELBO) {
                totalLogSumResp += wc_d(n) * (maxScore_n + log(sumResp_n));
            }
        } // end if statement branch for nnz > 1
    } // end for loop over tokens n
    return totalLogSumResp;
}

void sparseLocalStepSingleDoc_ActiveOnly(
        double* ElogLik_d_IN,
        double* wc_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int d,
        int D,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT,
        int doTrackELBO,
        double* elboVec_OUT,
        int numRestarts,
        int REVISE_FIRST,
        int REVISE_EVERY,
        int* rAcceptVec_IN,
        int* rTrialVec_IN,
        int verbose
        )
{
    nCoordAscentIterLP = max(nCoordAscentIterLP + max(0, initProbsToEbeta), 1);

    // Unpack inputs, treated as fixed constants
    ExtArr2D_d ElogLik_d (ElogLik_d_IN, N, K);
    ExtArr1D_d wc_d (wc_d_IN, N);
    ExtArr1D_d alphaEbeta (alphaEbeta_IN, K);

    // Unpack outputs
    ExtArr1D_d spResp_data (spResp_data_OUT, N * nnzPerRow);
    ExtArr1D_i spResp_colids (spResp_colids_OUT, N * nnzPerRow);
    ExtArr1D_d topicCount_d (topicCount_d_OUT, K);

    ExtArr1D_i numIterVec (numIterVec_OUT, D);
    ExtArr1D_d maxDiffVec (maxDiffVec_OUT, D);
    ExtArr1D_d elboVec (elboVec_OUT, nCoordAscentIterLP);

    ExtArr1D_i rAcceptVec (rAcceptVec_IN, 1);
    ExtArr1D_i rTrialVec (rTrialVec_IN, 1);

    // Temporary storage
    Arr1D_d ElogProb_d (K);
    Arr1D_d prevTopicCount_d (K);
    Arr1D_d logScores_n (K);
    Arr1D_d tempScores_n (K);

    Arr1D_i activeTopics_d (K);
    Arr1D_i spareActiveTopics_d (nnzPerRow);

    prevTopicCount_d.fill(-1);
    double maxDiff = N;
    int iter = 0;
    int Kactive = K;
    double ACTIVE_THR = 1e-9;
    double totalLogSumResp = 0.0;
    double sum_gammalnalphaEbeta = 0.0;
    if (doTrackELBO || numRestarts > 0) {
        for (int k = 0; k < K; k++) {
            sum_gammalnalphaEbeta += boost::math::lgamma(alphaEbeta(k));
        }
    }
    // Before any updates... All topics are active!
    for (int k = 0; k < K; k++) {
        activeTopics_d(k) = k;
    }
    if (verbose > 1) {
        printf("Initial topic counts \n");
        for (int k = 0; k < K; k++) {
            printf("%d:%6.1f ", k, topicCount_d(k));
        }
        printf("\n");
    }

    for (iter = 0; iter < nCoordAscentIterLP; iter++) {
        int doReviseActiveSet;
        if (iter >= REVISE_FIRST && Kactive <= nnzPerRow) {
            // Set of active docs is already very small,
            // so nothing big to gain from revising
            doReviseActiveSet = 0;
        } else if (iter < REVISE_FIRST || (iter - 1) % REVISE_EVERY == 0) {
            doReviseActiveSet = 1;
        } else {
            doReviseActiveSet = 0;
        }

        if (iter > 0 || initProbsToEbeta < 0) {
            if (doReviseActiveSet) {
                int newKactive = 0;
                int ia = 0; // index for spare inactive topics
                for (int a = 0; a < Kactive; a++) {
                    int k = activeTopics_d(a);
                    if (topicCount_d(k) > ACTIVE_THR) {
                        activeTopics_d(newKactive) = k;
                        prevTopicCount_d(k) = topicCount_d(k);
                        newKactive += 1;
                    } else if (newKactive < nnzPerRow - ia) {
                        spareActiveTopics_d(ia) = k;
                        ia += 1;
                    }
                }
                // If num topics above threshold is less than nnzPerRow,
                // We need to fill in with some spare empty topics.
                Kactive = newKactive;
                while (Kactive < nnzPerRow) {
                    int k = spareActiveTopics_d(Kactive - newKactive);
                    activeTopics_d(Kactive) = k;
                    Kactive++;
                }
            } else {
                for (int a = 0; a < Kactive; a++) {
                    int k = activeTopics_d(a);
                    prevTopicCount_d(k) = topicCount_d(k);
                }
            }
        }
        assert(Kactive >= nnzPerRow);
        assert(Kactive <= K);
        if (nnzPerRow == 1 || doReviseActiveSet) {
            totalLogSumResp = updateAssignments_ActiveOnly(
                ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
                topicCount_d, ElogProb_d,
                logScores_n, tempScores_n,
                spResp_data, spResp_colids,
                N, K, Kactive, nnzPerRow,
                iter, initProbsToEbeta, doTrackELBO
                );
        } else {
            totalLogSumResp = updateAssignments_FixPerTokenActiveSet(
                ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
                topicCount_d, ElogProb_d,
                logScores_n,
                spResp_data, spResp_colids,
                N, K, Kactive, nnzPerRow,
                doTrackELBO
                );
        }

        if (doTrackELBO) {
            elboVec(iter) = calcELBOForSingleDoc_V1(
                alphaEbeta, topicCount_d, ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
            /*
            double elboV2 = calcELBOForSingleDoc_V2(
                alphaEbeta, topicCount_d, spResp_data, spResp_colids,
                ElogLik_d, K, N, nnzPerRow);

            double elboV1 = calcELBOForSingleDoc_V1(
                alphaEbeta, topicCount_d, ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
            printf(" V1: %.6f\n V2: %.6f\n", elboV2, elboV1);
            */
        }

        if (verbose > 1) {
            printf("end of iter %3d Kactive %3d maxDiff %10.6f\n",
                iter, Kactive, maxDiff);
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%d:%6.1f ", k, topicCount_d(k));
            }
            printf("\n");
        }

        // END ITERATION. Decide whether to quit early
        if (iter > 0 && iter % 5 == 0) {
            double absDiff_k = 0.0;
            maxDiff = 0.0;
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                absDiff_k = abs(prevTopicCount_d(k) - topicCount_d(k));
                if (absDiff_k > maxDiff) {
                    maxDiff = absDiff_k;
                }
            }
            if (maxDiff <= convThrLP) {
                break;
            }
        }
    }

    // Figure out which topics are eligible for sparse restarts
    double smallThr = 1e-6;
    double curELBO = 0.0;
    for (int riter = 0; riter < numRestarts; riter++) {
        if (riter == 0) {
            totalLogSumResp = updateAssignments_ActiveOnly(
                ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
                topicCount_d, ElogProb_d,
                logScores_n, tempScores_n,
                spResp_data, spResp_colids,
                N, K, Kactive, nnzPerRow,
                0, 0, 1);
            curELBO = calcELBOForSingleDoc_V1(
                alphaEbeta, topicCount_d, ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
            // Remember the best-known topic-count vector!
            for (int k = 0; k < K; k++) {
                prevTopicCount_d(k) = topicCount_d(k);
            }
        }
        // SEARCH FOR SMALLEST TOPIC HAVE NOT YET TRIED YET
        int numAboveThr = 0;
        double minVal = N + 1.0; // topicCount_d must never have this value
        int minLoc = 0;
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            if (topicCount_d(k) > smallThr) {
                numAboveThr += 1;
                if (topicCount_d(k) < minVal) {
                    minVal = topicCount_d(k);
                    minLoc = k;
                }
            }
        }
        if (numAboveThr == 1) {
            break;
        }

        if (verbose) {
            printf("START: best known counts. ELBO=%.5e \n", curELBO);
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%02d:%06.2f ", k, topicCount_d(k));
            }
            printf("\n");
        }
        smallThr = minVal;
        topicCount_d(minLoc) = 0.0;
        if (verbose) {
            printf(
                "RESTART: Set index %d to zero (%d left)\n",
                minLoc, numAboveThr - 1);
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%02d:%06.2f ", k, topicCount_d(k));
            }
            printf("\n");
        }
        int NSTEP = 2;
        for (int step = 0; step < NSTEP; step++) {
            totalLogSumResp = updateAssignments_ActiveOnly(
                ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
                topicCount_d, ElogProb_d,
                logScores_n, tempScores_n,
                spResp_data, spResp_colids,
                N, K, Kactive, nnzPerRow,
                step, 0, step == NSTEP-1);
        }
        // If the change is small, abandon this
        double propELBO;
        if (abs(prevTopicCount_d(minLoc) -topicCount_d(minLoc)) < 1e-6) {
            propELBO = curELBO;
        } else {
            propELBO = calcELBOForSingleDoc_V1(
                alphaEbeta, topicCount_d, ElogProb_d, activeTopics_d,
                totalLogSumResp, sum_gammalnalphaEbeta, Kactive);
        }

        if (verbose) {
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                printf("%02d:%06.2f ", k, topicCount_d(k));
            }
            printf("\n");
            printf("propELBO % .6e\n", propELBO);
            printf(" curELBO % .6e\n", curELBO);
            if (propELBO > curELBO) {
                printf("beforeCount: %.6f\n", prevTopicCount_d(minLoc));
                printf(" afterCount: %.6f\n", topicCount_d(minLoc));
                printf("gainELBO % .6e  *** ACCEPTED \n", propELBO - curELBO);
            } else {
                printf("gainELBO % .6e      rejected \n", propELBO - curELBO);
            }
        }
        // If accepted, set current best doc-topic counts to latest proposal
        // Otherwise, reset the starting point for the next proposal.
        if (propELBO > curELBO) {
            curELBO = propELBO;
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                prevTopicCount_d(k) = topicCount_d(k);
            }
            rAcceptVec(0) += 1;
        } else {
            for (int ka = 0; ka < Kactive; ka++) {
                int k = activeTopics_d(ka);
                topicCount_d(k) = prevTopicCount_d(k);
            }
        }
        rTrialVec(0) += 1;
    }
    if (numRestarts > 0) {
        for (int ka = 0; ka < Kactive; ka++) {
            int k = activeTopics_d(ka);
            topicCount_d(k) = prevTopicCount_d(k);
        }
        // Final update! Make sure spResp reflects best topicCounts found
        updateAssignments_ActiveOnly(
            ElogLik_d, wc_d, alphaEbeta, activeTopics_d,
            topicCount_d, ElogProb_d,
            logScores_n, tempScores_n,
            spResp_data, spResp_colids,
            N, K, Kactive, nnzPerRow,
            0, 0, 0);
    }
    maxDiffVec(d) = maxDiff;
    numIterVec(d) = iter; // will already have +1 from last iter of for loop
}

void sparseLocalStepSingleDoc(
        double* ElogLik_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT,
        int d,
        int D,
        int* numIterVec_OUT,
        double* maxDiffVec_OUT
        )
{
    // Unpack inputs, treated as fixed constants
    ExtArr2D_d ElogLik_d (ElogLik_d_IN, N, K);
    ExtArr1D_d alphaEbeta (alphaEbeta_IN, K);
    // Unpack outputs
    ExtArr1D_d spResp_data (spResp_data_OUT, N * nnzPerRow);
    ExtArr1D_i spResp_colids (spResp_colids_OUT, N * nnzPerRow);
    ExtArr1D_d topicCount_d (topicCount_d_OUT, K);

    ExtArr1D_i numIterVec (numIterVec_OUT, D);
    ExtArr1D_d maxDiffVec (maxDiffVec_OUT, D);

    // Temporary storage
    Arr1D_d ElogProb_d (K);
    Arr1D_d prevTopicCount_d (K);
    Arr1D_d logScores_n (K);
    Arr1D_d tempScores_n (K);

    prevTopicCount_d.fill(-1);
    double maxDiff = N;
    int iter = 0;

    for (iter = 0; iter < nCoordAscentIterLP + initProbsToEbeta; iter++) {

        if (iter == 0 and initProbsToEbeta == 1) {
            for (int k = 0; k < K; k++) {
                ElogProb_d(k) = log(alphaEbeta(k));
            }
        } else {
            for (int k = 0; k < K; k++) {
                ElogProb_d(k) = boost::math::digamma(
                    topicCount_d(k) + alphaEbeta(k));
            }
        }

        topicCount_d.fill(0);
        // Step over each data atom
        for (int n = 0; n < N; n++) {
            int m = n * nnzPerRow;
            int argmax_n = 0;
            logScores_n(0) = ElogProb_d(0) + ElogLik_d(n,0);
            double maxScore_n = logScores_n(0);
            for (int k = 1; k < K; k++) {
                logScores_n(k) = ElogProb_d(k) + ElogLik_d(n,k);
                if (logScores_n(k) > maxScore_n) {
                    maxScore_n = logScores_n(k);
                    argmax_n = k;
                }
            }
            if (nnzPerRow == 1) {
                spResp_data(m) = 1.0;
                spResp_colids(m) = argmax_n;
                // Update topicCount_d
                topicCount_d(argmax_n) += 1.0;
            } else {
                // Find the top L entries in logScores_n
                // Copy current row over into a temp buffer
                std::copy(
                    logScores_n.data(),
                    logScores_n.data() + K,
                    tempScores_n.data());
                // Sort the data in the temp buffer (in place)
                std::nth_element(
                    tempScores_n.data(),
                    tempScores_n.data() + K - nnzPerRow,
                    tempScores_n.data() + K);
                // Walk thru this row and find the top "nnzPerRow" positions
                double pivotScore = tempScores_n(K - nnzPerRow);
                int nzk = 0;
                double sumResp_n = 0.0;
                for (int k = 0; k < K; k++) {
                    if (logScores_n(k) >= pivotScore) {
                        spResp_data(m + nzk) = \
                            exp(logScores_n(k) - maxScore_n);
                        spResp_colids(m + nzk) = k;
                        sumResp_n += spResp_data(m + nzk);
                        nzk += 1;
                    }
                }
                // Normalize for doc-topic counts
                for (int nzk = 0; nzk < nnzPerRow; nzk++) {
                    spResp_data(m + nzk) /= sumResp_n;
                    topicCount_d(spResp_colids(m + nzk)) += \
                        spResp_data(m + nzk);
                }
            }
        }
        // END ITERATION. Decide whether to quit early
        if (iter > 0 && iter % 5 == 0) {
            double absDiff_k = 0.0;
            maxDiff = 0.0;
            for (int k = 0; k < K; k++) {
                absDiff_k = abs(prevTopicCount_d(k) - topicCount_d(k));
                if (absDiff_k > maxDiff) {
                    maxDiff = absDiff_k;
                }
                prevTopicCount_d(k) = topicCount_d(k); // copy over
            }
            if (maxDiff <= convThrLP) {
                break;
            }
        }
    }
    maxDiffVec(d) = maxDiff;
    numIterVec(d) = iter + 1;
}

void sparseLocalStepSingleDocWithWordCounts(
        double* wordcounts_d_IN,
        double* ElogLik_d_IN,
        double* alphaEbeta_IN,
        int nnzPerRow,
        int N,
        int K,
        int nCoordAscentIterLP,
        double convThrLP,
        int initProbsToEbeta,
        double* topicCount_d_OUT,
        double* spResp_data_OUT,
        int* spResp_colids_OUT
        )
{
    // Unpack inputs, treated as fixed constants
    ExtArr1D_d wc_d (wordcounts_d_IN, N);
    ExtArr2D_d ElogLik_d (ElogLik_d_IN, N, K);
    ExtArr1D_d alphaEbeta (alphaEbeta_IN, K);
    // Unpack outputs
    ExtArr1D_d spResp_data (spResp_data_OUT, N * nnzPerRow);
    ExtArr1D_i spResp_colids (spResp_colids_OUT, N * nnzPerRow);
    ExtArr1D_d topicCount_d (topicCount_d_OUT, K);
    // Temporary storage
    VectorXd ElogProb_d (K);
    VectorXd prevTopicCount_d (K);
    VectorXd logScores_n (K);
    VectorXd tempScores_n (K);
    prevTopicCount_d.fill(-1);
    int iter = 0;
    double maxDiff = 0.0;
    for (iter = 0; iter < nCoordAscentIterLP + initProbsToEbeta; iter++) {

        if (iter == 0 and initProbsToEbeta == 1) {
            for (int k = 0; k < K; k++) {
                ElogProb_d(k) = log(alphaEbeta(k));
            }
        } else {
            for (int k = 0; k < K; k++) {
                ElogProb_d(k) = boost::math::digamma(
                    topicCount_d(k) + alphaEbeta(k));
            }
        }

        topicCount_d.fill(0);
        // Step over each data atom
        for (int n = 0; n < N; n++) {
            int m = n * nnzPerRow;
            int argmax_n = 0;
            logScores_n(0) = ElogProb_d(0) + ElogLik_d(n,0);
            double maxScore_n = logScores_n(0);
            for (int k = 1; k < K; k++) {
                logScores_n(k) = ElogProb_d(k) + ElogLik_d(n,k);
                if (logScores_n(k) > maxScore_n) {
                    maxScore_n = logScores_n(k);
                    argmax_n = k;
                }
            }
            if (nnzPerRow == 1) {
                spResp_data(m) = 1.0;
                spResp_colids(m) = argmax_n;
                // Update topicCount_d
                topicCount_d(argmax_n) += wc_d(n);
            } else {
                // Find the top L entries in logScores_n
                // Copy current row over into a temp buffer
                std::copy(
                    logScores_n.data(),
                    logScores_n.data() + K,
                    tempScores_n.data());
                // Sort the data in the temp buffer (in place)
                std::nth_element(
                    tempScores_n.data(),
                    tempScores_n.data() + K - nnzPerRow,
                    tempScores_n.data() + K);
                // Walk thru this row and find the top "nnzPerRow" positions
                double pivotScore = tempScores_n(K - nnzPerRow);
                int nzk = 0;
                double sumResp_n = 0.0;
                for (int k = 0; k < K; k++) {
                    if (logScores_n(k) >= pivotScore) {
                        spResp_data(m + nzk) = \
                            exp(logScores_n(k) - maxScore_n);
                        spResp_colids(m + nzk) = k;
                        sumResp_n += spResp_data(m + nzk);
                        nzk += 1;
                    }
                }
                // Normalize for doc-topic counts
                for (int nzk = 0; nzk < nnzPerRow; nzk++) {
                    spResp_data(m + nzk) /= sumResp_n;
                    topicCount_d(spResp_colids(m + nzk)) += \
                        wc_d(n) * spResp_data(m + nzk);
                }
            }
        }
        // END ITERATION. Decide whether to quit early
        if (iter > 0 && iter % 5 == 0) {
            double absDiff_k = 0.0;
            maxDiff = 0.0;
            for (int k = 0; k < K; k++) {
                absDiff_k = abs(prevTopicCount_d(k) - topicCount_d(k));
                if (absDiff_k > maxDiff) {
                    maxDiff = absDiff_k;
                }
                prevTopicCount_d(k) = topicCount_d(k); // copy over
            }
            if (maxDiff <= convThrLP) {
                break;
            }
        }
    }
}


#ifndef SKYLARK_BLOCKADMM_HPP
#define SKYLARK_BLOCKADMM_HPP

#include <El.hpp>
#include <skylark.hpp>
#include <cmath>
#include <boost/mpi.hpp>

#ifdef SKYLARK_HAVE_OPENMP
#include <omp.h>
#endif

#include "../utility/timer.hpp"

template <class InputType>
struct BlockADMMSolver {

    typedef typename skylark::utility::typer_t<InputType>::value_type value_type;

    typedef InputType data_matrix_t;
    typedef El::Matrix<value_type> feature_matrix_t;
    typedef El::Matrix<value_type> target_matrix_t;

    typedef skylark::sketch::sketch_transform_t<data_matrix_t, feature_matrix_t>
    feature_transform_t;
    typedef std::vector<const feature_transform_t *> feature_transform_array_t;

    // No feature transdeforms (aka just linear regression).
    BlockADMMSolver(const skylark::algorithms::loss_t<value_type>* loss,
        const skylark::algorithms::regularizer_t<value_type>* regularizer,
        double lambda, // regularization parameter
        int NumFeatures,
        int NumFeaturePartitions = 1);

    // Easy interface, aka kernel based.
    template<typename Kernel, typename MapTypeTag>
    BlockADMMSolver<InputType>(skylark::base::context_t& context,
        const skylark::algorithms::loss_t<value_type>* loss,
        const skylark::algorithms::regularizer_t<value_type>* regularizer,
        double lambda, // regularization parameter
        int NumFeatures,
        Kernel kernel,
        MapTypeTag tag,
        int NumFeaturePartitions = 1);

    // Easy interface, aka kernel based, with quasi-random features.
    template<typename Kernel>
    BlockADMMSolver<InputType>(skylark::base::context_t& context,
        const skylark::algorithms::loss_t<value_type>* loss,
        const skylark::algorithms::regularizer_t<value_type>* regularizer,
        double lambda, // regularization parameter
        int NumFeatures,
        Kernel kernel,
        skylark::ml::quasi_feature_transform_tag tag,
        int NumFeaturePartitions);

    // Guru interface.
    BlockADMMSolver<InputType>(const skylark::algorithms::loss_t<value_type>* loss,
        const skylark::algorithms::regularizer_t<value_type>* regularizer,
        const feature_transform_array_t& featureMaps,
        double lambda, // regularization parameter
        bool ScaleFeatureMaps = true);

    void set_nthreads(int NumThreads) { this->NumThreads = NumThreads; }
    void set_rho(double RHO) { this->RHO = RHO; }
    void set_maxiter(double MAXITER) { this->MAXITER = MAXITER; }
    void set_tol(double TOL) { this->TOL = TOL; }
    void set_cache_transform(bool CacheTransforms) {this->CacheTransforms = CacheTransforms;}

    ~BlockADMMSolver();

    void InitializeFactorizationCache();
    void InitializeTransformCache(int n);

    skylark::ml::hilbert_model_t* train(data_matrix_t& X,
        target_matrix_t& Y, data_matrix_t& Xv, target_matrix_t& Yv,
        bool regression, const boost::mpi::communicator& comm);

    int get_numfeatures() {return NumFeatures;}

    feature_transform_array_t& get_feature_maps() {return featureMaps;}

private:

    typedef El::Matrix<value_type> local_matrix_t;

    feature_transform_array_t featureMaps;
    int NumFeatures;
    int NumFeaturePartitions;
    const skylark::algorithms::loss_t<value_type>* loss;
    const skylark::algorithms::regularizer_t<value_type>* regularizer;
    std::vector<int> starts, finishes;
    bool ScaleFeatureMaps;
    bool OwnFeatureMaps;
    local_matrix_t **Cache;
    local_matrix_t **TransformCache;
    int NumThreads;

    double lambda;
    double RHO;
    int MAXITER;
    double TOL;

    bool CacheTransforms;
};

template <class InputType>
void BlockADMMSolver<InputType>::InitializeFactorizationCache() {
    Cache = new local_matrix_t*[NumFeaturePartitions];
    for(int j=0; j<NumFeaturePartitions; j++) {
        int start = starts[j];
        int finish = finishes[j];
        int sj = finish - start  + 1;
        Cache[j]  = new local_matrix_t(sj, sj);
    }
}

template <class InputType>
void BlockADMMSolver<InputType>::InitializeTransformCache(int n) {
    TransformCache = new local_matrix_t*[NumFeaturePartitions];
    for(int j=0; j<NumFeaturePartitions; j++) {
        int start = starts[j];
        int finish = finishes[j];
        int sj = finish - start  + 1;
        TransformCache[j]  = new local_matrix_t(sj, n);
    }
}

// No feature transforms (aka just linear regression).
template <class InputType>
BlockADMMSolver<InputType>::BlockADMMSolver(
        const skylark::algorithms::loss_t<value_type>* loss,
        const skylark::algorithms::regularizer_t<value_type>* regularizer,
        double lambda, // regularization parameter
        int NumFeatures,
        int NumFeaturePartitions) :
        NumFeatures(NumFeatures),
        NumFeaturePartitions(NumFeaturePartitions),
        loss(loss), regularizer(regularizer),
        starts(NumFeaturePartitions), finishes(NumFeaturePartitions),
        NumThreads(1), lambda(lambda), RHO(1.0), MAXITER(1000), TOL(0.1) {

    int cstart = 0, nf = NumFeatures, np = NumFeaturePartitions;
    for(int i = 0; i < NumFeaturePartitions; i++) {
        int sj = int(floor(double(nf) / np));
        starts[i] = cstart;
        finishes[i] = cstart + sj - 1;
        cstart += sj;
        nf -= sj;
        np--;
    }
    this->ScaleFeatureMaps = false;
    OwnFeatureMaps = false;
    InitializeFactorizationCache();
    CacheTransforms = false;
}

// Easy interface, aka kernel based.
template<class InputType>
template<typename Kernel, typename MapTypeTag>
BlockADMMSolver<InputType>::BlockADMMSolver(skylark::base::context_t& context,
    const skylark::algorithms::loss_t<value_type>* loss,
    const skylark::algorithms::regularizer_t<value_type>* regularizer,
    double lambda, // regularization parameter
    int NumFeatures,
    Kernel kernel,
    MapTypeTag tag,
    int NumFeaturePartitions) :
    featureMaps(NumFeaturePartitions),
    NumFeatures(NumFeatures), NumFeaturePartitions(NumFeaturePartitions),
    loss(loss), regularizer(regularizer),
    starts(NumFeaturePartitions), finishes(NumFeaturePartitions),
    NumThreads(1), lambda(lambda), RHO(1.0), MAXITER(1000), TOL(0.1) {

    int cstart = 0, nf = NumFeatures, np = NumFeaturePartitions;
    for(int i = 0; i < NumFeaturePartitions; i++) {
        int sj = int(floor(double(nf) / np));
        starts[i] = cstart;
        finishes[i] = cstart + sj - 1;
        cstart += sj;
        nf -= sj;
        np--;

        featureMaps[i] =
            kernel.template create_rft<data_matrix_t, local_matrix_t>(sj, tag, context);
    }
    this->ScaleFeatureMaps = true;
    OwnFeatureMaps = true;
    InitializeFactorizationCache();
    CacheTransforms = false;
}

// Easy interface, aka kernel based, with quasi-random features.
template<class InputType>
template<typename Kernel>
BlockADMMSolver<InputType>::BlockADMMSolver(skylark::base::context_t& context,
    const skylark::algorithms::loss_t<value_type>* loss,
    const skylark::algorithms::regularizer_t<value_type>* regularizer,
    double lambda, // regularization parameter
    int NumFeatures,
    Kernel kernel,
    skylark::ml::quasi_feature_transform_tag tag,
    int NumFeaturePartitions) :
    featureMaps(NumFeaturePartitions),
    NumFeatures(NumFeatures), NumFeaturePartitions(NumFeaturePartitions),
    loss(loss), regularizer(regularizer),
    starts(NumFeaturePartitions), finishes(NumFeaturePartitions),
    NumThreads(1), lambda(lambda), RHO(1.0), MAXITER(1000), TOL(0.1) {

    skylark::base::leaped_halton_sequence_t<value_type>
        qmcseq(kernel.qrft_sequence_dim()); // TODO size
    int cstart = 0, nf = NumFeatures, np = NumFeaturePartitions;
    for(int i = 0; i < NumFeaturePartitions; i++) {
        int sj = int(floor(double(nf) / np));
        starts[i] = cstart;
        finishes[i] = cstart + sj - 1;
        cstart += sj;
        nf -= sj;
        np--;

        featureMaps[i] =
            kernel.template create_qrft< data_matrix_t, local_matrix_t,
              skylark::base::leaped_halton_sequence_t>(sj, qmcseq,
                  starts[i], context);
    }
    this->ScaleFeatureMaps = true;
    OwnFeatureMaps = true;
    InitializeFactorizationCache();
    CacheTransforms = false;
}

// Guru interface
template <class InputType>
BlockADMMSolver<InputType>::BlockADMMSolver(
    const skylark::algorithms::loss_t<value_type>* loss,
    const skylark::algorithms::regularizer_t<value_type>* regularizer,
    const feature_transform_array_t &featureMaps,
    double lambda,
    bool ScaleFeatureMaps) :
    featureMaps(featureMaps),
    NumFeaturePartitions(featureMaps.size()),
    loss(loss), regularizer(regularizer),
    starts(NumFeaturePartitions), finishes(NumFeaturePartitions),
    NumThreads(1), lambda(lambda), RHO(1.0), MAXITER(1000), TOL(0.1)  {

    NumFeaturePartitions = featureMaps.size();
    NumFeatures = 0;
    for(int i = 0; i < NumFeaturePartitions; i++) {
        starts[i] = NumFeatures;
        finishes[i] = NumFeatures + featureMaps[i]->get_S() - 1;
        NumFeatures += featureMaps[i]->get_S();
    }
    this->ScaleFeatureMaps = ScaleFeatureMaps;
    OwnFeatureMaps = false;
    InitializeFactorizationCache();
    CacheTransforms = false;
}

template <class InputType>
BlockADMMSolver<InputType>::~BlockADMMSolver() {
    for(int i=0; i  < NumFeaturePartitions; i++) {
        delete Cache[i];
        if (OwnFeatureMaps)
            delete featureMaps[i];
    }
    delete[] Cache;
}

namespace internal {

template<typename T>
void GetSlice(El::Matrix<T> &X, El::Matrix<T> &Z,
    El::Int i, El::Int j, El::Int height, El::Int width) {

    El::View(Z, X, i, j, height, width);
}

template<typename T>
void GetSlice(skylark::base::sparse_matrix_t<T> &X, El::Matrix<T> &Z,
    El::Int i, El::Int j, El::Int height, El::Int width) {

    skylark::base::DenseSubmatrixCopy(X, Z, i, j, height, width);
}

}

template <class InputType>
skylark::ml::hilbert_model_t* BlockADMMSolver<InputType>::train(data_matrix_t& X, target_matrix_t& Y,
    data_matrix_t& Xv, target_matrix_t& Yv,
    bool regression, const boost::mpi::communicator& comm) {

    int rank = comm.rank();
    int size = comm.size();

    int P = size;

    int ni = skylark::base::Width(X);
    int d = skylark::base::Height(X);
    int targets = regression ? 1 : GetNumTargets(comm, Y);

    skylark::ml::hilbert_model_t* model =
        new skylark::ml::hilbert_model_t(featureMaps,
            ScaleFeatureMaps, NumFeatures, targets, regression);

    local_matrix_t Wbar;
    El::View(Wbar, model->get_coef());

    int k = Wbar.Width();

    // number of classes, targets - to generalize
    int D = NumFeatures;

    // exception: check if D = Wbar.Height();

    local_matrix_t O(k, ni);
    El::Zero(O);

    local_matrix_t Obar(k, ni);
    El::Zero(Obar);

    local_matrix_t nu(k, ni);
    El::Zero(nu);

    local_matrix_t W, mu, Wi, mu_ij, ZtObar_ij;

    if(rank==0) {
        El::Zeros(W,  D, k);
        El::Zeros(mu, D, k);
    }
    El::Zeros(Wi, D, k);
    El::Zeros(mu_ij, D, k);
    El::Zeros(ZtObar_ij, D, k);

    int iter = 0;

    value_type localloss = loss->evaluate(O, Y);
    value_type totalloss, accuracy, obj;

    int Dk = D*k;
    int nik  = ni*k;
    int start, finish, sj;

    boost::mpi::timer timer;

    local_matrix_t sum_o, del_o, wbar_output;
    El::Zeros(del_o, k, ni);
    local_matrix_t Yp(Yv.Height(), k);
    local_matrix_t Yp_labels(Yv.Height(), 1);

    if (CacheTransforms)
        InitializeTransformCache(ni);

    SKYLARK_TIMER_INITIALIZE(ITERATIONS_PROFILE);
    SKYLARK_TIMER_INITIALIZE(COMMUNICATION_PROFILE);
    SKYLARK_TIMER_INITIALIZE(TRANSFORM_PROFILE);
    SKYLARK_TIMER_INITIALIZE(ZTRANSFORM_PROFILE);
    SKYLARK_TIMER_INITIALIZE(ZMULT_PROFILE);
    SKYLARK_TIMER_INITIALIZE(PROXLOSS_PROFILE);
    SKYLARK_TIMER_INITIALIZE(BARRIER_PROFILE);
    SKYLARK_TIMER_INITIALIZE(PREDICTION_PROFILE);

    while(iter<MAXITER) {

        SKYLARK_TIMER_RESTART(ITERATIONS_PROFILE);

        iter++;

        SKYLARK_TIMER_RESTART(COMMUNICATION_PROFILE);
        broadcast(comm, Wbar.Buffer(), Dk, 0);

        SKYLARK_TIMER_ACCUMULATE(COMMUNICATION_PROFILE)

            // mu_ij = mu_ij - Wbar
            El::Axpy(-1.0, Wbar, mu_ij);

        // Obar = Obar - nu
        El::Axpy(-1.0, nu, Obar);

        SKYLARK_TIMER_RESTART(PROXLOSS_PROFILE);
        loss->proxoperator(Obar, 1.0/RHO, Y, O);
        SKYLARK_TIMER_ACCUMULATE(PROXLOSS_PROFILE);

        if(rank==0) {
            regularizer->proxoperator(Wbar, lambda/RHO, mu, W);
        }

        El::Zeros(sum_o, k, ni);
        El::Zeros(wbar_output, k, ni);

        int j;
        const feature_transform_t* featureMap;

        SKYLARK_TIMER_RESTART(TRANSFORM_PROFILE);

#       ifdef SKYLARK_HAVE_OPENMP
#       pragma omp parallel for if(NumThreads > 1) private(j, start, finish, sj, featureMap) num_threads(NumThreads)
#       endif
        for(j = 0; j < NumFeaturePartitions; j++) {
            start = starts[j];
            finish = finishes[j];
            sj = finish - start  + 1;

            local_matrix_t Z;

            // Get the Z matrix
            if (CacheTransforms && (iter > 1))
                El::View(Z, *TransformCache[j], 0, 0, sj, ni);
            else {
                if (featureMaps.size() > 0) {
                    featureMap = featureMaps[j];

                    SKYLARK_TIMER_RESTART(ZTRANSFORM_PROFILE);
                    Z.Resize(sj, ni); // TODO do we need this?
                    featureMap->apply(X, Z, skylark::sketch::columnwise_tag());
                    SKYLARK_TIMER_ACCUMULATE(ZTRANSFORM_PROFILE);

                    if (ScaleFeatureMaps)
                        El::Scale(sqrt(double(sj) / d), Z);
                } else
                    internal::GetSlice(X, Z, start, 0, sj, ni);
            }

            local_matrix_t tmp(sj, k);
            local_matrix_t rhs(sj, k);
            local_matrix_t o(k, ni);

            if(iter==1) {

                local_matrix_t Ones;
                El::Ones(Ones, sj, 1);
                El::Gemm(El::NORMAL, El::TRANSPOSE, 1.0, Z, Z, 0.0, *Cache[j]);
                El::UpdateDiagonal(*Cache[j], 1.0, Ones);
                El::Inverse(*Cache[j]);

                if (CacheTransforms)
                    *TransformCache[j] = Z;
            }

            El::View(tmp, Wbar, start, 0, sj, k); //tmp = Wbar[J,:]

            local_matrix_t wbar_tmp;
            El::Zeros(wbar_tmp, k, ni);

            if (NumThreads > 1) {
                El::Gemm(El::TRANSPOSE, El::NORMAL, 1.0, tmp, Z, 0.0, wbar_tmp);

#               ifdef SKYLARK_HAVE_OPENMP
#               pragma omp critical
#               endif
                El::Axpy(1.0, wbar_tmp, wbar_output);
            } else
                El::Gemm(El::TRANSPOSE, El::NORMAL, 1.0, tmp, Z, 1.0, wbar_output);

            rhs = tmp; //rhs = Wbar[J,:]
            El::View(tmp, mu_ij, start, 0, sj, k); //tmp = mu_ij[J,:]
            El::Axpy(-1.0, tmp, rhs); // rhs = rhs - mu_ij[J,:] = Wbar[J,:] - mu_ij[J,:]
            El::View(tmp, ZtObar_ij, start, 0, sj, k);
            El::Axpy(+1.0, tmp, rhs); // rhs = rhs + ZtObar_ij[J,:]

            SKYLARK_TIMER_RESTART(ZMULT_PROFILE);
            local_matrix_t dsum = del_o;
            El::Axpy(NumFeaturePartitions + 1.0, nu, dsum);
            El::Gemm(El::NORMAL, El::TRANSPOSE,
                1.0/(NumFeaturePartitions + 1.0), Z, dsum, 1.0, rhs); // rhs = rhs + z'*(1/(n+1) * del_o + nu)
            SKYLARK_TIMER_ACCUMULATE(ZMULT_PROFILE);

            El::View(tmp, Wi, start, 0, sj, k);
            El::Gemm(El::NORMAL, El::NORMAL, 1.0, *Cache[j], rhs, 0.0, tmp); // ]tmp = Wi[J,:] = Cache[j]*rhs

            SKYLARK_TIMER_RESTART(ZMULT_PROFILE);
            El::Gemm(El::TRANSPOSE, El::NORMAL, 1.0, tmp, Z, 0.0, o); // o = (z*tmp)' = (z*Wi[J,:])'
            SKYLARK_TIMER_ACCUMULATE(ZMULT_PROFILE);

            // mu_ij[JJ,:] = mu_ij[JJ,:] + Wi[JJ,:];
            El::View(tmp, mu_ij, start, 0, sj, k); //tmp = mu_ij[J,:]
            El::View(rhs, Wi, start, 0, sj, k);
            El::Axpy(+1.0, rhs, tmp);

            //ZtObar_ij[JJ,:] = numpy.dot(Z.T, o);
            El::View(tmp, ZtObar_ij, start, 0, sj, k);
            El::Gemm(El::NORMAL, El::TRANSPOSE, 1.0, Z, o, 0.0, tmp);

            //  sum_o += o
            if (NumThreads > 1) {
#               ifdef SKYLARK_HAVE_OPENMP
#               pragma omp critical
#               endif
                El::Axpy(1.0, o, sum_o);
            } else
                El::Axpy(1.0, o, sum_o);

            Z.Empty(); // TODO do we need this?
        }

        SKYLARK_TIMER_ACCUMULATE(TRANSFORM_PROFILE);

        localloss = 0.0 ;
        //  El::Zeros(o, ni, k);
        local_matrix_t o(k, ni);
        El::Zero(o);
        El::Scale(-1.0, sum_o);
        El::Axpy(+1.0, O, sum_o); // sum_o = O.Matrix - sum_o
        del_o = sum_o;

        SKYLARK_TIMER_RESTART(PREDICTION_PROFILE);
        if (skylark::base::Width(Xv) > 0) {
            El::Zero(Yp);
            El::Zero(Yp_labels);
            model->predict(Xv, Yp_labels, Yp, NumThreads);

            if (regression) {
                El::Axpy(-1.0, Yv, Yp);
                value_type localerr = std::pow(El::Nrm2(Yp), 2);
                value_type localnrm = std::pow(El::Nrm2(Yv), 2);
                value_type err, nrm;
                boost::mpi::reduce(comm, localerr, err,
                    std::plus<value_type>(), 0);
                boost::mpi::reduce(comm, localnrm, nrm,
                    std::plus<value_type>(), 0);

                if (comm.rank() == 0)
                    accuracy = std::sqrt(err / nrm);
            } else {
                El::Int correct = skylark::ml::classification_accuracy(Yv, Yp);
                El::Int totalcorrect, total;
                boost::mpi::reduce(comm, correct, totalcorrect,
                    std::plus<El::Int>(), 0);
                boost::mpi::reduce(comm, Yv.Height(), total,
                    std::plus<El::Int>(), 0);

                if(comm.rank() == 0)
                    accuracy =  totalcorrect * 100.0 / total;
            }
        }
        SKYLARK_TIMER_ACCUMULATE(PREDICTION_PROFILE);

        localloss += loss->evaluate(wbar_output, Y);

        SKYLARK_TIMER_RESTART(COMMUNICATION_PROFILE);
        reduce(comm, localloss, totalloss, std::plus<value_type>(), 0);
        SKYLARK_TIMER_ACCUMULATE(COMMUNICATION_PROFILE);

        if(rank == 0) {
            obj = totalloss + lambda * regularizer->evaluate(Wbar);

            if (skylark::base::Width(Xv) <=0) {
                std::cout << "iteration " << iter
                          << " objective " << obj
                          << " time " << timer.elapsed()
                          << " seconds" << std::endl;
            }
            else {
                std::cout << "iteration " << iter
                          << " objective " << obj
                          << " accuracy " << boost::format("%.2f") % accuracy
                          << " time " << timer.elapsed()
                          << " seconds" << std::endl;
            }
        }

        El::Copy(O, Obar);
        El::Scale(1.0/(NumFeaturePartitions+1.0), sum_o);
        El::Axpy(-1.0, sum_o, Obar);

        El::Axpy(+1.0, O, nu);
        El::Axpy(-1.0, Obar, nu);

        SKYLARK_TIMER_RESTART(COMMUNICATION_PROFILE);
        boost::mpi::reduce (comm,
            Wi.LockedBuffer(),
            Wi.MemorySize(),
            Wbar.Buffer(),
            std::plus<value_type>(),
            0);
        SKYLARK_TIMER_ACCUMULATE(COMMUNICATION_PROFILE);

        if(rank==0) {
            //Wbar = (Wisum + W)/(P+1)
            El::Axpy(1.0, W, Wbar);
            El::Scale(1.0/(P+1), Wbar);

            // mu = mu + W - Wbar;
            El::Axpy(+1.0, W, mu);
            El::Axpy(-1.0, Wbar, mu);
        }

        SKYLARK_TIMER_RESTART(BARRIER_PROFILE);
        comm.barrier();
        SKYLARK_TIMER_ACCUMULATE(BARRIER_PROFILE);

        SKYLARK_TIMER_ACCUMULATE(ITERATIONS_PROFILE);
    }

    SKYLARK_TIMER_PRINT(ITERATIONS_PROFILE, comm);
    SKYLARK_TIMER_PRINT(COMMUNICATION_PROFILE, comm);
    SKYLARK_TIMER_PRINT(TRANSFORM_PROFILE, comm);
    SKYLARK_TIMER_PRINT(ZTRANSFORM_PROFILE, comm);
    SKYLARK_TIMER_PRINT(ZMULT_PROFILE, comm);
    SKYLARK_TIMER_PRINT(PROXLOSS_PROFILE, comm);
    SKYLARK_TIMER_PRINT(BARRIER_PROFILE, comm);
    SKYLARK_TIMER_PRINT(PREDICTION_PROFILE, comm);

    return model;
}

#endif /* SKYLARK_BLOCKADDM_HPP */