#include "trajopt/collision_terms.hpp"
#include "trajopt/collision_checker.hpp"
#include "trajopt/rave_utils.hpp"
#include "trajopt/utils.hpp"
#include "sco/expr_vec_ops.hpp"
#include "sco/expr_ops.hpp"
#include "sco/sco_common.hpp"
#include <boost/foreach.hpp>
#include "utils/eigen_conversions.hpp"
#include "sco/modeling_utils.hpp"
#include "utils/stl_to_string.hpp"
#include "utils/logging.hpp"
#include <boost/functional/hash.hpp>
using namespace OpenRAVE;
using namespace sco;
using namespace util;
using namespace std;

namespace trajopt {

void CollisionsToDistances(const vector<Collision>& collisions, const Link2Int& m_link2ind,
    DblVec& dists) {
  // Note: this checking (that the links are in the list we care about) is probably unnecessary
  // since we're using LinksVsAll
  dists.clear();
  dists.reserve(collisions.size());
  BOOST_FOREACH(const Collision& col, collisions) {
    Link2Int::const_iterator itA = m_link2ind.find(col.linkA);
    Link2Int::const_iterator itB = m_link2ind.find(col.linkB);
    if (itA != m_link2ind.end() || itB != m_link2ind.end()) {
      dists.push_back(col.distance);
    }
  }
}

void CollisionsToDistanceExpressions(const vector<Collision>& collisions, Configuration& rad,
    const Link2Int& link2ind, const VarVector& vars, const DblVec& dofvals, vector<AffExpr>& exprs, bool isTimestep1) {

  exprs.clear();
  exprs.reserve(collisions.size());
  rad.SetDOFValues(dofvals); // since we'll be calculating jacobians
  BOOST_FOREACH(const Collision& col, collisions) {
    AffExpr dist(col.distance);
    Link2Int::const_iterator itA = link2ind.find(col.linkA);
    if (itA != link2ind.end()) {
      VectorXd dist_grad = toVector3d(col.normalB2A).transpose()*rad.PositionJacobian(itA->second, col.ptA);
      exprInc(dist, varDot(dist_grad, vars));
      exprInc(dist, -dist_grad.dot(toVectorXd(dofvals)));
    }
    Link2Int::const_iterator itB = link2ind.find(col.linkB);
    if (itB != link2ind.end()) {
      VectorXd dist_grad = -toVector3d(col.normalB2A).transpose()*rad.PositionJacobian(itB->second, (isTimestep1 && (col.cctype == CCType_Between)) ? col.ptB1 : col.ptB);
      exprInc(dist, varDot(dist_grad, vars));
      exprInc(dist, -dist_grad.dot(toVectorXd(dofvals)));
    }
    if (itA != link2ind.end() || itB != link2ind.end()) {
      exprs.push_back(dist);
    }
  }
  LOG_DEBUG("%ld distance expressions\n", exprs.size());
}

void CollisionsToDistanceExpressions(const vector<Collision>& collisions, Configuration& rad, const Link2Int& link2ind,
    const VarVector& vars0, const VarVector& vars1, const DblVec& vals0, const DblVec& vals1,
    vector<AffExpr>& exprs) {
  vector<AffExpr> exprs0, exprs1;
  CollisionsToDistanceExpressions(collisions, rad, link2ind, vars0, vals0, exprs0, false);
  CollisionsToDistanceExpressions(collisions, rad, link2ind, vars1, vals1, exprs1,true);

  exprs.resize(exprs0.size());
  for (int i=0; i < exprs0.size(); ++i) {
    exprScale(exprs0[i], (1-collisions[i].time));
    exprScale(exprs1[i], collisions[i].time);
    exprs[i] = AffExpr(0);
    exprInc(exprs[i], exprs0[i]);
    exprInc(exprs[i], exprs1[i]);
    cleanupAff(exprs[i]);
  }
}

inline size_t hash(const DblVec& x) {
  return boost::hash_range(x.begin(), x.end());
}

void CollisionEvaluator::GetCollisionsCached(const DblVec& x, vector<Collision>& collisions) {
  double key = hash(getDblVec(x, GetVars()));
  vector<Collision>* it = m_cache.get(key);
  if (it != NULL) {
    LOG_DEBUG("using cached collision check\n");
    collisions = *it;
  }
  else {
    LOG_DEBUG("not using cached collision check\n");
    CalcCollisions(x, collisions);
    m_cache.put(key, collisions);
  }
}

SingleTimestepCollisionEvaluator::SingleTimestepCollisionEvaluator(ConfigurationPtr rad, const VarVector& vars) :
  m_env(rad->GetEnv()),
  m_cc(CollisionChecker::GetOrCreate(*m_env)),
  m_rad(rad),
  m_vars(vars),
  m_link2ind(),
  m_links(),
  m_filterMask(-1) {
  vector<KinBody::LinkPtr> links;
  vector<int> inds;
  rad->GetAffectedLinks(m_links, true, inds);
  for (int i=0; i < m_links.size(); ++i) {
    m_link2ind[m_links[i].get()] = inds[i];
  }
  // TODO add argument
}

void SingleTimestepCollisionEvaluator::CalcCollisions(const DblVec& x, vector<Collision>& collisions) {
  DblVec dofvals = getDblVec(x, m_vars);
  m_rad->SetDOFValues(dofvals);
  m_cc->LinksVsAll(m_links, collisions, m_filterMask);
}

void SingleTimestepCollisionEvaluator::CalcDists(const DblVec& x, DblVec& dists) {
  vector<Collision> collisions;
  GetCollisionsCached(x, collisions);
  CollisionsToDistances(collisions, m_link2ind, dists);
}

void SingleTimestepCollisionEvaluator::CalcDistExpressions(const DblVec& x, vector<AffExpr>& exprs) {
  vector<Collision> collisions;
  GetCollisionsCached(x, collisions);
  DblVec dofvals = getDblVec(x, m_vars);
  CollisionsToDistanceExpressions(collisions, *m_rad, m_link2ind, m_vars, dofvals, exprs, false);
}

////////////////////////////////////////

CastCollisionEvaluator::CastCollisionEvaluator(ConfigurationPtr rad, const VarVector& vars0, const VarVector& vars1) :
  m_env(rad->GetEnv()),
  m_cc(CollisionChecker::GetOrCreate(*m_env)),
  m_rad(rad),
  m_vars0(vars0),
  m_vars1(vars1),
  m_link2ind(),
  m_links() {
  vector<KinBody::LinkPtr> links;
  vector<int> inds;
  rad->GetAffectedLinks(m_links, true, inds);
  for (int i=0; i < m_links.size(); ++i) {
    m_link2ind[m_links[i].get()] = inds[i];
  }
}

void CastCollisionEvaluator::CalcCollisions(const DblVec& x, vector<Collision>& collisions) {
  DblVec dofvals0 = getDblVec(x, m_vars0);
  DblVec dofvals1 = getDblVec(x, m_vars1);
  m_rad->SetDOFValues(dofvals0);
  m_cc->CastVsAll(*m_rad, m_links, dofvals0, dofvals1, collisions);
}
void CastCollisionEvaluator::CalcDistExpressions(const DblVec& x, vector<AffExpr>& exprs) {
  vector<Collision> collisions;
  GetCollisionsCached(x, collisions);
  DblVec dofvals0 = getDblVec(x, m_vars0);
  DblVec dofvals1 = getDblVec(x, m_vars1);
  CollisionsToDistanceExpressions(collisions, *m_rad, m_link2ind, m_vars0, m_vars1, dofvals0, dofvals1, exprs);
}
void CastCollisionEvaluator::CalcDists(const DblVec& x, DblVec& dists) {
  vector<Collision> collisions;
  GetCollisionsCached(x, collisions);
  CollisionsToDistances(collisions, m_link2ind, dists);
}

//////////////////////////////////////////

typedef OpenRAVE::RaveVector<float> RaveVectorf;

void PlotCollisions(const std::vector<Collision>& collisions, OR::EnvironmentBase& env, vector<OR::GraphHandlePtr>& handles, double safe_dist) {
  BOOST_FOREACH(const Collision& col, collisions) {
    RaveVectorf color;
    if (col.distance < 0) color = RaveVectorf(1,0,0,1);
    else if (col.distance < safe_dist) color = RaveVectorf(1,1,0,1);
    else color = RaveVectorf(0,1,0,1);

    if (col.cctype == CCType_Between) {
      // handles.push_back(env.drawarrow(col.ptB, col.ptB1, .002, RaveVectorf(0,0,0,1)));
    }
    OR::Vector ptB = (col.cctype == CCType_Between)  ? ((1-col.time)* col.ptB +col.time*col.ptB1) : col.ptB;
    // handles.push_back(env.drawarrow(col.ptA, ptB, .0025, color));
  }
}

CollisionCost::CollisionCost(double dist_pen, double coeff, ConfigurationPtr rad, const VarVector& vars) :
    Cost("collision"),
    m_calc(new SingleTimestepCollisionEvaluator(rad, vars)), m_dist_pen(dist_pen), m_coeff(coeff)
{}

CollisionCost::CollisionCost(double dist_pen, double coeff, ConfigurationPtr rad, const VarVector& vars0, const VarVector& vars1) :
    Cost("cast_collision"),
    m_calc(new CastCollisionEvaluator(rad, vars0, vars1)), m_dist_pen(dist_pen), m_coeff(coeff)
{}
ConvexObjectivePtr CollisionCost::convex(const vector<double>& x, Model* model) {
  ConvexObjectivePtr out(new ConvexObjective(model));
  vector<AffExpr> exprs;
  m_calc->CalcDistExpressions(x, exprs);
  for (int i=0; i < exprs.size(); ++i) {
    AffExpr viol = exprSub(AffExpr(m_dist_pen), exprs[i]);
    out->addHinge(viol, m_coeff);
  }
  return out;
}
double CollisionCost::value(const vector<double>& x) {
  DblVec dists;
  m_calc->CalcDists(x, dists);
  double out = 0;
  for (int i=0; i < dists.size(); ++i) {
    out += pospart(m_dist_pen - dists[i]) * m_coeff;
  }
  return out;
}

void CollisionCost::Plot(const DblVec& x, OR::EnvironmentBase& env, std::vector<OR::GraphHandlePtr>& handles) {
  vector<Collision> collisions;
  m_calc->GetCollisionsCached(x, collisions);
  PlotCollisions(collisions, env, handles, m_dist_pen);
}

// ALMOST EXACTLY COPIED FROM CollisionCost

CollisionConstraint::CollisionConstraint(double dist_pen, double coeff, ConfigurationPtr rad, const VarVector& vars) :
    m_calc(new SingleTimestepCollisionEvaluator(rad, vars)), m_dist_pen(dist_pen), m_coeff(coeff)
{
  name_="collision";
}

CollisionConstraint::CollisionConstraint(double dist_pen, double coeff, ConfigurationPtr rad, const VarVector& vars0, const VarVector& vars1) :
    m_calc(new CastCollisionEvaluator(rad, vars0, vars1)), m_dist_pen(dist_pen), m_coeff(coeff)
{
  name_="collision";
}
ConvexConstraintsPtr CollisionConstraint::convex(const vector<double>& x, Model* model) {
  ConvexConstraintsPtr out(new ConvexConstraints(model));
  vector<AffExpr> exprs;
  m_calc->CalcDistExpressions(x, exprs);
  for (int i=0; i < exprs.size(); ++i) {
    AffExpr viol = exprSub(AffExpr(m_dist_pen), exprs[i]);
    out->addIneqCnt(exprMult(viol,m_coeff));
  }
  return out;
}
DblVec CollisionConstraint::value(const vector<double>& x) {
  DblVec dists;
  m_calc->CalcDists(x, dists);
  DblVec out(dists.size());
  for (int i=0; i < dists.size(); ++i) {
    out[i] = pospart(m_dist_pen - dists[i]) * m_coeff;
  }
  return out;
}

/////////////////////////////////////////////////////////
// Risk-aware section starts here.
/////////////////////////////////////////////////////////

/**
 * Constructor for the collision chance constraint
 *
 * @param uncertain_body_names: the names of the bodies representing uncertain obstacles
 * @param location_covariances: the covariance matrices of the Gaussian uncertainty in the location of the
 *                              corresponding uncertain bodies. Must have the same length as uncertain_body_names
 * @param risk_tolerance: a double representing the maximum allowable risk at this time step.
 * @param precision: a double representing the desired precision tolerance in our risk estimates. True risk is guaranteed
 *                   not to exceed the estimate by more than precision/2.
 * @param rad: a pointer to the configuration of the robot and evironment.
 * @param vars: an array of variables to be used in constructing symbolic expressions for linearized risk.
 */
CollisionChanceConstraint::CollisionChanceConstraint(std::vector<std::string> uncertain_body_names,
                                                     std::vector<Matrix3d> location_covariances,
                                                     double risk_tolerance,
                                                     double required_precision,
                                                     double coeff,
                                                     double grad_scale_factor,
                                                     int numsteps,
                                                     ConfigurationPtr rad,
                                                     const VarArray& vars) :
    m_calc(new OverallCollisionRiskEvaluator(uncertain_body_names,
                                             location_covariances,
                                             required_precision,
                                             grad_scale_factor,
                                             numsteps,
                                             rad,
                                             vars)),
    m_vars(vars),
    m_uncertain_body_names(uncertain_body_names),
    m_location_covariances(location_covariances),
    m_risk_tolerance(risk_tolerance),
    m_coeff(coeff),
    m_numsteps(m_numsteps)
{
  name_="collision";
}

/**
 * Convexify the constraint with the given values for the degrees of freedom, yielding symbolic expressions in terms
 * of the given model
 *
 * @param x: a vector of doubles containing the current values of the degrees of freedom (i.e. \theta_0)
 * @param model: the model with which to construct the convexified expressions.
 *
 * @returns a pointer to the convexified constraint.
 */
ConvexConstraintsPtr CollisionChanceConstraint::convex(const vector<double>& x, Model* model) {
  // Set up the convexified constraint object
  ConvexConstraintsPtr out(new ConvexConstraints(model));

  // Get the linearized risk expressions
  vector<AffExpr> exprs;
  m_calc->CalcRiskExpressions(x, exprs);

  // Accumulate the risks into a single constraint
  AffExpr risk_violation(-1.0*m_risk_tolerance);
  for (int i=0; i < exprs.size(); ++i) {
    // Add this risk expression to the risk violation expression
    exprInc(risk_violation, exprs[i]);
  }
  // Add the accumulated constraint to the model
  out->addIneqCnt(exprMult(risk_violation, m_coeff)); // <= 0
  LOG_DEBUG("VIOLATION convex: %f", out->violation(x));
  LOG_DEBUG("Total risk: %f", out->violation(x) + m_risk_tolerance);
  return out;
}

/**
 * Evaluate the constraint with the given values for the degrees of freedom.
 *
 * @param x: a vector of doubles containing the current values of the degrees of freedom (i.e. \theta_0)
 *
 * @returns a vector of doubles with one element representing the total risk at this timestep:
 *
 * The constraint is (summed over all obstacles)
 *          risk_tolerance - risk >= 0
 * but we relax it using the hingle loss to be
 *          |risk - risk_tolerance|^+
 */
DblVec CollisionChanceConstraint::value(const vector<double>& x) {
  // Start by getting the risk estimates
  DblVec risks;
  m_calc->CalcRisks(x, risks);

  // Now compute the relaxed hinge-loss penalties
  DblVec out(1);
  double accumulator = 0.0;
  for (int i=0; i < risks.size(); ++i) {
    accumulator += risks[i];
  }
  out[0] = pospart(accumulator - m_risk_tolerance) * m_coeff;
  LOG_DEBUG("Total RISK: %.6f", accumulator);
  return out;
}

typedef OpenRAVE::RaveVector<float> RaveVectorf;

void SingleTimestepCollisionRiskEvaluator::PlotRisks(const std::vector<RiskQueryResult>& risk_query_results, OR::EnvironmentBase& env, vector<OR::GraphHandlePtr>& handles) {
  BOOST_FOREACH(const RiskQueryResult& risk_result, risk_query_results) {
    RaveVectorf color1 = RaveVectorf(0,0,1,1);

    // OR::Vector endpoint = OR::Vector(
    //   risk_result.ptRobotTwoShot[0] + risk_result.contact_normal_into_robot_two_shot(0),
    //   risk_result.ptRobotTwoShot[1] + risk_result.contact_normal_into_robot_two_shot(1),
    //   risk_result.ptRobotTwoShot[2] + risk_result.contact_normal_into_robot_two_shot(2));
    // handles.push_back(env.drawarrow(risk_result.ptRobotTwoShot, endpoint, .0025, color));

    RaveVectorf color2 = RaveVectorf(1,0,0,1);

    // OR::Vector endpoint2 = OR::Vector(
    //   risk_result.ptRobotOneShot[0] + risk_result.contact_normal_into_robot_one_shot(0),
    //   risk_result.ptRobotOneShot[1] + risk_result.contact_normal_into_robot_one_shot(1),
    //   risk_result.ptRobotOneShot[2] + risk_result.contact_normal_into_robot_one_shot(2));
    // handles.push_back(env.drawarrow(risk_result.ptRobotOneShot, endpoint2, .0025, color2));

    // Get the indices of the two robot links that the risk estimate collides with
    Link2Int::const_iterator itRobotOneShot = m_link2ind.find(risk_result.linkRobotOneShot);
    Link2Int::const_iterator itRobotTwoShot = m_link2ind.find(risk_result.linkRobotTwoShot);

    if (risk_result.epsilon > m_precision && itRobotOneShot != m_link2ind.end()) {
      Vector3d n_hat_one_shot = risk_result.contact_normal_into_robot_one_shot.normalized();

      VectorXd one_shot_jacobian = m_rad->PositionJacobian(itRobotOneShot->second, risk_result.ptRobotOneShot);
      VectorXd d_epsilon_d_theta_one_shot = risk_result.d_epsilon_dx_one_shot * n_hat_one_shot * n_hat_one_shot.transpose() * one_shot_jacobian;

      VectorXd scaled_nhat_oneshot = risk_result.d_epsilon_dx_one_shot;

      OR::Vector endpoint1 = OR::Vector(
        risk_result.ptRobotOneShot[0] - scaled_nhat_oneshot(0),
        risk_result.ptRobotOneShot[1] - scaled_nhat_oneshot(1),
        risk_result.ptRobotOneShot[2] - scaled_nhat_oneshot(2));

      // handles.push_back(env.drawarrow(risk_result.ptRobotOneShot, endpoint1, .0025, color1));

      // We need to check if we're in case 3 before calculating the gradient for the second shot
      VectorXd risk_grad;
      Vector3d n_hat_two_shot;
      if (itRobotTwoShot != m_link2ind.end()) {
        // make sure the contact unit vector is normalized
        Vector3d n_hat_two_shot = risk_result.contact_normal_into_robot_two_shot.normalized();

        VectorXd two_shot_jacobian = m_rad->PositionJacobian(itRobotTwoShot->second, risk_result.ptRobotTwoShot);
        VectorXd d_epsilon_d_theta_two_shot = risk_result.d_epsilon_dx_two_shot * n_hat_two_shot * n_hat_two_shot.transpose() * two_shot_jacobian;

        // Because the combined one- and two-shot risk estimate is epsilon = 0.5*(epsilon_1 + epsilon_2),
        // The combined gradient is the sum of half of each separate gradient
        risk_grad = 0.5 * d_epsilon_d_theta_one_shot + 0.5 * d_epsilon_d_theta_two_shot;

        VectorXd scaled_nhat_twoshot = risk_result.d_epsilon_dx_two_shot;
        OR::Vector endpoint2 = OR::Vector(
          risk_result.ptRobotTwoShot[0] - scaled_nhat_twoshot(0),
          risk_result.ptRobotTwoShot[1] - scaled_nhat_twoshot(1),
          risk_result.ptRobotTwoShot[2] - scaled_nhat_twoshot(2));
        // handles.push_back(env.drawarrow(risk_result.ptRobotTwoShot, endpoint2, .0025, color2));
      }
    }
  }
}

void CollisionChanceConstraint::Plot(const DblVec& x, OR::EnvironmentBase& env, std::vector<OR::GraphHandlePtr>& handles) {
  m_calc->PlotRisks(x, env, handles);
}

/**
 * Constructor for the SingleTimestepCollisionRiskEvaluator class, which is an object used to evaluate
 * the risk of collision between the robot and all uncertain obstacles at a single timestep.
 *
 * @param uncertain_body_names: the names of the bodies representing uncertain obstacles
 * @param location_covariances: the covariance matrices of the Gaussian uncertainty in the location of the
 *                              corresponding uncertain bodies. Must have the same length as uncertain_body_names
 * @param precision: a double representing the desired precision tolerance in our risk estimates. True risk is guaranteed
 *                   not to exceed the estimate by more than precision/2.
 * @param rad: a pointer to the configuration of the robot and evironment.
 * @param vars: a vector of variables to be used in constructing symbolic expressions for linearized risk.
 */
SingleTimestepCollisionRiskEvaluator::SingleTimestepCollisionRiskEvaluator(
    std::vector<std::string> uncertain_body_names,
    std::vector<Matrix3d> location_covariances,
    double precision,
    double grad_scale_factor,
    ConfigurationPtr rad,
    const VarVector& vars) :
  m_env(rad->GetEnv()),
  m_cc(CollisionChecker::GetOrCreate(*m_env)),
  m_rad(rad),
  m_vars(vars),
  m_link2ind(),
  m_link2covariance(),
  m_links(),
  m_filterMask(RobotFilter),
  m_precision(precision),
  m_grad_scale_factor(grad_scale_factor)
{
  // Find uncertain links by body name
  vector<int> inds;
  vector<Matrix3d> covariances;
  int body_index = 0;

  // Match links to covariances
  BOOST_FOREACH(std::string body_name, uncertain_body_names) {
    BOOST_FOREACH(const KinBody::LinkPtr& link, m_env->GetKinBody(body_name)->GetLinks()) {
      if (!link->GetGeometries().empty()) {
        m_links.push_back(link);
        inds.push_back(link->GetIndex());
        covariances.push_back(location_covariances[body_index]);
      }
    }

    ++body_index;
  }

  // For easy access, save the OpenRAVE index and covariance matrix of each link in a map
  for (int i=0; i < m_links.size(); ++i) {
    m_link2covariance[m_links[i].get()] = covariances[i];
  }

  // We also need to save a map from robot link indices to OpenRAVE indices,
  // so we can get the Jacobian of points on those links later.
  vector<KinBody::LinkPtr> links;
  inds.clear();
  rad->GetAffectedLinks(links, true, inds);
  for (int i=0; i < links.size(); ++i) {
    m_link2ind[links[i].get()] = inds[i];
  }
}

/**
 * Serves as a wrapper for the PerformRiskCheck function, caching results of previous queries to give
 * performance benefits.
 *
 * @param x: a vector of doubles containing the current values of the degrees of freedom (i.e. \theta_0)
 * @param risks: a vector of RiskQueryResult structs into which we will put the results of the risk query.
 *               These structs contain both the risk estimate itself and information for computing the gradient.
 *
 * @returns void, but risks will filled with the risk query results.
 */
void CollisionRiskEvaluator::GetRisksCached(const DblVec& x, vector<RiskQueryResult>& risks) {
  double key = hash(getDblVec(x, GetVars()));
  vector<RiskQueryResult>* it = m_cache.get(key);
  if (it != NULL) {
    LOG_DEBUG("using cached risk query result\n");
    risks = *it;
  }
  else {
    LOG_DEBUG("not using cached risk query result\n");
    PerformRiskCheck(x, risks);
    m_cache.put(key, risks);
  }
}

/**
 * Calls the ProximityAlert risk estimation via the collision checker to estimate the risks of collision
 * between the robot and uncertain obstacles.
 *
 * @param x: a vector of doubles containing the current values of the degrees of freedom (i.e. \theta_0)
 * @param risks: a vector of RiskQueryResult structs into which we will put the results of the risk query.
 *               These structs contain both the risk estimate itself and information for computing the gradient.
 *
 * @returns void, but risks will filled with the risk query results.
 */
void SingleTimestepCollisionRiskEvaluator::PerformRiskCheck(const DblVec& x, vector<RiskQueryResult>& risks) {
  DblVec dofvals = getDblVec(x, m_vars);
  m_rad->SetDOFValues(dofvals);
  m_cc->LinksVsAllRiskEstimate(m_links, risks, m_filterMask, m_link2covariance, m_precision);
}

/**
 * Compute the risks of collision between the robot and each uncertain obstacle. Will use cached values if
 * possible; otherwise, will use the ProximityAlert epsilon-shadow algorithm to estimate the risk.
 *
 * @param x: a vector of doubles containing the current values of the degrees of freedom (i.e. \theta_0)
 * @param risks: a vector of doubles into which we will put the risk estimates
 *
 * @returns void, but risks will be cleared and then filled with the risk estimates.
 */
void SingleTimestepCollisionRiskEvaluator::CalcRisks(const DblVec& x, DblVec& risks) {
  vector<RiskQueryResult> risk_query_results;
  GetRisksCached(x, risk_query_results);

  // Extract risk measures from risk_query_results
  risks.clear();
  BOOST_FOREACH(const RiskQueryResult& risk_result, risk_query_results) {
    risks.push_back(risk_result.epsilon);
  }
}

/**
 * Construct symbolic expressions for the linearized collision risks at this single timestep.
 *
 * Begins by computing the risk of collision between each uncertain obstacle and the robot at this specific
 * timestep, then for each obstacle, computes the gradient of the risk with respect to joint values and
 * adds a symbolic affine expression to exprs of the form:
 *
 *    \epsilon_0 + \nabla_\theta \epsilon (\theta - \theta_0)
 *
 * where \epsilon_0 is the risk of collision with the obstacle, and \theta_0 is the vector of joint values at
 * this timestep, \nabla_\theta \epsilon is the gradient of \epsilon w.r.t. the joint values, and \theta is a
 * symbolic vector of joint values.
 *
 * @param x: a vector of doubles containing the current values of the degrees of freedom (i.e. \theta_0)
 * @param exprs: a vector of symbolic affine expressions into which we will put the linearized risk expressions
 *
 * @returns void, but exprs will be cleared and then filled with the linearized risk expressions.
 */
void SingleTimestepCollisionRiskEvaluator::CalcRiskExpressions(const DblVec& x, vector<AffExpr>& exprs) {
  vector<RiskQueryResult> risk_query_results;
  GetRisksCached(x, risk_query_results);
  DblVec dofvals = getDblVec(x, m_vars);

  exprs.clear();
  exprs.reserve(risk_query_results.size());
  m_rad->SetDOFValues(dofvals); // since we'll be calculating jacobians

  // For each risk result, we need to compute the linearized expression for the collision risk
  // risk(theta) \approx risk_0 + gradient * (theta - theta_0)
  BOOST_FOREACH(const RiskQueryResult& risk_result, risk_query_results) {
    AffExpr risk(risk_result.epsilon);

    // Get the indices of the two robot links that the risk estimate collides with
    Link2Int::const_iterator itRobotOneShot = m_link2ind.find(risk_result.linkRobotOneShot);
    Link2Int::const_iterator itRobotTwoShot = m_link2ind.find(risk_result.linkRobotTwoShot);

    // If the risk is extremely small, that means that the largest epsilon shadow used to estimate
    // the risk did not touch the robot. In that case, the risk is constant at zero, since small motions
    // are assumed not to bring the robot "into risk" from zero-risk.
    // Because of this, we only add the linear terms if the risk is non-zero (within tolerance).
    // We can also skip this if both contact links are not valid
    if (risk_result.epsilon > m_precision && itRobotOneShot != m_link2ind.end()) {

      // There are three cases:
      //  1.) neither the one shot nor the two shot search ended in collision with the controlled links
      //  2.) only the one shot search ended in collision
      //  3.) both ended in collision
      // (collision may have occured with an uncontrolled link, like the robot base).
      //
      // If (1) occured, then neither risk_result.linkRobotOneShot nor risk_result.linkRobotTwo
      // are valid.
      //
      // If (2) occured, then risk_result.linkRobotOneShot is a valid link but
      // risk_result.linkRobotTwoShot is not, so we only add the gradient terms from the first shot
      //
      // If (3) occured, then both risk_result.linkRobotOneShot and risk_result.linkRobotTwoShot
      // will be valid links, so we add the gradient terms from both.

      // In either case 2 or case 3, we include the terms from the first shot, so we can calculate
      // the gradient: d_epsilon_d_theta (epsilon is the risk estimate and theta the joint angles).
      //
      // The explanation of this calculation relies on some knowledge of the proximity alert algorithms.
      // The ProximityAlert risk estimator gives us d_epsilon_d_x for both the one and two shot estimates.
      // In this context, x is the vector from the center of the epsilon-shadow ellipsoid to the contact point
      // with the robot, n is a vector into the robot normal to the contact with the shadow, and J is the Jacobian
      // of the contact point w.r.t. the joint angles. Then, we have
      //
      // d_epsilon_d_theta = d_epsilon_d_x * J
      //
      // First we need to make sure the contact unit vector is normalized (Note: this might not be needed anymore)
      Vector3d n_hat_one_shot = risk_result.contact_normal_into_robot_one_shot.normalized();

      DblMatrix one_shot_jacobian = m_rad->PositionJacobian(itRobotOneShot->second, risk_result.ptRobotOneShot);
      VectorXd d_epsilon_d_theta_one_shot = risk_result.d_epsilon_dx_one_shot * one_shot_jacobian;

      // We need to check if we're in case 3 before calculating the gradient for the second shot
      VectorXd risk_grad;
      if (itRobotTwoShot != m_link2ind.end()) {
        // make sure the contact unit vector is normalized
        Vector3d n_hat_two_shot = risk_result.contact_normal_into_robot_two_shot.normalized();

        DblMatrix two_shot_jacobian = m_rad->PositionJacobian(itRobotTwoShot->second, risk_result.ptRobotTwoShot);
        VectorXd d_epsilon_d_theta_two_shot = risk_result.d_epsilon_dx_two_shot * two_shot_jacobian;

        // Because the combined one- and two-shot risk estimate is epsilon = 0.5*(epsilon_1 + epsilon_2),
        // The combined gradient is the sum of half of each separate gradient
        risk_grad = 0.5 * d_epsilon_d_theta_one_shot + 0.5 * d_epsilon_d_theta_two_shot;
      } else {
        // Otherwise, we're in case 2 and the gradient is just the gradient from the first shot
        risk_grad = d_epsilon_d_theta_one_shot;
      }
      risk_grad *= m_grad_scale_factor;
      // LOG_DEBUG("Risk Gradient: %ld elements, %ld vars", risk_grad.size(), m_vars.size());
      // for (int j = 0; j < risk_grad.size(); j++) {
      //   LOG_DEBUG("\t %.8f", risk_grad(j));
      // }

      // Now we add risk_grad * (theta - theta_0) to the affine expression
      exprInc(risk, varDot(risk_grad, m_vars));
      exprInc(risk, -risk_grad.dot(toVectorXd(dofvals)));
    }

    // And return the risk expression by pushing the affine expression into the supplied vector.
    exprs.push_back(risk);
  }
  LOG_DEBUG("%ld risk expressions\n", exprs.size());
}

OverallCollisionRiskEvaluator::OverallCollisionRiskEvaluator(
    std::vector<std::string> uncertain_body_names,
    std::vector<Matrix3d> location_covariances,
    double precision,
    double grad_scale_factor,
    int numsteps,
    ConfigurationPtr rad,
    const VarArray& vars) :
  m_timestep_risk_evals(),
  m_env(rad->GetEnv()),
  m_cc(CollisionChecker::GetOrCreate(*m_env)),
  m_rad(rad),
  m_vars(vars),
  m_numsteps(numsteps),
  m_max_risk(false)
{
  // Create a single timestep risk evaluator for every timestep
  for (int i = 0; i < numsteps; ++i) {
    m_timestep_risk_evals.push_back(SingleTimestepCollisionRiskEvaluator(
      uncertain_body_names,
      location_covariances,
      precision,
      grad_scale_factor,
      rad,
      vars.row(i)));
  }
}

OverallCollisionRiskEvaluator::OverallCollisionRiskEvaluator(
    std::vector<std::string> uncertain_body_names,
    std::vector<Matrix3d> location_covariances,
    double precision,
    double grad_scale_factor,
    int numsteps,
    ConfigurationPtr rad,
    const VarArray& vars,
    bool use_max_risk) :
  m_timestep_risk_evals(),
  m_env(rad->GetEnv()),
  m_cc(CollisionChecker::GetOrCreate(*m_env)),
  m_rad(rad),
  m_vars(vars),
  m_numsteps(numsteps),
  m_max_risk(use_max_risk)
{
  // Create a single timestep risk evaluator for every timestep
  for (int i = 0; i < numsteps; ++i) {
    m_timestep_risk_evals.push_back(SingleTimestepCollisionRiskEvaluator(
      uncertain_body_names,
      location_covariances,
      precision,
      grad_scale_factor,
      rad,
      vars.row(i)));
  }
}

/**
 * Accumulate the risk at each timestep into a single measure of collision risk
 *
 * @param x: a vector of doubles containing the current values of the degrees of freedom (i.e. \theta_0)
 * @param risks: a vector of doubles into which we will put the risk estimates.
 *
 * @returns void, but risks will be cleared and then filled with the risk estimates from each timestep.
 */
void OverallCollisionRiskEvaluator::CalcRisks(const DblVec& x, DblVec& risks) {
  risks.clear();

  // There are two modes. If m_max_risk is true, then we only save the maximum risk
  // for each obstacle (across time). Otherwise, we save all risks for all obstacles.
  if (m_max_risk) {
    // Get the number of obstacles, and create a vector to store the maximum risks for each
    int num_obstacles = m_timestep_risk_evals[0].m_links.size();
    DblVec max_obstacle_risk;
    for (int i = 0; i < num_obstacles; i++) {
      max_obstacle_risk.push_back(0.0);
    }

    // Now loop through the timesteps, updating our maximum risk for each obstacle
    DblVec single_step_risks;
    BOOST_FOREACH(SingleTimestepCollisionRiskEvaluator& evaluator, m_timestep_risk_evals) {
      // Get risks for this timestep
      evaluator.CalcRisks(x, single_step_risks);

      // Save new obstacle risks if they are higher than the incumbent
      for (int i = 0; i < num_obstacles; i++) {
        if (single_step_risks[i] > max_obstacle_risk[i]) {
          max_obstacle_risk[i] = single_step_risks[i];
        }
      }
    }

    // Now save the maximum risks
    risks.insert(risks.begin(), max_obstacle_risk.begin(), max_obstacle_risk.end());
  } else {
    // Create a new vector to store the results of each timestep's risk
    // Has to be different since CalcRisks clears the given vector
    DblVec single_step_risks;
    BOOST_FOREACH(SingleTimestepCollisionRiskEvaluator& evaluator, m_timestep_risk_evals) {
      // Get risks for this timestep
      evaluator.CalcRisks(x, single_step_risks);

      // Save them by appending to risks
      risks.insert(risks.begin(), single_step_risks.begin(), single_step_risks.end());
    }
    // After this step, risks will be filled with the results of all risk queries
    // for all timesteps for each uncertain link.
  }
}

/**
 * Accumulate the risk expressions at each timestep into a single measure of collision risk
 *
 * @param x: a vector of doubles containing the current values of the degrees of freedom (i.e. \theta_0)
 * @param exprs: a vector of symbolic affine expressions into which we will put the linearized risk expressions
 *
 * @returns void, but exprs will be cleared and then filled with the risk estimate expressions from each timestep.
 */
void OverallCollisionRiskEvaluator::CalcRiskExpressions(const DblVec& x, vector<AffExpr>& exprs) {
  exprs.clear();

  // There are two modes. If m_max_risk is true, then we only save the maximum risk
  // for each obstacle (across time). Otherwise, we save all risks for all obstacles.
  if (m_max_risk) {
    // Get the number of obstacles, and create a vector to store the maximum risks for each
    int num_obstacles = m_timestep_risk_evals[0].m_links.size();
    vector<AffExpr> max_obstacle_risk;
    for (int i = 0; i < num_obstacles; i++) {
      max_obstacle_risk.push_back(AffExpr(0.0));
    }

    // Now loop through the timesteps, saving the expression associated with the maximum
    // risk
    vector<AffExpr> single_step_exprs;
    BOOST_FOREACH(SingleTimestepCollisionRiskEvaluator& evaluator, m_timestep_risk_evals) {
      // Get risk expressions for this timestep
      evaluator.CalcRiskExpressions(x, single_step_exprs);

      // Save new obstacle risks if they are higher than the incumbent
      for (int i = 0; i < num_obstacles; i++) {
        if (single_step_exprs[i].constant > max_obstacle_risk[i].constant) {
          max_obstacle_risk[i] = single_step_exprs[i];
        }
      }
    }

    // Now save the maximum risks
    exprs.insert(exprs.begin(), max_obstacle_risk.begin(), max_obstacle_risk.end());
  } else {
    // Create a new vector to store the results of each timestep's risk expressions
    // Has to be different since CalcRiskExpressions clears the given vector
    vector<AffExpr> single_step_exprs;
    BOOST_FOREACH(SingleTimestepCollisionRiskEvaluator& evaluator, m_timestep_risk_evals) {
      // Get risk expressions for this timestep
      evaluator.CalcRiskExpressions(x, single_step_exprs);

      // Save them by appending to exprs
      exprs.insert(exprs.begin(), single_step_exprs.begin(), single_step_exprs.end());
    }
    // After this step, exprs will be filled with the results of all risk queries
    // for all timesteps for each uncertain link.
  }
}

/**
 * Plot the risks at each timestep.
 */
void OverallCollisionRiskEvaluator::PlotRisks(const DblVec& x, OR::EnvironmentBase& env, vector<OR::GraphHandlePtr>& handles) {
  BOOST_FOREACH(SingleTimestepCollisionRiskEvaluator& evaluator, m_timestep_risk_evals) {
    vector<RiskQueryResult> risk_query_results;
    evaluator.GetRisksCached(x, risk_query_results);
    evaluator.PlotRisks(risk_query_results, env, handles);
  }
}

}


