solver.cc

/*
  LENS Optimizer: fifth place winner of the IOPDDL track, ASPLOS&EuroSys'25
  Contest.

  Members: Haodi Jiang*, Yitian Yang*, Ruwen Fan, Shiwei Gao, Shaoxun Zeng,
  Junrong Huang, Huajun Bai, Hao Guo and Youyou Lu

  Below is the original copyright information from the framework:

  Copyright 2024 Google LLC

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

  https://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

#include "solver.h"

#include <absl/numeric/int128.h>
#include <algorithm>
#include <cstdlib>
#include <iostream>
#include <mutex>
#include <optional>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <thread>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_join.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"

#include "iopddl.h"

namespace iopddl {

using EdgeList = std::vector<std::vector<std::tuple<NodeIdx, EdgeIdx, bool>>>;

// =================================
// Auxilary Data Structures:
// =================================

// ==================================================
// Optimized Usage Calculation: Segment Tree Based
// ==================================================
class UsageTree {
  struct Node {
    TotalUsage max_usage;
    TotalUsage lazy;
  };
  int64_t n;
  std::vector<Node> tree;
  void push_up(int64_t x) {
    tree[x].max_usage =
        std::max(tree[x * 2].max_usage, tree[x * 2 + 1].max_usage);
  }
  void push_down(int64_t x) {
    if (tree[x].lazy) {
      tree[x * 2].max_usage += tree[x].lazy;
      tree[x * 2 + 1].max_usage += tree[x].lazy;
      tree[x * 2].lazy += tree[x].lazy;
      tree[x * 2 + 1].lazy += tree[x].lazy;
      tree[x].lazy = 0;
    }
  }
  void build(int64_t x, int64_t l, int64_t r,
             const std::vector<TotalUsage> &v) {
    if (l == r) {
      tree[x].max_usage = v[l];
      return;
    }
    int64_t mid = (l + r) / 2;
    build(x * 2, l, mid, v);
    build(x * 2 + 1, mid + 1, r, v);
    push_up(x);
  }
  void update(int64_t x, int64_t l, int64_t r, int64_t ql, int64_t qr,
              TotalUsage v) {
    if (ql <= l && r <= qr) {
      tree[x].max_usage += v;
      tree[x].lazy += v;
      return;
    }
    push_down(x);
    int64_t mid = (l + r) / 2;
    if (ql <= mid)
      update(x * 2, l, mid, ql, qr, v);
    if (qr > mid)
      update(x * 2 + 1, mid + 1, r, ql, qr, v);
    push_up(x);
  }
  TotalUsage query(int64_t x, int64_t l, int64_t r, int64_t ql, int64_t qr) {
    if (ql <= l && r <= qr) {
      return tree[x].max_usage;
    }
    push_down(x);
    int64_t mid = (l + r) / 2;
    TotalUsage ret = 0;
    if (ql <= mid)
      ret = std::max(ret, query(x * 2, l, mid, ql, qr));
    if (qr > mid)
      ret = std::max(ret, query(x * 2 + 1, mid + 1, r, ql, qr));
    return ret;
  }

public:
  UsageTree(const std::vector<TotalUsage> &v) : n(v.size()), tree(4 * n) {
    build(1, 0, n - 1, v);
  }

  void update(int64_t l, int64_t r, TotalUsage v) {
    if (l < 0 || r >= n || l > r) {
      return;
    }
    update(1, 0, n - 1, l, r, v);
  }

  TotalUsage query(int64_t l, int64_t r) {
    if (l < 0 || r >= n || l > r) {
      return 0;
    }

    return query(1, 0, n - 1, l, r);
  }
  TotalUsage query() { return query(0, n - 1); }
};

// ================================================================
// General Interface for the Problem's structure,
// which allows fast usage update, cost and validaty maintainance,
// and supports auto-revert
// ================================================================
class ExtendedProblem {
  const Problem &problem;
  const EdgeList &edges;
  const std::vector<std::pair<TimeIdx, TimeIdx>> &new_interval;
  std::vector<StrategyIdx> &solution;
  UsageTree &usage_tree;
  std::stack<std::pair<NodeIdx, StrategyIdx>> history;
  TotalCost cost;
  Cost get_edge_cost(EdgeIdx eid, StrategyIdx u, StrategyIdx v,
                     bool reverse = false) {
    const auto &edge = problem.edges[eid];
    StrategyIdx strategy_idx = 0;
    if (reverse)
      std::swap(u, v);
    std::vector<StrategyIdx> p = {u, v};
    auto idx = p.begin();
    for (const NodeIdx node_idx : edge.nodes) {
      strategy_idx *= problem.nodes[node_idx].strategies.size();
      strategy_idx += *(idx++);
    }
    return edge.strategies[strategy_idx].cost;
  }

public:
  ExtendedProblem(const Problem &problem, const EdgeList &edges,
                  const std::vector<std::pair<TimeIdx, TimeIdx>> &new_interval,
                  Solution &solution, UsageTree &usage_tree, TotalCost cost)
      : problem(problem), edges(edges), new_interval(new_interval),
        solution(solution), usage_tree(usage_tree), cost(cost) {}
  void check() { return; }
  std::optional<TotalCost> get_cost() const {
    auto max_usage =
        problem.usage_limit ? *problem.usage_limit : absl::Int128Max();
    return (usage_tree.query() > max_usage) ? std::nullopt
                                            : std::make_optional(cost);
  }

  Solution get_solution() const { return solution; }
  StrategyIdx get_strategy(NodeIdx x) const { return solution[x]; }
  void update(NodeIdx x, StrategyIdx s, bool record = true) {
    if (record)
      history.push({x, solution[x]});
    const auto &node = problem.nodes[x];
    auto [begin, end] = new_interval[x];
    auto old_strategy = solution[x];
    solution[x] = s;
    auto max_usage =
        problem.usage_limit ? *problem.usage_limit : absl::Int128Max();
    usage_tree.update(begin, end - 1,
                      node.strategies[s].usage -
                          node.strategies[old_strategy].usage);

    for (auto [v, e, rev] : edges[x]) {
      auto old_cost = get_edge_cost(e, old_strategy, solution[v], rev);
      auto new_cost = get_edge_cost(e, s, solution[v], rev);
      cost += new_cost - old_cost;
    }
    cost += node.strategies[s].cost - node.strategies[old_strategy].cost;
  }
  void undo() {
    while (!history.empty()) {
      auto [x, s] = history.top();
      history.pop();
      update(x, s, false);
    }
  }
  void commit() { history = std::stack<std::pair<NodeIdx, StrategyIdx>>(); }
  ~ExtendedProblem() { undo(); }
};

// ================================================================
// Technique #1: Chained Restructuring
// A seperate module
// ================================================================
template <class UsageManager, class EdgeManager> struct ChainedRestruct {
  std::random_device rd;
  std::mt19937 gen{rd()};
  std::uniform_real_distribution<> dis{0.0, 1.0};

  const Problem &problem;
  const EdgeManager &edges;
  const std::vector<Interval> &interval;

  std::tuple<NodeIdx, StrategyIdx, TotalCost> *modify_queue;
  NodeIdx modify_queue_head = 0, modify_queue_tail = 0;
  bool *modify_flag;
  std::tuple<NodeIdx, StrategyIdx> *saved_modify_queue;
  NodeIdx saved_modify_queue_length;
  TotalCost saved_total_cost;

  NodeIdx *random_seq;
  bool *random_flag;
  NodeIdx random_i = 0;

  const TotalCost inf = 1000000000000000000;
  std::vector<NodeIdx> first_strategy_global_indices;
  std::vector<StrategyIdx> first_force_select_indices;
  std::vector<std::pair<NodeIdx, StrategyIdx>> force_selects;

  ChainedRestruct(const Problem &problem, const EdgeManager &edges,
                  const std::vector<Interval> &interval)
      : problem(problem), edges(edges), interval(interval) {
    modify_queue =
        new std::tuple<NodeIdx, StrategyIdx, TotalCost>[problem.nodes.size()];
    modify_flag = new bool[problem.nodes.size()]();
    saved_modify_queue =
        new std::tuple<NodeIdx, StrategyIdx>[problem.nodes.size()];

    random_seq = new NodeIdx[problem.nodes.size()];
    for (NodeIdx i = 0; i < problem.nodes.size(); ++i) {
      random_seq[i] = i;
    }
    std::shuffle(random_seq, random_seq + problem.nodes.size(), gen);
    random_flag = new bool[problem.nodes.size()]();

    for (NodeIdx i = 0; i < problem.nodes.size(); ++i) {
      first_strategy_global_indices.emplace_back(
          first_force_select_indices.size());
      auto &node = problem.nodes[i];
      for (StrategyIdx k = 0; k < node.strategies.size(); ++k) {
        first_force_select_indices.emplace_back(force_selects.size());
        for (EdgeIdx j = 0; j < get_n_edges(i); ++j) {
          auto &edge = problem.edges[get_edge_idx(i, j)];
          auto pair_node_idx =
              i == edge.nodes[0] ? edge.nodes[1] : edge.nodes[0];
          auto &pair_node = problem.nodes[pair_node_idx];
          StrategyIdx cnt = 0;
          StrategyIdx best_pair_strategy_idx;
          for (StrategyIdx l = 0; l < pair_node.strategies.size(); ++l) {
            if (get_edge_cost(edge, i, pair_node_idx, k, l) != inf) {
              cnt++;
              best_pair_strategy_idx = l;
            }
          }
          if (cnt == 1) {
            force_selects.emplace_back(pair_node_idx, best_pair_strategy_idx);
          }
        }
      }
    }
    first_strategy_global_indices.emplace_back(
        first_force_select_indices.size());
    first_force_select_indices.emplace_back(force_selects.size());
  }

  ~ChainedRestruct() {
    delete[] modify_queue;
    delete[] modify_flag;
    delete[] saved_modify_queue;

    delete[] random_seq;
    delete[] random_flag;
  }

  inline void add_usage(UsageManager &usage_manager, NodeIdx node_idx,
                        Usage usage) {
    if constexpr (std::is_same<UsageManager, std::vector<TotalUsage>>::value ==
                  true) {
      for (auto i = interval[node_idx].first; i < interval[node_idx].second;
           ++i) {
        usage_manager[i] += usage;
      }
    } else if constexpr (std::is_same<UsageManager, UsageTree>::value == true) {
      usage_manager.update(interval[node_idx].first,
                           interval[node_idx].second - 1, usage);
    }
  }

  inline EdgeIdx get_n_edges(NodeIdx node_idx) {
    if constexpr (std::is_same<EdgeManager, EdgeList>::value == true) {
      return edges[node_idx].size();
    }
    return 0;
  }

  inline EdgeIdx get_edge_idx(NodeIdx node_idx, EdgeIdx edge_idx) {
    if constexpr (std::is_same<EdgeManager, EdgeList>::value == true) {
      return std::get<1>(edges[node_idx][edge_idx]);
    }
    return -1;
  }

  inline Cost get_edge_cost(Solution &solution, const Edge &edge) {
    auto ret =
        edge.strategies[solution[edge.nodes[0]] *
                            problem.nodes[edge.nodes[1]].strategies.size() +
                        solution[edge.nodes[1]]]
            .cost;
    return ret;
  }

  inline Cost get_edge_cost(const Edge &edge, NodeIdx node0_idx,
                            NodeIdx node1_idx, StrategyIdx strategy0_idx,
                            StrategyIdx strategy1_idx) {
    return edge.nodes[0] == node0_idx
               ? edge.strategies[strategy0_idx * problem.nodes[node1_idx]
                                                     .strategies.size() +
                                 strategy1_idx]
                     .cost
               : edge.strategies[strategy1_idx * problem.nodes[node0_idx]
                                                     .strategies.size() +
                                 strategy0_idx]
                     .cost;
  }

  template <bool mark_modify>
  inline TotalCost get_node_total_cost(Solution &solution, NodeIdx node_idx) {
    auto &node = problem.nodes[node_idx];
    auto strategy_idx = solution[node_idx];
    TotalCost total_cost = node.strategies[strategy_idx].cost;
    for (EdgeIdx i = 0; i < get_n_edges(node_idx); ++i) {
      auto &edge = problem.edges[get_edge_idx(node_idx, i)];
      auto cost = get_edge_cost(solution, edge);
      total_cost += cost;

      if constexpr (mark_modify && false) {
        auto pair_node_idx =
            edge.nodes[0] != node_idx ? edge.nodes[0] : edge.nodes[1];
        auto &pair_node = problem.nodes[pair_node_idx];
        if (cost == inf && modify_flag[pair_node_idx] == false) {
          auto saved_strategy_idx = solution[pair_node_idx];
          for (StrategyIdx j = 0; j < pair_node.strategies.size(); ++j) {
            solution[pair_node_idx] = j;
            if (get_edge_cost(solution, edge) != inf) {
              modify_queue[modify_queue_tail++] = {pair_node_idx, j, -1};
              modify_flag[pair_node_idx] = true;
              break;
            }
          }
          solution[pair_node_idx] = saved_strategy_idx;
        }
      }
    }

    if constexpr (mark_modify && true) {
      auto strategy_global_idx =
          first_strategy_global_indices[node_idx] + solution[node_idx];
      for (auto i = first_force_select_indices[strategy_global_idx];
           i < first_force_select_indices[strategy_global_idx + 1]; ++i) {
        auto &[pair_node_idx, pair_strategy_idx] = force_selects[i];
        if (modify_flag[pair_node_idx] == false &&
            solution[pair_node_idx] != pair_strategy_idx) {
          modify_queue[modify_queue_tail++] = {pair_node_idx, pair_strategy_idx,
                                               -1};
          modify_flag[pair_node_idx] = true;
        }
      }
    }

    return total_cost;
  }

  bool try_modify(Solution &solution, TotalCost &total_cost,
                  UsageManager &usage_manager) {
    while (modify_queue_head < modify_queue_tail) {
      auto &[node_idx, saved_strategy_idx, saved_total_cost] =
          modify_queue[modify_queue_head++];
      auto &node = problem.nodes[node_idx];
      add_usage(usage_manager, node_idx,
                node.strategies[saved_strategy_idx].usage -
                    node.strategies[solution[node_idx]].usage);
      saved_total_cost = total_cost;
      total_cost -= get_node_total_cost<false>(solution, node_idx);
      std::swap(solution[node_idx], saved_strategy_idx);
      total_cost += get_node_total_cost<true>(solution, node_idx);
    }

    return check_modify(usage_manager);
  }

  void discard_modify(Solution &solution, TotalCost &total_cost,
                      UsageManager &usage_manager, NodeIdx step) {
    for (NodeIdx i = step; i < modify_queue_tail; ++i) {
      auto &[node_idx, saved_strategy_idx, saved_total_cost] = modify_queue[i];
      auto &node = problem.nodes[node_idx];
      add_usage(usage_manager, node_idx,
                node.strategies[saved_strategy_idx].usage -
                    node.strategies[solution[node_idx]].usage);
      solution[node_idx] = saved_strategy_idx;
      modify_flag[node_idx] = false;
    }
    total_cost = std::get<2>(modify_queue[step]);

    modify_queue_head = step;
    modify_queue_tail = step;
  }

  void stage_modify(Solution &solution, TotalCost &total_cost) {
    for (NodeIdx i = 0; i < modify_queue_tail; ++i) {
      auto &[node_idx, saved_strategy_idx, saved_total_cost] = modify_queue[i];
      saved_modify_queue[i] = {node_idx, solution[node_idx]};
    }
    saved_modify_queue_length = modify_queue_tail;
    saved_total_cost = total_cost;
  }

  void commit_modify(Solution &solution, TotalCost &total_cost,
                     UsageManager &usage_manager) {
    for (NodeIdx i = 0; i < saved_modify_queue_length; ++i) {
      auto &[node_idx, saved_strategy_idx] = saved_modify_queue[i];
      auto &node = problem.nodes[node_idx];
      add_usage(usage_manager, node_idx,
                node.strategies[saved_strategy_idx].usage -
                    node.strategies[solution[node_idx]].usage);
      solution[node_idx] = saved_strategy_idx;
    }
    total_cost = saved_total_cost;
  }

  bool change_one_node(Solution &solution, TotalCost &total_cost,
                       UsageManager &usage_manager, double temperature) {
    for (; random_i < problem.nodes.size() && random_flag[random_seq[random_i]];
         ++random_i)
      ;
    if (random_i == problem.nodes.size()) {
      std::shuffle(random_seq, random_seq + problem.nodes.size(), gen);
      std::memset(random_flag, 0, problem.nodes.size());
      random_i = 0;
    }
    auto node_idx = random_seq[random_i];

    auto &node = problem.nodes[node_idx];
    auto best_total_cost = absl::Int128Max();
    for (StrategyIdx i = 0; i < node.strategies.size(); ++i) {
      modify_queue[modify_queue_tail++] = {node_idx, i, -1};
      modify_flag[node_idx] = true;
      if (try_modify(solution, total_cost, usage_manager) == true &&
          total_cost < best_total_cost) {
        best_total_cost = total_cost;
        stage_modify(solution, total_cost);
      }

      for (NodeIdx j = 0; j < modify_queue_tail; ++j) {
        random_flag[std::get<0>(modify_queue[j])] = true;
      }

      discard_modify(solution, total_cost, usage_manager, 0);
    }

    if (best_total_cost < total_cost ||
        exp(double(total_cost - best_total_cost) / temperature) > dis(gen)) {
      commit_modify(solution, total_cost, usage_manager);
      return true;
    }
    return false;
  }

  inline void modify(Solution &solution, TotalCost &total_cost,
                     UsageManager &usage_manager, NodeIdx node_idx,
                     StrategyIdx strategy_idx) {
    auto &node = problem.nodes[node_idx];
    total_cost += node.strategies[strategy_idx].cost -
                  node.strategies[solution[node_idx]].cost;
    add_usage(usage_manager, node_idx,
              node.strategies[strategy_idx].usage -
                  node.strategies[solution[node_idx]].usage);
    for (EdgeIdx i = 0; i < get_n_edges(node_idx); ++i) {
      auto &edge = problem.edges[get_edge_idx(node_idx, i)];
      auto pair_node_idx =
          edge.nodes[0] != node_idx ? edge.nodes[0] : edge.nodes[1];
      total_cost += get_edge_cost(edge, node_idx, pair_node_idx, strategy_idx,
                                  solution[pair_node_idx]) -
                    get_edge_cost(edge, node_idx, pair_node_idx,
                                  solution[node_idx], solution[pair_node_idx]);
    }
  }

  inline void modify(Solution &solution, TotalCost &total_cost,
                     UsageManager &usage_manager) {
    while (modify_queue_head < modify_queue_tail) {
      auto &[node_idx, saved_strategy_idx, saved_total_cost] =
          modify_queue[modify_queue_head++];
      saved_total_cost = total_cost;
      modify(solution, total_cost, usage_manager, node_idx, saved_strategy_idx);
      std::swap(saved_strategy_idx, solution[node_idx]);

      auto strategy_global_idx =
          first_strategy_global_indices[node_idx] + solution[node_idx];
      for (auto i = first_force_select_indices[strategy_global_idx];
           i < first_force_select_indices[strategy_global_idx + 1]; ++i) {
        auto &[pair_node_idx, pair_strategy_idx] = force_selects[i];
        if (modify_flag[pair_node_idx] == false &&
            solution[pair_node_idx] != pair_strategy_idx) {
          modify_queue[modify_queue_tail++] = {pair_node_idx, pair_strategy_idx,
                                               -1};
          modify_flag[pair_node_idx] = true;
        }
      }
    }
  }

  inline bool check_modify(UsageManager &usage_manager) {
    if constexpr (std::is_same<UsageManager, std::vector<TotalUsage>>::value ==
                  true) {
      for (NodeIdx i = 0; i < modify_queue_tail; ++i) {
        auto &[node_idx, saved_strategy_idx, saved_total_cost] =
            modify_queue[i];
        for (auto j = interval[node_idx].first; j < interval[node_idx].second;
             ++j) {
          if (usage_manager[j] > *problem.usage_limit) {
            return false;
          }
        }
      }
      return true;
    } else if constexpr (std::is_same<UsageManager, UsageTree>::value == true) {
      return usage_manager.query() <= *problem.usage_limit;
    }
    return false;
  }

  inline bool change_some_nodes(Solution &solution, TotalCost &total_cost,
                                UsageManager &usage_manager, NodeIdx max_nodes,
                                double temperature) {
    std::vector<NodeIdx> node_indices;
    std::vector<StrategyIdx> strategy_indices;

    for (; random_i < problem.nodes.size() && random_flag[random_seq[random_i]];
         ++random_i)
      ;
    if (random_i == problem.nodes.size()) {
      std::shuffle(random_seq, random_seq + problem.nodes.size(), gen);
      std::memset(random_flag, 0, problem.nodes.size());
      random_i = 0;
    }
    node_indices.emplace_back(random_seq[random_i]);
    strategy_indices.emplace_back(-1);
    if (max_nodes == 2) {
      for (EdgeIdx i = 0; i < get_n_edges(node_indices[0]); ++i) {
        auto &edge = problem.edges[get_edge_idx(node_indices[0], i)];
        auto pair_node_idx =
            edge.nodes[0] != node_indices[0] ? edge.nodes[0] : edge.nodes[1];
        auto &pair_node = problem.nodes[pair_node_idx];
        auto flag = false;
        for (StrategyIdx j = 0; j < pair_node.strategies.size(); ++j) {
          if (get_edge_cost(edge, node_indices[0], pair_node_idx, 0, j) ==
              inf) {
            flag = true;
            break;
          }
        }
        if (flag == false) {
          node_indices.emplace_back(pair_node_idx);
          strategy_indices.emplace_back(-1);
          strategy_indices[0] = 0;
          break;
        }
      }
    }
    auto best_total_cost = absl::Int128Max();
    auto original_total_cost = total_cost;

    while (true) {
      NodeIdx i = node_indices.size() - 1;
      for (; i >= 0; --i) {
        if (strategy_indices[i] + 1 ==
            problem.nodes[node_indices[i]].strategies.size()) {
          strategy_indices[i] = 0;
        } else {
          ++strategy_indices[i];
          break;
        }
      }
      if (i < 0) {
        break;
      }
      for (NodeIdx j = 0; j < node_indices.size(); ++j) {
        if (modify_flag[node_indices[j]] == false) {
          modify_queue[modify_queue_tail++] = {node_indices[j],
                                               strategy_indices[j], -1};
          modify_flag[node_indices[j]] = true;
          modify(solution, total_cost, usage_manager);
        }
      }

      if (check_modify(usage_manager) == true && total_cost < best_total_cost &&
          total_cost != original_total_cost) {
        best_total_cost = total_cost;
        stage_modify(solution, total_cost);
      }
      for (NodeIdx j = 0; j < modify_queue_tail; ++j) { // TODO 优化？
        random_flag[std::get<0>(modify_queue[j])] = true;
      }
      discard_modify(solution, total_cost, usage_manager, 0);
    }

    if (best_total_cost < total_cost ||
        exp(double(total_cost - best_total_cost) / temperature * 100) >
            dis(gen)) {
      commit_modify(solution, total_cost, usage_manager);
      return true;
    }
    return false;
  }
};

absl::StatusOr<Solution> Solver::Solve(const Problem &problem,
                                       absl::Duration timeout) {
  const absl::Time start_time = absl::Now();
  std::optional<TotalCost> best_cost;
  std::optional<Solution> best_solution;

  // =================== Initialization ===========================

  // Find an inital solution
  Solution cur_sol_global;
  for (const auto &node : problem.nodes) {
    Strategy best;
    int idx = -1;
    for (int i = 0; i < (int)node.strategies.size(); i++) {
      const auto &strategy = node.strategies[i];
      if ((idx == -1) ||
          (strategy.usage < best.usage ||
           (strategy.usage == best.usage && strategy.cost < best.cost))) {
        idx = i;
        best = strategy;
      }
    }
    cur_sol_global.push_back(idx);
  }

  auto tmp_cost = Evaluate(problem, cur_sol_global);
  TotalCost cur_cost_global;
  if (!tmp_cost.ok()) {
    std::cout << "No Initial Solution!" << std::endl;
    return absl::NotFoundError("No solution found");
  } else {
    std::cout << "# Found initial solution ["
              << absl::StrJoin(cur_sol_global, ", ") << "] with cost "
              << *tmp_cost << std::endl;
    std::cerr << *tmp_cost << std::endl;
    best_cost = cur_cost_global = *tmp_cost;
    best_solution = cur_sol_global;
  }

  // Generate the Usage Table
  std::vector<Interval> new_interval;
  NodeIdx n = problem.nodes.size();
  std::vector<TotalUsage> cur_usage_global;

  {
    std::vector<TimeIdx> times;
    for (const auto &node : problem.nodes) {
      times.push_back(node.interval.first);
      times.push_back(node.interval.second);
    }
    std::sort(times.begin(), times.end());
    times.erase(std::unique(times.begin(), times.end()), times.end());

    for (const auto &node : problem.nodes) {
      auto [begin, end] = node.interval;
      begin =
          std::lower_bound(times.begin(), times.end(), begin) - times.begin();
      end = std::lower_bound(times.begin(), times.end(), end) - times.begin();

      new_interval.emplace_back(begin, end);
    }
    cur_usage_global.resize(times.size());
  }

  {
    for (NodeIdx i = 0; i < n; i++) {
      auto sel = cur_sol_global[i];
      auto usage = problem.nodes[i].strategies[sel].usage;
      auto [begin, end] = new_interval[i];
      for (auto j = begin; j < end; j++) {
        cur_usage_global[j] += usage;
      }
    }
  }

  // Helper function to get the edge cost
  auto get_edge_cost =
      [&problem = std::as_const(problem)](EdgeIdx eid, StrategyIdx u,
                                          StrategyIdx v, bool reverse = false) {
        const auto &edge = problem.edges[eid];
        StrategyIdx strategy_idx = 0;
        if (reverse)
          std::swap(u, v);
        std::vector<StrategyIdx> p = {u, v};
        auto idx = p.begin();
        for (const NodeIdx node_idx : edge.nodes) {
          strategy_idx *= problem.nodes[node_idx].strategies.size();
          strategy_idx += *(idx++);
        }
        return edge.strategies[strategy_idx].cost;
      };

  // Adj. List of the edges
  std::vector<std::vector<std::tuple<NodeIdx, EdgeIdx, bool>>> edges(n);
  for (EdgeIdx i = 0; i < problem.edges.size(); i++) {
    const auto &edge = problem.edges[i];
    if (edge.nodes.size() != 2) {
      std::cerr << "Invalid Input (not an edge)" << std::endl;
      return absl::InvalidArgumentError("# of node in an edge != 2");
    }
    auto u = edge.nodes[0];
    auto v = edge.nodes[1];
    if (u == v) {
      std::cerr << "Invalid Input (u == v)" << std::endl;
      return absl::InvalidArgumentError("u == v");
    }
    edges[u].emplace_back(v, i, 0);
    edges[v].emplace_back(u, i, 1);
  }
  for (NodeIdx i = 0; i < n; i++) {
    std::sort(edges[i].begin(), edges[i].end());
  }
  std::vector<std::pair<Cost, EdgeIdx>> edge_order_global;
  EdgeIdx last_edge = 0;
  {
    for (EdgeIdx i = 0; i < problem.edges.size(); i++) {
      auto u = problem.edges[i].nodes[0];
      auto v = problem.edges[i].nodes[1];
      auto c = get_edge_cost(i, cur_sol_global[u], cur_sol_global[v]);
      edge_order_global.emplace_back(c, i);
    }
    std::sort(edge_order_global.begin(), edge_order_global.end(),
              [](auto x, auto y) { return x.first > y.first; });
  }
  std::cerr << n << std::endl;

  std::cerr << absl::Now() - start_time << std::endl;

  Usage max_usage = problem.usage_limit ? *problem.usage_limit : INT64_MAX;

  std::cerr << "!!!!" << std::endl;
  std::cerr << absl::Now() - start_time << std::endl;

  std::mutex result_mutex;

  std::set<std::pair<NodeIdx, NodeIdx>> edge_set;
  for (const auto &e : problem.edges) {
    edge_set.emplace(e.nodes[0], e.nodes[1]);
    edge_set.emplace(e.nodes[1], e.nodes[1]);
  }

  const double c_inc_ratio_bfs = 2;
  const int c_max_ratio_bfs = 1e7;
  const int c_min_iter_bfs = 2;
  const int c_start_iter_bfs = 10;
  auto init_sol = cur_sol_global;
  auto init_cost = cur_cost_global;
  auto init_usage = cur_usage_global;

  // Get initial solution for reruns.
  auto gen_init_sol = [&problem = std::as_const(problem)](unsigned int &seed) {
    Solution sol;
    for (const auto &node : problem.nodes) {
      Strategy best;
      std::vector<NodeIdx> idx;
      for (int i = 0; i < (int)node.strategies.size(); i++) {
        const auto &strategy = node.strategies[i];
        if (idx.empty() || (strategy.usage < best.usage)) {
          idx.clear();
          idx.push_back(i);
          best = strategy;
        } else if (strategy.usage == best.usage) {
          idx.push_back(i);
        }
      }
      sol.push_back(idx[rand_r(&seed) % idx.size()]);
    }

    auto tmp_cost = Evaluate(problem, sol);
    return std::make_pair(sol, *tmp_cost);
  };
  std::cerr << cur_usage_global.size() << std::endl;
  std::cerr << edge_order_global.size() << std::endl;

  // Entrance Function for each thread.
  auto run = [=, &problem = std::as_const(problem), &result_mutex, &best_cost,
              &best_solution](int threadid) {
    int64_t iters = 0;
    auto last = absl::Now();
    auto edge_order = edge_order_global;
    EdgeIdx last_edge = 0;
    unsigned seed = 2025 + threadid;
    auto cur_sol = cur_sol_global;
    auto cur_cost = cur_cost_global;
    auto best_sol_local = cur_sol;
    auto best_cost_local = cur_cost;
    UsageTree cur_usage_tree(cur_usage_global);
    std::random_device rd;
    std::mt19937 g(rd());

    ChainedRestruct<UsageTree, EdgeList> chain_restruct(problem, edges,
                                                        new_interval);

    const int c_use_ratio_bfs = 5000;

    double f_init_temp = 50000;
    double f_degrade_fact = 1e-5;
    int f_use_ratio_bfs = c_use_ratio_bfs;
    int f_stop_iter_bfs = c_start_iter_bfs;

    double temperature = f_init_temp;

    auto recompute_usage = [&]() {
      cur_usage_tree = UsageTree(std::vector<TotalUsage>(n, 0));
      for (NodeIdx i = 0; i < n; i++) {
        auto sel = cur_sol[i];
        auto usage = problem.nodes[i].strategies[sel].usage;
        auto [begin, end] = new_interval[i];
        cur_usage_tree.update(begin, end - 1, usage);
      }
    };

    int64_t acc_iters = 0, acc_upds = 0;
    int64_t last_upds = 0, last_iters = 0;
    const int c_check_freq = 2;
    const int c_stop_thres = 2;
    auto last_check = absl::Now();
    auto last_cost = best_cost_local;
    std::cerr << threadid << "Ready" << std::endl;

    // ================== Main iteration Starts ==========================
    while (absl::Now() - start_time < timeout) {
      iters++;
      acc_iters++;
      const int X = 2;

      // ======================= Periodic checks & logging ===============
      if (absl::Now() - last > absl::Seconds(X)) {
        if (threadid == 0) {
          std::cerr << "[" << absl::Now() - start_time << "]" << std::endl;
          std::cerr << "Rate " << 1.0 * iters / X << "iters / s" << std::endl;
          iters = 0;
          last = absl::Now();

          TotalUsage musage = cur_usage_tree.query();
          std::cerr << "Usage: " << musage << " / " << max_usage << std::endl;
          std::cerr << "Temperature: " << temperature << std::endl;
          TotalCost x = 0;
          edge_order.clear();
          for (int i = 0; i < n; i++) {
            x += problem.nodes[i].strategies[cur_sol[i]].cost;
          }
          std::cerr << "Node Cost: " << x << std::endl;
          x = 0;
          for (int i = 0; i < problem.edges.size(); i++) {
            auto u = problem.edges[i].nodes[0];
            auto v = problem.edges[i].nodes[1];
            auto c = get_edge_cost(i, cur_sol[u], cur_sol[v]);
            x += c;
            edge_order.emplace_back(c, i);
          }
          std::sort(edge_order.begin(), edge_order.end(),
                    [](auto x, auto y) { return x.first > y.first; });
          last_edge = 0;
          std::cerr << "Edge Cost: " << x << std::endl;
        } else {
          edge_order.clear();
          for (int i = 0; i < problem.edges.size(); i++) {
            auto u = problem.edges[i].nodes[0];
            auto v = problem.edges[i].nodes[1];
            auto c = get_edge_cost(i, cur_sol[u], cur_sol[v]);
            edge_order.emplace_back(c, i);
          }
          std::sort(edge_order.begin(), edge_order.end(),
                    [](auto x, auto y) { return x.first > y.first; });
          last_edge = 0;
        }
      }

      // =================== Checking Reruns & update globals ============
      if (absl::Now() - last_check > absl::Seconds(c_check_freq) &&
          (absl::Now() - start_time < timeout - absl::Seconds(c_stop_thres))) {
        if ((acc_upds - last_upds) == 0 ||
            ((last_cost - best_cost_local) < best_cost_local / 1000 &&
             (acc_upds - last_upds) < 10)) {
          {
            std::lock_guard<std::mutex> g(result_mutex);
            if (best_cost_local < *best_cost) {
              best_cost = best_cost_local;
              best_solution = best_sol_local;
            }
          }
          std::cerr << threadid << " Restart " << best_cost_local << "!"
                    << std::endl;
          std::tie(best_sol_local, best_cost_local) =
              std::tie(cur_sol, cur_cost) = gen_init_sol(seed);
          temperature = f_init_temp;
          recompute_usage();
          last_check = absl::Now();

        } else {
          last_check = absl::Now();
        }
        last_upds = acc_upds;
        last_iters = acc_iters;
        last_cost = best_cost_local;
      }
      const Solution &solution = cur_sol;

      // =================== Technique #2: Greedy Restructuring =======
      if (rand_r(&seed) % f_use_ratio_bfs == 0) {
        std::vector<std::pair<TotalCost, NodeIdx>> total_costs;
        for (int i = 0; i < n; i++) {
          TotalCost c = 0;
          for (auto [v, e, rev] : edges[i]) {
            c += get_edge_cost(e, cur_sol[i], cur_sol[v], rev);
          }
          total_costs.emplace_back(c, i);
        }
        std::sort(total_costs.begin(), total_costs.end(),
                  [](auto x, auto y) { return x.first > y.first; });
        bool found = false;
        int i = 0;
        int upds = 0;
        int fails = 0;
        for (auto [cost, node_u] : total_costs) {
          i++;
          if (i > 1 && cost < best_cost_local / 1000)
            break;
          auto alters = edges[node_u];
          std::shuffle(alters.begin(), alters.end(), g);

          bool stop = false;
          for (StrategyIdx i = 0; i < problem.nodes[node_u].strategies.size();
               i++) {
            auto neo_ctx = ExtendedProblem(problem, edges, new_interval,
                                           cur_sol, cur_usage_tree, cur_cost);
            neo_ctx.update(node_u, i);
            std::queue<NodeIdx> q;
            std::set<NodeIdx> visited;
            q.push(node_u);
            visited.insert(node_u);
            auto iters = 0;
            auto last = cur_cost;
            while (!q.empty()) {
              auto cur = q.front();
              q.pop();
              auto s = neo_ctx.get_strategy(cur);
              for (auto [v, e, rev] : edges[cur]) {
                if (visited.count(v))
                  continue;
                StrategyIdx x = neo_ctx.get_strategy(v);
                for (StrategyIdx j = 0; j < problem.nodes[v].strategies.size();
                     j++) {
                  if (std::make_pair(get_edge_cost(e, s, j, rev),
                                     problem.nodes[v].strategies[j].usage) <
                      std::make_pair(get_edge_cost(e, s, x, rev),
                                     problem.nodes[v].strategies[x].usage)) {
                    x = j;
                  }
                }
                if (x != neo_ctx.get_strategy(v)) {
                  neo_ctx.update(v, x);
                  q.push(v);
                  visited.insert(v);
                }
              }

              auto result = neo_ctx.get_cost();
              if (result.has_value() && *result < best_cost_local) {
                best_cost_local = *result;
                best_sol_local = neo_ctx.get_solution();
                cur_cost = *result;
                neo_ctx.commit();
                found = true;
                ++upds;
                ++acc_upds;
              }
              if (!result.has_value())
                break;
              if (++iters > f_stop_iter_bfs && *result >= last)
                break;
              if (iters > 2 * f_stop_iter_bfs && *result > best_cost_local)
                break;
              last = *result;
              if (upds > 10 ||
                  absl::Now() - start_time > timeout - absl::Seconds(10)) {
                stop = true;
                break;
              }
            }

            if (stop)
              break;
          }
          if (stop || upds > 10 || i > 5)
            break;
        }
        if (upds == 0) {
          f_use_ratio_bfs = std::min(c_max_ratio_bfs,
                                     (int)(f_use_ratio_bfs * c_inc_ratio_bfs));
        } else {
          f_use_ratio_bfs = c_use_ratio_bfs;
          f_stop_iter_bfs = c_start_iter_bfs;
        }
      }

      // =================== Edge Refinement ==========================
      int x = rand_r(&seed) % 10;
      // Alternate one edge
      if (true || x == 0) {
        EdgeIdx eid = rand_r(&seed) % problem.edges.size();
        {
          if (last_edge < edge_order.size()) {
            auto [c, e] = edge_order[last_edge++];

            if (static_cast<TotalCost>(c) * n / 10 >= best_cost_local) {
              eid = e;
            }
          }
        }
        auto node_u = problem.edges[eid].nodes[0];
        auto node_v = problem.edges[eid].nodes[1];
        auto cur_u = cur_sol[node_u];
        auto cur_v = cur_sol[node_v];
        const auto &data_u = problem.nodes[node_u];
        const auto &data_v = problem.nodes[node_v];

        TotalUsage usage_u = 0, usage_v = 0, usage_uv = 0;
        auto [begin_u, end_u] = new_interval[node_u];
        auto [begin_v, end_v] = new_interval[node_v];
        auto get_usage = [&cur_usage_tree](TimeIdx begin_u, TimeIdx end_u,
                                           TimeIdx begin_v, TimeIdx end_v) {
          if (end_u <= begin_v || end_v <= begin_u) {
            return cur_usage_tree.query(begin_u, end_u - 1);
          } else if (begin_v <= begin_u && end_v >= end_u) {
            return TotalUsage(0);
          } else if (begin_v > begin_u && end_v < end_u) {

            return std::max(cur_usage_tree.query(begin_u, begin_v - 1),
                            cur_usage_tree.query(end_v, end_u - 1));
          } else if (begin_v <= begin_u) {
            return cur_usage_tree.query(end_v, end_u - 1);
          } else if (end_v >= end_u) {
            return cur_usage_tree.query(begin_u, begin_v - 1);
          }
          return TotalUsage(0);
        };
        usage_u = get_usage(begin_u, end_u, begin_v, end_v);
        usage_v = get_usage(begin_v, end_v, begin_u, end_u);
        usage_uv = cur_usage_tree.query(std::max(begin_u, begin_v),
                                        std::min(end_u, end_v) - 1);
        usage_uv = max_usage - usage_uv + data_u.strategies[cur_u].usage +
                   data_v.strategies[cur_v].usage;
        usage_u = max_usage - usage_u + data_u.strategies[cur_u].usage;
        usage_v = max_usage - usage_v + data_v.strategies[cur_v].usage;

        TotalCost cand_cost = absl::Int128Max();
        StrategyIdx sel_u = -1, sel_v = -1;
        for (StrategyIdx i = 0; i < data_u.strategies.size(); i++) {
          if (data_u.strategies[i].usage <= usage_u) {
            for (StrategyIdx j = 0; j < data_v.strategies.size(); j++) {
              if (data_v.strategies[j].usage <= usage_v &&
                  data_u.strategies[i].usage + data_v.strategies[j].usage <=
                      usage_uv) {
                TotalCost tot_delta =
                    get_edge_cost(eid, i, j) - get_edge_cost(eid, cur_u, cur_v);
                tot_delta +=
                    data_u.strategies[i].cost - data_u.strategies[cur_u].cost;
                tot_delta +=
                    data_v.strategies[j].cost - data_v.strategies[cur_v].cost;
                for (auto [w, e, rev] : edges[node_u]) {
                  if (e == eid)
                    continue;
                  if (w == node_v) {
                    tot_delta += get_edge_cost(e, i, j, rev) -
                                 get_edge_cost(e, cur_u, cur_v, rev);
                  } else {
                    tot_delta += get_edge_cost(e, i, cur_sol[w], rev) -
                                 get_edge_cost(e, cur_u, cur_sol[w], rev);
                  }
                }
                for (auto [w, e, rev] : edges[node_v]) {
                  if (w == node_u)
                    continue;
                  tot_delta += get_edge_cost(e, j, cur_sol[w], rev) -
                               get_edge_cost(e, cur_v, cur_sol[w], rev);
                  // if (debug) std::cerr << tot_delta << "!" << w << "\n";
                }
                if (sel_u == -1 || tot_delta < cand_cost) {
                  cand_cost = tot_delta;
                  sel_u = i, sel_v = j;
                }
              }
            }
          }
        }

        bool found = false;

        if (sel_u != -1) {
          bool accept =
              cand_cost < 0 ||
              rand_r(&seed) < RAND_MAX * exp(-double(cand_cost) / temperature);
          ;
          if (accept) {
            cur_sol[node_u] = sel_u;
            cur_sol[node_v] = sel_v;
            cur_cost += cand_cost;
            cur_usage_tree.update(begin_u, end_u - 1,
                                  data_u.strategies[sel_u].usage -
                                      data_u.strategies[cur_u].usage);
            cur_usage_tree.update(begin_v, end_v - 1,
                                  data_v.strategies[sel_v].usage -
                                      data_v.strategies[cur_v].usage);
            found = true;
          }
        }

        if (found) {
          if (cur_cost < best_cost_local) {
            best_cost_local = cur_cost;
            best_sol_local = cur_sol;
            ++acc_upds;
          }
        }
      }

      // =================== Invoke `Chain Resturcturing` edgewise ====
      if (true) {
        if (chain_restruct.change_some_nodes(cur_sol, cur_cost, cur_usage_tree,
                                             2, temperature)) {

          if (cur_cost < best_cost_local) {
            best_cost_local = cur_cost;
            best_sol_local = cur_sol;
            ++acc_upds;
          }
        }
      }

      // =================== Node Refinement ==========================
      if (true) {
        // Alternate one node
        NodeIdx idx = rand_r(&seed) % problem.nodes.size();

        // Find the maximum usage increase
        auto &cur_node = problem.nodes[idx];
        auto sel = cur_sol[idx];
        auto cur_node_usage = cur_node.strategies[sel].usage;
        auto cur_node_cost = cur_node.strategies[sel].cost;
        auto [begin, end] = new_interval[idx];

        // Find all strategies
        TotalCost max_inc_usage =
            begin == end ? INT64_MAX
                         : max_usage - cur_usage_tree.query(begin, end - 1);

        std::vector<std::pair<TotalCost, Usage>> costs;

        for (StrategyIdx i = 0; i < cur_node.strategies.size(); i++) {
          TotalCost cost = cur_node.strategies[i].cost;
          for (auto [v, eid, rev] : edges[idx]) {
            cost += get_edge_cost(eid, i, solution[v], rev);
          }
          costs.emplace_back(cost, cur_node.strategies[i].usage);
        }

        // Option 1: Select the min cost with Max usage < constraint
        StrategyIdx cand_sel = -1;
        TotalCost min_cost, cur_tot_cost = costs[sel].first;
        for (StrategyIdx i = 0; i < cur_node.strategies.size(); i++) {
          if (i == sel)
            continue;
          auto [cost, usage] = costs[i];
          if ((usage - cur_node_usage <= max_inc_usage) &&
              (cand_sel == -1 || cost < min_cost)) {
            cand_sel = i;
            min_cost = cost;
          }
        }

        bool found = false;
        if (cand_sel != -1) {
          bool accept =
              min_cost < cur_tot_cost ||
              rand_r(&seed) < RAND_MAX * exp(-double(min_cost - cur_tot_cost) /
                                             temperature);
          if (accept) {
            cur_sol[idx] = cand_sel;
            cur_cost += min_cost - cur_tot_cost;
            cur_usage_tree.update(begin, end - 1,
                                  costs[cand_sel].second - cur_node_usage);
            found = true;
          }
        }

        if (found) {
          if (cur_cost < best_cost_local) {
            best_cost_local = cur_cost;
            best_sol_local = cur_sol;
            ++acc_upds;
          }
        }
      }

      // Use Simulated Annealing to update
      temperature *= (1 - f_degrade_fact);
      if (temperature < 1) {
        temperature = f_init_temp;
      }
    }
    // Final update of the result
    {
      std::lock_guard<std::mutex> g(result_mutex);
      if (best_cost_local < *best_cost) {
        best_cost = best_cost_local;
        best_solution = best_sol_local;
      }
    }
    std::cerr << threadid << "End" << std::endl;
  };
  std::vector<std::thread> pool;
  for (int i = 1; i < 8; i++)
    pool.emplace_back(run, i);
  run(0);
  for (auto &t : pool)
    t.join();
  std::cerr << absl::Now() - start_time << std::endl;
  return *best_solution;
}

} // namespace iopddl