#include <trajopt_utils/macros.h>
TRAJOPT_IGNORE_WARNINGS_PUSH
#include <Eigen/Core>
TRAJOPT_IGNORE_WARNINGS_POP

#include <trajopt_sco/expr_ops.hpp>
#include <trajopt_sco/modeling_utils.hpp>
#include <trajopt_utils/eigen_conversions.hpp>

#include "trajopt/trajectory_costs.hpp"

namespace
{
/** @brief Returns the difference between each row of a matrixXd and the row before */
static Eigen::MatrixXd diffAxis0(const Eigen::MatrixXd& in)
{
  return in.middleRows(1, in.rows() - 1) - in.middleRows(0, in.rows() - 1);
}
}  // namespace

namespace trajopt
{
//////////// Joint cost functions /////////////////

//////////////////// Position /////////////////////
JointPosEqCost::JointPosEqCost(const VarArray& vars, const Eigen::VectorXd& coeffs, const Eigen::VectorXd& targets,
                               int& first_step, int& last_step)
  : Cost("JointPosEq"), vars_(vars), coeffs_(coeffs), targets_(targets), first_step_(first_step), last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // pos = x1 - targ
      sco::AffExpr pos;
      sco::exprInc(pos, sco::exprMult(vars(i, j), 1));
      sco::exprDec(pos, targets_[j]);
      // expr_ = coeff * vel^2
      sco::exprInc(expr_, sco::exprMult(sco::exprSquare(pos), coeffs_[j]));
    }
  }
}
double JointPosEqCost::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = (getTraj(xvec, vars_));
  // Takes diff b/n the subsequent rows and subtract each row by the targets_ vector
  Eigen::MatrixXd diff =
      (traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())).rowwise() - targets_.transpose();
  // Element-wise square it, multiply it by a diagonal matrix of coefficients, and sums output
  return (diff.array().square().matrix() * coeffs_.asDiagonal()).sum();
}
sco::ConvexObjectivePtr JointPosEqCost::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexObjectivePtr out(new sco::ConvexObjective(model));
  out->addQuadExpr(expr_);
  return out;
}

JointPosIneqCost::JointPosIneqCost(const VarArray& vars, const Eigen::VectorXd& coeffs, const Eigen::VectorXd& targets,
                                   const Eigen::VectorXd& upper_tols, const Eigen::VectorXd& lower_tols,
                                   int& first_step, int& last_step)
  : Cost("JointPosIneq")
  , vars_(vars)
  , coeffs_(coeffs)
  , upper_tols_(upper_tols)
  , lower_tols_(lower_tols)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      sco::AffExpr expr;
      sco::AffExpr expr_neg;
      // pos = x1 - targ
      sco::AffExpr pos;
      sco::exprInc(pos, sco::exprMult(vars(i, j), 1));
      sco::exprDec(pos, targets_[j]);
      sco::exprInc(expr, upper_tols_[j]);  // expr_ = upper_tol

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprDec(expr, pos);           // expr = upper_tol_- (vel - targets_)
      sco::exprScale(expr, -coeffs[j]);  // expr = - (upper_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr);

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprDec(expr_neg, pos);          // expr = lower_tol_- (vel - targets_)
      sco::exprScale(expr_neg, coeffs[j]);  // expr = (lower_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr_neg);
    }
  }
}

double JointPosIneqCost::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows to get velocity
  Eigen::MatrixXd pos = traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols());
  // Subtract targets to center about 0 and then subtract from tolerance
  Eigen::MatrixXd diff0 = (pos.rowwise() - targets_.transpose());
  Eigen::MatrixXd diff1 = (diff0.rowwise() - upper_tols_.transpose()) * coeffs_.asDiagonal();
  Eigen::MatrixXd diff2 = ((diff0 * -1).rowwise() + lower_tols_.transpose()) * coeffs_.asDiagonal();
  // Applies hinge, multiplies it by a diagonal matrix of coefficients, sums each corresponding value, and converts to
  // vector
  return diff1.cwiseMax(0).sum() + diff2.cwiseMax(0).sum();
}

sco::ConvexObjectivePtr JointPosIneqCost::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexObjectivePtr out(new sco::ConvexObjective(model));
  // Add hinge cost. Set the coefficient to 1 here since we include it in the AffExpr already
  // This is necessary since we want a seperate coefficient per joint
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addHinge(expr, 1);
  }
  return out;
}

JointPosEqConstraint::JointPosEqConstraint(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                           const Eigen::VectorXd& targets, int& first_step, int& last_step)
  : EqConstraint("JointPosEq")
  , vars_(vars)
  , coeffs_(coeffs)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // pos = x1 - targ
      sco::AffExpr pos;
      sco::exprInc(pos, sco::exprMult(vars(i, j), 1));
      sco::exprDec(pos, targets_[j]);
      // expr_ = coeff * vel - Not squared b/c QuadExpr cnt not yet supported (TODO)
      expr_vec_.push_back(sco::exprMult(pos, coeffs_[j]));
    }
  }
}

DblVec JointPosEqConstraint::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows and subtract each row by the targets_ vector
  Eigen::MatrixXd diff =
      (traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())).rowwise() - targets_.transpose();
  // Squares it, multiplies it by a diagonal matrix of coefficients, and converts to vector
  return util::toDblVec((diff.array().square()).matrix() * coeffs_.asDiagonal());
}
sco::ConvexConstraintsPtr JointPosEqConstraint::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexConstraintsPtr out(new sco::ConvexConstraints(model));
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addEqCnt(expr);
  }
  return out;
}

JointPosIneqConstraint::JointPosIneqConstraint(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                               const Eigen::VectorXd& targets, const Eigen::VectorXd& upper_tols,
                                               const Eigen::VectorXd& lower_tols, int& first_step, int& last_step)
  : IneqConstraint("JointPosIneq")
  , vars_(vars)
  , coeffs_(coeffs)
  , upper_tols_(upper_tols)
  , lower_tols_(lower_tols)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      sco::AffExpr expr;
      sco::AffExpr expr_neg;
      // pos = x1 - targ
      sco::AffExpr pos;
      sco::exprInc(pos, sco::exprMult(vars(i, j), 1));
      sco::exprDec(pos, targets_[j]);
      sco::exprInc(expr, upper_tols_[j]);  // expr_ = upper_tol

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprDec(expr, pos);           // expr = upper_tol_- (vel - targets_)
      sco::exprScale(expr, -coeffs[j]);  // expr = - (upper_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr);

      // Form lower limit expr = (upper_tol-(vel-targ))
      sco::exprDec(expr_neg, pos);          // expr = lower_tol_- (vel - targets_)
      sco::exprScale(expr_neg, coeffs[j]);  // expr = (lower_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr_neg);
    }
  }
}

DblVec JointPosIneqConstraint::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows to get velocity
  Eigen::MatrixXd pos = diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols()));
  // Subtract targets to center about 0 and then subtract from tolerance
  Eigen::MatrixXd diff0 = (pos.rowwise() - targets_.transpose());
  Eigen::MatrixXd diff1 = (diff0.rowwise() - upper_tols_.transpose()) * coeffs_.asDiagonal();
  Eigen::MatrixXd diff2 = ((diff0 * -1).rowwise() + lower_tols_.transpose()) * coeffs_.asDiagonal();
  // Applies hinge, multiplies it by a diagonal matrix of coefficients, sums each corresponding value, and converts to
  // vector
  Eigen::MatrixXd out(diff1.rows(), diff1.cols() + diff2.cols());
  out << diff1, diff2;
  return util::toDblVec(out);
}

sco::ConvexConstraintsPtr JointPosIneqConstraint::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexConstraintsPtr out(new sco::ConvexConstraints(model));
  // Add hinge cost. Set the coefficient to 1 here since we include it in the AffExpr already
  // This is necessary since we want a seperate coefficient per joint
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addIneqCnt(expr);
  }
  return out;
}

//////////////////// Velocity /////////////////////
JointVelEqCost::JointVelEqCost(const VarArray& vars, const Eigen::VectorXd& coeffs, const Eigen::VectorXd& targets,
                               int& first_step, int& last_step)
  : Cost("JointVelEq"), vars_(vars), coeffs_(coeffs), targets_(targets), first_step_(first_step), last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 1; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // vel = (x2 - x1) - targ
      sco::AffExpr vel;
      sco::exprInc(vel, sco::exprMult(vars(i, j), -1));
      sco::exprInc(vel, sco::exprMult(vars(i + 1, j), 1));
      exprDec(vel, targets_[j]);
      // expr_ = coeff * vel^2
      exprInc(expr_, exprMult(exprSquare(vel), coeffs_[j]));
    }
  }
}
double JointVelEqCost::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = (getTraj(xvec, vars_));
  // Takes diff b/n the subsequent rows and subtract each row by the targets_ vector
  Eigen::MatrixXd diff = (diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols()))).rowwise() -
                         targets_.transpose();
  // Element-wise square it, multiply it by a diagonal matrix of coefficients, and sums output
  return (diff.array().square().matrix() * coeffs_.asDiagonal()).sum();
}
sco::ConvexObjectivePtr JointVelEqCost::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexObjectivePtr out(new sco::ConvexObjective(model));
  out->addQuadExpr(expr_);
  return out;
}

