/**
 * @file environment.hpp
 * @author Licheng Wen (wenlc@zju.edu.cn)
 * @brief Environment class header
 * @date 2020-11-12
 *
 * @copyright Copyright (c) 2020
 *
 */

#pragma once
#include <ompl/base/State.h>
#include <ompl/base/spaces/DubinsStateSpace.h>
#include <ompl/base/spaces/ReedsSheppStateSpace.h>
#include <ompl/base/spaces/SE2StateSpace.h>

#include <boost/functional/hash.hpp>
#include <boost/geometry.hpp>
#include <boost/geometry/algorithms/intersection.hpp>
#include <boost/geometry/geometries/linestring.hpp>
#include <boost/geometry/geometries/point_xy.hpp>
#include <boost/heap/fibonacci_heap.hpp>
#include <unordered_map>
#include <unordered_set>

#include "neighbor.hpp"
#include "planresult.hpp"

namespace Constants {
// [m] --- The minimum turning radius of the vehicle
static float r = 3;
static float deltat = 6.75 * 6 / 180.0 * M_PI;
// [#] --- A movement cost penalty for turning (choosing non straight motion
// primitives)
static float penaltyTurning = 1.5;
// [#] --- A movement cost penalty for reversing (choosing motion primitives >
// 2)
static float penaltyReversing = 2.0;
// [#] --- A movement cost penalty for change of direction (changing from
// primitives < 3 to primitives > 2)
static float penaltyCOD = 2.0;
// map resolution
static float mapResolution = 2.0;
// change to set calcIndex resolution
static float xyResolution = r * deltat;
static float yawResolution = deltat;

// width of car
static float carWidth = 2.0;
// distance from rear to vehicle front end
static float LF = 2.0;
// distance from rear to vehicle back end
static float LB = 1.0;
// obstacle default radius
static float obsRadius = 1;
// least time to wait for constraint
static int constraintWaitTime = 2;

// R = 3, 6.75 DEG
std::vector<double> dyaw = {0, deltat, -deltat, 0, -deltat, deltat};
std::vector<double> dx = {r * deltat, r *sin(deltat),  r *sin(deltat),
                          -r *deltat, -r *sin(deltat), -r *sin(deltat)};
std::vector<double> dy = {0, -r *(1 - cos(deltat)), r *(1 - cos(deltat)),
                          0, -r *(1 - cos(deltat)), r *(1 - cos(deltat))};

static inline float normalizeHeadingRad(float t) {
  if (t < 0) {
    t = t - 2.f * M_PI * (int)(t / (2.f * M_PI));
    return 2.f * M_PI + t;
  }

  return t - 2.f * M_PI * (int)(t / (2.f * M_PI));
}
}  // namespace Constants

namespace libMultiRobotPlanning {

using libMultiRobotPlanning::Neighbor;
using libMultiRobotPlanning::PlanResult;
using namespace libMultiRobotPlanning;
typedef ompl::base::SE2StateSpace::StateType OmplState;
typedef boost::geometry::model::d2::point_xy<double> Point;
typedef boost::geometry::model::segment<Point> Segment;
/**
 * @brief  Environment class
 *
 * @tparam Location
 * @tparam State
 * @tparam Action
 * @tparam Cost
 * @tparam Conflict
 * @tparam Constraint
 * @tparam Constraints
 */
template <typename Location, typename State, typename Action, typename Cost,
          typename Conflict, typename Constraint, typename Constraints>
class Environment {
 public:
  Environment(size_t maxx, size_t maxy, std::unordered_set<Location> obstacles,
              std::multimap<int, State> dynamic_obstacles,
              std::vector<State> goals)
      : m_obstacles(std::move(obstacles)),
        m_dynamic_obstacles(std::move(dynamic_obstacles)),
        m_agentIdx(0),
        m_constraints(nullptr),
        m_lastGoalConstraint(-1),
        m_highLevelExpanded(0),
        m_lowLevelExpanded(0) {
    m_dimx = (int)maxx / Constants::mapResolution;
    m_dimy = (int)maxy / Constants::mapResolution;
    // std::cout << "env build " << m_dimx << " " << m_dimy << " "
    //           << m_obstacles.size() << std::endl;
    holonomic_cost_maps = std::vector<std::vector<std::vector<double>>>(
        goals.size(), std::vector<std::vector<double>>(
                          m_dimx, std::vector<double>(m_dimy, 0)));
    m_goals.clear();
    for (const auto &g : goals) {
      if (g.x < 0 || g.x > maxx || g.y < 0 || g.y > maxy) {
        std::cout << "\033[1m\033[31m Goal out of boundary, Fail to build "
                     "environment \033[0m\n";
        return;
      }
      m_goals.emplace_back(
          State(g.x, g.y, Constants::normalizeHeadingRad(g.yaw)));
    }
    updateCostmap();
  }

  Environment(const Environment &) = delete;
  Environment &operator=(const Environment &) = delete;

  /// High Level Environment functions
  bool getFirstConflict(
      const std::vector<PlanResult<State, Action, double>> &solution,
      Conflict &result) {
    int max_t = 0;
    for (const auto &sol : solution) {
      max_t = std::max<int>(max_t, sol.states.size() - 1);
    }
    for (int t = 0; t < max_t; ++t) {
      // check drive-drive collisions
      for (size_t i = 0; i < solution.size(); ++i) {
        State state1 = getState(i, solution, t);
        for (size_t j = i + 1; j < solution.size(); ++j) {
          State state2 = getState(j, solution, t);
          if (state1.agentCollision(state2)) {
            result.time = t;
            result.agent1 = i;
            result.agent2 = j;
            result.s1 = state1;
            result.s2 = state2;
            return true;
          }
        }
      }
    }
    return false;
  }

  void createConstraintsFromConflict(
      const Conflict &conflict, std::map<size_t, Constraints> &constraints) {
    Constraints c1;
    c1.constraints.emplace(
        Constraint(conflict.time, conflict.s2, conflict.agent2));
    constraints[conflict.agent1] = c1;
    Constraints c2;
    c2.constraints.emplace(
        Constraint(conflict.time, conflict.s1, conflict.agent1));
    constraints[conflict.agent2] = c2;
  }

  void onExpandHighLevelNode(int /*cost*/) {
    m_highLevelExpanded++;
    if (m_highLevelExpanded % 50 == 0)
      std::cout << "Now expand " << m_highLevelExpanded
                << " high level nodes.\n";
  }

  int highLevelExpanded() { return m_highLevelExpanded; }

  /// Low Level Environment functions
  void setLowLevelContext(size_t agentIdx, const Constraints *constraints) {
    assert(constraints);  // NOLINT
    m_agentIdx = agentIdx;
    m_constraints = constraints;
    m_lastGoalConstraint = -1;
    for (const auto &c : constraints->constraints) {
      if (m_goals[m_agentIdx].agentCollision(c.s)) {
        m_lastGoalConstraint = std::max(m_lastGoalConstraint, c.time);
      }
    }

    // std::cout << "Setting Lowlevel agent idx:" << agentIdx
    //           << " Constraints:" << constraints->constraints.size()
    //           << "  lastGoalConstraints:" << m_lastGoalConstraint <<
    //           std::endl;
  }

  int admissibleHeuristic(const State &s) {
    // non-holonomic-without-obstacles heuristic: use a Reeds-Shepp
    ompl::base::ReedsSheppStateSpace reedsSheppPath(Constants::r);
    OmplState *rsStart = (OmplState *)reedsSheppPath.allocState();
    OmplState *rsEnd = (OmplState *)reedsSheppPath.allocState();
    rsStart->setXY(s.x, s.y);
    rsStart->setYaw(s.yaw);
    rsEnd->setXY(m_goals[m_agentIdx].x, m_goals[m_agentIdx].y);
    rsEnd->setYaw(m_goals[m_agentIdx].yaw);
    double reedsSheppCost = reedsSheppPath.distance(rsStart, rsEnd);
    // std::cout << "ReedsShepps cost:" << reedsSheppCost << std::endl;
    // Euclidean distance
    double euclideanCost = sqrt(pow(m_goals[m_agentIdx].x - s.x, 2) +
                                pow(m_goals[m_agentIdx].y - s.y, 2));
    // std::cout << "Euclidean cost:" << euclideanCost << std::endl;
    // holonomic-with-obstacles heuristic
    double twoDoffset =
        sqrt(pow((s.x - (int)s.x) -
                     (m_goals[m_agentIdx].x - (int)m_goals[m_agentIdx].x),
                 2) +
             pow((s.y - (int)s.y) -
                     (m_goals[m_agentIdx].y - (int)m_goals[m_agentIdx].y),
                 2));
    double twoDCost =
        holonomic_cost_maps[m_agentIdx][(int)s.x / Constants::mapResolution]
                           [(int)s.y / Constants::mapResolution] -
        twoDoffset;
    // std::cout << "holonomic cost:" << twoDCost << std::endl;

    return std::max({reedsSheppCost, euclideanCost, twoDCost});
    return 0;
  }