// __BEGIN_LICENSE__
//  Copyright (c) 2006-2013, United States Government as represented by the
//  Administrator of the National Aeronautics and Space Administration. All
//  rights reserved.
//
//  The NASA Vision Workbench is licensed under the Apache License,
//  Version 2.0 (the "License"); you may not use this file except in
//  compliance with the License. You may obtain a copy of the License at
//  http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// __END_LICENSE__

/**
  Program to clean out low-intensity junk pixels from the edges of
  Sentinel-1 images.  If these pixels are not cleaned out, the
  water detection algorithm will almost certainly flag them as water
  and may corrupt the statistics used to find water in the rest of
  the image.
*/

#include <queue>
#include <vw/Image/ImageIO.h>
#include <vw/Math/Statistics.h>
#include <vw/tools/modis_utilities.h>
#include <vw/tools/modis_water_detection.h>

#include <boost/algorithm/string.hpp>
#include <boost/program_options.hpp>
namespace po = boost::program_options;

using namespace vw;

/// Applies a median filter to a vector with a window size
/// - TODO: Move to somewhere in VW, but first we need a place to put it.
template <typename T>
void window_median_filter(std::vector<T> const& v_in, std::vector<T> &v_out, const int width) {

  // Setup
  const int length = static_cast<int>(v_in.size());
  const int half_width = width / 2;
  v_out = v_in;

  // Handle small input sizes
  if ((width < 3) || (length < width))
    return;

  // Just copy the first few elements
  for (int i=0; i<half_width; ++i)
    v_out[i] = v_in[i];

  // Init a queue of neigboring values
  std::list<T> locals;
  for (int i=0; i<width-1; ++i)
    locals.push_back(v_in[i]);

  for (int i=half_width; i<length-half_width; ++i) {
    // Load next value into queue
    locals.push_back(v_in[i+half_width]);

    // Compute the median of the neigbors
    v_out[i] = math::median(locals);

    // Drop oldest value
    locals.pop_front();
  }
  // Just copy the last few elements
  for (int i=length-half_width; i<length; ++i) {
    v_out[i] = v_in[i];
  }
} // End median_filter

