/*
 * LPSolver.cpp
 * @brief:
 * @date: May 1, 2014
 * @author: Duy-Nguyen Ta
 */

#include <boost/foreach.hpp>
#include <boost/range/adaptor/map.hpp>
#include <gtsam/inference/Symbol.h>
#include <gtsam_unstable/linear/LPSolver.h>

using namespace std;
using namespace gtsam;

namespace gtsam {

/* ************************************************************************* */
void LPSolver::buildMetaInformation() {
  size_t firstVarIndex = 1; // Warning: lpsolve's column number starts from 1 !!!
  // Collect variables in objective function first
  BOOST_FOREACH(Key key, objectiveCoeffs_ | boost::adaptors::map_keys) {
    variableColumnNo_.insert(make_pair(key, firstVarIndex));
    variableDims_.insert(make_pair(key,objectiveCoeffs_.dim(key)));
    firstVarIndex += variableDims_[key];
    freeVars_.insert(key);
  }
  // Now collect remaining keys in constraints
  VariableIndex factorIndex(*constraints_);
  BOOST_FOREACH(Key key, factorIndex | boost::adaptors::map_keys) {
    if (!variableColumnNo_.count(key)) {
      JacobianFactor::shared_ptr jacobian = boost::dynamic_pointer_cast<JacobianFactor>
      (constraints_->at(*factorIndex[key].begin()));
      if (!jacobian || !jacobian->isConstrained()) {
        throw std::runtime_error("Invalid constrained graph!");
      }
      size_t dim = jacobian->getDim(jacobian->find(key));
      variableColumnNo_.insert(make_pair(key, firstVarIndex));
      variableDims_.insert(make_pair(key,dim));
      firstVarIndex += variableDims_[key];
      freeVars_.insert(key);
    }
  }
  // Collect the remaining keys in lowerBounds
  BOOST_FOREACH(Key key, lowerBounds_ | boost::adaptors::map_keys){
    if (!variableColumnNo_.count(key)) {
      variableColumnNo_.insert(make_pair(key, firstVarIndex));
      variableDims_.insert(make_pair(key,lowerBounds_.dim(key)));
      firstVarIndex += variableDims_[key];
    }
    freeVars_.erase(key);
  }
  // Collect the remaining keys in upperBounds
  BOOST_FOREACH(Key key, upperBounds_ | boost::adaptors::map_keys){
    if (!variableColumnNo_.count(key)) {
      variableColumnNo_.insert(make_pair(key, firstVarIndex));
      variableDims_.insert(make_pair(key,upperBounds_.dim(key)));
      firstVarIndex += variableDims_[key];
    }
    freeVars_.erase(key);
  }

  nrColumns_ = firstVarIndex - 1;
}

/* ************************************************************************* */
void LPSolver::addConstraints(const boost::shared_ptr<lprec>& lp,
    const JacobianFactor::shared_ptr& jacobian) const {
  if (!jacobian || !jacobian->isConstrained())
    throw std::runtime_error("LP only accepts constrained factors!");

  // Build column number from keys
  KeyVector keys = jacobian->keys();
  vector<int> columnNo = buildColumnNo(keys);

  // Add each row to lp one by one. TODO: is there a faster way?
  Vector sigmas = jacobian->get_model()->sigmas();
  Matrix A = jacobian->getA();
  Vector b = jacobian->getb();
  for (int i = 0; i<A.rows(); ++i) {
    // A.row(i).data() gives wrong result so have to make a copy
    // TODO: Why? Probably because lpsolve's add_constraintex modifies this raw buffer!!!
    Vector r = A.row(i);

    // WARNING: lpsolve's add_constraintex modifies the columnNo raw buffer!
    // so we have to make a new copy for every row!!!!!
    vector<int> columnNoCopy(columnNo);

    if (sigmas[i]>0) {
      cout << "Warning: Ignore Gaussian noise (sigma>0) in LP constraints!" << endl;
    }
    int constraintType = (sigmas[i]<0)?LE:EQ;
    if(!add_constraintex(lp.get(), columnNoCopy.size(), r.data(), columnNoCopy.data(),
        constraintType, b[i]))
      throw runtime_error("LP can't accept Gaussian noise!");
  }
}


/* ************************************************************************* */
void LPSolver::addBounds(const boost::shared_ptr<lprec>& lp) const {
  // Set lower bounds
  BOOST_FOREACH(Key key, lowerBounds_ | boost::adaptors::map_keys){
    Vector lb = lowerBounds_.at(key);
    for (size_t i = 0; i<lb.size(); ++i) {
      set_lowbo(lp.get(), variableColumnNo_.at(key)+i, lb[i]);
    }
  }

  // Set upper bounds
  BOOST_FOREACH(Key key, upperBounds_ | boost::adaptors::map_keys){
    Vector ub = upperBounds_.at(key);
    for (size_t i = 0; i<ub.size(); ++i) {
      set_upbo(lp.get(), variableColumnNo_.at(key)+i, ub[i]);
    }
  }

  // Set the rest as free variables
  BOOST_FOREACH(Key key, freeVars_) {
    for (size_t i = 0; i<variableDims_.at(key); ++i) {
      set_unbounded(lp.get(), variableColumnNo_.at(key)+i);
    }
  }
}

/* ************************************************************************* */
boost::shared_ptr<lprec> LPSolver::buildModel() const {
  boost::shared_ptr<lprec> lp(make_lp(0, nrColumns_));

  // Makes building the model faster if it is done rows by row
  set_add_rowmode(lp.get(), TRUE);

  // Add constraints
  BOOST_FOREACH(const GaussianFactor::shared_ptr& factor, *constraints_) {
    JacobianFactor::shared_ptr jacobian =
        boost::dynamic_pointer_cast<JacobianFactor>(factor);
    addConstraints(lp, jacobian);
  }

  // Add bounds
  addBounds(lp);

  // Rowmode should be turned off again when done building the model
  set_add_rowmode(lp.get(), FALSE);

  // Finally, the objective function from the objective coefficients
  KeyVector keys;
  BOOST_FOREACH(Key key, objectiveCoeffs_ | boost::adaptors::map_keys) {
    keys.push_back(key);
  }

  Vector f = objectiveCoeffs_.vector(keys);
  vector<int> columnNo = buildColumnNo(keys);

  if(!set_obj_fnex(lp.get(), f.size(), f.data(), columnNo.data()))
    throw std::runtime_error("lpsolve cannot set obj function!");

  // Set the object direction to minimize
  set_minim(lp.get());

  // Set lp's verbose
  set_verbose(lp.get(), CRITICAL);

  return lp;
}

/* ************************************************************************* */
VectorValues LPSolver::convertToVectorValues(REAL* row) const {
  VectorValues values;
  BOOST_FOREACH(Key key, variableColumnNo_ | boost::adaptors::map_keys) {
    // Warning: the columnNo starts from 1, but C's array index starts from 0!!
    Vector v = Eigen::Map<Eigen::VectorXd>(&row[variableColumnNo_.at(key)-1], variableDims_.at(key));
    values.insert(key, v);
  }
  return values;
}

/* ************************************************************************* */
VectorValues LPSolver::solve() const {
  static const bool debug = false;

  boost::shared_ptr<lprec> lp = buildModel();

  /* just out of curioucity, now show the model in lp format on screen */
  /* this only works if this is a console application. If not, use write_lp and a filename */
  if (debug) write_LP(lp.get(), stdout);

  int ret = ::solve(lp.get());
  if (ret != 0) {
    throw std::runtime_error(
        (boost::format( "lpsolve cannot find the optimal solution and terminates with %d error. "\
            "See lpsolve's solve() documentation for details.") % ret).str());
  }
  REAL* row = NULL;
  get_ptr_variables(lp.get(), &row);

  VectorValues solution = convertToVectorValues(row);

  return solution;
}

} /* namespace gtsam */