JointVelIneqCost::JointVelIneqCost(const VarArray& vars, const Eigen::VectorXd& coeffs, const Eigen::VectorXd& targets,
                                   const Eigen::VectorXd& upper_tols, const Eigen::VectorXd& lower_tols,
                                   int& first_step, int& last_step)
  : Cost("JointVelIneq")
  , vars_(vars)
  , coeffs_(coeffs)
  , upper_tols_(upper_tols)
  , lower_tols_(lower_tols)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 1; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // vel = (x2 - x1) - targ
      sco::AffExpr vel;
      sco::AffExpr expr;
      sco::AffExpr expr_neg;
      sco::exprInc(vel, sco::exprMult(vars(i, j), -1));
      sco::exprInc(vel, sco::exprMult(vars(i + 1, j), 1));
      sco::exprDec(vel, targets_[j]);  // offset to center about 0

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprInc(expr, upper_tols_[j]);  // expr_ = upper_tol
      sco::exprDec(expr, vel);             // expr = upper_tol_- (vel - targets_)
      sco::exprScale(expr, -coeffs[j]);    // expr = - (upper_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr);

      // Form lower limit expr =  (upper_tol-(vel-targ))
      sco::exprInc(expr_neg, lower_tols_[j]);  // expr_ = lower_tol_
      sco::exprDec(expr_neg, vel);             // expr = lower_tol_- (vel - targets_)
      sco::exprScale(expr_neg, coeffs[j]);     // expr = (lower_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr_neg);
    }
  }
}

double JointVelIneqCost::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows to get velocity
  Eigen::MatrixXd vel = diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols()));
  // Subtract targets_ to center about 0 and then subtract from tolerance
  Eigen::MatrixXd diff0 = (vel.rowwise() - targets_.transpose());
  Eigen::MatrixXd diff1 = (diff0.rowwise() - upper_tols_.transpose()) * coeffs_.asDiagonal();
  Eigen::MatrixXd diff2 = ((diff0 * -1).rowwise() + lower_tols_.transpose()) * coeffs_.asDiagonal();
  // Applies hinge, multiplies it by a diagonal matrix of coefficients, sums each corresponding value, and converts to
  // vector
  return diff1.cwiseMax(0).sum() + diff2.cwiseMax(0).sum();
}

sco::ConvexObjectivePtr JointVelIneqCost::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexObjectivePtr out(new sco::ConvexObjective(model));
  // Add hinge cost. Set the coefficient to 1 here since we include it in the AffExpr already
  // This is necessary since we want a seperate coefficient per joint
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addHinge(expr, 1);
  }
  return out;
}

JointVelEqConstraint::JointVelEqConstraint(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                           const Eigen::VectorXd& targets, int& first_step, int& last_step)
  : EqConstraint("JointVelEq")
  , vars_(vars)
  , coeffs_(coeffs)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 1; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // vel = (x2 - x1) - targ
      sco::AffExpr vel;
      sco::exprInc(vel, sco::exprMult(vars(i, j), -1));
      sco::exprInc(vel, sco::exprMult(vars(i + 1, j), 1));
      sco::exprDec(vel, targets_[j]);
      // expr_ = coeff * vel - Not squared b/c QuadExpr cnt not yet supported (TODO)
      expr_vec_.push_back(sco::exprMult(vel, coeffs_[j]));
    }
  }
}

DblVec JointVelEqConstraint::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows and subtract each row by the targets_ vector
  Eigen::MatrixXd diff = (diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols()))).rowwise() -
                         targets_.transpose();
  // Squares it, multiplies it by a diagonal matrix of coefficients, and converts to vector
  return util::toDblVec((diff.array().square()).matrix() * coeffs_.asDiagonal());
}
sco::ConvexConstraintsPtr JointVelEqConstraint::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexConstraintsPtr out(new sco::ConvexConstraints(model));
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addEqCnt(expr);
  }
  return out;
}

JointVelIneqConstraint::JointVelIneqConstraint(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                               const Eigen::VectorXd& targets, const Eigen::VectorXd& upper_tols,
                                               const Eigen::VectorXd& lower_tols, int& first_step, int& last_step)
  : IneqConstraint("JointVelIneq")
  , vars_(vars)
  , coeffs_(coeffs)
  , upper_tols_(upper_tols)
  , lower_tols_(lower_tols)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 1; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // vel = (x2 - x1) - targ
      sco::AffExpr vel;
      sco::AffExpr expr;
      sco::AffExpr expr_neg;
      sco::exprInc(vel, sco::exprMult(vars(i, j), -1));
      sco::exprInc(vel, sco::exprMult(vars(i + 1, j), 1));
      sco::exprDec(vel, targets_[j]);  // offset to center about 0

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprInc(expr, upper_tols_[j]);  // expr_ = upper_tol
      sco::exprDec(expr, vel);             // expr = upper_tol_- (vel - targets_)
      sco::exprScale(expr, -coeffs[j]);    // expr = - (upper_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr);

      // Form lower limit expr = (lower_tol-(vel-targ))
      sco::exprInc(expr_neg, lower_tols_[j]);  // expr_ = lower_tol_
      sco::exprDec(expr_neg, vel);             // expr = lower_tol_- (vel - targets_)
      sco::exprScale(expr_neg, coeffs[j]);     // expr = (lower_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr_neg);
    }
  }
}

DblVec JointVelIneqConstraint::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows to get velocity
  Eigen::MatrixXd vel = diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols()));
  // Subtract targets_ to center about 0 and then subtract from tolerance
  Eigen::MatrixXd diff0 = (vel.rowwise() - targets_.transpose());
  Eigen::MatrixXd diff1 = (diff0.rowwise() - upper_tols_.transpose()) * coeffs_.asDiagonal();
  Eigen::MatrixXd diff2 = ((diff0 * -1).rowwise() + lower_tols_.transpose()) * coeffs_.asDiagonal();
  // Applies hinge, multiplies it by a diagonal matrix of coefficients, sums each corresponding value, and converts to
  // vector
  Eigen::MatrixXd out(diff1.rows(), diff1.cols() + diff2.cols());
  out << diff1, diff2;
  return util::toDblVec(out.cwiseMax(0));
}

sco::ConvexConstraintsPtr JointVelIneqConstraint::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexConstraintsPtr out(new sco::ConvexConstraints(model));
  // Add hinge cost. Set the coefficient to 1 here since we include it in the AffExpr already
  // This is necessary since we want a seperate coefficient per joint
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addIneqCnt(expr);
  }
  return out;
}

//////////////////// Acceleration /////////////////////
JointAccEqCost::JointAccEqCost(const VarArray& vars, const Eigen::VectorXd& coeffs, const Eigen::VectorXd& targets,
                               int& first_step, int& last_step)
  : Cost("JointAccEq"), vars_(vars), coeffs_(coeffs), targets_(targets), first_step_(first_step), last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 2; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // acc = (x3 - 2*x2 + x1) - targ
      sco::AffExpr acc;
      sco::exprInc(acc, sco::exprMult(vars(i, j), 1.0));
      sco::exprInc(acc, sco::exprMult(vars(i + 1, j), -2.0));
      sco::exprInc(acc, sco::exprMult(vars(i + 2, j), 1.0));

      sco::exprDec(acc, targets_[j]);
      // expr_ = coeff * acc^2
      sco::exprInc(expr_, sco::exprMult(sco::exprSquare(acc), coeffs_[j]));
    }
  }
}
double JointAccEqCost::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = (getTraj(xvec, vars_));
  // Takes diff b/n the subsequent rows and subtract each row by the targets_ vector
  Eigen::MatrixXd diff =
      (diffAxis0(diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())))).rowwise() -
      targets_.transpose();
  // Element-wise square it, multiply it by a diagonal matrix of coefficients, and sums output
  return (diff.array().square().matrix() * coeffs_.asDiagonal()).sum();
}
sco::ConvexObjectivePtr JointAccEqCost::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexObjectivePtr out(new sco::ConvexObjective(model));
  out->addQuadExpr(expr_);
  return out;
}

JointAccIneqCost::JointAccIneqCost(const VarArray& vars, const Eigen::VectorXd& coeffs, const Eigen::VectorXd& targets,
                                   const Eigen::VectorXd& upper_tols, const Eigen::VectorXd& lower_tols,
                                   int& first_step, int& last_step)
  : Cost("JointAccIneq")
  , vars_(vars)
  , coeffs_(coeffs)
  , upper_tols_(upper_tols)
  , lower_tols_(lower_tols)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 2; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // acc = (x3 - 2*x2 + x1) - targ
      sco::AffExpr acc;
      sco::AffExpr expr;
      sco::AffExpr expr_neg;
      sco::exprInc(acc, sco::exprMult(vars(i, j), 1.0));
      sco::exprInc(acc, sco::exprMult(vars(i + 1, j), -2.0));
      sco::exprInc(acc, sco::exprMult(vars(i + 2, j), 1.0));
      sco::exprDec(acc, targets_[j]);  // offset to center about 0

      // Form upper limit expr = - (upper_tol-(acc-targ))
      sco::exprInc(expr, upper_tols_[j]);  // expr_ = upper_tol
      sco::exprDec(expr, acc);             // expr = upper_tol_- (vel - targets_)
      sco::exprScale(expr, -coeffs[j]);    // expr = - (upper_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr);

      // Form lower limit expr = (lower_tol-(acc-targ))
      sco::exprInc(expr_neg, lower_tols_[j]);  // expr_ = lower_tol_
      sco::exprDec(expr_neg, acc);             // expr = lower_tol_- (vel - targets_)
      sco::exprScale(expr_neg, coeffs[j]);     // expr = (lower_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr_neg);
    }
  }
}