/// Returns the percentage of values equal to a given value
/// - T is a type that can be accessed with standard iterators.
template <typename T>
double percent_equal(T const& values, double equal_to) {

  double count = 0;
  typename T::const_iterator iter;
  for (iter = values.begin(); iter != values.end(); ++iter) {
    if (static_cast<double>(*iter) == equal_to)
      count += 1.0;
  }
  return count / static_cast<double>(values.size());
}

/// Class to locate the best position (if any) of a jump from low pixel values to higher pixel values.
/// - This is a somewhat arbitrary decision tuned for the available sentinel-1 data.
/// - This class works by traversing inwards from the image border and maintaining three
///   sets of pixels.  The leading list is filled up first, then the buffer list, then
///   the trailing list.  The lead and buffer lists are capped at a fixed size, the trailing
///   list has no size limit.  The statistics of the lead and trailing buffers are compared
///   to try and identify the most likely border of real image pixels.
class JumpFinder {
public:

  /// Constructor
  /// - lead_size is the number of pixels it in the leading pixel set.
  JumpFinder(int lead_size)
    : m_lead_size(static_cast<size_t>(lead_size)),
      m_best_position(0), m_best_score(15), m_trail_zero_count(0), m_trail_sum(0) {
  }

  /// Add a new value and return true if a jump was detected
  bool add_value(double value) {

    const int    BUFFER_WIDTH = 1; // Keep this many numbers out of the statistics
    const double CLOSE_ZERO   = 8; // Treat values this small as if they are zero

    // Add new value to the leading list
    m_lead_list.push_front(value);

    if (m_lead_list.size() > m_lead_size) {
      // Move back of lead list to the trailing list
      m_buffer_list.push_front(m_lead_list.back());
      m_lead_list.pop_back();

      if (m_buffer_list.size() > BUFFER_WIDTH) {
        double trail_val = m_buffer_list.back();
        if (trail_val <= CLOSE_ZERO)
          ++m_trail_zero_count;
        m_trail_list.push_front(trail_val);
        m_buffer_list.pop_back();
        m_trail_sum += trail_val;
      }
    }

    if (m_trail_list.empty()) // Nothing to do until lead buffer is filled
      return false;

    // Compute statistics
    double mean_trail=0, mean_lead=0, dev_trail=0;
    double trail_count  = static_cast<double>(m_trail_list.size());
    double percent_zero = m_trail_zero_count / trail_count;
    mean_lead = math::mean(m_lead_list);
    if (percent_zero < 1.0) { // Save time on big blank regions
      mean_trail = m_trail_sum / trail_count;
      dev_trail  = math::standard_deviation(m_trail_list, mean_trail);
    }

    const double MAX_DEV_TRAIL    = 5;    // After STD_DEV of trail portion is higher, not a good jump
    const double MIN_PERCENT_ZERO = 0.30; // Must be this percent or higher zeroes in trail for a jump
    const double DEV_TRAIL_BREAK  = 20;   // Quit if the trail deviation exceeds this threshold.

    // Evaluate the value of an edge located at the start of the trail region
    double score = (mean_lead - mean_trail); // Valid pixels should be much brighter than border junk
    if ((dev_trail > MAX_DEV_TRAIL) || (percent_zero < MIN_PERCENT_ZERO))
      score = 0;
    //std::cout << "Lead  stats: " << mean_lead  << ", " << dev_lead  << std::endl;
    //std::cout << "Trail stats: " << mean_trail << ", " << dev_trail
    //          << ", size = " << m_trail_list.size() << ", percent = " << percent_zero<< std::endl;
    if (score > m_best_score) { // Record the best edge score to date
      m_best_score    = score;
      m_best_position = get_jump_index();
    }

    // If one of these things is true, very unlikely to see the edge beyond this point.
    return ((percent_zero < MIN_PERCENT_ZERO) || (dev_trail > DEV_TRAIL_BREAK));
  }

