#include "routing.h"
#include "ortools/constraint_solver/routing.h"
#include "ortools/constraint_solver/routing_enums.pb.h"
#include "ortools/constraint_solver/routing_index_manager.h"
#include "ortools/constraint_solver/routing_parameters.h"

using namespace operations_research;

// Aux struct and functions.
struct InternalParameters {
  InternalParameters()
      : numSolutionsPerRun(1), numRuns(1), minNumTransectsPerRun(5),
        stop([] { return false; }) {}
  std::size_t numSolutionsPerRun;
  std::size_t numRuns;
  std::size_t minNumTransectsPerRun;
  std::function<bool(void)> stop;
  mutable std::string errorString;
};

bool getRoute(const FPolygon &area, const LineStringArray &transects,
              std::vector<Solution> &solutionVector,
              const InternalParameters &par = InternalParameters());


bool getRoute(const FPolygon &area, const LineStringArray &transects,
              std::vector<Solution> &solutionVector,
              const InternalParameters &par) {

#ifdef SNAKE_SHOW_TIME
  auto start = std::chrono::high_resolution_clock::now();
#endif
  //================================================================
  // Create routing model.
  //================================================================
  // Use integer polygons to increase numerical robustness.
  // Convert area;
  IPolygon intArea;
  for (const auto &v : area.outer()) {
    auto p = float2Int(v);
    intArea.outer().push_back(p);
  }
  for (const auto &ring : area.inners()) {
    IRing intRing;
    for (const auto &v : ring) {
      auto p = float2Int(v);
      intRing.push_back(p);
    }
    intArea.inners().push_back(std::move(intRing));
  }

  // Helper classes.
  struct VirtualNode {
    VirtualNode(std::size_t f, std::size_t t) : fromIndex(f), toIndex(t) {}
    std::size_t fromIndex; // index for leaving node
    std::size_t toIndex;   // index for entering node
  };
  struct NodeToTransect {
    NodeToTransect(std::size_t i, bool r) : transectsIndex(i), reversed(r) {}
    std::size_t transectsIndex; // transects index
    bool reversed;              // transect reversed?
  };
  // Create vertex and node list
  std::vector<IPoint> vertices;
  std::vector<std::pair<std::size_t, std::size_t>> disjointNodes;
  std::vector<VirtualNode> nodeList;
  std::vector<NodeToTransect> nodeToTransectList;
  for (std::size_t i = 0; i < transects.size(); ++i) {
    const auto &t = transects[i];
    // Copy line edges only.
    if (t.size() == 1 || i == 0) {
      auto p = float2Int(t.back());
      vertices.push_back(p);
      nodeToTransectList.emplace_back(i, false);
      auto idx = vertices.size() - 1;
      nodeList.emplace_back(idx, idx);
    } else if (t.size() > 1) {
      auto p1 = float2Int(t.front());
      auto p2 = float2Int(t.back());
      vertices.push_back(p1);
      vertices.push_back(p2);
      nodeToTransectList.emplace_back(i, false);
      nodeToTransectList.emplace_back(i, true);
      auto fromIdx = vertices.size() - 1;
      auto toIdx = fromIdx - 1;
      nodeList.emplace_back(fromIdx, toIdx);
      nodeList.emplace_back(toIdx, fromIdx);
      disjointNodes.emplace_back(toIdx, fromIdx);
    } else { // transect empty
      std::cout << "ignoring empty transect with index " << i << std::endl;
    }
  }
#ifdef SNAKE_DEBUG
  // Print.
  std::cout << "nodeToTransectList:" << std::endl;
  std::cout << "node:transectIndex:reversed" << std::endl;
  std::size_t c = 0;
  for (const auto &n2t : nodeToTransectList) {
    std::cout << c++ << ":" << n2t.transectsIndex << ":" << n2t.reversed
              << std::endl;
  }
  std::cout << "nodeList:" << std::endl;
  std::cout << "node:fromIndex:toIndex" << std::endl;
  c = 0;
  for (const auto &n : nodeList) {
    std::cout << c++ << ":" << n.fromIndex << ":" << n.toIndex << std::endl;
  }
  std::cout << "disjoint nodes:" << std::endl;
  std::cout << "number:nodes" << std::endl;
  c = 0;
  for (const auto &d : disjointNodes) {
    std::cout << c++ << ":" << d.first << "," << d.second << std::endl;
  }
#endif

  // Add polygon vertices.
  for (auto &v : intArea.outer()) {
    vertices.push_back(v);
  }
  for (auto &ring : intArea.inners()) {
    for (auto &v : ring) {
      vertices.push_back(v);
    }
  }

  // Create connection graph (inf == no connection between vertices).
  // Note: graph is not symmetric.
  auto n = vertices.size();
  // Matrix must be double since integers don't have infinity and nan
  Matrix<double> connectionGraph(n, n);
  for (std::size_t i = 0; i < n; ++i) {
    auto &fromVertex = vertices[i];
    for (std::size_t j = 0; j < n; ++j) {
      auto &toVertex = vertices[j];
      ILineString line{fromVertex, toVertex};
      if (bg::covered_by(line, intArea)) {
        connectionGraph(i, j) = bg::length(line);
      } else {
        connectionGraph(i, j) = std::numeric_limits<double>::infinity();
      }
    }
  }
#ifdef SNAKE_DEBUG
  std::cout << "connection grah:" << std::endl;
  std::cout << connectionGraph << std::endl;
#endif

  // Create distance matrix.
  auto distLambda = [&connectionGraph](std::size_t i, std::size_t j) -> double {
    return connectionGraph(i, j);
  };
  auto nNodes = nodeList.size();
  Matrix<IntType> distanceMatrix(nNodes, nNodes);
  for (std::size_t i = 0; i < nNodes; ++i) {
    distanceMatrix(i, i) = 0;
    for (std::size_t j = i + 1; j < nNodes; ++j) {
      auto dist = connectionGraph(i, j);
      if (std::isinf(dist)) {
        std::vector<std::size_t> route;
        if (!dijkstraAlgorithm(n, i, j, route, dist, distLambda)) {
          std::stringstream ss;
          ss << "Distance matrix calculation failed. connection graph: "
             << connectionGraph << std::endl;
          ss << "area: " << bg::wkt(area) << std::endl;
          ss << "transects:" << std::endl;
          for (const auto &t : transects) {

            ss << bg::wkt(t) << std::endl;
          }

          par.errorString = ss.str();
          return false;
        }
        (void)route;
      }
      distanceMatrix(i, j) = dist;
      distanceMatrix(j, i) = dist;
    }
  }
#ifdef SNAKE_DEBUG
  std::cout << "distance matrix:" << std::endl;
  std::cout << distanceMatrix << std::endl;
#endif

  // Create (asymmetric) routing matrix.
  Matrix<IntType> routingMatrix(nNodes, nNodes);
  for (std::size_t i = 0; i < nNodes; ++i) {
    auto fromNode = nodeList[i];
    for (std::size_t j = 0; j < nNodes; ++j) {
      auto toNode = nodeList[j];
      routingMatrix(i, j) = distanceMatrix(fromNode.fromIndex, toNode.toIndex);
    }
  }
  // Insert max for disjoint nodes.
  for (const auto &d : disjointNodes) {
    auto i = d.first;
    auto j = d.second;
    routingMatrix(i, j) = std::numeric_limits<IntType>::max();
    routingMatrix(j, i) = std::numeric_limits<IntType>::max();
  }
#ifdef SNAKE_DEBUG
  std::cout << "routing matrix:" << std::endl;
  std::cout << routingMatrix << std::endl;
#endif

  // Create Routing Index Manager.
  auto minNumTransectsPerRun =
      std::max<std::size_t>(1, par.minNumTransectsPerRun);
  auto maxRuns = std::max<std::size_t>(
      1, std::floor(double(transects.size()) / minNumTransectsPerRun));
  auto numRuns = std::max<std::size_t>(1, par.numRuns);
  numRuns = std::min<std::size_t>(numRuns, maxRuns);

  RoutingIndexManager::NodeIndex depot(0);
  //  std::vector<RoutingIndexManager::NodeIndex> depots(numRuns, depot);
  //  RoutingIndexManager manager(nNodes, numRuns, depots, depots);
  RoutingIndexManager manager(nNodes, numRuns, depot);

  // Create Routing Model.
  RoutingModel routing(manager);
  // Create and register a transit callback.
  const int transitCallbackIndex = routing.RegisterTransitCallback(
      [&routingMatrix, &manager](int64 from_index, int64 to_index) -> int64 {
        // Convert from routing variable Index to distance matrix NodeIndex.
        auto from_node = manager.IndexToNode(from_index).value();
        auto to_node = manager.IndexToNode(to_index).value();
        return routingMatrix(from_node, to_node);
      });
  // Define cost of each arc.
  routing.SetArcCostEvaluatorOfAllVehicles(transitCallbackIndex);
  // Add distance dimension.
  if (numRuns > 1) {
    routing.AddDimension(transitCallbackIndex, 0, 300000000,
                         true, // start cumul to zero
                         "Distance");
    routing.GetMutableDimension("Distance")
        ->SetGlobalSpanCostCoefficient(100000000);
  }

  // Define disjunctions.
#ifdef SNAKE_DEBUG
  std::cout << "disjunctions:" << std::endl;
#endif
  for (const auto &d : disjointNodes) {
    auto i = d.first;
    auto j = d.second;
#ifdef SNAKE_DEBUG
    std::cout << i << "," << j << std::endl;
#endif
    auto idx0 = manager.NodeToIndex(RoutingIndexManager::NodeIndex(i));
    auto idx1 = manager.NodeToIndex(RoutingIndexManager::NodeIndex(j));
    std::vector<int64> disj{idx0, idx1};
    routing.AddDisjunction(disj, -1 /*force cardinality*/, 1 /*cardinality*/);
  }

  // Set first solution heuristic.
  auto searchParameters = DefaultRoutingSearchParameters();
  searchParameters.set_first_solution_strategy(
      FirstSolutionStrategy::PATH_CHEAPEST_ARC);
  // Number of solutions.
  auto numSolutionsPerRun = std::max<std::size_t>(1, par.numSolutionsPerRun);
  searchParameters.set_number_of_solutions_to_collect(numSolutionsPerRun);
  // Set costume limit.
  auto *solver = routing.solver();
  auto *limit = solver->MakeCustomLimit(par.stop);
  routing.AddSearchMonitor(limit);
#ifdef SNAKE_SHOW_TIME
  auto delta = std::chrono::duration_cast<std::chrono::milliseconds>(
      std::chrono::high_resolution_clock::now() - start);
  cout << "create routing model: " << delta.count() << " ms" << endl;
#endif

  //================================================================
  // Solve model.
  //================================================================
#ifdef SNAKE_SHOW_TIME
  start = std::chrono::high_resolution_clock::now();
#endif
  auto pSolutions = std::make_unique<std::vector<const Assignment *>>();
  (void)routing.SolveWithParameters(searchParameters, pSolutions.get());
#ifdef SNAKE_SHOW_TIME
  delta = std::chrono::duration_cast<std::chrono::milliseconds>(
      std::chrono::high_resolution_clock::now() - start);
  cout << "solve routing model: " << delta.count() << " ms" << endl;
#endif
  if (par.stop()) {
    par.errorString = "User terminated.";
    return false;
  }
  if (pSolutions->size() == 0) {
    std::stringstream ss;
    ss << "No solution found." << std::endl;
    par.errorString = ss.str();
    return false;
  }

  //================================================================
  // Construc route.
  //================================================================
#ifdef SNAKE_SHOW_TIME
  start = std::chrono::high_resolution_clock::now();
#endif
  long long counter = -1;
  // Note: route number 0 corresponds to the best route which is the last entry
  // of *pSolutions.
  for (auto solution = pSolutions->end() - 1; solution >= pSolutions->begin();
       --solution) {
    ++counter;
    if (!*solution || (*solution)->Size() <= 1) {
      std::stringstream ss;
      ss << par.errorString << "Solution " << counter << "invalid."
         << std::endl;
      par.errorString = ss.str();
      continue;
    }
    // Iterate over all routes.
    Solution routeVector;
    for (std::size_t vehicle = 0; vehicle < numRuns; ++vehicle) {
      if (!routing.IsVehicleUsed(**solution, vehicle))
        continue;

      // Create index list.
      auto index = routing.Start(vehicle);
      std::vector<size_t> route_idx;
      route_idx.push_back(manager.IndexToNode(index).value());
      while (!routing.IsEnd(index)) {
        index = (*solution)->Value(routing.NextVar(index));
        route_idx.push_back(manager.IndexToNode(index).value());
      }

#ifdef SNAKE_DEBUG
      // Print route.
      std::cout << "route " << counter
                << " route_idx.size() = " << route_idx.size() << std::endl;
      std::cout << "route: ";
      for (const auto &idx : route_idx) {
        std::cout << idx << ", ";
      }
      std::cout << std::endl;
#endif
      if (route_idx.size() < 2) {
        std::stringstream ss;
        ss << par.errorString
           << "Error while assembling route (solution = " << counter
           << ", run = " << vehicle << ")." << std::endl;
        par.errorString = ss.str();
        continue;
      }

      // Assemble route.
      Route r;
      auto &path = r.path;
      auto &info = r.info;
      for (size_t i = 0; i < route_idx.size() - 1; ++i) {
        size_t nodeIndex0 = route_idx[i];
        size_t nodeIndex1 = route_idx[i + 1];
        const auto &n2t0 = nodeToTransectList[nodeIndex0];
        info.emplace_back(n2t0.transectsIndex, n2t0.reversed);
        // Copy transect to route.
        const auto &t = transects[n2t0.transectsIndex];
        if (n2t0.reversed) { // transect reversal needed?
          for (auto it = t.end() - 1; it > t.begin(); --it) {
            path.push_back(*it);
          }
        } else {
          for (auto it = t.begin(); it < t.end() - 1; ++it) {
            path.push_back(*it);
          }
        }
        // Connect transects.
        std::vector<size_t> idxList;
        if (!shortestPathFromGraph(connectionGraph,
                                   nodeList[nodeIndex0].fromIndex,
                                   nodeList[nodeIndex1].toIndex, idxList)) {
          std::stringstream ss;
          ss << par.errorString
             << "Error while assembling route (solution = " << counter
             << ", run = " << vehicle << ")." << std::endl;
          par.errorString = ss.str();
          continue;
        }
        if (i != route_idx.size() - 2) {
          idxList.pop_back();
        }
        for (auto idx : idxList) {
          auto p = int2Float(vertices[idx]);
          path.push_back(p);
        }
      }
      // Append last transect info.
      const auto &n2t0 = nodeToTransectList.back();
      info.emplace_back(n2t0.transectsIndex, n2t0.reversed);

      if (path.size() < 2 || info.size() < 2) {
        std::stringstream ss;
        ss << par.errorString << "Route empty (solution = " << counter
           << ", run = " << vehicle << ")." << std::endl;
        par.errorString = ss.str();
        continue;
      }

      routeVector.push_back(std::move(r));
    }
    if (routeVector.size() > 0) {
      solutionVector.push_back(std::move(routeVector));
    } else {
      std::stringstream ss;
      ss << par.errorString << "Solution " << counter << " empty." << std::endl;
      par.errorString = ss.str();
    }
  }
#ifdef SNAKE_SHOW_TIME
  delta = std::chrono::duration_cast<std::chrono::milliseconds>(
      std::chrono::high_resolution_clock::now() - start);
  cout << "reconstruct route: " << delta.count() << " ms" << endl;
#endif

  if (solutionVector.size() > 0) {
    return true;
  } else {
    return false;
  }
}