double JointAccIneqCost::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows to get velocity
  Eigen::MatrixXd acc = diffAxis0(diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())));
  // Subtract targets_ to center about 0 and then subtract from tolerance
  Eigen::MatrixXd diff0 = (acc.rowwise() - targets_.transpose());
  Eigen::MatrixXd diff1 = (diff0.rowwise() - upper_tols_.transpose()) * coeffs_.asDiagonal();
  Eigen::MatrixXd diff2 = ((diff0 * -1).rowwise() + lower_tols_.transpose()) * coeffs_.asDiagonal();
  // Applies hinge, multiplies it by a diagonal matrix of coefficients, sums each corresponding value, and converts to
  // vector
  return diff1.cwiseMax(0).sum() + diff2.cwiseMax(0).sum();
}

sco::ConvexObjectivePtr JointAccIneqCost::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexObjectivePtr out(new sco::ConvexObjective(model));
  // Add hinge cost. Set the coefficient to 1 here since we include it in the AffExpr already
  // This is necessary since we want a seperate coefficient per joint
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addHinge(expr, 1);
  }
  return out;
}

JointAccEqConstraint::JointAccEqConstraint(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                           const Eigen::VectorXd& targets, int& first_step, int& last_step)
  : EqConstraint("JointAccEq")
  , vars_(vars)
  , coeffs_(coeffs)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 2; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // acc = (x3 - 2*x2 + x1) - targ
      sco::AffExpr acc;
      sco::exprInc(acc, sco::exprMult(vars(i, j), 1.0));
      sco::exprInc(acc, sco::exprMult(vars(i + 1, j), -2.0));
      sco::exprInc(acc, sco::exprMult(vars(i + 2, j), 1.0));

      sco::exprDec(acc, targets_[j]);  // offset to center about 0
      // expr_ = coeff * vel - Not squared b/c QuadExpr cnt not yet supported (TODO)
      expr_vec_.push_back(sco::exprMult(acc, coeffs_[j]));
    }
  }
}

DblVec JointAccEqConstraint::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows and subtract each row by the targets_ vector
  Eigen::MatrixXd diff =
      (diffAxis0(diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())))).rowwise() -
      targets_.transpose();
  // Squares it, multiplies it by a diagonal matrix of coefficients, and converts to vector
  return util::toDblVec((diff.array().square()).matrix() * coeffs_.asDiagonal());
}
sco::ConvexConstraintsPtr JointAccEqConstraint::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexConstraintsPtr out(new sco::ConvexConstraints(model));
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addEqCnt(expr);
  }
  return out;
}

JointAccIneqConstraint::JointAccIneqConstraint(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                               const Eigen::VectorXd& targets, const Eigen::VectorXd& upper_tols,
                                               const Eigen::VectorXd& lower_tols, int& first_step, int& last_step)
  : IneqConstraint("JointAccIneq")
  , vars_(vars)
  , coeffs_(coeffs)
  , upper_tols_(upper_tols)
  , lower_tols_(lower_tols)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  // Form upper limit expr = - (upper_tol-(vel-targ))
  for (int i = first_step_; i <= last_step_ - 2; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      // acc = (x3 - 2*x2 + x1) - targ
      sco::AffExpr acc;
      sco::AffExpr expr;
      sco::AffExpr expr_neg;
      sco::exprInc(acc, sco::exprMult(vars(i, j), 1.0));
      sco::exprInc(acc, sco::exprMult(vars(i + 1, j), -2.0));
      sco::exprInc(acc, sco::exprMult(vars(i + 2, j), 1.0));
      sco::exprDec(acc, targets_[j]);  // offset to center about 0

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprInc(expr, upper_tols_[j]);  // expr_ = upper_tol
      sco::exprDec(expr, acc);             // expr = upper_tol_- (vel - targets_)
      sco::exprScale(expr, -coeffs[j]);    // expr = - (upper_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr);

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprInc(expr_neg, lower_tols_[j]);  // expr_ = lower_tol_
      sco::exprDec(expr_neg, acc);             // expr = lower_tol_- (vel - targets_)
      sco::exprScale(expr_neg, coeffs[j]);     // expr = (lower_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr_neg);
    }
  }
}

DblVec JointAccIneqConstraint::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows to get velocity
  Eigen::MatrixXd acc = diffAxis0(diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())));
  // Subtract targets_ to center about 0 and then subtract from tolerance
  Eigen::MatrixXd diff0 = (acc.rowwise() - targets_.transpose());
  Eigen::MatrixXd diff1 = (diff0.rowwise() - upper_tols_.transpose()) * coeffs_.asDiagonal();
  Eigen::MatrixXd diff2 = ((diff0 * -1).rowwise() + lower_tols_.transpose()) * coeffs_.asDiagonal();
  // Applies hinge, multiplies it by a diagonal matrix of coefficients, sums each corresponding value, and converts to
  // vector
  Eigen::MatrixXd out(diff1.rows(), diff1.cols() + diff2.cols());
  out << diff1, diff2;
  return util::toDblVec(out.cwiseMax(0));
}

sco::ConvexConstraintsPtr JointAccIneqConstraint::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexConstraintsPtr out(new sco::ConvexConstraints(model));
  // Add hinge cost. Set the coefficient to 1 here since we include it in the AffExpr already
  // This is necessary since we want a seperate coefficient per joint
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addIneqCnt(expr);
  }
  return out;
}

//////////////////// Jerk /////////////////////
JointJerkEqCost::JointJerkEqCost(const VarArray& vars, const Eigen::VectorXd& coeffs, const Eigen::VectorXd& targets,
                                 int& first_step, int& last_step)
  : Cost("JointJerkEq"), vars_(vars), coeffs_(coeffs), targets_(targets), first_step_(first_step), last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 4; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      sco::AffExpr jerk;
      sco::exprInc(jerk, sco::exprMult(vars(i, j), -1.0 / 2.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 1, j), 1.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 2, j), 0.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 3, j), -1.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 4, j), 1.0 / 2.0));

      sco::exprDec(jerk, targets_[j]);
      // expr_ = coeff * jerk^2
      sco::exprInc(expr_, sco::exprMult(sco::exprSquare(jerk), coeffs_[j]));
    }
  }
}
double JointJerkEqCost::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = (getTraj(xvec, vars_));
  // Takes diff b/n the subsequent rows and subtract each row by the targets_ vector
  Eigen::MatrixXd diff =
      (diffAxis0(diffAxis0(diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())))))
          .rowwise() -
      targets_.transpose();
  // Element-wise square it, multiply it by a diagonal matrix of coefficients, and sums output
  return (diff.array().square().matrix() * coeffs_.asDiagonal()).sum();
}
sco::ConvexObjectivePtr JointJerkEqCost::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexObjectivePtr out(new sco::ConvexObjective(model));
  out->addQuadExpr(expr_);
  return out;
}

JointJerkIneqCost::JointJerkIneqCost(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                     const Eigen::VectorXd& targets, const Eigen::VectorXd& upper_tols,
                                     const Eigen::VectorXd& lower_tols, int& first_step, int& last_step)
  : Cost("JointJerkIneq")
  , vars_(vars)
  , coeffs_(coeffs)
  , upper_tols_(upper_tols)
  , lower_tols_(lower_tols)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 4; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      sco::AffExpr jerk;
      sco::AffExpr expr;
      sco::AffExpr expr_neg;
      sco::exprInc(jerk, sco::exprMult(vars(i, j), -1.0 / 2.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 1, j), 1.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 2, j), 0.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 3, j), -1.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 4, j), 1.0 / 2.0));
      sco::exprDec(jerk, targets_[j]);

      // Form upper limit expr = - (upper_tol-(jerk-targ))
      sco::exprInc(expr, upper_tols_[j]);  // expr_ = upper_tol
      sco::exprDec(expr, jerk);            // expr = upper_tol_- (vel - targets_)
      sco::exprScale(expr, -coeffs[j]);    // expr = - (upper_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr);

      // Form lower limit expr = (lower_tol-(acc-targ))
      sco::exprInc(expr_neg, lower_tols_[j]);  // expr_ = lower_tol_
      sco::exprDec(expr_neg, jerk);            // expr = lower_tol_- (vel - targets_)
      sco::exprScale(expr_neg, coeffs[j]);     // expr = (lower_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr_neg);
    }
  }
}