  /// Return the index at which the jump occurs.
  int get_jump_index() const {
    return m_trail_list.size() + 1; // Advance into the buffer pixel
  }

  int get_best_jump() const {
    //printf("Best score = %lf and position %d\n", m_best_score, m_best_position);
    return m_best_position;
  }

private:
  size_t m_lead_size; ///< Max size of the lead list
  std::list<double> m_lead_list;   ///< Values at the leading part of the window
  std::list<double> m_trail_list;  ///< Values at the trailing part of the window
  std::list<double> m_buffer_list; ///< Values in the middle part of the window.
  int    m_best_position;    ///< Most likely jump position to date.
  double m_best_score;       ///< Score at m_best_position
  double m_trail_zero_count; ///< Number of zero pixels in the trailing buffer
  double m_trail_sum;        ///< Sum of all pixels in the trailing buffer

}; // End class JumpFinder

/// View class which attempts to zero out the border noise on each border of the image.
/// - Because this uses the standard VW tile implementation and does not share information
///   between tiles, the maximum border size that can be handled is roughly the tile size.
/// - The JumpFinder class is used to find the border in each row/column and a median filter
///   is applied to those results to remove outliers.
template <class ImageT>
class Sentinel1CleanBordersView : public ImageViewBase<Sentinel1CleanBordersView<ImageT> > {

public: // Definitions

  typedef uint16     pixel_type;
  typedef pixel_type result_type;

private: // Variables

  ImageT const& m_input_image;

public: // Functions

  // Constructor
  Sentinel1CleanBordersView( ImageT  const& input_image)
                  : m_input_image(input_image){}

  inline int32 cols  () const { return m_input_image.cols(); }
  inline int32 rows  () const { return m_input_image.rows(); }
  inline int32 planes() const { return 1; }

  inline result_type operator()( int32 i, int32 j, int32 p=0 ) const
  { return 0; } // NOT IMPLEMENTED!

  typedef ProceduralPixelAccessor<Sentinel1CleanBordersView<ImageT> > pixel_accessor;
  inline pixel_accessor origin() const { return pixel_accessor( *this, 0, 0 ); }

  // Determin which edges of the image this tile is on
  void check_bbox_edges(BBox2i const& bbox, bool &left, bool &right, bool &top, bool &bottom) const {
    left   = (bbox.min().x() == 0);
    top    = (bbox.min().y() == 0);
    right  = (bbox.max().x() == cols());
    bottom = (bbox.max().y() == rows());
  }