  bool isSolution(
      const State &state, double gscore,
      std::unordered_map<State, std::tuple<State, Action, double, double>,
                         std::hash<State>> &_camefrom) {
    double goal_distance =
        sqrt(pow(state.x - getGoal().x, 2) + pow(state.y - getGoal().y, 2));
    if (goal_distance > 3 * (Constants::LB + Constants::LF)) return false;
    ompl::base::ReedsSheppStateSpace reedsSheppSpace(Constants::r);
    OmplState *rsStart = (OmplState *)reedsSheppSpace.allocState();
    OmplState *rsEnd = (OmplState *)reedsSheppSpace.allocState();
    rsStart->setXY(state.x, state.y);
    rsStart->setYaw(-state.yaw);
    rsEnd->setXY(getGoal().x, getGoal().y);
    rsEnd->setYaw(-getGoal().yaw);
    ompl::base::ReedsSheppStateSpace::ReedsSheppPath reedsShepppath =
        reedsSheppSpace.reedsShepp(rsStart, rsEnd);

    std::vector<State> path;
    std::unordered_map<State, std::tuple<State, Action, double, double>,
                       std::hash<State>>
        cameFrom;
    cameFrom.clear();
    path.emplace_back(state);
    for (auto pathidx = 0; pathidx < 5; pathidx++) {
      if (fabs(reedsShepppath.length_[pathidx]) < 1e-6) continue;
      double deltat = 0, dx = 0, act = 0, cost = 0;
      switch (reedsShepppath.type_[pathidx]) {
        case 0:  // RS_NOP
          continue;
          break;
        case 1:  // RS_LEFT
          deltat = -reedsShepppath.length_[pathidx];
          dx = Constants::r * sin(-deltat);
          // dy = Constants::r * (1 - cos(-deltat));
          act = 2;
          cost = reedsShepppath.length_[pathidx] * Constants::r *
                 Constants::penaltyTurning;
          break;
        case 2:  // RS_STRAIGHT
          deltat = 0;
          dx = reedsShepppath.length_[pathidx] * Constants::r;
          // dy = 0;
          act = 0;
          cost = dx;
          break;
        case 3:  // RS_RIGHT
          deltat = reedsShepppath.length_[pathidx];
          dx = Constants::r * sin(deltat);
          // dy = -Constants::r * (1 - cos(deltat));
          act = 1;
          cost = reedsShepppath.length_[pathidx] * Constants::r *
                 Constants::penaltyTurning;
          break;
        default:
          std::cout << "\033[1m\033[31m"
                    << "Warning: Receive unknown ReedsSheppPath type"
                    << "\033[0m\n";
          break;
      }
      if (cost < 0) {
        cost = -cost * Constants::penaltyReversing;
        act = act + 3;
      }
      State s = path.back();
      std::vector<std::pair<State, double>> next_path;
      if (generatePath(s, act, deltat, dx, next_path)) {
        for (auto iter = next_path.begin(); iter != next_path.end(); iter++) {
          State next_s = iter->first;
          gscore += iter->second;
          if (!(next_s == path.back())) {
            cameFrom.insert(std::make_pair<>(
                next_s,
                std::make_tuple<>(path.back(), act, iter->second, gscore)));
          }
          path.emplace_back(next_s);
        }
      } else {
        return false;
      }
    }

    if (path.back().time <= m_lastGoalConstraint) {
      return false;
    }

    m_goals[m_agentIdx] = path.back();

    _camefrom.insert(cameFrom.begin(), cameFrom.end());
    return true;
  }

  void getNeighbors(const State &s, Action action,
                    std::vector<Neighbor<State, Action, double>> &neighbors) {
    neighbors.clear();
    double g = Constants::dx[0];
    for (Action act = 0; act < 6; act++) {  // has 6 directions for Reeds-Shepp
      double xSucc, ySucc, yawSucc;
      g = Constants::dx[0];
      xSucc = s.x + Constants::dx[act] * cos(-s.yaw) -
              Constants::dy[act] * sin(-s.yaw);
      ySucc = s.y + Constants::dx[act] * sin(-s.yaw) +
              Constants::dy[act] * cos(-s.yaw);
      yawSucc = Constants::normalizeHeadingRad(s.yaw + Constants::dyaw[act]);
      // if (act != action) {  // penalize turning
      //   g = g * Constants::penaltyTurning;
      //   if (act >= 3)  // penalize change of direction
      //     g = g * Constants::penaltyCOD;
      // }
      // if (act > 3) {  // backwards
      //   g = g * Constants::penaltyReversing;
      // }
      if (act % 3 != 0) {  // penalize turning
        g = g * Constants::penaltyTurning;
      }
      if ((act < 3 && action >= 3) || (action < 3 && act >= 3)) {
        // penalize change of direction
        g = g * Constants::penaltyCOD;
      }
      if (act >= 3) {  // backwards
        g = g * Constants::penaltyReversing;
      }
      State tempState(xSucc, ySucc, yawSucc, s.time + 1);
      if (stateValid(tempState)) {
        neighbors.emplace_back(
            Neighbor<State, Action, double>(tempState, act, g));
      }
    }
    // wait
    g = Constants::dx[0];
    State tempState(s.x, s.y, s.yaw, s.time + 1);
    if (stateValid(tempState)) {
      neighbors.emplace_back(Neighbor<State, Action, double>(tempState, 6, g));
    }
  }
  State getGoal() { return m_goals[m_agentIdx]; }
  uint64_t calcIndex(const State &s) {
    return (uint64_t)s.time * (2 * M_PI / Constants::deltat) *
               (m_dimx / Constants::xyResolution) *
               (m_dimy / Constants::xyResolution) +
           (uint64_t)(Constants::normalizeHeadingRad(s.yaw) /
                      Constants::yawResolution) *
               (m_dimx / Constants::xyResolution) *
               (m_dimy / Constants::xyResolution) +
           (uint64_t)(s.y / Constants::xyResolution) *
               (m_dimx / Constants::xyResolution) +
           (uint64_t)(s.x / Constants::xyResolution);
  }

  void onExpandLowLevelNode(const State & /*s*/, int /*fScore*/,
                            int /*gScore*/) {
    m_lowLevelExpanded++;
  }

  int lowLevelExpanded() const { return m_lowLevelExpanded; }

  bool startAndGoalValid(const std::vector<State> &m_starts, const size_t iter,
                         const int batchsize) {
    assert(m_goals.size() == m_starts.size());
    for (size_t i = 0; i < m_goals.size(); i++)
      for (size_t j = i + 1; j < m_goals.size(); j++) {
        if (m_goals[i].agentCollision(m_goals[j])) {
          std::cout << "ERROR: Goal point of " << i + iter * batchsize << " & "
                    << j + iter * batchsize << " collide!\n";
          return false;
        }
        if (m_starts[i].agentCollision(m_starts[j])) {
          std::cout << "ERROR: Start point of " << i + iter * batchsize << " & "
                    << j + iter * batchsize << " collide!\n";
          return false;
        }
      }
    return true;
  }

 private:
  State getState(size_t agentIdx,
                 const std::vector<PlanResult<State, Action, double>> &solution,
                 size_t t) {
    assert(agentIdx < solution.size());
    if (t < solution[agentIdx].states.size()) {
      return solution[agentIdx].states[t].first;
    }
    assert(!solution[agentIdx].states.empty());
    return solution[agentIdx].states.back().first;
  }

  bool stateValid(const State &s) {
    double x_ind = s.x / Constants::mapResolution;
    double y_ind = s.y / Constants::mapResolution;
    if (x_ind < 0 || x_ind >= m_dimx || y_ind < 0 || y_ind >= m_dimy)
      return false;

    for (auto it = m_obstacles.begin(); it != m_obstacles.end(); it++) {
      if (s.obsCollision(*it)) return false;
    }

    auto it = m_dynamic_obstacles.equal_range(s.time);
    for (auto itr = it.first; itr != it.second; ++itr) {
      if (s.agentCollision(itr->second)) return false;
    }
    auto itlow = m_dynamic_obstacles.lower_bound(-s.time);
    auto itup = m_dynamic_obstacles.upper_bound(-1);
    for (auto it = itlow; it != itup; ++it)
      if (s.agentCollision(it->second)) return false;

    for (auto it = m_constraints->constraints.begin();
         it != m_constraints->constraints.end(); it++) {
      if (!it->satisfyConstraint(s)) return false;
    }

    return true;
  }

 private:
  struct compare_node {
    bool operator()(const std::pair<State, double> &n1,
                    const std::pair<State, double> &n2) const {
      return (n1.second > n2.second);
    }
  };
  void updateCostmap() {
    boost::heap::fibonacci_heap<std::pair<State, double>,
                                boost::heap::compare<compare_node>>
        heap;

    std::set<std::pair<int, int>> temp_obs_set;
    for (auto it = m_obstacles.begin(); it != m_obstacles.end(); it++) {
      temp_obs_set.insert(
          std::make_pair((int)it->x / Constants::mapResolution,
                         (int)it->y / Constants::mapResolution));
    }

    for (size_t idx = 0; idx < m_goals.size(); idx++) {
      heap.clear();
      int goal_x = (int)m_goals[idx].x / Constants::mapResolution;
      int goal_y = (int)m_goals[idx].y / Constants::mapResolution;
      heap.push(std::make_pair(State(goal_x, goal_y, 0), 0));

      while (!heap.empty()) {
        std::pair<State, double> node = heap.top();
        heap.pop();

        int x = node.first.x;
        int y = node.first.y;
        for (int dx = -1; dx <= 1; dx++)
          for (int dy = -1; dy <= 1; dy++) {
            if (dx == 0 && dy == 0) continue;
            int new_x = x + dx;
            int new_y = y + dy;
            if (new_x == goal_x && new_y == goal_y) continue;
            if (new_x >= 0 && new_x < m_dimx && new_y >= 0 && new_y < m_dimy &&
                holonomic_cost_maps[idx][new_x][new_y] == 0 &&
                temp_obs_set.find(std::make_pair(new_x, new_y)) ==
                    temp_obs_set.end()) {
              holonomic_cost_maps[idx][new_x][new_y] =
                  holonomic_cost_maps[idx][x][y] +
                  sqrt(pow(dx * Constants::mapResolution, 2) +
                       pow(dy * Constants::mapResolution, 2));
              heap.push(std::make_pair(State(new_x, new_y, 0),
                                       holonomic_cost_maps[idx][new_x][new_y]));
            }
          }
      }
    }

    // for (size_t idx = 0; idx < m_goals.size(); idx++) {
    //   std::cout << "---------Cost Map -------Agent: " << idx
    //             << "------------\n";
    //   for (size_t i = 0; i < m_dimx; i++) {
    //     for (size_t j = 0; j < m_dimy; j++)
    //       std::cout << holonomic_cost_maps[idx][i][j] << "\t";
    //     std::cout << std::endl;
    //   }
    // }
  }