double JointJerkIneqCost::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows to get velocity
  Eigen::MatrixXd jerk =
      diffAxis0(diffAxis0(diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols()))));
  // Subtract targets_ to center about 0 and then subtract from tolerance
  Eigen::MatrixXd diff0 = (jerk.rowwise() - targets_.transpose());
  Eigen::MatrixXd diff1 = (diff0.rowwise() - upper_tols_.transpose()) * coeffs_.asDiagonal();
  Eigen::MatrixXd diff2 = ((diff0 * -1).rowwise() + lower_tols_.transpose()) * coeffs_.asDiagonal();
  // Applies hinge, multiplies it by a diagonal matrix of coefficients, sums each corresponding value, and converts to
  // vector
  return diff1.cwiseMax(0).sum() + diff2.cwiseMax(0).sum();
}

sco::ConvexObjectivePtr JointJerkIneqCost::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexObjectivePtr out(new sco::ConvexObjective(model));
  // Add hinge cost. Set the coefficient to 1 here since we include it in the AffExpr already
  // This is necessary since we want a seperate coefficient per joint
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addHinge(expr, 1);
  }
  return out;
}

JointJerkEqConstraint::JointJerkEqConstraint(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                             const Eigen::VectorXd& targets, int& first_step, int& last_step)
  : EqConstraint("JointJerkEq")
  , vars_(vars)
  , coeffs_(coeffs)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 4; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      sco::AffExpr jerk;
      sco::exprInc(jerk, sco::exprMult(vars(i, j), -1.0 / 2.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 1, j), 1.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 2, j), 0.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 3, j), -1.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 4, j), 1.0 / 2.0));

      sco::exprDec(jerk, targets_[j]);  // offset to center about 0
      // expr_ = coeff * jerk - Not squared b/c QuadExpr cnt not yet supported (TODO)
      expr_vec_.push_back(sco::exprMult(jerk, coeffs_[j]));
    }
  }
}

DblVec JointJerkEqConstraint::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows and subtract each row by the targets_ vector
  Eigen::MatrixXd diff =
      (diffAxis0(diffAxis0(diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())))))
          .rowwise() -
      targets_.transpose();
  // Squares it, multiplies it by a diagonal matrix of coefficients, and converts to vector
  return util::toDblVec((diff.array().square()).matrix() * coeffs_.asDiagonal());
}
sco::ConvexConstraintsPtr JointJerkEqConstraint::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexConstraintsPtr out(new sco::ConvexConstraints(model));
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addEqCnt(expr);
  }
  return out;
}

JointJerkIneqConstraint::JointJerkIneqConstraint(const VarArray& vars, const Eigen::VectorXd& coeffs,
                                                 const Eigen::VectorXd& targets, const Eigen::VectorXd& upper_tols,
                                                 const Eigen::VectorXd& lower_tols, int& first_step, int& last_step)
  : IneqConstraint("JointJerkIneq")
  , vars_(vars)
  , coeffs_(coeffs)
  , upper_tols_(upper_tols)
  , lower_tols_(lower_tols)
  , targets_(targets)
  , first_step_(first_step)
  , last_step_(last_step)
{
  for (int i = first_step_; i <= last_step_ - 4; ++i)
  {
    for (int j = 0; j < vars.cols(); ++j)
    {
      sco::AffExpr jerk;
      sco::AffExpr expr;
      sco::AffExpr expr_neg;
      sco::exprInc(jerk, sco::exprMult(vars(i, j), -1.0 / 2.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 1, j), 1.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 2, j), 0.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 3, j), -1.0));
      sco::exprInc(jerk, sco::exprMult(vars(i + 4, j), 1.0 / 2.0));
      sco::exprDec(jerk, targets_[j]);  // offset to center about 0

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprInc(expr, upper_tols_[j]);  // expr_ = upper_tol
      sco::exprDec(expr, jerk);            // expr = upper_tol_- (vel - targets_)
      sco::exprScale(expr, -coeffs[j]);    // expr = - (upper_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr);

      // Form upper limit expr = - (upper_tol-(vel-targ))
      sco::exprInc(expr_neg, lower_tols_[j]);  // expr_ = lower_tol_
      sco::exprDec(expr_neg, jerk);            // expr = lower_tol_- (vel - targets_)
      sco::exprScale(expr_neg, coeffs[j]);     // expr = (lower_tol_- (vel - targets_)) * coeffs_
      expr_vec_.push_back(expr_neg);
    }
  }
}

DblVec JointJerkIneqConstraint::value(const DblVec& xvec)
{
  // Convert vector from optimization to trajectory
  Eigen::MatrixXd traj = getTraj(xvec, vars_);
  // Takes diff b/n the subsequent rows to get velocity
  Eigen::MatrixXd acc = diffAxis0(diffAxis0(traj.block(first_step_, 0, last_step_ - first_step_ + 1, traj.cols())));
  // Subtract targets_ to center about 0 and then subtract from tolerance
  Eigen::MatrixXd diff0 = (acc.rowwise() - targets_.transpose());
  Eigen::MatrixXd diff1 = (diff0.rowwise() - upper_tols_.transpose()) * coeffs_.asDiagonal();
  Eigen::MatrixXd diff2 = ((diff0 * -1).rowwise() + lower_tols_.transpose()) * coeffs_.asDiagonal();
  // Applies hinge, multiplies it by a diagonal matrix of coefficients, sums each corresponding value, and converts to
  // vector
  Eigen::MatrixXd out(diff1.rows(), diff1.cols() + diff2.cols());
  out << diff1, diff2;
  return util::toDblVec(out.cwiseMax(0));
}

sco::ConvexConstraintsPtr JointJerkIneqConstraint::convex(const DblVec& /*x*/, sco::Model* model)
{
  sco::ConvexConstraintsPtr out(new sco::ConvexConstraints(model));
  // Add hinge cost. Set the coefficient to 1 here since we include it in the AffExpr already
  // This is necessary since we want a seperate coefficient per joint
  for (sco::AffExpr& expr : expr_vec_)
  {
    out->addIneqCnt(expr);
  }
  return out;
}

}  // namespace trajopt


//==============================================================================
//         Copyright 2003 - 2012   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2012   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_LINALG_FUNCTIONS_GENERAL_FUNM_HPP_INCLUDED
#define NT2_LINALG_FUNCTIONS_GENERAL_FUNM_HPP_INCLUDED

#include <nt2/linalg/functions/funm.hpp>
#include <nt2/core/container/table/table.hpp>
#include <nt2/core/container/dsl/size.hpp>
#include <nt2/sdk/complex/meta/is_complex.hpp>
#include <nt2/include/constants/zero.hpp>
#include <nt2/include/functions/abs.hpp>
#include <nt2/include/functions/conj.hpp>
#include <nt2/include/functions/colon.hpp>
#include <nt2/include/functions/diag_of.hpp>
#include <nt2/include/functions/divides.hpp>
#include <nt2/include/functions/eye.hpp>
#include <nt2/include/functions/eyeminus.hpp>
#include <nt2/include/functions/factorial.hpp>
#include <nt2/include/functions/from_diag.hpp>
#include <nt2/include/functions/isdiagonal.hpp>
#include <nt2/include/functions/is_finite.hpp>
#include <nt2/include/functions/length.hpp>
#include <nt2/include/functions/linsolve.hpp>
#include <nt2/include/functions/max.hpp>
#include <nt2/include/functions/globalmean.hpp>
#include <nt2/include/functions/mtimes.hpp>
#include <nt2/include/functions/multiplies.hpp>
#include <nt2/include/functions/mnorminf.hpp>
#include <nt2/include/functions/oneplus.hpp>
#include <nt2/include/functions/ones.hpp>
#include <nt2/include/functions/schur.hpp>
#include <nt2/include/functions/size.hpp>
#include <nt2/include/functions/tie.hpp>
#include <nt2/include/functions/transpose.hpp>
#include <nt2/include/functions/triu.hpp>
#include <nt2/include/functions/zeros.hpp>
#include <nt2/include/functions/extent.hpp>
#include <nt2/include/functions/issquare.hpp>
#include <boost/assert.hpp>

// there are many optimisations/ameliorations to be done
// -- passing parameters to functor
// -- fine tuning of types
// -- some code can be shared with logm

namespace nt2
{
  namespace ext
  {
    BOOST_DISPATCH_IMPLEMENT  ( funm_, tag::cpu_
                              , (A0)(A1)
                              , (unspecified_<A0>)
                                (scalar_<unspecified_<A1> >)
                              )
    {
      typedef A1 result_type;
      NT2_FUNCTOR_CALL(2)
      {
        return a0(a1, 0);
      }
    };

    // funm tag only used for functor/matrix
    template<class Domain, int N, class Expr>
    struct size_of<tag::funm_,Domain,N,Expr>
    {
      typedef _2D                                                      result_type;
      BOOST_FORCEINLINE result_type operator()(Expr& e) const
      {
        BOOST_ASSERT_MSG( issquare(boost::proto::child_c<1>(e)),
                          "funm needs a functor and scalar or a ");
        return nt2::extent(boost::proto::child_c<1>(e));
      }
    };