  // This function does most of the work
  typedef CropView<ImageView<result_type> > prerasterize_type;
  inline prerasterize_type prerasterize( BBox2i const& bbox ) const {

    ImageView<result_type> output_tile = crop(m_input_image, bbox);

    // Check if this is a border tile.  If not, no processing will happen.
    bool left_edge, right_edge, top_edge, bottom_edge;
    check_bbox_edges(bbox, left_edge, right_edge, top_edge, bottom_edge);

    const unsigned short NODATA_VALUE        = 0; // Output nodata flag
    const int            MEDIAN_FILTER_WIDTH = 9; // Width of median filter applied to edge detect results
    const int            JUMP_WIDTH          = 8; // Width of lead region in the JumpFinder class

    // For each line coming in from the border, choose the best break point.
    int num_rows = bbox.height();
    int num_cols = bbox.width();

    if (left_edge) {
      std::vector<int> breaks(num_rows), breaks_filtered;
      for (int r=0; r<num_rows; ++r) { // For each row in this tile
        breaks[r] = 0; // Default if no jump found
        JumpFinder edge_detector(JUMP_WIDTH);
        // Keep iterating to the right until we detect the start of valid pixels
        for (int c=0; c<num_cols; ++c) { // Move along the row
          int value = output_tile(c,r);
          if (edge_detector.add_value(value))
            break;
        } // End col loop
        breaks[r] = edge_detector.get_best_jump();
      } // End row loop

      // Smooth out the borders
      window_median_filter(breaks, breaks_filtered, MEDIAN_FILTER_WIDTH);

      // Update the output image
      for (int r=0; r<num_rows; ++r) { // For each row in this tile
        for (int c=0; c<=breaks_filtered[r]; ++c) // Move along the row
          output_tile(c,r) = NODATA_VALUE;
      }
    } // End left edge case

    if (right_edge) {
      std::vector<int> breaks(num_rows), breaks_filtered;
      for (int r=0; r<num_rows; ++r) { // For each row in this tile
        //std::cout << "Row = " << r+bbox.min().y() << std::endl;
        breaks[r] = num_cols-1; // Default if no jump found
        JumpFinder edge_detector(JUMP_WIDTH);
        // Keep iterating to the right until we detect the start of valid pixels
        for (int c=num_cols-1; c>0; --c) { // Move along the row
          int value = output_tile(c,r);
          if (edge_detector.add_value(value))
            break;
        } // End col loop
        breaks[r] = num_cols-1 - edge_detector.get_best_jump();
      } // End row loop

      // Smooth out the borders
      window_median_filter(breaks, breaks_filtered, MEDIAN_FILTER_WIDTH);

      // Update the output image
      for (int r=0; r<num_rows; ++r) { // For each row in this tile
        for (int c=num_cols-1; c>=breaks_filtered[r]; --c) // Move along the row
          output_tile(c,r) = NODATA_VALUE;
      }
    } // End right edge case

    if (top_edge) {
      std::vector<int> breaks(num_cols), breaks_filtered;
      for (int c=0; c<num_cols; ++c) { // For each column in this tile
        breaks[c] = 0; // Default if no jump found
        JumpFinder edge_detector(JUMP_WIDTH);
        // Keep iterating to the right until we detect the start of valid pixels
        for (int r=0; r<num_rows; ++r) { // Move along the column
          int value = output_tile(c,r);
          if (edge_detector.add_value(value))
            break;
        } // End row loop
        breaks[c] = edge_detector.get_best_jump();
        //std::cout << "breaks[c] = " << breaks[c] << std::endl;
      } // End col loop

      // Smooth out the borders
      window_median_filter(breaks, breaks_filtered, MEDIAN_FILTER_WIDTH);

      // Update the output image
      for (int c=0; c<num_cols; ++c) { // For each column in this tile
        //std::cout << "breaks_filtered[c] = " << breaks_filtered[c] << std::endl;
        for (int r=0; r<breaks_filtered[c]; ++r) // Move along the column
          output_tile(c,r) = NODATA_VALUE;
      }
    } // End top edge case

    if (bottom_edge) {
      std::vector<int> breaks(num_cols), breaks_filtered;
      for (int c=0; c<num_cols; ++c) { // For each column in this tile
        breaks[c] = num_rows-1; // Default if no jump found
        JumpFinder edge_detector(JUMP_WIDTH);
        // Keep iterating to the right until we detect the start of valid pixels
        for (int r=num_rows-1; r>0; --r) { // Move along the column
          int value = output_tile(c,r);
          if (edge_detector.add_value(value))
            break;
        } // End row loop
        breaks[c] = num_rows-1 - edge_detector.get_best_jump();
      } // End col loop

      // Smooth out the borders
      window_median_filter(breaks, breaks_filtered, MEDIAN_FILTER_WIDTH);

      // Update the output image
      for (int c=0; c<num_cols; ++c) { // For each column in this tile
        for (int r=num_rows-1; r>breaks_filtered[c]; --r) // Move along the column
          output_tile(c,r) = NODATA_VALUE;
      }
    } // End bottom edge case

    // Return the tile we created with fake borders to make it look the size of the entire output image
    return prerasterize_type(output_tile, -bbox.min().x(), -bbox.min().y(), cols(), rows() );

  } // End prerasterize function

 template <class DestT>
 inline void rasterize( DestT const& dest, BBox2i const& bbox ) const {
   vw::rasterize( prerasterize(bbox), dest, bbox );
 }
}; // End class Sentinel1CleanBordersView

/// Helper function
template <class T>
Sentinel1CleanBordersView<T> clean_sentinel1_borders(T const& input) {
  return Sentinel1CleanBordersView<T>(input);
}

/// Compute an optimal tile size close to the input tile size
/// - The purpose of this function is to avoid a potential issue caused
///   by processing the image in tiles.  The edge cleaning is limited to
///   the size of one tile and there is the potential for truncated partial
///   tiles at the right and bottom edges of the images.
/// - This function slightly changes the tile size to ensure that any
///   truncated tiles are at least close to a full tile size and will have
///   adequate edge cleaning range.
int compute_new_tile_size(int input_size, int image_size) {

  // Define the safe range of percent size of edge tiles
  double MIN_TILE_PERCENTAGE = 0.75;
  double MAX_TILE_PERCENTAGE = 0.99;

  // The current tile usage.  Last tile is the size of a tile at the right
  //  or bottom of the tiling grid.
  double num_tiles_used       = (double)image_size / (double)input_size;
  double last_tile_percentage = num_tiles_used - floor(num_tiles_used);

  // If the input is in the acceptable range, keep it!
  if ((last_tile_percentage >= MIN_TILE_PERCENTAGE) && (last_tile_percentage <= MAX_TILE_PERCENTAGE))
    return input_size;

  //printf("Adjusting input tile size:\n");
  //printf("num_tiles = %lf, input tile size: %d\n", num_tiles_used, input_size);

  // Otherwise the tile size should be adjusted.

  // Decide how many tiles to shoot for
  double target_num_tiles = floor(num_tiles_used);
  if (last_tile_percentage > 0.75)
    target_num_tiles = ceil(num_tiles_used);

  // Compute new tile size, rounding up to ensure we don't end up with a tile sliver.
  int tile_size = ceil((double)image_size / target_num_tiles);

  //double first_num_tiles = (double)image_size / (double)tile_size;
  //printf("first_num_tiles = %lf, Computed new tile size: %d\n", first_num_tiles, tile_size);

  // GDAL block write sizes must be a multiple to 16 so if the input value is
  //  not a multiple of 16 increase it until it is.  This should not have a significant
  //  effect on the size of the border tiles.
  const int TILE_MULTIPLE = 16;
  if (tile_size % TILE_MULTIPLE != 0)
    tile_size = ((tile_size / TILE_MULTIPLE) + 1) * TILE_MULTIPLE;

  //double new_num_tiles = (double)image_size / (double)tile_size;
  //printf("new_num_tiles = %lf, Computed new tile size: %d\n", new_num_tiles, tile_size);

  return tile_size;
}

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

  std::string input_file;
  std::string output_path;
  int num_threads = 0;
  int tile_size   = 2048;

  /// Handle the command line parameters
  cartography::GdalWriteOptions write_options;

  po::options_description general_options("Clean up the borders of a Sentinel1 radar image.\n\nGeneral Options");
  general_options.add_options()
    ("output-path,o",    po::value<std::string>(&output_path), "The output file path")
    ("num-threads",      po::value<int>(&num_threads)->default_value(0),
                         "Number of threads to use for writing")
    ("tile-size",        po::value<int>(&tile_size)->default_value(2048),
                         "This is the farthest in that bad edges can be corrected.")
    ("help,h", "Display this help message");

  po::options_description hidden_options("");
  hidden_options.add_options()
    ("input-file", po::value<std::string>(&input_file));

  po::options_description options("Allowed Options");
  options.add(general_options).add(hidden_options);

  po::positional_options_description p;
  p.add("input-file", -1);

  std::ostringstream usage;
  usage << "Usage: clean_sentinel1_borders [options] <image file>" <<std::endl << std::endl;
  usage << general_options << std::endl;

  po::variables_map vm;
  try {
    po::store( po::command_line_parser( argc, argv ).options(options).positional(p).run(), vm );
    po::notify( vm );
  } catch (const po::error& e) {
    std::cout << "An error occurred while parsing command line arguments.\n";
    std::cout << "\t" << e.what() << "\n\n";
    std::cout << usage.str();
    return 1;
  }

  if( vm.count("help") ) {
    std::cout << usage.str();
    return 0;
  }

  if( vm.count("output-path") < 1 ) {
    std::cerr << "Error: must specify the output file!" << std::endl << std::endl;
    std::cout << usage.str();
    return 1;
  }

  // Modify tile size to make sure good coverage of the borders.
  ImageFormat input_format = image_format(input_file);
  int tile_size_h = compute_new_tile_size(tile_size, input_format.cols);
  int tile_size_v = compute_new_tile_size(tile_size, input_format.rows);

  // Load tile size and thread count into VW internals
  write_options.raster_tile_size = Vector2i(tile_size_h, tile_size_v);
  vw_settings().set_default_tile_size(tile_size);
  if (num_threads > 0) {
    write_options.num_threads = num_threads;
    vw_settings().set_default_num_threads(write_options.num_threads);
  }

  // Generate the output image
  const double OUTPUT_NODATA = 0;
  std::remove(output_path.c_str());
  block_write_gdal_image(output_path,
                         clean_sentinel1_borders(DiskImageView<uint16>(input_file)),
                         OUTPUT_NODATA, write_options);

}


// $Id: matrix_BOOST.cpp 27906 2007-04-27 11:50:53Z wmeeusse $
// Copyright (C) 2002 Klaas Gadeyne <first dot last at gmail dot com>

//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation; either version 2.1 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
//

#include "../config.h"
#ifdef __MATRIXWRAPPER_BOOST__

#include "matrix_BOOST.h"
#include "vector_BOOST.h"

#include <boost/numeric/ublas/matrix_proxy.hpp>

using namespace std;

// Passing the constructor arguments...
MyMatrix::Matrix() : BoostMatrix() {}
MyMatrix::Matrix(int num_rows, int num_cols) : BoostMatrix(num_rows,
							   num_cols){}

// Destructor
MyMatrix::~Matrix(){}

// Copy constructor
MyMatrix::Matrix(const MyMatrix& a) : BoostMatrix(a){}

// ill-designed
MyMatrix::Matrix(const BoostMatrix & a) : BoostMatrix(a){}

MyMatrix::Matrix(int num_rows,const RowVector& v):BoostMatrix(num_rows,v.size()){
  BoostMatrix & m = *this;
  for(unsigned int i=0;i<num_rows;i++)
    row(m,i).assign(v);
}

MyRowVector MyMatrix::operator[](unsigned int i) const{
  return this->rowCopy(i);
}

// Size/Capacity
unsigned int MyMatrix::size() const { return this->size1();}
unsigned int MyMatrix::capacity() const { return this->size1();}

// Number of Rows/Cols
unsigned int MyMatrix::rows() const { return this->size1();}
unsigned int MyMatrix::columns() const { return this->size2();}