  bool generatePath(State startState, int act, double deltaSteer,
                    double deltaLength,
                    std::vector<std::pair<State, double>> &result) {
    double xSucc, ySucc, yawSucc, dx, dy, dyaw, ratio;
    result.emplace_back(std::make_pair<>(startState, 0));
    if (act == 0 || act == 3) {
      for (size_t i = 0; i < (size_t)(deltaLength / Constants::dx[act]); i++) {
        State s = result.back().first;
        xSucc = s.x + Constants::dx[act] * cos(-s.yaw) -
                Constants::dy[act] * sin(-s.yaw);
        ySucc = s.y + Constants::dx[act] * sin(-s.yaw) +
                Constants::dy[act] * cos(-s.yaw);
        yawSucc = Constants::normalizeHeadingRad(s.yaw + Constants::dyaw[act]);
        State nextState(xSucc, ySucc, yawSucc, result.back().first.time + 1);
        if (!stateValid(nextState)) return false;
        result.emplace_back(std::make_pair<>(nextState, Constants::dx[0]));
      }
      ratio = (deltaLength -
               (int)(deltaLength / Constants::dx[act]) * Constants::dx[act]) /
              Constants::dx[act];
      dyaw = 0;
      dx = ratio * Constants::dx[act];
      dy = 0;
    } else {
      for (size_t i = 0; i < (size_t)(deltaSteer / Constants::dyaw[act]); i++) {
        State s = result.back().first;
        xSucc = s.x + Constants::dx[act] * cos(-s.yaw) -
                Constants::dy[act] * sin(-s.yaw);
        ySucc = s.y + Constants::dx[act] * sin(-s.yaw) +
                Constants::dy[act] * cos(-s.yaw);
        yawSucc = Constants::normalizeHeadingRad(s.yaw + Constants::dyaw[act]);
        State nextState(xSucc, ySucc, yawSucc, result.back().first.time + 1);
        if (!stateValid(nextState)) return false;
        result.emplace_back(std::make_pair<>(
            nextState, Constants::dx[0] * Constants::penaltyTurning));
      }
      ratio = (deltaSteer - (int)(deltaSteer / Constants::dyaw[act]) *
                                Constants::dyaw[act]) /
              Constants::dyaw[act];
      dyaw = ratio * Constants::dyaw[act];
      dx = Constants::r * sin(dyaw);
      dy = -Constants::r * (1 - cos(dyaw));
      if (act == 2 || act == 5) {
        dx = -dx;
        dy = -dy;
      }
    }
    State s = result.back().first;
    xSucc = s.x + dx * cos(-s.yaw) - dy * sin(-s.yaw);
    ySucc = s.y + dx * sin(-s.yaw) + dy * cos(-s.yaw);
    yawSucc = Constants::normalizeHeadingRad(s.yaw + dyaw);
    // std::cout << m_agentIdx << " ratio::" << ratio << std::endl;
    State nextState(xSucc, ySucc, yawSucc, result.back().first.time + 1);
    if (!stateValid(nextState)) return false;
    result.emplace_back(std::make_pair<>(nextState, ratio * Constants::dx[0]));

    // std::cout << "Have generate " << result.size() << " path segments:\n\t";
    // for (auto iter = result.begin(); iter != result.end(); iter++)
    //   std::cout << iter->first << ":" << iter->second << "->";
    // std::cout << std::endl;

    return true;
  }

 private:
  int m_dimx;
  int m_dimy;
  std::vector<std::vector<std::vector<double>>> holonomic_cost_maps;
  std::unordered_set<Location> m_obstacles;
  std::multimap<int, State> m_dynamic_obstacles;
  std::vector<State> m_goals;
  // std::vector< std::vector<int> > m_heuristic;
  std::vector<double> m_vel_limit;
  size_t m_agentIdx;
  const Constraints *m_constraints;
  int m_lastGoalConstraint;
  int m_highLevelExpanded;
  int m_lowLevelExpanded;
};

}  // namespace libMultiRobotPlanning


// Copyright (c) 2019, The Monero And Italo Project
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification, are
// permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this list of
//    conditions and the following disclaimer.
//
// 2. 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.
//
// 3. Neither the name of the copyright holder nor the names of its contributors may be
//    used to endorse or promote products derived from this software without specific
//    prior written permission.
//
// 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 HOLDER 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 <boost/multiprecision/cpp_int.hpp>

void div128_64(uint64_t dividend_hi, uint64_t dividend_lo, uint64_t divisor, uint64_t* quotient_hi, uint64_t *quotient_lo, uint64_t *remainder_hi, uint64_t *remainder_lo)
{
  typedef boost::multiprecision::uint128_t uint128_t;

  uint128_t dividend = dividend_hi;
  dividend <<= 64;
  dividend |= dividend_lo;

  uint128_t q, r;
  divide_qr(dividend, uint128_t(divisor), q, r);

  *quotient_hi = ((q >> 64) & 0xffffffffffffffffull).convert_to<uint64_t>();
  *quotient_lo = (q & 0xffffffffffffffffull).convert_to<uint64_t>();
  if (remainder_hi)
    *remainder_hi = ((r >> 64) & 0xffffffffffffffffull).convert_to<uint64_t>();
  if (remainder_lo)
    *remainder_lo = (r & 0xffffffffffffffffull).convert_to<uint64_t>();
}


//  (C) Copyright Jeremy William Murphy 2016.

//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_INTEGER_COMMON_FACTOR_RT_HPP
#define BOOST_INTEGER_COMMON_FACTOR_RT_HPP

#include <boost/assert.hpp>
#include <boost/core/enable_if.hpp>

#include <boost/config.hpp>  // for BOOST_NESTED_TEMPLATE, etc.
#include <boost/limits.hpp>  // for std::numeric_limits
#include <climits>           // for CHAR_MIN
#include <boost/detail/workaround.hpp>
#include <iterator>
#include <algorithm>
#include <limits>
#ifndef BOOST_NO_CXX11_HDR_TYPE_TRAITS
#include <type_traits>
#endif
#ifdef BOOST_NO_CXX11_HDR_FUNCTIONAL
#include <functional>
#endif

#if ((defined(BOOST_MSVC) && (BOOST_MSVC >= 1600)) || (defined(__clang__) && defined(__c2__)) || (defined(BOOST_INTEL) && defined(_MSC_VER))) && (defined(_M_IX86) || defined(_M_X64))
#include <intrin.h>
#endif

#ifdef BOOST_MSVC
#pragma warning(push)
#pragma warning(disable:4127 4244)  // Conditional expression is constant
#endif

#if !defined(BOOST_NO_CXX11_HDR_TYPE_TRAITS) && !defined(BOOST_NO_CXX11_NOEXCEPT)
#define BOOST_GCD_NOEXCEPT(T) noexcept(std::is_arithmetic<T>::value)
#else
#define BOOST_GCD_NOEXCEPT(T)
#endif

namespace boost {

   template <class I>
   class rational;

   namespace integer {

      namespace gcd_detail{

         //
         // some helper functions which really should be constexpr already, but sadly aren't:
         //
#ifndef BOOST_NO_CXX14_CONSTEXPR
         template <class T>
         inline constexpr T constexpr_min(T const& a, T const& b) BOOST_GCD_NOEXCEPT(T)
         {
            return a < b ? a : b;
         }
         template <class T>
         inline constexpr auto constexpr_swap(T&a, T& b) BOOST_GCD_NOEXCEPT(T) -> decltype(a.swap(b))
         {
            return a.swap(b);
         }
         template <class T, class U>
         inline constexpr void constexpr_swap(T&a, U& b...) BOOST_GCD_NOEXCEPT(T)
         {
            T t(static_cast<T&&>(a));
            a = static_cast<T&&>(b);
            b = static_cast<T&&>(t);
         }
#else
         template <class T>
         inline T constexpr_min(T const& a, T const& b) BOOST_GCD_NOEXCEPT(T)
         {
            return a < b ? a : b;
         }
         template <class T>
         inline void constexpr_swap(T&a, T& b) BOOST_GCD_NOEXCEPT(T)
         {
            using std::swap;
            swap(a, b);
         }
#endif

      template <class T, bool a =
#ifndef BOOST_NO_CXX11_HDR_TYPE_TRAITS
         std::is_unsigned<T>::value ||
#endif
         (std::numeric_limits<T>::is_specialized && !std::numeric_limits<T>::is_signed)>
      struct gcd_traits_abs_defaults
      {
         inline static BOOST_CXX14_CONSTEXPR const T& abs(const T& val) BOOST_GCD_NOEXCEPT(T) { return val; }
      };
      template <class T>
      struct gcd_traits_abs_defaults<T, false>
      {
         inline static T BOOST_CXX14_CONSTEXPR abs(const T& val) BOOST_GCD_NOEXCEPT(T)
         {
            // This sucks, but std::abs is not constexpr :(
            return val < T(0) ? -val : val;
         }
      };

      enum method_type
      {
         method_euclid = 0,
         method_binary = 1,
         method_mixed = 2
      };

      struct any_convert
      {
         template <class T>
         any_convert(const T&);
      };

      struct unlikely_size
      {
         char buf[9973];
      };

      unlikely_size operator <<= (any_convert, any_convert);
      unlikely_size operator >>= (any_convert, any_convert);

      template <class T>
      struct gcd_traits_defaults : public gcd_traits_abs_defaults<T>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(T& val) BOOST_GCD_NOEXCEPT(T)
         {
            unsigned r = 0;
            while(0 == (val & 1u))
            {
#ifdef _MSC_VER  // VC++ can't handle operator >>= in constexpr code for some reason
               val = val >> 1;
#else
               val >>= 1;
#endif
               ++r;
            }
            return r;
         }
         inline static BOOST_CXX14_CONSTEXPR bool less(const T& a, const T& b) BOOST_GCD_NOEXCEPT(T)
         {
            return a < b;
         }

         static T& get_value();

#ifndef BOOST_NO_SFINAE
         static const bool has_operator_left_shift_equal = sizeof(get_value() <<= 2) != sizeof(unlikely_size);
         static const bool has_operator_right_shift_equal = sizeof(get_value() >>= 2) != sizeof(unlikely_size);
#else
         static const bool has_operator_left_shift_equal = true;
         static const bool has_operator_right_shift_equal = true;
#endif
         static const method_type method = std::numeric_limits<T>::is_specialized && std::numeric_limits<T>::is_integer && has_operator_left_shift_equal && has_operator_right_shift_equal ? method_mixed : method_euclid;
      };
      //
      // Default gcd_traits just inherits from defaults:
      //
      template <class T>
      struct gcd_traits : public gcd_traits_defaults<T> {};