  BOOST_DISPATCH_IMPLEMENT  ( run_assign_, tag::cpu_
                            , (A0)(A1)
                            , ((ast_<A0, nt2::container::domain>))
                              ((node_<A1, nt2::tag::funm_, boost::mpl::long_<2>, nt2::container::domain>))
                            )
  {
      typedef void                                                    result_type;
      typedef typename A0::value_type                                  value_type;
      typedef typename meta::as_real<value_type>::type                     r_type;
      typedef typename meta::as_complex<r_type>::type                   cplx_type;
      typedef typename meta::as_integer<r_type>::type                      i_type;
      typedef nt2::table<value_type >                                       tab_t;
      typedef nt2::table<r_type>                                           btab_t;
      typedef table<cplx_type>                                             ctab_t;
      typedef table<i_type>                                                itab_t;

      BOOST_FORCEINLINE result_type operator()(const A0& a0, const A1& a1) const
      {
        compute_funm(boost::proto::value(boost::proto::child_c<0>(a1)), boost::proto::child_c<1>(a1), a0);
      }
    private:
      template <class F, class T >
      BOOST_FORCEINLINE static void compute_funm(const F& f, const T& a0, A0& res)
      {
         r_type tol = nt2::Eps<r_type>();
         uint32_t maxterms = 250;
         //u, t and r are complex arrays
         res.resize(extent(a0));
         ctab_t u, t;
         nt2::tie(u, t) = schur(a0, nt2::cmplx_);
         if (isdiagonal(t))
         {
           t = nt2::from_diag(f(nt2::diag_of(t), 0));
           BOOST_AUTO_TPL(r, nt2::mtimes(u, nt2::mtimes(t, nt2::trans(nt2::conj(u)))));
           transtype(res, r, typename nt2::meta::is_complex<value_type>::type());
           return;
         }
         else
         {
           ctab_t r = zeros(extent(a0), meta::as_<cplx_type>()); //is it necessary ?
           r_type delta = 0.1;
           itab_t  ord(nt2::of_size(2u, nt2::size(a0,1)));
           blocking(diag_of(t),delta, ord);
           uint32_t lord = nt2::size(ord, 2);
           ctab_t ca0 = a0;

           for(uint32_t col=1; col <= lord ; ++col)
           {
             BOOST_AUTO_TPL(j, nt2::_(ord(1, col), ord(2, col)));

             itab_t terms(nt2::of_size(1, lord));
             uint32_t nj =  length(j);
             ctab_t rj(nt2::of_size(nj, nj));
             terms(col) = funm_atom(f, a0(j, j), tol, maxterms, rj);
             r(j, j) = rj;
             for(uint32_t row=col-1; row >= 1; --row)
             {
               BOOST_AUTO_TPL(i, nt2::_(ord(1, row), ord(2, row)));
               if (length(i) == 1 && length(j) == 1)
               {
                 size_t ii = i(1), jj = j(1);
                 BOOST_AUTO_TPL(k, nt2::_(ii+1, jj-1));
                 cplx_type temp = ca0(ii,jj)*(r(ii,ii) - r(jj,jj));
                 if (!isempty(k)) temp += mtimes(r(ii,k), ca0(k,jj)) - mtimes(ca0(ii,k), r(k,jj));
                 r(ii,jj) = temp/(ca0(ii,ii)-ca0(jj,jj));
               }
               else
               {
                 itab_t k(nt2::of_size(1, 0));
                 for(uint32_t l = row+1; l < col; ++l)
                 {
                   itab_t k1 = horzcat(k, nt2::_(ord(1, l), ord(2, l)));
                   k = k1;
                 }
                 ctab_t rhs =  mtimes(r(i,i), ca0(i,j)) - mtimes(ca0(i,j), r(j,j));
                 if(!isempty(k)) rhs += mtimes(r(i,k), ca0(k,j)) -  mtimes(ca0(i,k), r(k,j));
                 r(i,j) = sylv_tri(ca0(i,i),-ca0(j,j),rhs);
               }
             }
             ctab_t z =  mtimes(mtimes(u, r), trans(conj(u)));
             transtype(res, z, typename nt2::meta::is_complex<value_type>::type());
           }
         }
       }

      template <class F, class X, class N> static inline
        cplx_type feval(const F & f, const X& x,  const N& k)
      {
        ctab_t xx = x;
        return f(xx, k);
      }

      template < class F, class D1, class D2> static inline
      int32_t funm_atom(const F& f, const D1& t,
                        const r_type& tol,
                        const uint32_t& maxterms,
                        D2& rj)
      {
        // Function of triangular matrix with nearly constant diagonal.
        //   n_terms = funm_atom(f, a0, tol, maxterms, fa0)
        //   evaluates function f at the upper triangular matrix t,
        //   where t has nearly constant diagonal.
        //   a taylor series is used, taking at most maxterms terms.
        //   the function represented by f must have a taylor series with an
        //   infinite radius of convergence.
        //   f(x,k) must return the k'th derivative of
        //   the function represented by fun evaluated at the vector x.
        //   tol is a convergence tolerance for the taylor series, defaulting to eps.
        //   if prnt ~= 0 information is printed on the convergence of the
        //   taylor series evaluation.
        //   n_terms is the number of terms taken in the taylor series.
        //   n_terms  = -1 signals lack of convergence.

        uint32_t n = length(t);
        if (n == 1)
        {
          rj = f(t, 0);
          return 1;
        }
        cplx_type lambda{nt2::globalmean(diag_of(t)),0};
        rj = feval(f, lambda,0)*eye(n, meta::as_<cplx_type>());
        btab_t f_deriv_max = zeros(maxterms+n-1,1, meta::as_<r_type>());
        ctab_t nn = t - lambda*eye(n, meta::as_<cplx_type>());
        r_type mu = mnorminf(linsolve(eyeminus(nt2::abs(triu(t,1))), nt2::ones(n,1, meta::as_<r_type>())));
        ctab_t p = nn;
        uint32_t max_d = 1;
        for(uint32_t k = 1; k <= maxterms; ++k)
        {
          cplx_type fval = feval(f, lambda,k);
          BOOST_ASSERT_MSG(is_finite(fval), "funm infinite derivative");
          ctab_t rj_old = rj;
          rj+= p*fval;
          r_type  rel_diff = nt2::mnorminf(rj - rj_old)/(tol+mnorminf(rj_old));
          p = mtimes(p, nn)/r_type(oneplus(k));
          if (rel_diff <= tol)
          {
            // Approximate the maximum of derivatives in convex set containing
            // eigenvalues by maximum of derivatives at eigenvalues.
            for(uint32_t j = max_d; j <=  k+n-1; ++j)
            {
              f_deriv_max(j) = nt2::mnorminf(f(nt2::diag_of(t),j));
            }
            max_d = k+n;
            r_type omega = nt2::Zero<r_type>();
            for(uint32_t j = 0; j <= n-1; ++j)
            {
              omega = nt2::max(omega,f_deriv_max(k+j)/nt2::factorial(j));
            }
            r_type trunc = nt2::mnorminf(p)*mu*omega;
            if (trunc <= tol*mnorminf(rj))
            {
              return k+1;
            }
          }
        }
        return -1;
      }

      template < class T1, class T2 >
      BOOST_FORCEINLINE static void transtype(T1& r, T2& z, boost::mpl::false_ const &)
      {
        r =  real(z);
      }

      template < class T1, class T2 >
        BOOST_FORCEINLINE static void transtype(T1& r, T2& z, boost::mpl::true_ const &)
      {
        r =  z;
      }
      template < class D> static inline void blocking(const D& a0, const r_type& delta, itab_t & ord)
      {
        uint32_t n = nt2::size(a0, 1);
        uint32_t j = 1;
        ord(1, 1) = 1;
        for(uint32_t i=2; i <= n; ++i)
        {
          if (nt2::abs(a0(i-1)-a0(i)) >= delta)
          {
            ord(2, j) = i-1;
            ord(1, ++j) = i;
          }
        }
        ord(2, j) = n;
        ord.resize(nt2::of_size(2u, j));
      }
      template < class T, class U, class B>
        static inline ctab_t sylv_tri(const T& t,const U& u, const B& b)
      {

        // sylv_tri    solve triangular sylvester equation.
        //    x = sylv_tri(t,u,b) solves the sylvester equation
        //  t*x + x*u = b, where t and u are square upper triangular matrices.

        uint32_t m = length(t);
        uint32_t n = length(u);
        ctab_t x = zeros(m,n, nt2::meta::as_<cplx_type>());
        for(uint32_t i = 1;  i <= n; ++i)
        {
          ctab_t bb = b(nt2::_,i);
          BOOST_AUTO_TPL(ii, nt2::_(1u, i-1));
          if(!isempty(ii)) bb -= mtimes(x(nt2::_,ii),u(ii,i));
          x(nt2::_,i) = linsolve(t + u(i,i)*eye(m, meta::as_<cplx_type>()),
                                 bb);
        }
        return x;
     }
    };

  }
}

#endif


/*
 * cuCTProjectionOperator.cpp
 *
 *  Created on: Feb 18, 2015
 *      Author: u051747
 */

#include "CTProjectionOperator.h"
#include <boost/math/constants/constants.hpp>
#include "ct_projection.h"
#include "hoNDArray_math.h"

using namespace boost::math::float_constants;
namespace Gadgetron {

    template<template<class> class ARRAY>
    CTProjectionOperator<ARRAY>::CTProjectionOperator() {
        // TODO Auto-generated constructor stub

        samples_per_pixel_ = 1.5f;

    }

    template<template<class> class ARRAY>
    CTProjectionOperator<ARRAY>::~CTProjectionOperator() {
        // TODO Auto-generated destructor stub
    }