// MATRIX - SCALAR operators
MyMatrix& MyMatrix::operator+= (double a)
{
  BoostMatrix & op1 = *this;
  op1 += boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a);
  return (MyMatrix&)op1;
}

MyMatrix& MyMatrix::operator-= (double a)
{
  BoostMatrix & op1 = (*this);
  op1 -= boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a);
  return (MyMatrix&) op1;
}

MyMatrix& MyMatrix::operator*= (double a)
{
  BoostMatrix & op1 = (*this);
  op1 *= a;
  return *this;
}

MyMatrix& MyMatrix::operator/= (double a)
{
  BoostMatrix & op1 = (*this);
  op1 /= a;
  return (MyMatrix&) op1;
}

MyMatrix MyMatrix::operator+ (double a) const
{
  return (MyMatrix)(((BoostMatrix)(*this)) + boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a));
}

MyMatrix MyMatrix::operator- (double a) const
{
  return (MyMatrix)(((BoostMatrix)(*this)) - boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a));
}

MyMatrix MyMatrix::operator* (double a) const
{
  const BoostMatrix& op1 = (*this);
  return (MyMatrix) (op1 *  a);
}

MyMatrix MyMatrix::operator/ (double a) const
{
  const BoostMatrix& op1 = (*this);
  return (MyMatrix) (op1 /  a);
}

MyMatrix&
MyMatrix::operator =(const MySymmetricMatrix& a)
{
  *this =(MyMatrix) a;

  return *this;
}

// MATRIX - MATRIX Operators
MyMatrix MyMatrix::operator- (const MyMatrix& a) const
{
  const BoostMatrix& op1 = *this;
  const BoostMatrix& op2 = a;

  return (MyMatrix)(op1 - op2);
}

MyMatrix MyMatrix::operator+ (const MyMatrix& a) const
{
  const BoostMatrix& op1 = *this;
  const BoostMatrix& op2 = a;

  return (MyMatrix)(op1 + op2);
}

MyMatrix MyMatrix::operator* (const MyMatrix& a) const
{
  const BoostMatrix& op1 = *this;
  const BoostMatrix& op2 = a;

  return (MyMatrix) prod(op1,op2);
}

MyMatrix & MyMatrix::operator+= (const MyMatrix& a)
{
  BoostMatrix & op1 = (*this);
  const BoostMatrix & op2 = a;
  op1 += op2;
  return (MyMatrix &) op1;
}

MyMatrix & MyMatrix::operator-= (const MyMatrix& a)
{
  BoostMatrix & op1 = (*this);
  const BoostMatrix & op2 = a;
  op1 -= op2;
  return (MyMatrix &) op1;
}

// MATRIX - VECTOR Operators
MyColumnVector MyMatrix::operator* (const MyColumnVector &b) const
{
  const BoostMatrix& op1 = (*this);
  return (MyColumnVector) prod(op1, ((const BoostColumnVector&)b));
}

double& MyMatrix::operator()(unsigned int a, unsigned int b)
{
  BoostMatrix & op1 = (*this);
  return op1(a-1,b-1);
}

double MyMatrix::operator()(unsigned int a, unsigned int b) const
{
  BoostMatrix  op1(*this);
  return op1(a-1,b-1);
}

bool MyMatrix::operator==(const MyMatrix& a) const
{
  if (this->rows() != a.rows()) return false;
  if (this->columns() != a.columns()) return false;
  return(norm_inf((BoostMatrix)(*this)-(BoostMatrix)a) == 0);
}

// Set all elements equal to a
MyMatrix&
 MyMatrix::operator=(double a)
{
  *this = (MyMatrix)boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a);

  return *this;
}

MyRowVector MyMatrix::rowCopy(unsigned int r) const
{
  unsigned int cols = columns();
  BoostRowVector temp(cols);
  for (unsigned int i=0; i<cols; i++)
    temp(i) = (*this)(r,i+1);
  return (MyRowVector) temp;
}

MyColumnVector MyMatrix::columnCopy(unsigned int c) const
{
  unsigned int ro = rows();
  BoostColumnVector temp(ro);
  for (unsigned int i=0; i<ro; i++)
    temp(i) = (*this)(i+1,c);
  return (MyColumnVector) temp;
}

MyMatrix MyMatrix::transpose() const
{
  const BoostMatrix &op1 = (*this);
  return (MyMatrix) trans(op1);
}

double MyMatrix::determinant() const
{
  unsigned int r = this->rows();
  assert(r == this->columns());
  double result = 1.0;
  const BoostMatrix& A = (*this);
  switch (r)
  {
        case 1:
            return A(0,0);
        case 2:
            return ( ( A(0,0) * A(1,1)) - ( A(1,0) * A(0,1)) );
        default:
            BoostMatrix LU(r,r);
            boost::numeric::ublas::permutation_matrix<> ndx(r);
            noalias(LU) = A;
            int res = lu_factorize(LU,ndx);
            assert(res == 0);

            int s = 1;
            for (boost::numeric::ublas::matrix<double>::size_type i=0;i<LU.size1();++i) {
              result *= LU(i,i);
              if (ndx(i)!=i) s = -s;
            }
            return result*s;
  }
}

MyMatrix MyMatrix::inverse() const
{
  unsigned int r = this->rows();
  assert(r == this->columns());
  const BoostMatrix& A = (*this);
  BoostMatrix Ai(r,r);
  switch (r)
  {
     case 1:
     {
        Ai(0,0) = 1/A(0,0);
        break;
     }
     case 2:
     {
       double det = A(0,0)*A(1,1)-A(0,1)*A(1,0);
       Ai(0,0) = A(1,1)/det;
       Ai(1,1) = A(0,0)/det;
       Ai(0,1) = -A(0,1)/det;
       Ai(1,0) = -A(1,0)/det;
       break;
     }
     default:
     {
       BoostMatrix LU(r,r);
       boost::numeric::ublas::permutation_matrix<> ndx(r);
       noalias(LU) = A;
       int res = lu_factorize(LU,ndx);
       assert(res == 0);
       noalias(Ai) = boost::numeric::ublas::identity_matrix<double>(r);
       lu_substitute(LU,ndx,Ai);
       break;
     }
  }
  return Ai;
}

int
MyMatrix::convertToSymmetricMatrix(MySymmetricMatrix& sym)
{
  // test if matrix is square matrix
  assert(this->rows() == this->columns());

  // if necessairy, resize sym
  // only check cols or rows. Symmetric matrix is square.
  if ( sym.rows() != this->rows() )
    sym = MySymmetricMatrix(this->rows());

  // copy elements
  for ( unsigned int i=0; i<this->rows(); i++ )
    for ( unsigned int j=0; j<=i; j++ )
      sym(i+1,j+1) = (*this)(i+1,j+1);
  return 0;
}

void
MyMatrix::resize(unsigned int i, unsigned int j, bool copy, bool initialize)
{
  BoostMatrix & temp = (BoostMatrix &) (*this);
  temp.resize(i,j,copy);
}

// get sub matrix
MyMatrix MyMatrix::sub(int i_start, int i_end, int j_start , int j_end) const
{
  MyMatrix submatrix(i_end-i_start+1, j_end-j_start+1);
  for (int i=i_start; i<=i_end; i++)
    for (int j=j_start; j<=j_end; j++)
      submatrix(i-i_start+1,j-j_start+1) = (*this)(i,j);

  return submatrix;
}

/////////////////////////////
// CLASS SYMMETRIC MATRIX  //
/////////////////////////////

MySymmetricMatrix::SymmetricMatrix() : BoostSymmetricMatrix() {}
MySymmetricMatrix::SymmetricMatrix(int n) : BoostSymmetricMatrix(n) {}
MySymmetricMatrix::SymmetricMatrix(int num_rows,const RowVector& v):BoostSymmetricMatrix(num_rows,v.size()){
  BoostSymmetricMatrix & m = *this;
  for(unsigned int i=0;i<num_rows;i++)
    row(m,i).assign(v);
}

MyRowVector MySymmetricMatrix::operator[](unsigned int i) const{
  return this->rowCopy(i);
}

// Copy constructor
MySymmetricMatrix::SymmetricMatrix(const SymmetricMatrix& a) : BoostSymmetricMatrix(a){}
MySymmetricMatrix::SymmetricMatrix(const BoostSymmetricMatrix & a) : BoostSymmetricMatrix(a){}

// Destructor
MySymmetricMatrix::~SymmetricMatrix(){}

// Size/Capacity
unsigned int MySymmetricMatrix::size() const { return this->size1();}
unsigned int MySymmetricMatrix::capacity() const { return this->size1();}

// Ask Number of Rows and Columns
unsigned int MySymmetricMatrix::rows() const { return this->size1();}
unsigned int MySymmetricMatrix::columns() const { return this->size2();}

MyRowVector MySymmetricMatrix::rowCopy(unsigned int r) const
{

  unsigned int cols = columns();
  BoostRowVector temp(cols);
  for (unsigned int i=0; i<cols; i++)
    temp(i) = (*this)(r,i+1);
  return (MyRowVector) temp;
}

MySymmetricMatrix MySymmetricMatrix::transpose() const {return (*this);}

MySymmetricMatrix MySymmetricMatrix::inverse() const
{
  unsigned int r = this->rows();
  assert(r == this->columns());
  const BoostMatrix& A = (*this);
  BoostMatrix Ai(r,r);
  switch (r)
  {
     case 1:
     {
        Ai(0,0) = 1/A(0,0);
        break;
     }
     case 2:
     {
       double det = A(0,0)*A(1,1)-A(0,1)*A(1,0);
       Ai(0,0) = A(1,1)/det;
       Ai(1,1) = A(0,0)/det;
       Ai(0,1) = -A(0,1)/det;
       Ai(1,0) = -A(1,0)/det;
       break;
     }
     default:
     {
       BoostMatrix LU(r,r);
       boost::numeric::ublas::permutation_matrix<> ndx(r);
       noalias(LU) = A;
       int res = lu_factorize(LU,ndx);
       assert(res == 0);
       noalias(Ai) = boost::numeric::ublas::identity_matrix<double>(r);
       lu_substitute(LU,ndx,Ai);
       break;
     }
  }

  return MySymmetricMatrix(Ai);
}

double MySymmetricMatrix::determinant() const
{
  unsigned int r = this->rows();
  assert(r == this->columns());
  const BoostMatrix& A = (*this);
  switch (r)
  {
     case 1:
     {
        return A(0,0);
     }
     case 2:
     {
       return A(0,0)*A(1,1)-A(0,1)*A(1,0);
     }
     default:
     {
        BoostMatrix LU(r,r);
        boost::numeric::ublas::permutation_matrix<> ndx(r);
        noalias(LU) = A;
        int res = lu_factorize(LU,ndx);
        assert(res == 0);

        double result = 1.0;
        int s = 1;
        for (boost::numeric::ublas::matrix<double>::size_type i=0;i<LU.size1();++i) {
          result *= LU(i,i);
          if (ndx(i)!=i) s = -s;
        }
        return result*s;
     }
  }
}

// Set all elements equal to a
MySymmetricMatrix& MySymmetricMatrix::operator=(const double a)
{
  *this = (MySymmetricMatrix)boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a);

  return *this;
}

// SYMMETRICMATRIX - SCALAR operators
MySymmetricMatrix& MySymmetricMatrix::operator +=(double a)
{
  BoostSymmetricMatrix & op1 = *this;
  op1 += boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a);
  return (MySymmetricMatrix&)op1;
}

MySymmetricMatrix& MySymmetricMatrix::operator -=(double a)
{
  BoostSymmetricMatrix & op1 = *this;
  op1 -= boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a);
  return (MySymmetricMatrix&)op1;
}

MySymmetricMatrix& MySymmetricMatrix::operator *=(double b)
{
  BoostSymmetricMatrix & op1 = (*this);
  op1 *= b;
  return (MySymmetricMatrix&) op1;
}