      //
      // Some platforms have fast bitscan operations, that allow us to implement
      // make_odd much more efficiently, unfortunately we can't use these if we want
      // the functions to be constexpr as the compiler intrinsics aren't constexpr.
      //
#if defined(BOOST_NO_CXX14_CONSTEXPR) && ((defined(BOOST_MSVC) && (BOOST_MSVC >= 1600)) || (defined(__clang__) && defined(__c2__)) || (defined(BOOST_INTEL) && defined(_MSC_VER))) && (defined(_M_IX86) || defined(_M_X64))
#pragma intrinsic(_BitScanForward,)
      template <>
      struct gcd_traits<unsigned long> : public gcd_traits_defaults<unsigned long>
      {
         BOOST_FORCEINLINE static unsigned find_lsb(unsigned long val) BOOST_NOEXCEPT
         {
            unsigned long result;
            _BitScanForward(&result, val);
            return result;
         }
         BOOST_FORCEINLINE static unsigned make_odd(unsigned long& val) BOOST_NOEXCEPT
         {
            unsigned result = find_lsb(val);
            val >>= result;
            return result;
         }
      };

#ifdef _M_X64
#pragma intrinsic(_BitScanForward64)
      template <>
      struct gcd_traits<unsigned __int64> : public gcd_traits_defaults<unsigned __int64>
      {
         BOOST_FORCEINLINE static unsigned find_lsb(unsigned __int64 mask) BOOST_NOEXCEPT
         {
            unsigned long result;
            _BitScanForward64(&result, mask);
            return result;
         }
         BOOST_FORCEINLINE static unsigned make_odd(unsigned __int64& val) BOOST_NOEXCEPT
         {
            unsigned result = find_lsb(val);
            val >>= result;
            return result;
         }
      };
#endif
      //
      // Other integer type are trivial adaptations of the above,
      // this works for signed types too, as by the time these functions
      // are called, all values are > 0.
      //
      template <> struct gcd_traits<long> : public gcd_traits_defaults<long>
      { BOOST_FORCEINLINE static unsigned make_odd(long& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
      template <> struct gcd_traits<unsigned int> : public gcd_traits_defaults<unsigned int>
      { BOOST_FORCEINLINE static unsigned make_odd(unsigned int& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
      template <> struct gcd_traits<int> : public gcd_traits_defaults<int>
      { BOOST_FORCEINLINE static unsigned make_odd(int& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
      template <> struct gcd_traits<unsigned short> : public gcd_traits_defaults<unsigned short>
      { BOOST_FORCEINLINE static unsigned make_odd(unsigned short& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
      template <> struct gcd_traits<short> : public gcd_traits_defaults<short>
      { BOOST_FORCEINLINE static unsigned make_odd(short& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
      template <> struct gcd_traits<unsigned char> : public gcd_traits_defaults<unsigned char>
      { BOOST_FORCEINLINE static unsigned make_odd(unsigned char& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
      template <> struct gcd_traits<signed char> : public gcd_traits_defaults<signed char>
      { BOOST_FORCEINLINE static signed make_odd(signed char& val)BOOST_NOEXCEPT{ signed result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
      template <> struct gcd_traits<char> : public gcd_traits_defaults<char>
      { BOOST_FORCEINLINE static unsigned make_odd(char& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
#ifndef BOOST_NO_INTRINSIC_WCHAR_T
      template <> struct gcd_traits<wchar_t> : public gcd_traits_defaults<wchar_t>
      { BOOST_FORCEINLINE static unsigned make_odd(wchar_t& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; } };
#endif
#ifdef _M_X64
      template <> struct gcd_traits<__int64> : public gcd_traits_defaults<__int64>
      { BOOST_FORCEINLINE static unsigned make_odd(__int64& val)BOOST_NOEXCEPT{ unsigned result = gcd_traits<unsigned __int64>::find_lsb(val); val >>= result; return result; } };
#endif

#elif defined(BOOST_GCC) || defined(__clang__) || (defined(BOOST_INTEL) && defined(__GNUC__))

      template <>
      struct gcd_traits<unsigned> : public gcd_traits_defaults<unsigned>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned find_lsb(unsigned mask)BOOST_NOEXCEPT
         {
            return __builtin_ctz(mask);
         }
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(unsigned& val)BOOST_NOEXCEPT
         {
            unsigned result = find_lsb(val);
            val >>= result;
            return result;
         }
      };
      template <>
      struct gcd_traits<unsigned long> : public gcd_traits_defaults<unsigned long>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned find_lsb(unsigned long mask)BOOST_NOEXCEPT
         {
            return __builtin_ctzl(mask);
         }
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(unsigned long& val)BOOST_NOEXCEPT
         {
            unsigned result = find_lsb(val);
            val >>= result;
            return result;
         }
      };
      template <>
      struct gcd_traits<boost::ulong_long_type> : public gcd_traits_defaults<boost::ulong_long_type>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned find_lsb(boost::ulong_long_type mask)BOOST_NOEXCEPT
         {
            return __builtin_ctzll(mask);
         }
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(boost::ulong_long_type& val)BOOST_NOEXCEPT
         {
            unsigned result = find_lsb(val);
            val >>= result;
            return result;
         }
      };
      //
      // Other integer type are trivial adaptations of the above,
      // this works for signed types too, as by the time these functions
      // are called, all values are > 0.
      //
      template <> struct gcd_traits<boost::long_long_type> : public gcd_traits_defaults<boost::long_long_type>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(boost::long_long_type& val)BOOST_NOEXCEPT { unsigned result = gcd_traits<boost::ulong_long_type>::find_lsb(val); val >>= result; return result; }
      };
      template <> struct gcd_traits<long> : public gcd_traits_defaults<long>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(long& val)BOOST_NOEXCEPT { unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; }
      };
      template <> struct gcd_traits<int> : public gcd_traits_defaults<int>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(int& val)BOOST_NOEXCEPT { unsigned result = gcd_traits<unsigned long>::find_lsb(val); val >>= result; return result; }
      };
      template <> struct gcd_traits<unsigned short> : public gcd_traits_defaults<unsigned short>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(unsigned short& val)BOOST_NOEXCEPT { unsigned result = gcd_traits<unsigned>::find_lsb(val); val >>= result; return result; }
      };
      template <> struct gcd_traits<short> : public gcd_traits_defaults<short>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(short& val)BOOST_NOEXCEPT { unsigned result = gcd_traits<unsigned>::find_lsb(val); val >>= result; return result; }
      };
      template <> struct gcd_traits<unsigned char> : public gcd_traits_defaults<unsigned char>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(unsigned char& val)BOOST_NOEXCEPT { unsigned result = gcd_traits<unsigned>::find_lsb(val); val >>= result; return result; }
      };
      template <> struct gcd_traits<signed char> : public gcd_traits_defaults<signed char>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR signed make_odd(signed char& val)BOOST_NOEXCEPT { signed result = gcd_traits<unsigned>::find_lsb(val); val >>= result; return result; }
      };
      template <> struct gcd_traits<char> : public gcd_traits_defaults<char>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(char& val)BOOST_NOEXCEPT { unsigned result = gcd_traits<unsigned>::find_lsb(val); val >>= result; return result; }
      };
#ifndef BOOST_NO_INTRINSIC_WCHAR_T
      template <> struct gcd_traits<wchar_t> : public gcd_traits_defaults<wchar_t>
      {
         BOOST_FORCEINLINE static BOOST_CXX14_CONSTEXPR unsigned make_odd(wchar_t& val)BOOST_NOEXCEPT { unsigned result = gcd_traits<unsigned>::find_lsb(val); val >>= result; return result; }
      };
#endif
#endif
   //
   // The Mixed Binary Euclid Algorithm
   // Sidi Mohamed Sedjelmaci
   // Electronic Notes in Discrete Mathematics 35 (2009) 169-176
   //
   template <class T>
   BOOST_CXX14_CONSTEXPR T mixed_binary_gcd(T u, T v) BOOST_GCD_NOEXCEPT(T)
   {
      if(gcd_traits<T>::less(u, v))
         constexpr_swap(u, v);

      unsigned shifts = 0;

      if(u == T(0))
         return v;
      if(v == T(0))
         return u;

      shifts = constexpr_min(gcd_traits<T>::make_odd(u), gcd_traits<T>::make_odd(v));

      while(gcd_traits<T>::less(1, v))
      {
         u %= v;
         v -= u;
         if(u == T(0))
            return v << shifts;
         if(v == T(0))
            return u << shifts;
         gcd_traits<T>::make_odd(u);
         gcd_traits<T>::make_odd(v);
         if(gcd_traits<T>::less(u, v))
            constexpr_swap(u, v);
      }
      return (v == 1 ? v : u) << shifts;
   }

    /** Stein gcd (aka 'binary gcd')
     *
     * From Mathematics to Generic Programming, Alexander Stepanov, Daniel Rose
     */
    template <typename SteinDomain>
    BOOST_CXX14_CONSTEXPR SteinDomain Stein_gcd(SteinDomain m, SteinDomain n) BOOST_GCD_NOEXCEPT(SteinDomain)
    {
        BOOST_ASSERT(m >= 0);
        BOOST_ASSERT(n >= 0);
        if (m == SteinDomain(0))
            return n;
        if (n == SteinDomain(0))
            return m;
        // m > 0 && n > 0
        int d_m = gcd_traits<SteinDomain>::make_odd(m);
        int d_n = gcd_traits<SteinDomain>::make_odd(n);
        // odd(m) && odd(n)
        while (m != n)
        {
            if (n > m)
               constexpr_swap(n, m);
            m -= n;
            gcd_traits<SteinDomain>::make_odd(m);
        }
        // m == n
        m <<= constexpr_min(d_m, d_n);
        return m;
    }

    /** Euclidean algorithm
     *
     * From Mathematics to Generic Programming, Alexander Stepanov, Daniel Rose
     *
     */
    template <typename EuclideanDomain>
    inline BOOST_CXX14_CONSTEXPR EuclideanDomain Euclid_gcd(EuclideanDomain a, EuclideanDomain b) BOOST_GCD_NOEXCEPT(EuclideanDomain)
    {
        while (b != EuclideanDomain(0))
        {
            a %= b;
            constexpr_swap(a, b);
        }
        return a;
    }

    template <typename T>
    inline BOOST_CXX14_CONSTEXPR BOOST_DEDUCED_TYPENAME enable_if_c<gcd_traits<T>::method == method_mixed, T>::type
       optimal_gcd_select(T const &a, T const &b) BOOST_GCD_NOEXCEPT(T)
    {
       return gcd_detail::mixed_binary_gcd(a, b);
    }

    template <typename T>
    inline BOOST_CXX14_CONSTEXPR BOOST_DEDUCED_TYPENAME enable_if_c<gcd_traits<T>::method == method_binary, T>::type
       optimal_gcd_select(T const &a, T const &b) BOOST_GCD_NOEXCEPT(T)
    {
       return gcd_detail::Stein_gcd(a, b);
    }

    template <typename T>
    inline BOOST_CXX14_CONSTEXPR BOOST_DEDUCED_TYPENAME enable_if_c<gcd_traits<T>::method == method_euclid, T>::type
       optimal_gcd_select(T const &a, T const &b) BOOST_GCD_NOEXCEPT(T)
    {
       return gcd_detail::Euclid_gcd(a, b);
    }

    template <class T>
    inline BOOST_CXX14_CONSTEXPR T lcm_imp(const T& a, const T& b) BOOST_GCD_NOEXCEPT(T)
    {
       T temp = boost::integer::gcd_detail::optimal_gcd_select(a, b);
#if BOOST_WORKAROUND(BOOST_GCC_VERSION, < 40500)
       return (temp != T(0)) ? T(a / temp * b) : T(0);
#else
       return temp != T(0) ? T(a / temp * b) : T(0);
#endif
    }

} // namespace detail

template <typename Integer>
inline BOOST_CXX14_CONSTEXPR Integer gcd(Integer const &a, Integer const &b) BOOST_GCD_NOEXCEPT(Integer)
{
    if(a == (std::numeric_limits<Integer>::min)())
       return a == static_cast<Integer>(0) ? gcd_detail::gcd_traits<Integer>::abs(b) : boost::integer::gcd(static_cast<Integer>(a % b), b);
    else if (b == (std::numeric_limits<Integer>::min)())
       return b == static_cast<Integer>(0) ? gcd_detail::gcd_traits<Integer>::abs(a) : boost::integer::gcd(a, static_cast<Integer>(b % a));
    return gcd_detail::optimal_gcd_select(static_cast<Integer>(gcd_detail::gcd_traits<Integer>::abs(a)), static_cast<Integer>(gcd_detail::gcd_traits<Integer>::abs(b)));
}

template <typename Integer>
inline BOOST_CXX14_CONSTEXPR Integer lcm(Integer const &a, Integer const &b) BOOST_GCD_NOEXCEPT(Integer)
{
   return gcd_detail::lcm_imp(static_cast<Integer>(gcd_detail::gcd_traits<Integer>::abs(a)), static_cast<Integer>(gcd_detail::gcd_traits<Integer>::abs(b)));
}
#ifndef BOOST_NO_CXX11_VARIADIC_TEMPLATES
//
// This looks slightly odd, but the variadic forms must have 3 or more arguments, and the variadic argument pack may be empty.
// This matters not at all for most compilers, but Oracle C++ selects the wrong overload in the 2-arg case unless we do this.
//
template <typename Integer, typename... Args>
inline BOOST_CXX14_CONSTEXPR Integer gcd(Integer const &a, Integer const &b, const Integer& c, Args const&... args) BOOST_GCD_NOEXCEPT(Integer)
{
   Integer t = gcd(b, c, args...);
   return t == 1 ? 1 : gcd(a, t);
}

template <typename Integer, typename... Args>
inline BOOST_CXX14_CONSTEXPR Integer lcm(Integer const &a, Integer const &b, Integer const& c, Args const&... args) BOOST_GCD_NOEXCEPT(Integer)
{
   return lcm(a, lcm(b, c, args...));
}
#endif
//
// Special handling for rationals:
//
template <typename Integer>
inline typename boost::enable_if_c<std::numeric_limits<Integer>::is_specialized, boost::rational<Integer> >::type gcd(boost::rational<Integer> const &a, boost::rational<Integer> const &b)
{
   return boost::rational<Integer>(static_cast<Integer>(gcd(a.numerator(), b.numerator())), static_cast<Integer>(lcm(a.denominator(), b.denominator())));
}

template <typename Integer>
inline typename boost::enable_if_c<std::numeric_limits<Integer>::is_specialized, boost::rational<Integer> >::type lcm(boost::rational<Integer> const &a, boost::rational<Integer> const &b)
{
   return boost::rational<Integer>(static_cast<Integer>(lcm(a.numerator(), b.numerator())), static_cast<Integer>(gcd(a.denominator(), b.denominator())));
}
/**
 * Knuth, The Art of Computer Programming: Volume 2, Third edition, 1998
 * Chapter 4.5.2, Algorithm C: Greatest common divisor of n integers.
 *
 * Knuth counts down from n to zero but we naturally go from first to last.
 * We also return the termination position because it might be useful to know.
 *
 * Partly by quirk, partly by design, this algorithm is defined for n = 1,
 * because the gcd of {x} is x. It is not defined for n = 0.
 *
 * @tparam  I   Input iterator.
 * @return  The gcd of the range and the iterator position at termination.
 */
template <typename I>
std::pair<typename std::iterator_traits<I>::value_type, I>
gcd_range(I first, I last) BOOST_GCD_NOEXCEPT(I)
{
    BOOST_ASSERT(first != last);
    typedef typename std::iterator_traits<I>::value_type T;

    T d = *first++;
    while (d != T(1) && first != last)
    {
        d = gcd(d, *first);
        first++;
    }
    return std::make_pair(d, first);
}
template <typename I>
std::pair<typename std::iterator_traits<I>::value_type, I>
lcm_range(I first, I last) BOOST_GCD_NOEXCEPT(I)
{
    BOOST_ASSERT(first != last);
    typedef typename std::iterator_traits<I>::value_type T;

    T d = *first++;
    while (d != T(1) && first != last)
    {
        d = lcm(d, *first);
        first++;
    }
    return std::make_pair(d, first);
}

template < typename IntegerType >
class gcd_evaluator
#ifdef BOOST_NO_CXX11_HDR_FUNCTIONAL
   : public std::binary_function<IntegerType, IntegerType, IntegerType>
#endif
{
public:
#ifndef BOOST_NO_CXX11_HDR_FUNCTIONAL
   typedef IntegerType first_argument_type;
   typedef IntegerType second_argument_type;
   typedef IntegerType result_type;
#endif
   IntegerType operator()(IntegerType const &a, IntegerType const &b)const
   {
      return boost::integer::gcd(a, b);
   }
};

template < typename IntegerType >
class lcm_evaluator
#ifdef BOOST_NO_CXX11_HDR_FUNCTIONAL
   : public std::binary_function<IntegerType, IntegerType, IntegerType>
#endif
{
public:
#ifndef BOOST_NO_CXX11_HDR_FUNCTIONAL
   typedef IntegerType first_argument_type;
   typedef IntegerType second_argument_type;
   typedef IntegerType result_type;
#endif
   IntegerType operator()(IntegerType const &a, IntegerType const &b)const
   {
      return boost::integer::lcm(a, b);
   }
};

}  // namespace integer
}  // namespace boost

#ifdef BOOST_MSVC
#pragma warning(pop)
#endif

#endif  // BOOST_INTEGER_COMMON_FACTOR_RT_HPP


/*    Copyright (c) 2010-2019, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 */

#include <boost/make_shared.hpp>
#include <boost/shared_ptr.hpp>

#include <Eigen/Core>

#include <tudat/astro/basic_astro/physicalConstants.h>
#include <tudat/basics/testMacros.h>
#include <tudat/math/basic/mathematicalConstants.h>
#include "tudat/astro/basic_astro/unitConversions.h"
#include <tudat/astro/basic_astro/orbitalElementConversions.h>

#include <tudat/io/basicInputOutput.h>
#include <tudat/io/applicationOutput.h>

#include "tudat/simulation/propagation_setup/propagationPatchedConicFullProblem.h"
#include "tudat/simulation/propagation_setup/propagationLambertTargeterFullProblem.h"

#include "tudat/astro/ephemerides/approximatePlanetPositions.h"
#include "tudat/astro/trajectory_design/trajectory.h"
#include "tudat/astro/trajectory_design/exportTrajectory.h"
#include "tudat/astro/trajectory_design/planetTrajectory.h"
#include <tudat/simulation/simulation.h>

int main( )
{
    using namespace tudat;
    using namespace tudat::input_output;
    using namespace tudat::input_output::parsed_data_vector_utilities;
    using namespace tudat::transfer_trajectories;

    ///  Characteristics of the interplanetary trajectory

    // Define central body of the trajectory and body to be propagated.
    std::vector< std::string > centralBody;
    centralBody.push_back( "Sun" );
    std::string bodyToPropagate = "spacecraft";

    // Specify the number and types of legs and type of legs.
    int numberOfLegs = 5;
    std::vector< TransferLegType > legTypeVector( numberOfLegs );
    legTypeVector[ 0 ] = mga1DsmVelocity_Departure;
    legTypeVector[ 1 ] = mga1DsmVelocity_Swingby;
    legTypeVector[ 2 ] = mga1DsmVelocity_Swingby;
    legTypeVector[ 3 ] = mga1DsmVelocity_Swingby;
    legTypeVector[ 4 ] = capture;

    // Name of the bodies involved in the trajectory
    std::vector< std::string > transferBodyTrajectory;
    transferBodyTrajectory.push_back("Earth");
    transferBodyTrajectory.push_back("Earth");
    transferBodyTrajectory.push_back("Venus");
    transferBodyTrajectory.push_back("Venus");
    transferBodyTrajectory.push_back("Mercury");

    // Create variable vector.
    std::vector< double > variableVector;

    // Add the time of flight and start epoch.
    variableVector.push_back( 1171.64503236 * physical_constants::JULIAN_DAY);
    variableVector.push_back( 399.999999715 * physical_constants::JULIAN_DAY);
    variableVector.push_back( 178.372255301 * physical_constants::JULIAN_DAY);
    variableVector.push_back( 299.223139512 * physical_constants::JULIAN_DAY);
    variableVector.push_back( 180.510754824 * physical_constants::JULIAN_DAY);
    variableVector.push_back( 1.0 ); // The capture time is irrelevant for the final leg.

    // Add the additional variables.
    // 1st leg.
    variableVector.push_back( 0.234594654679 );
    variableVector.push_back( 1408.99421278 );
    variableVector.push_back( 0.37992647165 * 2.0 * 3.14159265358979 );
    variableVector.push_back( std::acos(  2.0 * 0.498004040298 - 1. ) - 3.14159265358979 / 2.0 );
    // 2nd leg.
    variableVector.push_back( 0.0964769387134 );
    variableVector.push_back( 1.35077257078 );
    variableVector.push_back( 1.80629232251 * 6.378e6 );
    variableVector.push_back( 0.0 );
    // 3rd leg.
    variableVector.push_back( 0.829948744508);
    variableVector.push_back( 1.09554368115 );
    variableVector.push_back( 3.04129845698 * 6.052e6 );
    variableVector.push_back( 0.0 );
    // 4th leg.
    variableVector.push_back( 0.317174785637 );
    variableVector.push_back( 1.34317576594 );
    variableVector.push_back( 1.10000000891 * 6.052e6 );
    variableVector.push_back( 0.0 );

    // Create minimum pericenter radii vector
    std::vector< double > minimumPericenterRadii;
    minimumPericenterRadii.push_back( TUDAT_NAN ); minimumPericenterRadii.push_back( TUDAT_NAN ); minimumPericenterRadii.push_back( TUDAT_NAN );
    minimumPericenterRadii.push_back( TUDAT_NAN ); minimumPericenterRadii.push_back( TUDAT_NAN );

    // Create departure and capture variables.
    std::vector< double > semiMajorAxes;
    semiMajorAxes.push_back( std::numeric_limits< double >::infinity( ) ); semiMajorAxes.push_back( std::numeric_limits< double >::infinity( ) );
    std::vector< double > eccentricities;
    eccentricities.push_back( 0.0 ); eccentricities.push_back( 0.0 );

    // Define integrator settings.
    double initialTime = 0.0;
    double fixedStepSize = 1000.0;
    std::shared_ptr< numerical_integrators::IntegratorSettings< double > > integratorSettings =
            std::make_shared < numerical_integrators::IntegratorSettings < > > (
                numerical_integrators::rungeKutta4, initialTime, fixedStepSize);

    /// Ideal case

    // Create system of bodies.
    std::vector< double > gravitationalParametersTransferBodies;
    gravitationalParametersTransferBodies.push_back( 3.9860119e14 );
    gravitationalParametersTransferBodies.push_back( 3.9860119e14 );
    gravitationalParametersTransferBodies.push_back( 3.24860e14 );
    gravitationalParametersTransferBodies.push_back( 3.24860e14 );
    gravitationalParametersTransferBodies.push_back( 2.2321e13 );

    std::vector< ephemerides::EphemerisPointer > ephemerisVectorTransferBodies;
    ephemerisVectorTransferBodies.push_back( std::make_shared< ephemerides::ApproximatePlanetPositions>(
                ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData::earthMoonBarycenter) ) ;
    ephemerisVectorTransferBodies.push_back( std::make_shared< ephemerides::ApproximatePlanetPositions>(
                ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData::earthMoonBarycenter) ) ;
    ephemerisVectorTransferBodies.push_back( std::make_shared< ephemerides::ApproximatePlanetPositions>(
                ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData::venus) );
    ephemerisVectorTransferBodies.push_back( std::make_shared< ephemerides::ApproximatePlanetPositions>(
                ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData::venus) );
    ephemerisVectorTransferBodies.push_back( std::make_shared< ephemerides::ApproximatePlanetPositions>(
                ephemerides::ApproximatePlanetPositionsBase::BodiesWithEphemerisData::mercury) );

    simulation_setup::SystemOfBodies bodies = propagators::setupBodyMapFromUserDefinedEphemeridesForPatchedConicsTrajectory(
            centralBody[0], bodyToPropagate, transferBodyTrajectory,
            ephemerisVectorTransferBodies, gravitationalParametersTransferBodies, "ECLIPJ2000" );

    // Create acceleration map.
    std::vector< basic_astrodynamics::AccelerationMap > accelerationMap = propagators::setupAccelerationMapPatchedConicsTrajectory(
                transferBodyTrajectory.size(), centralBody[0], bodyToPropagate, bodies);

    // Calculate the patched conics solution and the propagation results of the associated full dynamics problem for each leg.
    std::map< int, std::map< double, Eigen::Vector6d > > patchedConicsTrajectory;
    std::map< int, std::map< double, Eigen::Vector6d > > fullProblemTrajectory;
    std::map< int, std::map< double, Eigen::VectorXd > > dependentVariablesTrajectory;

    propagators::fullPropagationPatchedConicsTrajectory(
                bodies, accelerationMap, transferBodyTrajectory, centralBody[0], bodyToPropagate,
            legTypeVector, variableVector, minimumPericenterRadii, semiMajorAxes, eccentricities, integratorSettings,
            patchedConicsTrajectory, fullProblemTrajectory, dependentVariablesTrajectory, false );

    for( auto itr : patchedConicsTrajectory )
    {
        std::cout << "Leg " << itr.first << "\n\n";
        std::cout << "Departure: " << fullProblemTrajectory[ itr.first ].begin( )->second -
                  patchedConicsTrajectory[ itr.first ].begin( )->second<< "\n\n";
        std::cout << "Arrival: " << fullProblemTrajectory[ itr.first ].rbegin( )->second -
                     patchedConicsTrajectory[ itr.first ].rbegin( )->second << "\n\n";
    }

    /// Perturbed case

    // Define accelerations acting on the spacecraft.
    std::map< std::string, std::vector< std::shared_ptr< simulation_setup::AccelerationSettings > > > bodyToPropagateAccelerations;
    bodyToPropagateAccelerations[ "Sun" ].push_back(std::make_shared< simulation_setup::AccelerationSettings >(
                                                                basic_astrodynamics::central_gravity ) );
    bodyToPropagateAccelerations[ "Earth" ].push_back(std::make_shared< simulation_setup::AccelerationSettings >(
                                                                basic_astrodynamics::central_gravity ) );
    bodyToPropagateAccelerations[ "Venus" ].push_back(std::make_shared< simulation_setup::AccelerationSettings >(
                                                                basic_astrodynamics::central_gravity ) );
    bodyToPropagateAccelerations[ "Mercury" ].push_back(std::make_shared< simulation_setup::AccelerationSettings >(
                                                                basic_astrodynamics::central_gravity ) );
    simulation_setup::SelectedAccelerationMap accelerationMapPerturbedCase;
    accelerationMapPerturbedCase[ "spacecraft" ] = bodyToPropagateAccelerations;

    // Create the acceleration map.
    basic_astrodynamics::AccelerationMap accelerationModelMapPerturbedCase = simulation_setup::createAccelerationModelsMap(
                bodies, accelerationMapPerturbedCase, {bodyToPropagate}, {centralBody} );
    std::vector<  basic_astrodynamics::AccelerationMap > accelerationMapVectorPerturbedCase;
    for (int i = 0 ; i < numberOfLegs ; i++){
        accelerationMapVectorPerturbedCase.push_back( accelerationModelMapPerturbedCase );
    }

    // Calculate the patched conics trajectory and propagate the full dynamics problem jointly.
    std::map< int, std::map< double, Eigen::Vector6d > > patchedConicsTrajectoryPerturbedCase;
    std::map< int, std::map< double, Eigen::Vector6d > > fullProblemTrajectoryPerturbedCase;
    std::map< int, std::map< double, Eigen::VectorXd > > dependentVariablesPerturbedCase;

    propagators::fullPropagationPatchedConicsTrajectory( bodies, accelerationMapVectorPerturbedCase,
            transferBodyTrajectory, centralBody[0], bodyToPropagate, legTypeVector, variableVector, minimumPericenterRadii,
            semiMajorAxes, eccentricities, integratorSettings, patchedConicsTrajectoryPerturbedCase,
            fullProblemTrajectoryPerturbedCase, dependentVariablesPerturbedCase, true );

    for( auto itr : patchedConicsTrajectoryPerturbedCase )
    {
        std::cout << "Leg " << itr.first << "\n\n";
        std::cout << "Departure: " << fullProblemTrajectoryPerturbedCase[ itr.first ].begin( )->second -
                  patchedConicsTrajectoryPerturbedCase[ itr.first ].begin( )->second<< "\n\n";
        std::cout << "Arrival: " << fullProblemTrajectoryPerturbedCase[ itr.first ].rbegin( )->second -
                     patchedConicsTrajectoryPerturbedCase[ itr.first ].rbegin( )->second << "\n\n";
    }

    /// Outputs

    for( std::map< int, std::map< double, Eigen::Vector6d > >::iterator itr = patchedConicsTrajectory.begin( );
         itr != patchedConicsTrajectory.end( ); itr++ ){

        input_output::writeDataMapToTextFile( fullProblemTrajectory[itr->first],
                                              "fullProblemInterplanetaryTrajectory_0_leg_" + std::to_string(itr->first) + ".dat",
                                              tudat_applications::getOutputPath( ),
                                              "",
                                              std::numeric_limits< double >::digits10,
                                              std::numeric_limits< double >::digits10,
                                              "," );

        input_output::writeDataMapToTextFile( fullProblemTrajectoryPerturbedCase[itr->first],
                                              "fullProblemInterplanetaryTrajectory_1_leg_" + std::to_string(itr->first) + ".dat",
                                              tudat_applications::getOutputPath( ),
                                              "",
                                              std::numeric_limits< double >::digits10,
                                              std::numeric_limits< double >::digits10,
                                              "," );

        input_output::writeDataMapToTextFile( patchedConicsTrajectory[itr->first],
                                              "patchedConicsInterplanetaryTrajectory_0_leg_" + std::to_string(itr->first) + ".dat",
                                              tudat_applications::getOutputPath( ),
                                              "",
                                              std::numeric_limits< double >::digits10,
                                              std::numeric_limits< double >::digits10,
                                              "," );

        input_output::writeDataMapToTextFile( patchedConicsTrajectoryPerturbedCase[itr->first],
                                              "patchedConicsInterplanetaryTrajectory_1_leg_" + std::to_string(itr->first) + ".dat",
                                              tudat_applications::getOutputPath( ),
                                              "",
                                              std::numeric_limits< double >::digits10,
                                              std::numeric_limits< double >::digits10,
                                              "," );

    }

    // Final statement.
    // The exit code EXIT_SUCCESS indicates that the program was successfully executed.
    return EXIT_SUCCESS;
}


#include <iostream>
#include <cstdlib>
#include <cmath>
#include <boost/pending/disjoint_sets.hpp>
#include <vector>
#include <queue>
#include <map>
#include <cstdio>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <fstream>

using namespace std;

template <class T>
class AffinityGraphCompare{
    private:
        const T * mEdgeWeightArray;
    public:
        AffinityGraphCompare(const T * EdgeWeightArray){
            mEdgeWeightArray = EdgeWeightArray;
        }
        bool operator() (const int& ind1, const int& ind2) const {
            return (mEdgeWeightArray[ind1] > mEdgeWeightArray[ind2]);
        }
};

/*
 * Compute the MALIS loss function and its derivative wrt the affinity graph
 * MAXIMUM spanning tree
 * Author: Srini Turaga (sturaga@mit.edu)
 * All rights reserved
 */
void preCompute(const uint64_t* conn_dims, const int32_t* nhood_data, const uint64_t* nhood_dims,
        uint64_t* pre_ve, uint64_t* pre_prodDims, int32_t* pre_nHood){

  uint64_t conn_num_dims = nhood_dims[1]+1; //4
  // nVert stores (x * y * z)
  uint64_t nVert = 1;
  for (int64_t i = 1; i < conn_num_dims; ++i) {
      nVert *= conn_dims[i];
  }
  pre_ve[0] = nVert;
  /* get number of edges */
  uint64_t nEdge = 0;
  // Loop over #edges
  for (int d = 0, i = 0; d < conn_dims[0]; ++d) {
    // Loop over Z
    for (int z = 0; z < conn_dims[1]; ++z) {
      // Loop over Y
      for (int y = 0; y < conn_dims[2]; ++y) {
        // Loop over X
        for (int x = 0; x < conn_dims[3]; ++x, ++i) {
          // Out-of-bounds check:
          if (!((z + nhood_data[d * nhood_dims[1] + 0] < 0)
              ||(z + nhood_data[d * nhood_dims[1] + 0] >= conn_dims[1])
              ||(y + nhood_data[d * nhood_dims[1] + 1] < 0)
              ||(y + nhood_data[d * nhood_dims[1] + 1] >= conn_dims[2])
              ||(x + nhood_data[d * nhood_dims[1] + 2] < 0)
              ||(x + nhood_data[d * nhood_dims[1] + 2] >= conn_dims[3]))) {
              ++nEdge;
          }
        }
      }
    }
  }
  pre_ve[1] = nEdge;

  // prodDims stores x, x*y, x*y*z offsets
  pre_prodDims[conn_num_dims - 2] = 1;
  for (int64_t i = 1; i < conn_num_dims - 1; ++i) {
      pre_prodDims[conn_num_dims - 2 - i] = pre_prodDims[conn_num_dims - 1 - i]
                                      * conn_dims[conn_num_dims - i];
  }

  /* convert n-d offset vectors into linear array offset scalars */
  // nHood is a vector of size #edges
    for (int64_t i = 0; i < nhood_dims[0]; ++i) {
        pre_nHood[i] = 0;
        for (int64_t j = 0; j < nhood_dims[1]; ++j) {
            pre_nHood[i] += nhood_data[j + i * nhood_dims[1]] * pre_prodDims[j];
        }
    }
}

const std::vector<int64_t> & initPqueue(const int64_t nEdge,
               const uint64_t* conn_dims, const int32_t* nhood_data, const uint64_t* nhood_dims ){
    static std::vector<int64_t> pqueue;
    if(pqueue.empty()){
        // initialization
        pqueue.resize(nEdge);
        int j = 0;
          // Loop over #edges
          for (int d = 0, i = 0; d < conn_dims[0]; ++d) {
            // Loop over Z
            for (int z = 0; z < conn_dims[1]; ++z) {
              // Loop over Y
              for (int y = 0; y < conn_dims[2]; ++y) {
                // Loop over X
                for (int x = 0; x < conn_dims[3]; ++x, ++i) {
                  // Out-of-bounds check:
                  if (!((z + nhood_data[d * nhood_dims[1] + 0] < 0)
                      ||(z + nhood_data[d * nhood_dims[1] + 0] >= conn_dims[1])
                      ||(y + nhood_data[d * nhood_dims[1] + 1] < 0)
                      ||(y + nhood_data[d * nhood_dims[1] + 1] >= conn_dims[2])
                      ||(x + nhood_data[d * nhood_dims[1] + 2] < 0)
                      ||(x + nhood_data[d * nhood_dims[1] + 2] >= conn_dims[3]))) {
                        pqueue[j++] = i;
                  }
                }
              }
            }
          }
        }
    return pqueue;
}

void malis_loss_weights_cpp_both(const uint64_t* seg,
               const uint64_t* conn_dims, const int32_t* nhood_data, const uint64_t* nhood_dims,
               const float* edgeWeight,
               float* nPairPerEdge,
               const uint64_t* pre_ve, const uint64_t* pre_prodDims, const int32_t* pre_nHood,
               const int pos, const float weight_opt){

    /* Disjoint sets and sparse overlap vectors */
    const uint64_t nVert = pre_ve[0];
    const uint64_t nEdge = pre_ve[1];
    uint64_t conn_num_dims = nhood_dims[1]+1;

    vector<map<uint64_t,uint64_t> > overlap(nVert);
    vector<uint64_t> rank(nVert);
    vector<uint64_t> parent(nVert);
    boost::disjoint_sets<uint64_t*, uint64_t*> dsets(&rank[0],&parent[0]);
    std::map<int64_t, int64_t> segSizes;
    uint64_t nLabeledVert = 0;
    uint64_t nPairPos = 0;
    // int mm=0;
    for (int i=0; i<nVert; ++i){
        dsets.make_set(i);
        if (0!=seg[i]) {
            overlap[i].insert(pair<uint64_t,uint64_t>(seg[i],1));
            ++nLabeledVert;
            ++segSizes[seg[i]];
            nPairPos += (segSizes[seg[i]] - 1); //add pair with previous seg label
            //if(seg[i]>mm){mm=seg[i];}
        }
    }

    // for (int i=0; i<=mm; ++i){std::cout<<segSizes[i]<<",";}std::cout<<std::endl;
    float nPairPosNormInv;
    if (pos==1){
        nPairPosNormInv = 1.0/(float)nPairPos;
    }else{
        nPairPosNormInv = 1.0/(nLabeledVert*(nLabeledVert-1)*0.5-nPairPos);
    }
    //std::cout<<nLabeledVert<<"hh: "<<nLabeledVert*(nLabeledVert-1)*0.5<<","<<nPairPos<<std::endl;
    if(weight_opt<=1){
        if (pos==1){
            nPairPosNormInv *= weight_opt;
        }else{
            nPairPosNormInv *= (1-weight_opt);
        }
    }
    //std::cout<<"norm:"<<nPairPosNormInv<<"\n";

    /* 2. Sort all the edges in increasing order of weight */
    const std::vector<int64_t> Pqueue1 = initPqueue(nEdge, conn_dims, nhood_data, nhood_dims);
    static std::vector<int64_t> Pqueue2(nEdge);
    std::copy(Pqueue1.begin(), Pqueue1.end(), Pqueue2.begin());

    sort( Pqueue2.begin(), Pqueue2.end(), AffinityGraphCompare<float>( edgeWeight ) );

    /* Start MST */
    uint64_t minEdge, e, v1, v2;
    uint64_t set1, set2;
    float nPair = 0;
    map<uint64_t,uint64_t>::iterator it1, it2;

    //std::ofstream db;
    //if(pos==1){ db.open("/n/coxfs01/donglai/malis_trans/db/mst-pt-pos.txt");}
    float ll=0;
    /* Start Kruskal's */
    for (uint64_t i = 0; i < Pqueue2.size(); ++i ) {
        minEdge = Pqueue2[i];
        // std::cout <<"do:"<<i<<","<<minEdge<< "\n";
        e =  minEdge / nVert;
        // v1: node at edge beginning
        v1 = minEdge % nVert;
        // v2: neighborhood node at edge e
        v2 = v1 + pre_nHood[e];
        set1 = dsets.find_set(v1);
        set2 = dsets.find_set(v2);
        if (set1!=set2){
            dsets.link(set1, set2);
            /* compute the number of pairs merged by this MST edge */
            for (it1 = overlap[set1].begin();
                    it1 != overlap[set1].end(); ++it1) {
                for (it2 = overlap[set2].begin();
                        it2 != overlap[set2].end(); ++it2) {
                    nPair = (float)(it1->second * it2->second);
                    if ((pos==1 && it1->first == it2->first) || (pos==0 && it1->first != it2->first)) { //pos
                        nPairPerEdge[minEdge] += nPairPosNormInv * nPair;
                        //if (pos==1){db <<minEdge<<","<<edgeWeight[minEdge]<<","<<v1<<","<<v2<<","<<nPair << "\n";}
                        // if (pos==1){ll+=nPairPosNormInv *(1-edgeWeight[minEdge])*(1-edgeWeight[minEdge])*nPair;db <<minEdge<<","<<ll << ","<<edgeWeight[minEdge]<<"," <<nPairPosNormInv << "\n";}
                    }
                }
            }
            /* move the pixel bags of the non-representative to the representative */
            if (dsets.find_set(set1) == set2) // make set1 the rep to keep and set2 the rep to empty
                swap(set1,set2);
            it2 = overlap[set2].begin();
            while (it2 != overlap[set2].end()) {
                it1 = overlap[set1].find(it2->first);
                if (it1 == overlap[set1].end()) {
                    overlap[set1].insert(pair<uint64_t,uint64_t>(it2->first,it2->second));
                } else {
                    it1->second += it2->second;
                }
                overlap[set2].erase(it2++);
            }
        } // end link
    } // end while
    // if(pos==1){ db.close();}
}

// for a given segmentation map (if not all 0), there are two cases:
// 1. only positive edges: single segmentation label
// 2. both positive and negative edges
void malis_loss_weights_cpp_pre(const uint64_t* seg,
               const uint64_t* conn_dims, const int32_t* nhood_data, const uint64_t* nhood_dims,
               const float* edgeWeight,
               float* nPairPerEdge,
               const uint64_t* pre_ve, const uint64_t* pre_prodDims, const int32_t* pre_nHood, const int pos){

    /* Disjoint sets and sparse overlap vectors */
    const int64_t nVert = pre_ve[0];
    const int64_t nEdge = pre_ve[1];
    uint64_t conn_num_dims = nhood_dims[1]+1;

    vector<map<uint64_t,uint64_t> > overlap(nVert);
    vector<uint64_t> rank(nVert);
    vector<uint64_t> parent(nVert);
    boost::disjoint_sets<uint64_t*, uint64_t*> dsets(&rank[0],&parent[0]);
    std::map<int64_t, int64_t> segSizes;
    int nLabeledVert = 0;
    int nPairPos = 0;
    for (int i=0; i<nVert; ++i){
        dsets.make_set(i);
        if (0!=seg[i]) {
            overlap[i].insert(pair<uint64_t,uint64_t>(seg[i],1));
            ++nLabeledVert;
            ++segSizes[seg[i]];
            nPairPos += (segSizes[seg[i]] - 1); //add pair with previous seg label
        }
    }

    float nPairPosNormInv;
    if (pos==1){
        nPairPosNormInv = 1.0/(float)nPairPos;
    }else{
        nPairPosNormInv = 1.0/(nLabeledVert*(nLabeledVert-1)*0.5-nPairPos);
    }
    //std::cout<<"norm:"<<nPairPosNormInv<<"\n";

    /* 2. Sort all the edges in increasing order of weight */
    const std::vector<int64_t> Pqueue1 = initPqueue(nEdge, conn_dims, nhood_data, nhood_dims);
    static std::vector<int64_t> Pqueue2(nEdge);
    std::copy(Pqueue1.begin(), Pqueue1.end(), Pqueue2.begin());

    sort( Pqueue2.begin(), Pqueue2.end(), AffinityGraphCompare<float>( edgeWeight ) );

    /* Start MST */
    int minEdge, e, v1, v2;
    uint64_t set1, set2;
    float nPair = 0;
    map<uint64_t,uint64_t>::iterator it1, it2;

    //std::ofstream db;
    //if(pos==1){ db.open("/n/coxfs01/donglai/malis_trans/db/mst-pt-pos.txt");}
    float ll=0;
    /* Start Kruskal's */
    for (unsigned int i = 0; i < Pqueue2.size(); ++i ) {
        minEdge = Pqueue2[i];
        // std::cout <<"do:"<<i<<","<<minEdge<< "\n";
        e =  minEdge / nVert;
        // v1: node at edge beginning
        v1 = minEdge % nVert;
        // v2: neighborhood node at edge e
        v2 = v1 + pre_nHood[e];
        set1 = dsets.find_set(v1);
        set2 = dsets.find_set(v2);
        if (set1!=set2){
            dsets.link(set1, set2);
            /* compute the number of pairs merged by this MST edge */
            for (it1 = overlap[set1].begin();
                    it1 != overlap[set1].end(); ++it1) {
                for (it2 = overlap[set2].begin();
                        it2 != overlap[set2].end(); ++it2) {
                    nPair = (float)(it1->second * it2->second);
                    if ((pos==1 && it1->first == it2->first) || (pos==0 && it1->first != it2->first)) { //pos
                        nPairPerEdge[minEdge] += nPairPosNormInv * nPair;
                        //if (pos==1){db <<minEdge<<","<<edgeWeight[minEdge]<<","<<v1<<","<<v2<<","<<nPair << "\n";}
                        // if (pos==1){ll+=nPairPosNormInv *(1-edgeWeight[minEdge])*(1-edgeWeight[minEdge])*nPair;db <<minEdge<<","<<ll << ","<<edgeWeight[minEdge]<<"," <<nPairPosNormInv << "\n";}
                    }
                }
            }
            /* move the pixel bags of the non-representative to the representative */
            if (dsets.find_set(set1) == set2) // make set1 the rep to keep and set2 the rep to empty
                swap(set1,set2);
            it2 = overlap[set2].begin();
            while (it2 != overlap[set2].end()) {
                it1 = overlap[set1].find(it2->first);
                if (it1 == overlap[set1].end()) {
                    overlap[set1].insert(pair<uint64_t,uint64_t>(it2->first,it2->second));
                } else {
                    it1->second += it2->second;
                }
                overlap[set2].erase(it2++);
            }
        } // end link
    } // end while
    // if(pos==1){ db.close();}
}

void malis_loss_weights_cpp(const uint64_t* seg,
               const uint64_t* conn_dims, const int32_t* nhood_data, const uint64_t* nhood_dims,
               const float* edgeWeight,
               const int pos,
               float* nPairPerEdge){

  uint64_t conn_num_dims = nhood_dims[1]+1;
  /* Cache for speed to access neighbors */
  // nVert stores (x * y * z)
  int64_t nVert = 1;
  for (int64_t i = 1; i < conn_num_dims; ++i) {
      nVert *= conn_dims[i];
  }

  // prodDims stores x, x*y, x*y*z offsets
  std::vector<int64_t> prodDims(conn_num_dims - 1);
  prodDims[conn_num_dims - 2] = 1;
  for (int64_t i = 1; i < conn_num_dims - 1; ++i) {
      prodDims[conn_num_dims - 2 - i] = prodDims[conn_num_dims - 1 - i]
                                      * conn_dims[conn_num_dims - i];
  }

  /* convert n-d offset vectors into linear array offset scalars */
  // nHood is a vector of size #edges
    std::vector<int32_t> nHood(nhood_dims[0]);
    for (int64_t i = 0; i < nhood_dims[0]; ++i) {
        nHood[i] = 0;
        for (int64_t j = 0; j < nhood_dims[1]; ++j) {
            nHood[i] += (int32_t) nhood_data[j + i * nhood_dims[1]] * prodDims[j];
        }
    }

    /* Disjoint sets and sparse overlap vectors */
    vector<map<uint64_t,uint64_t> > overlap(nVert);
    vector<uint64_t> rank(nVert);
    vector<uint64_t> parent(nVert);
    boost::disjoint_sets<uint64_t*, uint64_t*> dsets(&rank[0],&parent[0]);
    std::map<int64_t, int64_t> segSizes;
    int nLabeledVert = 0;
    int nPairPos = 0;
    for (int i=0; i<nVert; ++i){
        dsets.make_set(i);
        if (0!=seg[i]) {
            overlap[i].insert(pair<uint64_t,uint64_t>(seg[i],1));
            ++nLabeledVert;
            ++segSizes[seg[i]];
            nPairPos += (segSizes[seg[i]] - 1); //add pair with previous seg label
        }
    }

    int nPairTot = (nLabeledVert * (nLabeledVert - 1)) / 2;
    int nPairNeg = nPairTot - nPairPos;
    float nPairNormR;
    if(pos==1){
        nPairNormR = (float)nPairPos;
    }else{
        nPairNormR = (float)nPairNeg;
    }

    if (nPairNormR ==0){
        return;
    }else{
        nPairNormR = 1.0/nPairNormR;
    }
    /* get number of edges */
  int nEdge = 0;
  // Loop over #edges
  for (int d = 0, i = 0; d < conn_dims[0]; ++d) {
    // Loop over Z
    for (int z = 0; z < conn_dims[1]; ++z) {
      // Loop over Y
      for (int y = 0; y < conn_dims[2]; ++y) {
        // Loop over X
        for (int x = 0; x < conn_dims[3]; ++x, ++i) {
          // Out-of-bounds check:
          if (!((z + nhood_data[d * nhood_dims[1] + 0] < 0)
              ||(z + nhood_data[d * nhood_dims[1] + 0] >= conn_dims[1])
              ||(y + nhood_data[d * nhood_dims[1] + 1] < 0)
              ||(y + nhood_data[d * nhood_dims[1] + 1] >= conn_dims[2])
              ||(x + nhood_data[d * nhood_dims[1] + 2] < 0)
              ||(x + nhood_data[d * nhood_dims[1] + 2] >= conn_dims[3]))) {
              ++nEdge;
          }
        }
      }
    }
  }
    /* Sort all the edges in increasing order of weight */
    std::vector< int > pqueue( nEdge );
    int j = 0;
      // Loop over #edges
      for (int d = 0, i = 0; d < conn_dims[0]; ++d) {
        // Loop over Z
        for (int z = 0; z < conn_dims[1]; ++z) {
          // Loop over Y
          for (int y = 0; y < conn_dims[2]; ++y) {
            // Loop over X
            for (int x = 0; x < conn_dims[3]; ++x, ++i) {
              // Out-of-bounds check:
              if (!((z + nhood_data[d * nhood_dims[1] + 0] < 0)
                  ||(z + nhood_data[d * nhood_dims[1] + 0] >= conn_dims[1])
                  ||(y + nhood_data[d * nhood_dims[1] + 1] < 0)
                  ||(y + nhood_data[d * nhood_dims[1] + 1] >= conn_dims[2])
                  ||(x + nhood_data[d * nhood_dims[1] + 2] < 0)
                  ||(x + nhood_data[d * nhood_dims[1] + 2] >= conn_dims[3]))) {
                    pqueue[j++] = i;
              }
            }
          }
        }
      }
    pqueue.resize(j);
    sort( pqueue.begin(), pqueue.end(), AffinityGraphCompare<float>( edgeWeight ) );

    /* Start MST */
    int minEdge, e, v1, v2;
    uint64_t set1, set2;
    float nPair = 0;
    map<uint64_t,uint64_t>::iterator it1, it2;

    /* Start Kruskal's */
    // std::ofstream db;
    // db.open("/n/coxfs01/donglai/malis_trans/db/mst-pt.txt");

    //std::cout<<"hi:"<<std::endl;for (unsigned int i = 0; i < pqueue.size(); ++i ) {std::cout<<i<<",";}
    for (unsigned int i = 0; i < pqueue.size(); ++i ) {
        minEdge = pqueue[i];
        e =  minEdge / nVert;
        // v1: node at edge beginning
        v1 = minEdge % nVert;
        // v2: neighborhood node at edge e
        v2 = v1 + nHood[e];

        set1 = dsets.find_set(v1);
        set2 = dsets.find_set(v2);

        if (set1!=set2){
            dsets.link(set1, set2);

            /* compute the number of pairs merged by this MST edge */
            for (it1 = overlap[set1].begin();
                    it1 != overlap[set1].end(); ++it1) {
                for (it2 = overlap[set2].begin();
                        it2 != overlap[set2].end(); ++it2) {

                    nPair = (float)(it1->second * it2->second);

                    if (pos>0 && (it1->first == it2->first)) {
                        nPairPerEdge[minEdge] += nPairNormR * nPair;
                        //db <<minEdge<<","<<edgeWeight[minEdge]<<","<<v1<<","<<v2<<","<<nPair << "\n";
                    } else if ((pos==0) && (it1->first != it2->first)) {
                        nPairPerEdge[minEdge] += nPairNormR * nPair;
                    }
                }
            }

            /* move the pixel bags of the non-representative to the representative */
            if (dsets.find_set(set1) == set2) // make set1 the rep to keep and set2 the rep to empty
                swap(set1,set2);

            it2 = overlap[set2].begin();
            while (it2 != overlap[set2].end()) {
                it1 = overlap[set1].find(it2->first);
                if (it1 == overlap[set1].end()) {
                    overlap[set1].insert(pair<uint64_t,uint64_t>(it2->first,it2->second));
                } else {
                    it1->second += it2->second;
                }
                overlap[set2].erase(it2++);
            }
        } // end link

    } // end while
    //db.close();
}