    template<template<class> class ARRAY>
    void CTProjectionOperator<ARRAY>::mult_M(ARRAY<float> *input,
                                             ARRAY<float> *output, bool accumulate) {

        auto dims = *input->get_dimensions();
        std::vector<size_t> dims3d(dims);

        if (dims3d.size() == 4) dims3d.pop_back();
        auto outdims = *output->get_dimensions();
        std::vector<size_t> outbindims(outdims);
        float *input_ptr = input->get_data_ptr();
        ARRAY<float> *tmp_out = output;
        if (accumulate)
            tmp_out = new ARRAY<float>(output->get_dimensions());

        float *output_ptr = tmp_out->get_data_ptr();
        for (int bin = 0; bin < binning->get_number_of_bins(); bin++) {
            //Check for empty bins
            if (binning->get_bin(bin).size() == 0)
                continue;
            ARRAY<float> input_view(dims3d, input_ptr);
            outbindims.back() = detector_focal_cyls[bin].size();
            auto output_view = ARRAY<float>(outbindims, output_ptr);

            ct_forwards_projection(&output_view, &input_view, detector_focal_cyls[bin], focal_offset_cyls[bin],
                                   central_elements[bin], is_dims_in_mm, ps_spacing, ADD, samples_per_pixel_, false);
            //conebeam_forwards_projection(output_view2.get(),&input_view,angles[bin],offsets[bin],samples_per_pixel_,is_dims_in_mm_,acquisition_->get_geometry()->get_FOV(),acquisition_->get_geometry()->get_SDD(),acquisition_->get_geometry()->get_SAD(),accumulate);

            input_ptr += input_view.get_number_of_elements();
            output_ptr += output_view.get_number_of_elements();

        }
        if (weights)
            *tmp_out *= *weights;

        if (accumulate) {
            *output += *tmp_out;
            delete tmp_out;
        }

    }

    template<template<class> class ARRAY>
    void CTProjectionOperator<ARRAY>::mult_MH(ARRAY<float> *input,
                                              ARRAY<float> *output, bool accumulate) {

        ARRAY<float> *tmp_in = input;
        if (weights) {
            tmp_in = new ARRAY<float>(input);
            *tmp_in *= *weights;
        }
        auto dims = *output->get_dimensions();
        std::vector<size_t> dims3d = dims;
        if (dims3d.size() == 4) dims3d.pop_back();

        auto indims = *input->get_dimensions();
        std::vector<size_t> inbindims(indims);
        float *input_ptr = tmp_in->get_data_ptr();
        float *output_ptr = output->get_data_ptr();
        if (!accumulate)
            clear(output);
        for (int bin = 0; bin < binning->get_number_of_bins(); bin++) {
            //Check for empty bins
            if (binning->get_bin(bin).size() == 0)
                continue;

            ARRAY<float> output_view(dims3d, output_ptr);
            inbindims.back() = detector_focal_cyls[bin].size();
            auto input_view = ARRAY<float>(inbindims, input_ptr);

            vector_td<int, 3> is_dims_in_pixels{dims3d[0], dims3d[1], dims3d[2]};

            ct_backwards_projection(&input_view, &output_view, detector_focal_cyls[bin], focal_offset_cyls[bin],
                                    central_elements[bin], proj_indices[bin], is_dims_in_mm, ps_spacing, ADD,
                                    accumulate);
            input_ptr += input_view.get_number_of_elements();
            output_ptr += output_view.get_number_of_elements();
        }

        if (weights) {
            delete tmp_in;
        }

    }

    template<template<class> class ARRAY>
    void Gadgetron::CTProjectionOperator<ARRAY>::setup(boost::shared_ptr<CT_acquisition> acquisition,
                                                       floatd3 is_dims_in_mm) {
        std::vector<unsigned int> bins(acquisition->geometry.detectorFocalCenterAngularPosition.size());
        std::iota(bins.begin(), bins.end(), 0);
        auto tmp_binning = boost::make_shared<CBCT_binning>(std::vector<std::vector<unsigned int>>(1, bins));
        this->setup(acquisition, tmp_binning, is_dims_in_mm);
    }

    template<template<class> class ARRAY>
    std::vector<intd2> Gadgetron::CTProjectionOperator<ARRAY>::calculate_slice_indices(CT_acquisition &acquisition) {

        floatd2 detectorSize = acquisition.geometry.detectorSize;
        auto &centralElements = acquisition.geometry.detectorCentralElement;
        auto &axialPosition = acquisition.geometry.detectorFocalCenterAxialPosition;

        auto image_dims = *this->get_domain_dimensions();
        auto proj_dims = *this->get_codomain_dimensions();

        std::vector<intd2> slice_indices(image_dims[2]);
        std::fill(slice_indices.begin(), slice_indices.end(), intd2(0, 0));
        std::vector<float> start_point(centralElements.size() + 1);

        for (int i = 0; i < centralElements.size(); i++) {
            start_point[i] = axialPosition[i] + (centralElements[i][1] - float(proj_dims[1]) / 2) * detectorSize[1] -
                             detectorSize[1] * proj_dims[1] / 2;
            //std::cout << "Start point " << start_point[i] << std::endl;
        }
        start_point.back() = 2 * start_point[centralElements.size() - 1] - start_point[centralElements.size() - 2];

        int projection_start = 0;
        int projection_stop = 0;
        for (int i = 0; i < slice_indices.size(); i++) {
            float slice_start = is_dims_in_mm[2] / image_dims[2] * (i - 0.5f) - is_dims_in_mm[2] / 2;
            float slice_stop = is_dims_in_mm[2] / image_dims[2] * (i + 0.5f) - is_dims_in_mm[2] / 2;
            //std::cout << "Slice start " << slice_start << " slice end " << slice_stop << std::endl;
            while (slice_start > (start_point[projection_start] + detectorSize[1] * proj_dims[1])) {
                projection_start++;
                if (projection_start >= centralElements.size()) {
                    projection_start = centralElements.size();
                    break;
                }
            }
            while (slice_stop > start_point[projection_stop]) {

                projection_stop++;
                if (projection_stop >= centralElements.size()) {
                    projection_stop = centralElements.size();
                    break;
                }
            }

            slice_indices[i][0] = projection_start;
            slice_indices[i][1] = projection_stop;
            //std::cout << "Projection start stop " << projection_start << " " << projection_stop << " start point " << start_point[projection_start] << " " << start_point[projection_stop] << " " << detectorSize[1] << std::endl;
        }

        std::cout << "Indices size " << slice_indices.size() << std::endl;
        return slice_indices;

    }

    template<template<class> class ARRAY>
    void Gadgetron::CTProjectionOperator<ARRAY>::setup(boost::shared_ptr<CT_acquisition> acquisition,
                                                       boost::shared_ptr<ARRAY<float>> weights, floatd3 is_dims_in_mm) {
        this->weights = weights;
        setup(acquisition, is_dims_in_mm);
    }

    template<template<class> class ARRAY>
    void Gadgetron::CTProjectionOperator<ARRAY>::setup(boost::shared_ptr<CT_acquisition> acquisition,
                                                       boost::shared_ptr<CBCT_binning> binning, floatd3 is_dims_in_mm) {
        this->binning = binning;
        auto bins = binning->get_bins();
        this->is_dims_in_mm = is_dims_in_mm;

        //Variables needed for differewnt
        detector_focal_cyls = std::vector<std::vector<floatd3>>(bins.size());
        focal_offset_cyls = std::vector<std::vector<floatd3>>(bins.size());
        central_elements = std::vector<std::vector<floatd2>>(bins.size());
        proj_indices = std::vector<std::vector<intd2>>(bins.size());

        proj_indices[0] = calculate_slice_indices(*acquisition);
        CT_geometry &geometry = acquisition->geometry;
        ps_spacing = acquisition->geometry.detectorSize;
        ADD = acquisition->geometry.constantRadialDistance[0];

        if (bins.size() != 1) throw std::runtime_error("CT reconstruction does not fully support 4D data yet");
        for (size_t b = 0; b < bins.size(); b++) {
            for (auto i : bins[b]) {
                detector_focal_cyls[b].emplace_back(geometry.detectorFocalCenterAngularPosition[i],
                                                    geometry.detectorFocalRadialDistance[i],
                                                    geometry.detectorFocalCenterAxialPosition[i]);

                focal_offset_cyls[b].emplace_back(geometry.sourceAngularPositionShift[i],
                                                  geometry.sourceRadialDistanceShift[i],
                                                  geometry.sourceAxialPositionShift[i]);

                central_elements[b].push_back(geometry.detectorCentralElement[i]);
                //proj_indices[b].push_back(all_proj_indices[i]);
            }

        }
    }

    template
    class CTProjectionOperator<cuNDArray>;

    template
    class CTProjectionOperator<hoCuNDArray>;

} /* namespace Gadgetron */


 /*
  Copyright (c) 2016, Bart Vermeulen
  All rights reserved.

  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are
  met:

      * Redistributions of source code must retain the above copyright
        notice, this list of conditions and the following disclaimer.
      * Redistributions in binary form must reproduce the above copyright
        notice, this list of conditions and the following disclaimer in
        the documentation and/or other materials provided with the distribution

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  POSSIBILITY OF SUCH DAMAGE.
*/