MySymmetricMatrix& MySymmetricMatrix::operator /=(double b)
{
  BoostSymmetricMatrix & op1 = (*this);
  op1 /= b;
  return (MySymmetricMatrix&) op1;
}

MySymmetricMatrix MySymmetricMatrix::operator +(double a) const
{
  return (MySymmetricMatrix)(((BoostSymmetricMatrix)(*this)) + boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a));
}

MySymmetricMatrix MySymmetricMatrix::operator -(double a) const
{
  return (MySymmetricMatrix)(((BoostSymmetricMatrix)(*this)) - boost::numeric::ublas::scalar_matrix<double>(rows(),columns(),a));
}

MySymmetricMatrix MySymmetricMatrix::operator *(double b) const
{
 const BoostSymmetricMatrix& op1 = (*this);
  return (MySymmetricMatrix) (op1 *  b);
}

MySymmetricMatrix MySymmetricMatrix::operator /(double b) const
{
  const BoostSymmetricMatrix& op1 = (*this);
  return (MySymmetricMatrix) (op1 /  b);
}

// SYMMETRICMATRIX - MATRIX operators
MyMatrix& MySymmetricMatrix::operator +=(const MyMatrix& a)
{
  BoostSymmetricMatrix & op1 = (*this);
  op1 += a;
  return (MyMatrix &) op1;
}

MyMatrix& MySymmetricMatrix::operator -=(const MyMatrix& a)
{
  BoostSymmetricMatrix & op1 = (*this);
  op1 -= a;
  return (MyMatrix &) op1;
}

MyMatrix MySymmetricMatrix::operator+ (const MyMatrix &a) const
{
  const BoostSymmetricMatrix& op1 = *this;
  const BoostMatrix& op2 = a;

  return (MyMatrix) (op1 + op2);
}

MyMatrix MySymmetricMatrix::operator- (const MyMatrix &a) const
{
  const BoostSymmetricMatrix& op1 = *this;
  const BoostMatrix& op2 = a;

  return (MyMatrix) (op1 - op2);
}

MyMatrix MySymmetricMatrix::operator* (const MyMatrix &a) const
{
  const BoostSymmetricMatrix& op1 = *this;
  const BoostMatrix& op2 = a;

  return (MyMatrix) prod(op1, op2);
}

// SYMMETRICMATRIX - SYMMETRICMATRIX operators
MySymmetricMatrix& MySymmetricMatrix::operator +=(const MySymmetricMatrix& a)
{
  BoostSymmetricMatrix & op1 = (*this);
  const BoostSymmetricMatrix & op2 = a;
  op1 += op2;
  return (MySymmetricMatrix &) op1;
}

MySymmetricMatrix& MySymmetricMatrix::operator -=(const MySymmetricMatrix& a)
{
  BoostSymmetricMatrix & op1 = (*this);
  const BoostSymmetricMatrix & op2 = a;
  op1 -= op2;
  return (MySymmetricMatrix &) op1;
}

MySymmetricMatrix MySymmetricMatrix::operator+ (const MySymmetricMatrix &a) const
{
  const BoostSymmetricMatrix& op1 = *this;
  const BoostSymmetricMatrix& op2 = a;

  return (MySymmetricMatrix) (op1 + op2);
}

MySymmetricMatrix MySymmetricMatrix::operator- (const MySymmetricMatrix &a) const
{
  const BoostSymmetricMatrix& op1 = *this;
  const BoostSymmetricMatrix& op2 = a;

  return (MySymmetricMatrix) (op1 - op2);
}

MyMatrix MySymmetricMatrix::operator* (const MySymmetricMatrix &a) const
{
  const BoostSymmetricMatrix& op1 = *this;
  const BoostSymmetricMatrix& op2 = a;

  return (MyMatrix) prod(op1, op2);
}

MyColumnVector MySymmetricMatrix::operator* (const MyColumnVector &b) const
{
  const BoostSymmetricMatrix& op1 = (BoostSymmetricMatrix) *this;
  return (MyColumnVector) prod(op1, ((const BoostColumnVector&)b));
}

void MySymmetricMatrix::multiply (const MyColumnVector &b, MyColumnVector &result) const
{
  const BoostSymmetricMatrix& op1 = (BoostSymmetricMatrix) *this;
  result = (MyColumnVector) prod(op1, ((const BoostColumnVector&)b));
}

MyMatrix MySymmetricMatrix::sub(int i_start, int i_end, int j_start , int j_end) const
{
  MyMatrix submatrix(i_end-i_start+1, j_end-j_start+1);
  for (int i=i_start; i<=i_end; i++)
    for (int j=j_start; j<=j_end; j++)
      submatrix(i-i_start+1,j-j_start+1) = (*this)(i,j);

  return submatrix;
}

double& MySymmetricMatrix::operator()(unsigned int a, unsigned int b)
{
  BoostSymmetricMatrix & op1 = (*this);
  return op1(a-1,b-1);
}

double MySymmetricMatrix::operator()(unsigned int a, unsigned int b) const
{
  BoostSymmetricMatrix op1(*this);
  return op1(a-1,b-1);
}

bool MySymmetricMatrix::operator==(const MySymmetricMatrix& a) const
{
  if (this->rows() != a.rows()) return false;
  if (this->columns() != a.columns()) return false;
  return(norm_inf((BoostSymmetricMatrix)(*this)-(BoostSymmetricMatrix)a) == 0);
}

void
MySymmetricMatrix::resize(unsigned int i, bool copy, bool initialize)
{
  BoostSymmetricMatrix & temp = (BoostSymmetricMatrix &) (*this);
  temp.resize(i, copy);
}

#endif


//
//  Copyright © 2015 Claus Christmann <hcc |ä| gatech.edu>.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

#ifndef TYPEDEFS_H
#define TYPEDEFS_H

#include <Eigen/Dense>

namespace AMG {

  /** \brief A tupel of 3 values, representing a set of coordinates.
   *
   * The purpose of this type is to allow faster computations by accessing the
   * underlying Eigen::Vector3d type inside functions and algorithms that have
   * sorted out the FrameOfReference issues, i.e. all utilized entities are
   * expressed in the same FrameOfReference.
   */
  using CoordinateTupel = Eigen::Vector3d;

  /** \brief A tupel of 4 values, representing a set of quaternions.
   *
   * Quaternions are normally defined as either q = w + ix + jy + kz or
   * q = q0 + iq1 + jq2 + kq3.
   *
   * Eigen, the underlying matrix library has an own type for quaternions:
   * Eigen::Quaternion::Coefficients. However, those coefficients are
   * internally stored in an unconeventional manner, which could lead to
   * confusion. Eigen stores the quaternion tupel as [x,y,z,w], which could lead
   * to an ordering permutation when used in the "flight mechanics way" as
   *
   * \code
   * double q0=Eigen::Quaternion::Coefficient[3]; // w, the real part
   * double q1=Eigen::Quaternion::Coefficient[0]; // x, the 1st imaginary part
   * double q2=Eigen::Quaternion::Coefficient[1]; // y, the 2nd imaginary part
   * double q3=Eigen::Quaternion::Coefficient[2]; // z, the 3rd imaginary part
   * \endcode
   *
   * In order to avoid this, an  AMG::QuaternionTupel is defined as the
   * [q0,q1,q2,q3] tupel of quaternions:
   *
   * \code
   * double q0=AMG::QuaternionTupel[0]; // w, the real part
   * double q1=AMG::QuaternionTupel[1]; // x, the 1st imaginary part
   * double q2=AMG::QuaternionTupel[2]; // y, the 2nd imaginary part
   * double q3=AMG::QuaternionTupel[3]; // z, the 3rd imaginary part
   * \endcode
   *
   * \warning As a result of this
   *  AMG::QuaternionTupel != Eigen::Quaternion::Coefficients ! (It is a
   *  permutation of it.)
   */
  using QuaternionTupel = Eigen::Vector4d ;

  /** \brief A 3-tupel holding 3-2-1 Euler angles.
   *
   * This is a 3-tupel holding the roll/phi, pitch/theta, and yaw/psi
   * Euler-angles.
   *
   * \note The tupel does not keep track of the used units, i.e. degrees or
   *  radians! (That's a todo...)
   *
   * \code
   * AMG::EulerAngleTupel euler;
   *
   * double phi   = euler[0]; // roll angle
   * double theta = euler[1]; // pitch angle
   * double psi   = euler[2]; // yaw angle
   * \endcode
   *
   * \todo Make this a template taking Units::degree or Units::radian
   */
  using EulerAngleTupel = Eigen::Vector3d;

}

#endif // TYPEDEFS_H


/**
 * @file onepass_plus_impl.hpp
 * @author Leonardo Arcari (leonardo1.arcari@gmail.com)
 * @version 1.0.0
 * @date 2018-10-28
 *
 * @copyright Copyright (c) 2018 Leonardo Arcari
 *
 * MIT License
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 */

#ifndef ONEPASS_PLUS_IMPL_HPP
#define ONEPASS_PLUS_IMPL_HPP

#include <boost/graph/graph_concepts.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/properties.hpp>

#include <arlib/details/arlib_utils.hpp>
#include <arlib/type_traits.hpp>

#include <cassert>
#include <iostream>
#include <memory>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace arlib {
/**
 * Implementations details of kSPwLO algorithms
 */
namespace details {

//===----------------------------------------------------------------------===//
//                    OnePass+ algorithm support classes
//===----------------------------------------------------------------------===//

/**
 * A label for a node in the graph to keep track of its exploration
 *        state.
 *
 * A label tracks the path from the source to the node @c n it's attached, the
 * similarity of the path <tt>p(s -> n)</tt> wrt to the alternative paths
 * computed so far and the time step the similarities were updated (i.e. the
 * number of alternative paths against which the similarities are currently
 * computed)
 *
 * Labels can either be @c head labels if they are attached to source node, or
 * have a predecessor label, i.e. they are attached to a node @c n such that
 * there exist a label attached to a node @c n' and an edge <tt>(n', n)</tt> in
 * Graph.
 *
 * @tparam Graph A Boost::Graph
 * @tparam Vertex A vertex of Graph.
 * @tparam Length The value type of an edge weight of Graph.
 */
template <typename Graph, typename Length, typename Edge = edge_of_t<Graph>,
          typename Vertex = vertex_of_t<Graph>>
class OnePassLabel {
private:
  Vertex node;
  Length length;
  Length lower_bound;
  OnePassLabel *previous;
  std::vector<double> similarity_map;
  int k;
  int checked_at_step;

public:
  /**
   * size_type trait of similarity vector
   */
  using similarity_map_size_type = typename decltype(similarity_map)::size_type;

  /**
   * iterator trait of similarity vector
   */
  using similarity_map_iterator_type =
      typename decltype(similarity_map)::iterator;

  /**
   * Length
   */
  using length_type = Length;

  /**
   * Vertex
   */
  using vertex_type = Vertex;

  /**
   * Construct a new OnePass Label object with a predecessor label.
   *
   * @param node Node to attach this label to.
   * @param length Distance of this node from source following the path from
   *        this label to the source's one.
   * @param lower_bound AStar heuristic of the distance of @p node from target.
   * @param previous Predecessor label.
   * @param k Number of k alternative paths to compute.
   * @param checked_at_step The current time step (i.e. the number of
   *        alternative paths currently computed)
   */
  OnePassLabel(Vertex node, Length length, Length lower_bound,
               OnePassLabel *previous, int k, int checked_at_step)
      : node{node}, length{length},
        lower_bound{lower_bound}, previous{previous},
        similarity_map(k, 0), k{k}, checked_at_step{checked_at_step} {}