#include "Point.hpp"
#include <CGAL/basic.h>
#include <CGAL/Search_traits.h>
#include <CGAL/Orthogonal_incremental_neighbor_search.h>
#include <CGAL/iterator.h>
#include "mex.h"
#include <math.h>
#include <algorithm>
#include <iterator>
#include <Eigen/Dense>
#include <thread>
#include <functional>
#include <future>
#include <chrono>

using namespace std;

extern void main_();

typedef CGAL::Search_traits<double, Point, Point::Cit, Construct_coord_iterator> Traits; // search traits
typedef CGAL::Orthogonal_incremental_neighbor_search<Traits> K_inc_neighbor_search; // incremental searcher
typedef K_inc_neighbor_search::Tree Tree; // The tree
typedef K_inc_neighbor_search::iterator Piterator; //The search iterator
typedef K_inc_neighbor_search::Point_with_transformed_distance P_with_dist; //Return type of searcher
struct Point_rw_not_zero{
	bool operator() (const Piterator& it){
		return it->first.rw()==0; //functor to exclude points with rw == 0
	}
};
typedef CGAL::Filter_iterator< Piterator, Point_rw_not_zero > P_pos_rw_it; // construct filtered iterator

void loess(vector< Point > &,const vector< Point >&, vector< double > &, mwSize, mwSize, mwSize, mwSize);

typedef std::vector< P_with_dist >::const_iterator regpoints_iit;
typedef std::back_insert_iterator< std::vector< double > > weights_oit;
void triCube(regpoints_iit, regpoints_iit, weights_oit);

void biCube(Tree const & tree, const vector< double > &);

void localFit( Tree const &, vector< Point > const & qp,mwSize q, vector< double > &, mwSize, mwSize, double &);
double median(vector< double >);

void mexFunction(int nlhs,mxArray *plhs[],int nrhs, const mxArray * prhs[]){
   	mwSize nin,nd,nout,q,niter=0, order=1, nthreads=thread::hardware_concurrency() ;
	double * span=0, *x,*v,*xi,*vi;
	if (nrhs < 4)
		mexErrMsgIdAndTxt("loess:nargin","At least 4 arguments are required");
	// Input checking: x, v, xi
	nin=mxGetM( prhs[0] );
	nd=mxGetN( prhs[0] );
	nout=mxGetM( prhs[2] );
	if ( !( mxGetNumberOfDimensions( prhs[0] )==2 && mxGetNumberOfDimensions( prhs[1] )==2 && mxGetNumberOfDimensions( prhs[2] ) ==2 ) )
		mexErrMsgIdAndTxt("loess:notArray","Inputs should not have more than two dimensions");
	if ( nin != static_cast<mwSize>( mxGetM( prhs[1] ) ) )
		mexErrMsgIdAndTxt("loess:nin","Second input (values) should have same number of rows as first input (locations)");
	if ( nd != static_cast<mwSize>( mxGetN(prhs[2] ) ) )
		mexErrMsgIdAndTxt("loess:nd","Third input (query points) should have same number of columns as first input (locations)");
	if ( mxGetN( prhs[1] ) != 1)
		mexErrMsgIdAndTxt("loess:notColumnV","Second input should be a column vector");
	if ( !( mxIsDouble( prhs[0] ) && mxIsDouble( prhs[1] ) && mxIsDouble( prhs[2] ) && mxIsDouble( prhs[3] )  ) )
		mexErrMsgIdAndTxt("loess:isdouble","All input arguments should be of type Double");
	if ( mxGetNumberOfElements( prhs[3] ) != 1 )
		mexErrMsgIdAndTxt("loess:spanNotScalar","Fourth argument should be a scalar");
	// if ( nin < min_dat )
	// 	mexErrMsgIdAndTxt("loess:smallInput","At least 3 rows required in first input");
	//Read required inputs
	span = mxGetPr( prhs[3] );
	x = mxGetPr( prhs[0] );
	v = mxGetPr( prhs[1] );
	xi = mxGetPr( prhs [2] );

	if (nrhs > 4){
		if ( mxGetNumberOfElements( prhs[4] ) != 1)
			mexErrMsgIdAndTxt("loess:niterNotScalar","Fifth argument should be a scalar");
		if ( floor( * mxGetPr( prhs[4] )) != * mxGetPr( prhs[4] ) )
			mexErrMsgIdAndTxt("loess:niterNotInteger","Fifth argument should be an integer");
		if ( * mxGetPr( prhs[4] ) < 0 )
			mexErrMsgIdAndTxt("loess:niterNegative","Fifth argument should be positive");
		niter=static_cast< mwSize >(* mxGetPr( prhs[4] ) );
	}

	if (nrhs > 5){
		if ( floor( * mxGetPr( prhs[5] )) != * mxGetPr( prhs[5] ) )
			mexErrMsgIdAndTxt("loess:orderNotInteger","Sixth argument should be an integer");
		if ( mxGetNumberOfElements( prhs[5] ) != 1 )
			mexErrMsgIdAndTxt("loess:orderNotScalar","Sixth argument should be a scalar");
		if ( * mxGetPr( prhs[5] ) !=1 &&  * mxGetPr( prhs[5] ) !=2)
			mexErrMsgIdAndTxt("loess:orderNotOneOrTwo","Sixth argument should be equal to one or two");
		order = static_cast< mwSize >(* mxGetPr( prhs[5] ) );
	}

	if (nrhs > 6){
		if ( floor( * mxGetPr( prhs[6] )) != * mxGetPr( prhs[6] ) )
			mexErrMsgIdAndTxt("loess:orderNotInteger","Seventh argument (nthreads) should be an integer");
		if ( mxGetNumberOfElements( prhs[6] ) != 1 )
			mexErrMsgIdAndTxt("loess:orderNotScalar","Seventh argument (nthreads) should be a scalar");
		if ( * mxGetPr( prhs[6] ) < 1 ){
			mexWarnMsgIdAndTxt("loess:orderNotOneOrTwo","Using one thread instead!");
			nthreads=1;
		}
		else if ( * mxGetPr( prhs[6] ) > thread::hardware_concurrency() ){
			mexWarnMsgIdAndTxt("loess:orderNotOneOrTwo","Using number of available threads instead!");
			nthreads = thread::hardware_concurrency();
		}
		else
			nthreads = static_cast< unsigned int >(* mxGetPr( prhs[6] ) );
	}

	//Make output variable
	plhs[0]=mxCreateDoubleMatrix(nout,1, mxREAL);
	vi=mxGetPr( plhs[0] );

	//Transfor vars into vectors
	vector< Point > inpoints, outpoints;
	vector< double > valsout;
	//inpoints.resize(nin,Point(nd,0));
	outpoints.resize(nout,Point(nd,0));
	valsout.resize(nout,std::numeric_limits<double>::quiet_NaN() );
	for (mwSize cin=0; cin < nin; cin++){
		if (!std::isfinite(v[cin]))
			continue;
		Point tmpPoint(nd,0);
		tmpPoint.val(v[cin]);
		bool point_is_finite=true;
		for (mwSize cd=0; (cd < nd) & point_is_finite; cd++){
			if (!std::isfinite(x[cd*nin+cin])){
				point_is_finite=false;
				continue;
			}
			tmpPoint[cd]=(x[cd*nin+cin]);
		}
		if (point_is_finite)
			inpoints.push_back(tmpPoint);
	}
	for (mwSize co=0; co<nout; co++){
		for (mwSize cd=0; cd < nd; cd++) {
			outpoints[co][cd]=xi[cd*nout+co];
		}
	}

	//Make span
	if (*span > 1)
		q=static_cast<mwSize>(round(*span));
	else
		q=static_cast<mwSize>(round(*span * static_cast<double> (nin)));

	q=max(static_cast<mwSize>(3),min(nin,q));

	//Perform computation
	loess(inpoints, outpoints, valsout, q, niter, order, nthreads);

	// Copy output
	for (mwSize co=0; co<nout; co++){
		vi[co]=valsout[co];
	}
}

void loess(vector< Point > & inpoints,const vector < Point > & outpoints, vector <double> & valsout, mwSize q, mwSize niter, mwSize order, mwSize nthreads){
	// Build search tree
	Tree tree(inpoints.begin(), inpoints.end());

	// Perform regression on input data
	double prog_lf(0);
	vector<double> vals_reg=vector<double>(inpoints.size(),0); // holds regression results at input locations
	double frac_riter = double(inpoints.size())/double(inpoints.size()*niter+outpoints.size());
	double frac_interp = double(outpoints.size())/double(inpoints.size()*niter+outpoints.size());
	double prog(0);
	std::cout << std::fixed << std::setw(10) << std::setprecision(2);
	for (mwSize citer=0; citer < niter; citer++) {// robust iterations
	    prog_lf=0;
	    std::future< void > f = std::async(std::launch::async, localFit, std::ref(tree), std::ref(inpoints),q,std::ref(vals_reg), order, nthreads, std::ref(prog_lf));
		std::future_status status;
	    do {
	    	status = f.wait_for(std::chrono::seconds(1));
	    	prog = (double(citer) + prog_lf) * frac_riter;
                //mexPrintf("\r%10.2f%%", prog*100);
	    	//std::cout << "\r" << prog*100 << "%" << std::flush;
	    } while (status != std::future_status::ready);
	    biCube(tree,vals_reg); // Compute robust weights (This is not included in computation of progress)
	}

	// Perform regression on output points
	prog_lf = 0;
	std::future< void > f = std::async(std::launch::async, localFit, std::ref(tree), std::ref(outpoints),q,std::ref(valsout),order, nthreads, std::ref(prog_lf));
	std::future_status status;
    do {
    	status = f.wait_for(std::chrono::seconds(1));
    	prog = prog_lf * frac_interp + double(niter) * frac_riter;
        //mexPrintf("\r%10.2f%%", prog*100);
    	//std::cout << "\r" << prog*100 << "%" << std::flush;
    } while (status != std::future_status::ready);
    //std::cout << "\rDone.      " << std::endl;
}

void triCube(regpoints_iit pd_begin, regpoints_iit pd_end,  weights_oit w_it) {
	// Computes regression weights
	double arg=0;
	for (auto pd_it=pd_begin; pd_it != pd_end; pd_it++){
		arg=pd_it->second / (pd_end-1)->second;
		*w_it = pd_it->first.rw() * ( (arg < 1) ? pow(1 - pow(arg,1.5),1.5) : 0); // This is tricube, but distances are given squared, and we want to use the square of the weights in weighted regression hence twice the 1.5 exponent
		++w_it;
	}
}

void biCube(Tree const & tree, const vector<double> & vals_reg ){
	// biCube function for the computation of robust weights from residuals of the fit
	double arg=0;
	// calculate residuals
	vector<double> res=vector<double>(tree.size(),0);
	for (std::size_t cp=0; cp<static_cast<std::size_t>(tree.size()); cp++)
		res[cp]=abs((tree.begin()+cp)->val()-vals_reg[cp]);
	double sixmres=6*median(res);
	auto tit=tree.begin();
	for (std::size_t cp=0; cp<static_cast<std::size_t>(tree.size()); cp++){
		arg=res[cp]/sixmres;
		const_cast<Point *> (&(*tit))->rw( (arg < 1) ? 1-pow(arg,2) : 0 ); //const_cast shouldn't do any harm here. Second square is removed since we use square root of weights in regression
		tit++;
	}
}

void parFit(const Tree & tree, vector< Point >::const_iterator qp_begin, vector< Point >::const_iterator qp_end, vector< double >::iterator val_beg, mwSize q, mwSize n, mwSize order, double & prog){
// Performs the actual local regression
// Input:
// 		Tree:      the spatial search tree (not modified)
// 		qp_begin:  begin iterator of query points (not modified)
// 		qp_end:	   end iterator of query points (not modified)
// 		val_begin: begin iterator of values (modified to hold output of regression)
// 		q:         number of points for regression (not modified)
// 		n:		   number of terms in regression (not modified)
// 		order:	   order of regression (not modified)
// 		prog:	   to keep track of progress, between 0 and 1 (modified)
	// Initialize variable
	mwSize ndims=qp_begin->dims();      // number of dimensions

	// Search for N-nearest neighbors and perform regression
	for (auto qp_it = qp_begin; qp_it != qp_end; qp_it++){ // loop over all query points
		bool point_is_finite=true;
		for (auto c_it = qp_it->begin(); c_it != qp_it->end(); c_it++)
			if (!std::isfinite(*c_it))
				point_is_finite=false;
		if (!point_is_finite){
			++val_beg;
			continue;
		}
		K_inc_neighbor_search ins (tree,*qp_it); //Create incremental searcher
		P_pos_rw_it it(ins.end(), Point_rw_not_zero(), ins.begin()), end(ins.end(), Point_rw_not_zero()); //filtered iterator to exclude non-finite values and nan-values
		vector< P_with_dist > regpoints; // holds nearest neighbor search results

		// Copy N-nearest points to current query point
		for (mwSize cc=0; cc < q && it!=end; cc++){
			regpoints.push_back(*it);// store filtered nearest neighbors
			it++;
		}

		if (regpoints.size() < n){
			++val_beg;
			continue;
		}
		// Compute weights for regression
		vector< double > w; 				// hold regression weights
		triCube(regpoints.cbegin(), regpoints.cend(), std::back_inserter(w));

		// Eigen matrices for regression
		Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> A(regpoints.size(),n); // Regression matrix
		Eigen::Matrix<double, Eigen::Dynamic, 1> x(n,1);		      // Regression result vector
		Eigen::Matrix<double, Eigen::Dynamic, 1> b(regpoints.size(),1);              // Known values for regression
		// Fill EIGEN matrices and perform regression
		for (mwSize cc=0; cc < regpoints.size(); cc++) { // loop over points to regress
			A(cc,0)=w[cc]; // Term 1: Intercept, i.e. 1
			b(cc,0)=regpoints[cc].first.val()*w[cc]; // Known value
			for (mwSize cd=0; cd < ndims ; cd++)  // Regression terms (1st order)
				A(cc,cd+1)=w[cc]*(regpoints[cc].first[cd]-(*qp_it)[cd]); // Terms are centered coordinates
			if (order==2){ // Quadratic terms (2nd order), i.e. all cross products
				mwSize cpos=ndims+1; // Term number in matrix (column number)
				for (mwSize cd1=0; cd1<ndims; cd1++) // Loop over dimension
					for (mwSize cd2=cd1; cd2 < ndims; cd2++ ) // Loop from outer loop dimension to number of dimensions
						A(cc, cpos++)=w[cc]*((regpoints[cc].first[cd1]-(*qp_it)[cd1]) * (regpoints[cc].first[cd2]-(*qp_it)[cd2])); // Assign cross-products (centered and weighed)
			}
		}
		x=A.bdcSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(b); // Perform weighted least squares regression. Note the weights get squared in here
		* val_beg = x(0,0); // store result
		++val_beg;	// Increase pointer to output values
		prog = double(std::distance(qp_begin,qp_it))/double(std::distance(qp_begin,qp_end)); // keep track of progress
	}
	prog = 1;
}

void localFit(const Tree & tree, const vector< Point > & qp, mwSize q, vector<double> & val, mwSize order, mwSize nthreads, double & prog) {
	// This function controls the local fitting by calling the function parFit in separate computational threads.
	// Inputs are:
	// 	tree:     The Spatial search tree (not modified)
	// 	qp:       Query points (not modified)
	// 	q:	      Number of points in regression (not modified)
	// 	val:      Estimated function value at query points (modified)
	// 	order:    Order of the regression (not modified)
	// 	nthreads: Number of computational threads (not modified)
	// 	prog:	  Keeps track of progress (modified)

	// Compute number of input points and dimensions
	mwSize nin=qp.size();
	mwSize ndims=qp.begin()->dims();

	// Compute number of terms in regression
	// 	Linear terms
	mwSize n=ndims+1; // number of dimensions plus one
	//  Quadratic terms
	if (order == 2)
		for (mwSize cd=1; cd < ndims + 1; cd++)
			n += cd; // cross-products (1 + 2 + ... + ndims)

	mwSize usedthreads = nthreads>nin ? nin : nthreads; // Reduce the number of threads when there are less query points than threads
	// Compute number of points in each each computational thread
	vector< mwSize > n_in_thread(usedthreads,nin/usedthreads); // Number of threads is equal to the integral division of number of points and threads
	for (mwSize cr=0; cr < nin % usedthreads; cr++)         // Remaining point (modulo) are spread over the threads
		n_in_thread[cr]++;

	// Start computation in threads asynchronously
	vector<double> prog_th(usedthreads,0); 			 // Vector holding progress (between 0 and 1) of each computation thread
	std::vector< std::future < void > > all_futures; // Vector with one feature per thread
	mwSize total=0; 								 // Keeps track of points passed to previous threads
	for (mwSize cth = 0; cth<usedthreads; cth++){       // Loop to launch computations asynchronously
		all_futures.push_back( std::async( std::launch::async, parFit, std::ref(tree), qp.begin()+total, qp.begin()+total+n_in_thread[cth], val.begin()+total, q, n ,order, std::ref(prog_th[cth]) ) ); // Start thread
		total += n_in_thread[cth];	// Keep track of points already passed to the function
	}

	// Check status of threads every second
	std::future_status status;
	for (mwSize cth = 0; cth < usedthreads; cth++){ // Check status of each thread
		do {
			status = all_futures[cth].wait_for(std::chrono::seconds(1)); // Wait for 1 second, or for thread ending (whichever comes first)
			prog = std::accumulate(prog_th.begin(),prog_th.end(),0.0)/double(usedthreads); // Compute current progress of threads
		} while (status != std::future_status::ready); // Repeat loop until thread ended successfully
	}

} // end of localFit

double median(vector<double> v) // vector copied to avoid modifications on it
{
  if(v.empty()) {
    return 0.0;
  }
  auto n = v.size() / 2;
  nth_element(v.begin(), v.begin()+n, v.end());
  auto med = v[n];
  if(!(v.size() & 1)) { //If the set size is even
    auto max_it = max_element(v.begin(), v.begin()+n);
    med = (*max_it + med) / 2.0;
  }
  return med;
}