  /**
   * Construct a new OnePass Label object with no predecessor (i.e. a @c
   * head label)
   *
   * @param node Node to attach this label to.
   * @param length Distance of this node from source following the path from
   *        this label to the source's one.
   * @param lower_bound AStar heuristic of the distance of @p node from target.
   * @param k Number of k alternative paths to compute.
   * @param checked_at_step The current time step (i.e. the number of
   *        alternative paths currently computed)
   */
  OnePassLabel(Vertex node, Length length, Length lower_bound, int k,
               int checked_at_step)
      : node{node}, length{length}, lower_bound{lower_bound}, previous{nullptr},
        similarity_map(k, 0), k{k}, checked_at_step{checked_at_step} {}

  std::vector<Edge> get_path(Graph const &G) const {
    using namespace boost;
    auto edge_set = std::vector<Edge>{};

    auto v = node;
    auto prev = previous;
    while (prev != nullptr) {
      auto u = prev->node;
      auto [e, is_ok] = edge(u, v, G);
      assert(is_ok &&
             "[arlib::details::OnePassLabel::get_path] Edge not found.");
      edge_set.push_back(e);

      // Shift label pointer back
      v = u;
      prev = prev->previous;
    }

    return edge_set;
  }

  bool is_path_acyclic(Vertex begin) const {
    bool is_acyclic = true;
    auto current = this;
    while (current != nullptr) {
      if (current->node == begin) {
        is_acyclic = false;
        break;
      }
      current = current->previous;
    }
    return is_acyclic;
  }

  /**
   * @param kth The kth alternative path index.
   * @return A reference to the similarity of path <tt>p(source, n)</tt> wrt to
   *         @p kth alternative path.
   */
  double &get_similarity_with(int kth) { return similarity_map.at(kth); }

  /**
   * @param kth The kth alternative path index.
   * @return A const-reference to the similarity of path <tt>p(source, n)</tt>
   *         wrt to @p kth alternative path.
   */
  const double &get_similarity_with(int kth) const {
    return similarity_map.at(kth);
  }

  /**
   * @return the number of alternative paths for which there exists a similarity
   *         measure for, for this label.
   */
  similarity_map_size_type get_num_paths() const {
    return similarity_map.size();
  }

  /**
   * @return a copy of the similarity vector wrt the alternative paths.
   */
  std::vector<double> get_similarity_map() const {
    return std::vector<double>(std::begin(similarity_map),
                               std::end(similarity_map));
  }

  /**
   * Copies the similarity values from [@p first, @p last) iterators into
   *        label's similarity vector.
   *
   * @param first begin iterator.
   * @param last past-to-end iterator.
   */
  void set_similarities(similarity_map_iterator_type first,
                        similarity_map_iterator_type last) {
    std::copy(first, last, std::begin(similarity_map));
  }

  /**
   * @return the node this label is attached to.
   */
  Vertex get_node() const { return node; }

  /**
   * @return the distance of this node from source following the path from
   *         this label to the source's one.
   */
  length_type get_length() const { return length; }

  /**
   * @return AStar heuristic of the distance of the node from target.
   */
  length_type get_lower_bound() const { return lower_bound; }

  /**
   * @return Number of k alternative paths to compute.
   */
  int num_paths_k() const { return k; }

  /**
   * @return the time step the similarities were updated (i.e. the
   *         number of alternative paths against which the similarities are
   *         currently computed).
   */
  int last_check() const { return checked_at_step; }

  /**
   * @param currentStep the current time step.
   * @return true if @c last_check() < @p currentStep.
   * @return false otherwise.
   */
  bool is_outdated(int currentStep) const {
    return checked_at_step < currentStep;
  }

  /**
   * Set the time step of similarities update.
   *
   * @param step the time step.
   */
  void set_last_check(int step) {
    assert(step > 0);
    checked_at_step = step;
  }

private:
  template <typename Graph2, typename Length2>
  friend std::ostream &operator<<(std::ostream &os,
                                  const OnePassLabel<Graph2, Length2> &label);
};

template <typename Graph2, typename Length2>
std::ostream &operator<<(std::ostream &os,
                         const OnePassLabel<Graph2, Length2> &label) {
  os << "(node = " << label.node << ", length = " << label.length
     << ", lower_bound = " << label.lower_bound << ", k = " << label.k
     << ", checked_at_step = " << label.checked_at_step
     << ", similarities = [ ";
  for (auto sim : label.similarity_map) {
    os << sim << " ";
  }

  os << "])";
  return os;
}

/**
 * A conventient container for labels dominance checking.
 *
 * When a new label @c l' for a node @c n is created we must check it against
 * Lemma 2: Let @c P_LO the set of computed alternative paths so far. If all the
 * labels already existing for node @c n have a similarity with @c P_LO that is
 * less than the similarity of @c l' with @c P_LO, then @c l' is dominated and
 * should be pruned.
 *
 * @tparam Graph A Boost::Graph
 * @tparam Vertex A vertex of Graph.
 */
template <typename Graph, typename Length, typename Vertex = vertex_of_t<Graph>>
class SkylineContainer {
public:
  using Label = OnePassLabel<Graph, Length>;
  using LabelPtr = Label *;

  /**
   * Inserts a label into the skyline. Before insertion, you should check
   * if the skyline already dominates @p label. See dominates().
   *
   * @param label A label to add to the skyline.
   */
  void insert(Label *label) {
    auto node_n = label->get_node();
    // If node_n is new
    if (!contains(node_n)) {
      // Initialize the labels for node n with 'label'
      auto labels_for_n = std::vector<LabelPtr>{label};
      container.insert(std::make_pair(node_n, labels_for_n));
    } else {
      // If not, just add the label to node_n's list of labels
      container[node_n].push_back(label);
    }
  }

  /**
   * @param node A node of the graph.
   * @return true if the skyline contains @p node.
   * @return false otherwise.
   */
  bool contains(Vertex node) const {
    return container.find(node) != std::end(container);
  }

  /**
   * Check if the skyline dominates @p label. Refer to SkylineContainer
   * description for a more clear explanation.
   *
   * @param label A label to check if it's dominated by the skyline.
   * @return true if the skyline dominates @p label.
   * @return false otherwise.
   */
  bool dominates(const Label &label) {

    auto node_n = label.get_node();
    // If node_n is not in the skyline, then we have no labels to check
    if (!contains(node_n)) {
      return false;
    }

    // For each of the paths p' we have similarity measure in 'label' we
    // check if label.sim(p') is less than at least one of the labels in the
    // skyline. If so, the skyline does NOT dominates 'label'.
    for (const auto label_p : container[node_n]) {
      LabelPtr tmpLabel = label_p;
      bool skyline_dominates_label = true;

      for (int i = 0; i < static_cast<int>(label.get_num_paths()); ++i) {
        if (label.get_similarity_with(i) < tmpLabel->get_similarity_with(i)) {
          skyline_dominates_label = false;
          break;
        }
      }

      if (skyline_dominates_label) {
        return true;
      }
    }

    return false;
  }

  /**
   * @return the number of labels stored in the skyline.
   */
  int num_labels() const {
    int nb_labels = 0;

    for (auto it = std::begin(container); it != std::end(container); ++it) {
      nb_labels += it->second.size();
    }

    return nb_labels;
  }

private:
  std::unordered_map<Vertex, std::vector<LabelPtr>, boost::hash<Vertex>>
      container{};
};

/**
 * A Comparator functor to compare two labels in an A* fashion.
 *
 * A lower bound for the distance of the label from the target is used to decide
 * the ordering. The label with the lowest lower bound is the smallest.
 *
 * @tparam Graph A Boost::Graph
 */
template <typename Graph, typename Length> struct OnePassPlusASComparator {
  using LabelPtr = OnePassLabel<Graph, Length> *;

  bool operator()(LabelPtr lhs, LabelPtr rhs) const {
    return lhs->get_lower_bound() > rhs->get_lower_bound();
  }
};

//===----------------------------------------------------------------------===//
//                      OnePass+ algorithm routines
//===----------------------------------------------------------------------===//

template <typename Graph, typename EdgeMap,
          typename resPathIndex = typename EdgeMap::mapped_type::size_type>
void update_res_edges(const Graph &candidate, const Graph &graph,
                      EdgeMap &resEdges, resPathIndex paths_count) {
  using mapped_type = typename EdgeMap::mapped_type;
  for (auto ei = edges(candidate).first; ei != edges(candidate).second; ++ei) {
    auto edge_in_g =
        edge(source(*ei, candidate), target(*ei, candidate), graph).first;

    auto search = resEdges.find(edge_in_g);
    if (search != std::end(resEdges)) { // edge found
      search->second.push_back(paths_count - 1);
    } else { // new edge, add it to resEdges
      resEdges.insert(std::make_pair(edge_in_g, mapped_type{paths_count - 1}));
    }
  }
}

template <typename EdgeMap, typename Edge,
          typename resPathIndex = typename EdgeMap::mapped_type::size_type>
void update_res_edges(const std::vector<Edge> &candidate, EdgeMap &resEdges,
                      resPathIndex paths_count) {
  using mapped_type = typename EdgeMap::mapped_type;
  for (auto &e : candidate) {

    auto search = resEdges.find(e);
    if (search != std::end(resEdges)) { // edge found
      search->second.push_back(paths_count - 1);
    } else { // new edge, add it to resEdges
      resEdges.insert(std::make_pair(e, mapped_type{paths_count - 1}));
    }
  }
}

template <typename Label, typename Graph, typename EdgesMap, typename PathsMap,
          typename WeightMap>
bool update_label_similarity(Label &label, const Graph &G,
                             const EdgesMap &resEdges, const PathsMap &resPaths,
                             WeightMap &weight, double theta,
                             std::size_t step) {
  using namespace boost;
  bool below_sim_threshold = true;
  auto tmpPath = label.get_path(G);
  for (auto &e : tmpPath) {
    auto search = resEdges.find(e);
    // if tmpPath share an edge with any k-th shortest path, update the
    // overlapping factor
    if (search != std::end(resEdges)) {
      for (auto index : search->second) {
        if (static_cast<int>(index) > label.last_check() && index < step) {
          label.get_similarity_with(index) += weight[e];

          // Check Lemma 1. The similarity between the candidate path and
          // all the other k-shortest-paths must be less then theta
          auto const &alt_path = resPaths[index];
          auto alt_len = compute_length_from_edges(alt_path.begin(),
                                                   alt_path.end(), weight);
          if (label.get_similarity_with(index) / alt_len > theta) {
            below_sim_threshold = false;
            break;
          }
        }
      }
    }
  }
  return below_sim_threshold;
}

template <typename Label, typename Vertex = typename Label::Vertex,
          typename length_type = typename Label::length_type>
std::unique_ptr<Label> expand_path(Label *label, Vertex node,
                                   length_type node_lower_bound,
                                   length_type edge_weight, int step) {
  auto tmpLength = label->get_length() + edge_weight;
  auto tmpLowerBound = tmpLength + node_lower_bound;
  return std::make_unique<Label>(node, tmpLength, tmpLowerBound, label,
                                 label->num_paths_k(), step);
}

template <typename EdgesMap, typename PathsList, typename WeightMap,
          typename Edge = typename EdgesMap::key_type>
bool is_below_sim_threshold(const Edge &c_edge,
                            std::vector<double> &similarity_map, double theta,
                            const EdgesMap &resEdges, const PathsList &resPaths,
                            const WeightMap &weight) {
  auto search = resEdges.find(c_edge);
  if (search != std::end(resEdges)) {
    auto &res_paths_with_c_edge = search->second;
    for (auto index : res_paths_with_c_edge) {
      similarity_map[index] += weight[c_edge];
      auto const &alt_path = resPaths[index];
      auto alt_len =
          compute_length_from_edges(alt_path.begin(), alt_path.end(), weight);
      auto similarity = similarity_map[index] / alt_len;
      if (similarity > theta) {
        return false;
      }
    }
  }
  return true;
}
} // namespace details
} // namespace arlib

#endif