main.cpp

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <csignal>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <functional>
#include <iostream>
#include <limits>
#include <memory>
#include <numeric>
#include <optional>
#include <random>
#include <sstream>
#include <stdexcept>
#include <string>
#include <string_view>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#if defined(__unix__) || defined(__APPLE__)
#include <sys/resource.h>
#endif

struct SolverParameters {
  uint64_t seed = 42;
  size_t default_max_generations = 1'000'000'000ULL;
  size_t swap_stagnation_threshold = 1'000;
  double swap_probability = 0.1;
  size_t perturbation_stagnation_threshold = 5'000;
  double perturbation_probability = 0.1;

  size_t interrupted_fitness_penalty = std::numeric_limits<size_t>::max();
  size_t initial_best_fitness = std::numeric_limits<size_t>::max();

  double output_reserve_seconds = 4.0;

  size_t component_cache_memory_mib = 256;
};

volatile std::sig_atomic_t running = 1;
void signal_handler(int) { running = 0; }

std::chrono::steady_clock::time_point internal_deadline{};
bool internal_deadline_enabled = false;
bool internal_deadline_reached = false;

void configure_internal_deadline(
    std::chrono::steady_clock::time_point process_start, double timeout_seconds,
    double reserve_seconds) {
  if (timeout_seconds <= 0.0) {
    internal_deadline_enabled = false;
    return;
  }
  const double effective_reserve =
      std::min(reserve_seconds, timeout_seconds * 0.20);
  const double usable_seconds =
      std::max(0.0, timeout_seconds - effective_reserve);
  internal_deadline =
      process_start +
      std::chrono::duration_cast<std::chrono::steady_clock::duration>(
          std::chrono::duration<double>(usable_seconds));
  internal_deadline_enabled = true;
}

inline bool check_internal_deadline_now() {
  if (!internal_deadline_enabled || internal_deadline_reached) {
    return !internal_deadline_reached;
  }
  if (std::chrono::steady_clock::now() >= internal_deadline) {
    internal_deadline_reached = true;
    running = 0;
    return false;
  }
  return true;
}

inline bool keep_running() {
  if (running == 0) {
    return false;
  }
  // Avoid reading the clock on every inner-loop check.
  static uint32_t checks = 0;
  if (internal_deadline_enabled && ((++checks & 0x3ffU) == 0U)) {
    return check_internal_deadline_now();
  }
  return true;
}

class ExactComponentCompatibilityCache {
 public:
  struct Value {
    bool structures_match = false;
    uint32_t canonicalization_calls = 0;
    const size_t* lcas = nullptr;
  };

  ExactComponentCompatibilityCache(size_t memory_mib, size_t tree_count)
      : tree_count_(tree_count),
        budget_bytes_(static_cast<uint64_t>(memory_mib) * 1024ULL * 1024ULL) {
    if (memory_mib == 0) {
      throw std::runtime_error("component_cache_memory_mib must be positive");
    }
    if (tree_count_ != 2) {
      throw std::runtime_error(
          "component cache requires exactly two input trees");
    }

    constexpr uint64_t kExpectedLabelsPerSlot = 8;
    const uint64_t expected_payload_values =
        kExpectedLabelsPerSlot + static_cast<uint64_t>(tree_count_);
    const uint64_t bytes_per_planned_slot =
        sizeof(Slot) + expected_payload_values * sizeof(size_t);
    uint64_t raw_capacity = budget_bytes_ / bytes_per_planned_slot;
    if (raw_capacity < 16) {
      throw std::runtime_error("component cache memory budget is too small");
    }

    slot_capacity_ = 1;
    while (slot_capacity_ <= raw_capacity / 2) {
      slot_capacity_ *= 2;
    }
    const uint64_t slot_bytes =
        static_cast<uint64_t>(slot_capacity_) * sizeof(Slot);
    if (slot_bytes >= budget_bytes_) {
      throw std::runtime_error(
          "component cache memory budget cannot hold payload storage");
    }

    payload_capacity_ =
        static_cast<size_t>((budget_bytes_ - slot_bytes) / sizeof(size_t));
    max_entries_ = std::max<size_t>(1, slot_capacity_ * 7 / 10);
    slots_ = std::make_unique<Slot[]>(slot_capacity_);
    payload_ = std::make_unique<size_t[]>(payload_capacity_);
  }

  std::optional<Value> lookup(const std::vector<size_t>& labels) const {
    const uint64_t hash = hash_labels(labels);
    size_t index = static_cast<size_t>(hash) & (slot_capacity_ - 1);
    for (size_t probe = 0; probe < slot_capacity_; ++probe) {
      const Slot& slot = slots_[index];
      if (slot.hash == 0) {
        return std::nullopt;
      }
      if (slot.hash == hash) {
        if (equals(slot, labels)) {
          return Value{slot.structures_match != 0, slot.canonicalization_calls,
                       payload_.get() + slot.payload_offset + slot.label_count};
        }
      }
      index = (index + 1) & (slot_capacity_ - 1);
    }

    return std::nullopt;
  }

  void insert(const std::vector<size_t>& labels,
              const std::vector<size_t>& lcas, Value value) {
    if (lcas.size() != tree_count_) {
      throw std::runtime_error(
          "component cache received an invalid LCA vector");
    }
    const size_t required_values = labels.size() + tree_count_;
    if (entries_ >= max_entries_ ||
        required_values > payload_capacity_ - used_payload_values_) {
      return;
    }
    if (labels.size() > std::numeric_limits<uint32_t>::max() ||
        used_payload_values_ > std::numeric_limits<uint32_t>::max()) {
      return;
    }

    const uint64_t hash = hash_labels(labels);
    size_t index = static_cast<size_t>(hash) & (slot_capacity_ - 1);
    for (size_t probe = 0; probe < slot_capacity_; ++probe) {
      Slot& slot = slots_[index];
      if (slot.hash == 0) {
        const uint32_t offset = static_cast<uint32_t>(used_payload_values_);
        std::copy(labels.begin(), labels.end(),
                  payload_.get() + used_payload_values_);
        used_payload_values_ += labels.size();
        std::copy(lcas.begin(), lcas.end(),
                  payload_.get() + used_payload_values_);
        used_payload_values_ += lcas.size();

        slot.hash = hash;
        slot.payload_offset = offset;
        slot.label_count = static_cast<uint32_t>(labels.size());
        slot.canonicalization_calls = value.canonicalization_calls;
        if (value.structures_match) {
          slot.structures_match = 1U;
        } else {
          slot.structures_match = 0U;
        }
        ++entries_;
        return;
      }

      if (slot.hash == hash && equals(slot, labels)) {
        return;
      }
      index = (index + 1) & (slot_capacity_ - 1);
    }
  }

 private:
  struct Slot {
    uint64_t hash = 0;
    uint32_t payload_offset = 0;
    uint32_t label_count = 0;
    uint32_t canonicalization_calls = 0;
    uint8_t structures_match = 0;
    std::array<uint8_t, 3> padding{};
  };
  static uint64_t mix64(uint64_t value) {
    value ^= value >> 30;
    value *= 0xbf58476d1ce4e5b9ULL;
    value ^= value >> 27;
    value *= 0x94d049bb133111ebULL;
    value ^= value >> 31;
    return value;
  }

  static uint64_t hash_labels(const std::vector<size_t>& labels) {
    uint64_t hash = mix64(labels.size() + 0x9e3779b97f4a7c15ULL);
    for (size_t label : labels) {
      hash = mix64(hash ^
                   mix64(static_cast<uint64_t>(label) + 0x9e3779b97f4a7c15ULL));
    }
    if (hash == 0) {
      return 1;
    }
    return hash;
  }

  bool equals(const Slot& slot, const std::vector<size_t>& labels) const {
    if (slot.label_count != labels.size()) {
      return false;
    }
    return std::equal(labels.begin(), labels.end(),
                      payload_.get() + slot.payload_offset);
  }

  size_t tree_count_ = 0;
  uint64_t budget_bytes_ = 0;
  size_t slot_capacity_ = 0;
  size_t payload_capacity_ = 0;
  size_t max_entries_ = 0;
  size_t entries_ = 0;
  size_t used_payload_values_ = 0;
  std::unique_ptr<Slot[]> slots_;
  std::unique_ptr<size_t[]> payload_;
};
struct CanonicalKey {
  enum class Kind : uint8_t { Leaf, Branch };

  Kind kind = Kind::Leaf;
  size_t first = 0;
  size_t second = 0;

  bool operator==(const CanonicalKey&) const = default;
};

struct CanonicalKeyHash {
  size_t operator()(const CanonicalKey& key) const {
    size_t hash = static_cast<size_t>(key.kind);
    hash ^= key.first + 0x9e3779b97f4a7c15ULL + (hash << 6) + (hash >> 2);
    hash ^= key.second + 0x9e3779b97f4a7c15ULL + (hash << 6) + (hash >> 2);
    return hash;
  }
};

class StructuralInterner {
 public:
  explicit StructuralInterner(size_t reserve_hint = 0) {
    if (reserve_hint > 0) {
      ids_.reserve(reserve_hint);
    }
  }

  size_t leaf(size_t label) {
    return intern({CanonicalKey::Kind::Leaf, label, 0});
  }

  size_t branch(size_t left, size_t right) {
    if (left > right) {
      std::swap(left, right);
    }
    return intern({CanonicalKey::Kind::Branch, left, right});
  }

 private:
  size_t intern(const CanonicalKey& key) {
    auto [iterator, inserted] = ids_.try_emplace(key, next_id_);
    if (inserted) {
      ++next_id_;
    }
    return iterator->second;
  }

  std::unordered_map<CanonicalKey, size_t, CanonicalKeyHash> ids_;
  size_t next_id_ = 1;  // 0 denotes an empty restricted subtree.
};

struct Node {
  std::optional<size_t> label;
  std::vector<size_t> children;
  std::optional<size_t> parent;
  std::optional<size_t> edge_idx;
};

class Tree {
 public:
  std::vector<Node> nodes;
  size_t root = 0;
  size_t n_edges = 0;
  size_t max_leaf_label = 0;
  std::unordered_map<size_t, size_t> leaf_to_idx;
  std::vector<size_t> leaf_index_by_label;
  std::vector<size_t> depths;
  std::vector<std::vector<size_t>> up;
  std::vector<size_t> subtree_leaf_counts;
  std::vector<size_t> preorder_index;
  std::vector<size_t> subtree_preorder_end;

  size_t get_subtree_leaf_count(size_t node_idx) const {
    return subtree_leaf_counts[node_idx];
  }

  // Parse Newick iteratively and ignore branch lengths and internal
  // annotations.
  static Tree parse_newick(std::string text) {
    auto trim = [](std::string& value) {
      const size_t first = value.find_first_not_of(" \t\n\r");
      if (first == std::string::npos) {
        value.clear();
        return;
      }
      const size_t last = value.find_last_not_of(" \t\n\r");
      value = value.substr(first, last - first + 1);
    };

    trim(text);
    if (!text.empty() && text.back() == ';') {
      text.pop_back();
    }
    trim(text);
    if (text.empty()) {
      throw std::runtime_error("Empty Newick tree");
    }

    std::vector<Node> parsed_nodes;
    std::vector<size_t> parent_stack;
    std::optional<size_t> parsed_root;

    auto attach = [&](size_t node_index) {
      if (parent_stack.empty()) {
        if (parsed_root.has_value()) {
          throw std::runtime_error("Multiple roots in Newick tree");
        }
        parsed_root = node_index;
      } else {
        const size_t parent = parent_stack.back();
        parsed_nodes[parent].children.push_back(node_index);
        parsed_nodes[node_index].parent = parent;
      }
    };

    size_t i = 0;
    while (i < text.size()) {
      if (!keep_running()) {
        break;
      }
      const char current = text[i];
      if (current == ' ' || current == '\t' || current == '\n' ||
          current == '\r' || current == ',') {
        ++i;
        continue;
      }
      if (current == '(') {
        const size_t index = parsed_nodes.size();
        parsed_nodes.push_back({std::nullopt, {}, std::nullopt, std::nullopt});
        attach(index);
        parent_stack.push_back(index);
        ++i;
        continue;
      }
      if (current == ')') {
        if (parent_stack.empty()) {
          throw std::runtime_error("Unmatched ')' in Newick tree");
        }
        parent_stack.pop_back();
        ++i;
        // Ignore an optional internal-node annotation or branch length.
        while (i < text.size() && text[i] != ',' && text[i] != ')' &&
               text[i] != ';') {
          ++i;
        }
        continue;
      }
      if (current == ';') {
        ++i;
        continue;
      }

      const size_t label_start = i;
      while (i < text.size() && text[i] != ':' && text[i] != ',' &&
             text[i] != ')' && text[i] != '(' && text[i] != ';' &&
             text[i] != ' ' && text[i] != '\t' && text[i] != '\n' &&
             text[i] != '\r') {
        ++i;
      }
      if (i == label_start) {
        throw std::runtime_error("Malformed leaf token in Newick tree");
      }
      const std::string label_text = text.substr(label_start, i - label_start);
      size_t consumed = 0;
      const size_t label = std::stoull(label_text, &consumed);
      if (consumed != label_text.size()) {
        throw std::runtime_error("Non-numeric leaf label: " + label_text);
      }

      const size_t leaf_index = parsed_nodes.size();
      parsed_nodes.push_back({label, {}, std::nullopt, std::nullopt});
      attach(leaf_index);

      while (i < text.size() && (text[i] == ' ' || text[i] == '\t' ||
                                 text[i] == '\n' || text[i] == '\r')) {
        ++i;
      }
      if (i < text.size() && text[i] == ':') {
        ++i;
        while (i < text.size() && text[i] != ',' && text[i] != ')' &&
               text[i] != ';') {
          ++i;
        }
      }
    }

    if (!keep_running() && !parsed_root.has_value()) {
      throw std::runtime_error("Interrupted before a tree could be parsed");
    }
    if (!parent_stack.empty()) {
      throw std::runtime_error("Unclosed '(' in Newick tree");
    }
    if (!parsed_root.has_value()) {
      throw std::runtime_error("Newick tree contains no root");
    }

    Tree tree;
    tree.nodes = std::move(parsed_nodes);
    tree.root = *parsed_root;
    tree.n_edges = tree.assign_edge_indices();
    tree.precompute();
    return tree;
  }

  size_t assign_edge_indices() {
    if (nodes.empty() || root >= nodes.size()) {
      throw std::runtime_error("Invalid tree root");
    }
    for (auto& node : nodes) {
      node.parent.reset();
      node.edge_idx.reset();
    }

    size_t edge_index = 0;
    std::vector<size_t> stack = {root};
    std::vector<uint8_t> seen(nodes.size(), 0);
    seen[root] = 1;
    while (!stack.empty()) {
      const size_t node_index = stack.back();
      stack.pop_back();
      if (node_index != root) {
        nodes[node_index].edge_idx = edge_index++;
      }
      for (auto iterator = nodes[node_index].children.rbegin();
           iterator != nodes[node_index].children.rend(); ++iterator) {
        const size_t child = *iterator;
        if (child >= nodes.size()) {
          throw std::runtime_error("Tree contains an invalid child index");
        }
        if (seen[child]) {
          throw std::runtime_error("Tree is cyclic or has multiple parents");
        }
        seen[child] = 1;
        nodes[child].parent = node_index;
        stack.push_back(child);
      }
    }
    if (std::count(seen.begin(), seen.end(), uint8_t{1}) !=
        static_cast<std::ptrdiff_t>(nodes.size())) {
      throw std::runtime_error("Tree contains unreachable nodes");
    }
    nodes[root].parent.reset();
    nodes[root].edge_idx.reset();
    return edge_index;
  }

  void precompute() {
    const size_t node_count = nodes.size();
    if (node_count == 0) {
      throw std::runtime_error("Cannot precompute empty tree");
    }

    leaf_to_idx.clear();
    max_leaf_label = 0;
    for (size_t i = 0; i < node_count; ++i) {
      if (!nodes[i].label.has_value()) {
        continue;
      }
      const size_t label = *nodes[i].label;
      if (label == 0) {
        throw std::runtime_error("Leaf labels must be positive");
      }
      if (!leaf_to_idx.emplace(label, i).second) {
        throw std::runtime_error("Duplicate leaf label: " +
                                 std::to_string(label));
      }
      max_leaf_label = std::max(max_leaf_label, label);
    }

    const size_t missing = std::numeric_limits<size_t>::max();
    leaf_index_by_label.assign(max_leaf_label + 1, missing);
    for (const auto& [label, node_index] : leaf_to_idx) {
      leaf_index_by_label[label] = node_index;
    }

    depths.assign(node_count, 0);
    size_t log_n = 1;
    while (log_n < std::numeric_limits<size_t>::digits &&
           (size_t{1} << log_n) <= node_count) {
      ++log_n;
    }
    up.assign(node_count, std::vector<size_t>(log_n, root));

    std::vector<size_t> stack = {root};
    std::vector<size_t> preorder;
    preorder.reserve(node_count);
    while (!stack.empty()) {
      const size_t node_index = stack.back();
      stack.pop_back();
      preorder.push_back(node_index);
      for (auto iterator = nodes[node_index].children.rbegin();
           iterator != nodes[node_index].children.rend(); ++iterator) {
        const size_t child = *iterator;
        depths[child] = depths[node_index] + 1;
        up[child][0] = node_index;
        for (size_t level = 1; level < log_n; ++level) {
          up[child][level] = up[up[child][level - 1]][level - 1];
        }
        stack.push_back(child);
      }
    }

    preorder_index.assign(node_count, 0);
    for (size_t position = 0; position < preorder.size(); ++position) {
      preorder_index[preorder[position]] = position;
    }

    subtree_leaf_counts.assign(node_count, 0);
    std::vector<size_t> subtree_node_counts(node_count, 1);
    for (auto iterator = preorder.rbegin(); iterator != preorder.rend();
         ++iterator) {
      const size_t node_index = *iterator;
      if (nodes[node_index].label.has_value()) {
        subtree_leaf_counts[node_index] = 1;
      } else {
        size_t leaf_count = 0;
        size_t node_total = 1;
        for (size_t child : nodes[node_index].children) {
          leaf_count += subtree_leaf_counts[child];
          node_total += subtree_node_counts[child];
        }
        subtree_leaf_counts[node_index] = leaf_count;
        subtree_node_counts[node_index] = node_total;
      }
    }
    subtree_preorder_end.assign(node_count, 0);
    for (size_t node_index = 0; node_index < node_count; ++node_index) {
      subtree_preorder_end[node_index] =
          preorder_index[node_index] + subtree_node_counts[node_index];
    }
  }

  bool is_ancestor(size_t ancestor, size_t descendant) const {
    return preorder_index[ancestor] <= preorder_index[descendant] &&
           preorder_index[descendant] < subtree_preorder_end[ancestor];
  }

  // Compute exact full-subtree canonical IDs for reduction checks.
  std::vector<size_t> full_canonical_ids(StructuralInterner& interner) const {
    std::vector<size_t> order(nodes.size());
    std::iota(order.begin(), order.end(), 0);
    std::sort(order.begin(), order.end(), [&](size_t lhs, size_t rhs) {
      return preorder_index[lhs] > preorder_index[rhs];
    });

    std::vector<size_t> ids(nodes.size(), 0);
    for (size_t node_index : order) {
      const Node& node = nodes[node_index];
      if (node.label.has_value()) {
        ids[node_index] = interner.leaf(*node.label);
        continue;
      }
      if (node.children.size() != 2) {
        throw std::runtime_error("Full canonicalization requires binary trees");
      }
      ids[node_index] =
          interner.branch(ids[node.children[0]], ids[node.children[1]]);
    }
    return ids;
  }

  std::vector<size_t> get_subtree_leaf_labels(size_t node_index) const {
    std::vector<size_t> labels;
    labels.reserve(get_subtree_leaf_count(node_index));
    std::vector<size_t> stack = {node_index};
    while (!stack.empty()) {
      const size_t current = stack.back();
      stack.pop_back();
      const Node& node = nodes[current];
      if (node.label.has_value()) {
        labels.push_back(*node.label);
      } else {
        for (auto iterator = node.children.rbegin();
             iterator != node.children.rend(); ++iterator) {
          stack.push_back(*iterator);
        }
      }
    }
    return labels;
  }

  size_t get_lca(size_t first, size_t second) const {
    if (depths[first] < depths[second]) {
      std::swap(first, second);
    }
    size_t difference = depths[first] - depths[second];
    const size_t log_n = up[0].size();
    for (size_t level = 0; level < log_n; ++level) {
      if ((difference >> level) & 1U) {
        first = up[first][level];
      }
    }
    if (first == second) {
      return first;
    }
    for (size_t level = log_n; level-- > 0;) {
      if (up[first][level] != up[second][level]) {
        first = up[first][level];
        second = up[second][level];
      }
    }
    return up[first][0];
  }

  size_t get_lca_of_set(const std::vector<size_t>& leaves) const {
    if (leaves.empty()) {
      throw std::runtime_error("LCA of empty leaf set");
    }
    size_t current = find_leaf_idx(leaves.front());
    for (size_t i = 1; i < leaves.size(); ++i) {
      current = get_lca(current, find_leaf_idx(leaves[i]));
    }
    return current;
  }

  size_t find_leaf_idx(size_t label) const {
    const size_t missing = std::numeric_limits<size_t>::max();
    if (label >= leaf_index_by_label.size() ||
        leaf_index_by_label[label] == missing) {
      throw std::runtime_error("Leaf missing: " + std::to_string(label));
    }
    return leaf_index_by_label[label];
  }

  static Tree compact(std::vector<Node> source_nodes, size_t source_root) {
    if (source_root >= source_nodes.size()) {
      throw std::runtime_error("Invalid compact root");
    }

    std::vector<size_t> order;
    std::vector<size_t> stack = {source_root};
    while (!stack.empty()) {
      const size_t old_index = stack.back();
      stack.pop_back();
      order.push_back(old_index);
      for (auto iterator = source_nodes[old_index].children.rbegin();
           iterator != source_nodes[old_index].children.rend(); ++iterator) {
        stack.push_back(*iterator);
      }
    }

    std::vector<size_t> old_to_new(source_nodes.size(),
                                   std::numeric_limits<size_t>::max());
    std::vector<Node> compact_nodes;
    compact_nodes.reserve(order.size());
    for (size_t old_index : order) {
      old_to_new[old_index] = compact_nodes.size();
      Node copy = source_nodes[old_index];
      copy.parent.reset();
      copy.edge_idx.reset();
      compact_nodes.push_back(std::move(copy));
    }
    for (Node& node : compact_nodes) {
      for (size_t& child : node.children) {
        if (child >= old_to_new.size() ||
            old_to_new[child] == std::numeric_limits<size_t>::max()) {
          throw std::runtime_error("Compaction retained an unreachable child");
        }
        child = old_to_new[child];
      }
    }

    Tree tree;
    tree.nodes = std::move(compact_nodes);
    tree.root = 0;
    tree.n_edges = tree.assign_edge_indices();
    tree.precompute();
    return tree;
  }

  Tree prune_and_compact(const std::vector<size_t>& to_remove) const {
    std::vector<uint8_t> removed(nodes.size(), 0);
    for (size_t label : to_remove) {
      auto iterator = leaf_to_idx.find(label);
      if (iterator != leaf_to_idx.end()) {
        removed[iterator->second] = 1;
      }
    }

    std::vector<size_t> preorder;
    std::vector<size_t> stack = {root};
    while (!stack.empty()) {
      const size_t node_index = stack.back();
      stack.pop_back();
      preorder.push_back(node_index);
      for (size_t child : nodes[node_index].children) {
        stack.push_back(child);
      }
    }

    std::vector<std::optional<size_t>> mapped(nodes.size());
    std::vector<Node> new_nodes;
    new_nodes.reserve(nodes.size());
    for (auto iterator = preorder.rbegin(); iterator != preorder.rend();
         ++iterator) {
      const size_t old_index = *iterator;
      const Node& old_node = nodes[old_index];
      if (old_node.label.has_value()) {
        if (removed[old_index]) {
          continue;
        }
        const size_t new_index = new_nodes.size();
        new_nodes.push_back({*old_node.label, {}, std::nullopt, std::nullopt});
        mapped[old_index] = new_index;
        continue;
      }

      std::vector<size_t> kept_children;
      kept_children.reserve(old_node.children.size());
      for (size_t old_child : old_node.children) {
        if (mapped[old_child].has_value()) {
          kept_children.push_back(*mapped[old_child]);
        }
      }
      if (kept_children.empty()) {
        continue;
      }
      if (kept_children.size() == 1) {
        mapped[old_index] = kept_children.front();
        continue;
      }

      const size_t new_index = new_nodes.size();
      new_nodes.push_back(
          {std::nullopt, kept_children, std::nullopt, std::nullopt});
      mapped[old_index] = new_index;
    }

    if (!mapped[root].has_value()) {
      throw std::runtime_error("Tree empty after pruning");
    }

    Tree tree;
    tree.nodes = std::move(new_nodes);
    tree.root = *mapped[root];
    tree.n_edges = tree.assign_edge_indices();
    tree.precompute();
    return tree;
  }

  void fill_leaf_components(
      const std::vector<uint8_t>& bits, std::vector<size_t>& component_map,
      std::vector<std::pair<size_t, size_t>>& stack) const {
    if (component_map.size() < max_leaf_label + 1) {
      component_map.resize(max_leaf_label + 1);
    }
    stack.clear();
    stack.emplace_back(root, 0);
    size_t next_component = 1;
    while (!stack.empty()) {
      const auto [node_index, component] = stack.back();
      stack.pop_back();
      const Node& node = nodes[node_index];
      if (node.label.has_value()) {
        component_map[*node.label] = component;
        continue;
      }
      for (size_t child : node.children) {
        const std::optional<size_t> edge = nodes[child].edge_idx;
        const bool is_cut = edge.has_value() && bits[*edge] != 0;
        size_t child_component = component;
        if (is_cut) {
          child_component = next_component;
          ++next_component;
        }
        stack.emplace_back(child, child_component);
      }
    }
  }

  std::vector<size_t> get_leaf_components(
      const std::vector<uint8_t>& bits) const {
    std::vector<size_t> component_map(max_leaf_label + 1, 0);
    std::vector<std::pair<size_t, size_t>> stack;
    stack.reserve(nodes.size());
    fill_leaf_components(bits, component_map, stack);
    return component_map;
  }
};

// Reusable storage for exact induced-tree canonicalization.
struct InducedCanonicalScratch {
  std::vector<size_t> selected_nodes;
  std::vector<size_t> virtual_nodes;
  std::vector<size_t> parent_index;
  std::vector<size_t> first_child;
  std::vector<size_t> second_child;
  std::vector<uint8_t> child_count;
  std::vector<size_t> canonical_ids;
  std::vector<size_t> stack;
};

struct InducedCanonicalResult {
  size_t canonical_id = 0;
  size_t lca_node = 0;
};

InducedCanonicalResult canonical_induced_component_id(
    const Tree& tree, const std::vector<size_t>& leaf_labels,
    StructuralInterner& interner, InducedCanonicalScratch& scratch) {
  if (leaf_labels.empty()) {
    throw std::runtime_error("Cannot canonicalize an empty component");
  }
  if (leaf_labels.size() == 1) {
    const size_t leaf_node = tree.find_leaf_idx(leaf_labels.front());
    return {interner.leaf(leaf_labels.front()), leaf_node};
  }

  scratch.selected_nodes.clear();
  scratch.selected_nodes.reserve(leaf_labels.size());
  for (size_t label : leaf_labels) {
    scratch.selected_nodes.push_back(tree.find_leaf_idx(label));
  }

  auto by_preorder = [&](size_t lhs, size_t rhs) {
    return tree.preorder_index[lhs] < tree.preorder_index[rhs];
  };
  std::sort(scratch.selected_nodes.begin(), scratch.selected_nodes.end(),
            by_preorder);
  if (std::adjacent_find(scratch.selected_nodes.begin(),
                         scratch.selected_nodes.end()) !=
      scratch.selected_nodes.end()) {
    throw std::runtime_error("Component contains duplicate leaf labels");
  }

  scratch.virtual_nodes.assign(scratch.selected_nodes.begin(),
                               scratch.selected_nodes.end());
  scratch.virtual_nodes.reserve(scratch.selected_nodes.size() * 2);
  for (size_t i = 1; i < scratch.selected_nodes.size(); ++i) {
    scratch.virtual_nodes.push_back(
        tree.get_lca(scratch.selected_nodes[i - 1], scratch.selected_nodes[i]));
  }
  std::sort(scratch.virtual_nodes.begin(), scratch.virtual_nodes.end(),
            by_preorder);
  scratch.virtual_nodes.erase(
      std::unique(scratch.virtual_nodes.begin(), scratch.virtual_nodes.end()),
      scratch.virtual_nodes.end());

  const size_t node_count = scratch.virtual_nodes.size();
  const size_t none = std::numeric_limits<size_t>::max();
  scratch.parent_index.assign(node_count, none);
  scratch.first_child.assign(node_count, none);
  scratch.second_child.assign(node_count, none);
  scratch.child_count.assign(node_count, 0);
  scratch.canonical_ids.assign(node_count, 0);
  scratch.stack.clear();
  scratch.stack.reserve(node_count);

  size_t root_index = none;
  for (size_t index = 0; index < node_count; ++index) {
    const size_t node = scratch.virtual_nodes[index];
    while (
        !scratch.stack.empty() &&
        !tree.is_ancestor(scratch.virtual_nodes[scratch.stack.back()], node)) {
      scratch.stack.pop_back();
    }

    if (scratch.stack.empty()) {
      if (root_index != none) {
        throw std::runtime_error(
            "Virtual canonicalization produced multiple roots");
      }
      root_index = index;
    } else {
      const size_t parent = scratch.stack.back();
      scratch.parent_index[index] = parent;
      if (scratch.child_count[parent] == 0) {
        scratch.first_child[parent] = index;
      } else if (scratch.child_count[parent] == 1) {
        scratch.second_child[parent] = index;
      } else {
        throw std::runtime_error("Restricted component is not binary");
      }
      ++scratch.child_count[parent];
    }
    scratch.stack.push_back(index);
  }

  if (root_index == none) {
    throw std::runtime_error("Virtual canonicalization produced no root");
  }

  for (size_t index = node_count; index-- > 0;) {
    const Node& node = tree.nodes[scratch.virtual_nodes[index]];
    if (node.label.has_value()) {
      scratch.canonical_ids[index] = interner.leaf(*node.label);
      continue;
    }

    const uint8_t count = scratch.child_count[index];
    if (count == 0) {
      throw std::runtime_error(
          "Virtual canonicalization contains an empty internal node");
    }
    if (count == 1) {
      scratch.canonical_ids[index] =
          scratch.canonical_ids[scratch.first_child[index]];
      continue;
    }
    scratch.canonical_ids[index] =
        interner.branch(scratch.canonical_ids[scratch.first_child[index]],
                        scratch.canonical_ids[scratch.second_child[index]]);
  }

  return {scratch.canonical_ids[root_index], scratch.virtual_nodes[root_index]};
}

struct ComponentRecord {
  size_t component_id = 0;
  std::vector<size_t> leaves;
  std::vector<size_t> lcas;
  bool structures_match = true;
};

// Epoch-stamped owner arrays for repeated penalty evaluations.
struct StampedOwnerScratch {
  std::vector<size_t> owner;
  std::vector<uint64_t> stamp;
  uint64_t epoch = 0;

  void begin(size_t node_count) {
    if (owner.size() != node_count) {
      owner.assign(node_count, 0);
      stamp.assign(node_count, 0);
      epoch = 1;
      return;
    }
    ++epoch;
    if (epoch == 0) {
      std::fill(stamp.begin(), stamp.end(), 0);
      epoch = 1;
    }
  }

  bool conflict_or_claim(size_t node, size_t component_id) {
    if (stamp[node] != epoch) {
      stamp[node] = epoch;
      owner[node] = component_id;
      return false;
    }
    return owner[node] != component_id;
  }
};

class AgreementForestFitness {
 public:
  AgreementForestFitness(const std::vector<Tree>& trees,
                         size_t original_leaf_count,
                         const SolverParameters& params)
      : trees_(&trees),
        original_leaf_count_(original_leaf_count),
        params_(&params),
        canonical_scratch_(trees.size()),
        owner_scratch_(trees.size()) {
    const Tree& reference_tree = trees.front();
    leaf_iteration_order_.reserve(reference_tree.leaf_to_idx.size());
    for (const auto& [label, node_index] : reference_tree.leaf_to_idx) {
      (void)node_index;
      leaf_iteration_order_.push_back(label);
    }

    component_stack_.reserve(reference_tree.nodes.size());
    component_order_.reserve(reference_tree.leaf_to_idx.size());
    component_records_.reserve(reference_tree.leaf_to_idx.size());
    component_id_to_record_.resize(reference_tree.n_edges + 1);
    component_id_stamp_.assign(reference_tree.n_edges + 1, 0);
    for (size_t tree_index = 0; tree_index < trees.size(); ++tree_index) {
      owner_scratch_[tree_index].owner.resize(trees[tree_index].nodes.size());
      owner_scratch_[tree_index].stamp.assign(trees[tree_index].nodes.size(),
                                              0);
    }

    component_cache_ = std::make_unique<ExactComponentCompatibilityCache>(
        params.component_cache_memory_mib, trees.size());
  }

  AgreementForestFitness(const AgreementForestFitness&) = delete;
  AgreementForestFitness& operator=(const AgreementForestFitness&) = delete;
  AgreementForestFitness(AgreementForestFitness&&) noexcept = default;
  AgreementForestFitness& operator=(AgreementForestFitness&&) noexcept =
      default;

  const std::vector<Tree>& trees() const { return *trees_; }
  size_t original_leaf_count() const { return original_leaf_count_; }

  std::pair<size_t, size_t> calculate(const std::vector<uint8_t>& bits) const {
    if (!keep_running()) {
      return interrupted_result();
    }

    {
      extract_components(bits);
    }

    size_t structural_penalty = 0;
    {
      std::optional<StructuralInterner> interner;
      auto get_interner = [&]() -> StructuralInterner& {
        if (!interner.has_value()) {
          interner.emplace(trees_->front().nodes.size());
        }
        return *interner;
      };

      for (const auto& [component_id, record_index] : component_order_) {
        (void)component_id;
        if (!keep_running()) {
          return interrupted_result();
        }
        ComponentRecord& component = component_records_[record_index];
        prepare_component(component, get_interner);
        if (!component.structures_match) {
          structural_penalty += component.leaves.size();
        }
      }
    }

    std::optional<size_t> dynamic_penalty;
    {
      dynamic_penalty = calculate_penalty();
    }
    if (!dynamic_penalty.has_value()) {
      return interrupted_result();
    }

    const std::pair<size_t, size_t> result = {
        component_order_.size(), structural_penalty + *dynamic_penalty};
    record_completed_result(result);
    return result;
  }

 private:
  struct CompatibilityResult {
    bool structures_match = false;
    uint32_t canonicalization_calls = 0;
  };

  void extract_components(const std::vector<uint8_t>& bits) const {
    component_order_.clear();
    active_component_count_ = 0;
    ++component_id_epoch_;
    if (component_id_epoch_ == 0) {
      std::fill(component_id_stamp_.begin(), component_id_stamp_.end(), 0);
      component_id_epoch_ = 1;
    }

    const Tree& tree = trees_->front();
    component_stack_.clear();
    component_stack_.emplace_back(tree.root, 0);
    size_t next_component = 1;

    while (!component_stack_.empty()) {
      const auto [node_index, component_id] = component_stack_.back();
      component_stack_.pop_back();
      const Node& node = tree.nodes[node_index];
      if (node.label.has_value()) {
        add_leaf_to_component(component_id, *node.label);
        continue;
      }
      for (size_t child : node.children) {
        const std::optional<size_t> edge = tree.nodes[child].edge_idx;
        const bool is_cut = edge.has_value() && bits[*edge] != 0;
        size_t child_component = component_id;
        if (is_cut) {
          child_component = next_component;
          ++next_component;
        }
        component_stack_.emplace_back(child, child_component);
      }
    }
  }

  void add_leaf_to_component(size_t component_id, size_t leaf_label) const {
    size_t record_index = 0;
    if (component_id_stamp_[component_id] != component_id_epoch_) {
      component_id_stamp_[component_id] = component_id_epoch_;
      record_index = active_component_count_;
      component_id_to_record_[component_id] = record_index;
      component_order_.emplace_back(component_id, record_index);

      if (active_component_count_ == component_records_.size()) {
        component_records_.emplace_back();
      }
      ComponentRecord& component = component_records_[record_index];
      component.component_id = component_id;
      component.leaves.clear();
      component.lcas.resize(trees_->size());
      component.structures_match = true;
      ++active_component_count_;
    } else {
      record_index = component_id_to_record_[component_id];
    }
    component_records_[record_index].leaves.push_back(leaf_label);
  }

  template <typename GetInterner>
  void prepare_component(ComponentRecord& component,
                         GetInterner&& get_interner) const {
    if (component.leaves.size() == 1) {
      const size_t label = component.leaves.front();
      for (size_t tree_index = 0; tree_index < trees_->size(); ++tree_index) {
        component.lcas[tree_index] =
            trees_->at(tree_index).find_leaf_idx(label);
      }
      component.structures_match = true;
      return;
    }

    std::optional<ExactComponentCompatibilityCache::Value> cached;
    if (component_cache_) {
      cached = component_cache_->lookup(component.leaves);
    }

    if (cached.has_value()) {
      component.structures_match = cached->structures_match;
      std::copy_n(cached->lcas, trees_->size(), component.lcas.begin());
      return;
    }

    const CompatibilityResult compatibility =
        calculate_component_compatibility(component, get_interner());
    component.structures_match = compatibility.structures_match;
    if (component_cache_) {
      component_cache_->insert(component.leaves, component.lcas,
                               {compatibility.structures_match,
                                compatibility.canonicalization_calls, nullptr});
    }
  }

  CompatibilityResult calculate_component_compatibility(
      ComponentRecord& component, StructuralInterner& interner) const {
    bool structures_match = true;
    size_t reference_id = 0;
    for (size_t tree_index = 0; tree_index < trees_->size(); ++tree_index) {
      const InducedCanonicalResult result = canonical_induced_component_id(
          trees_->at(tree_index), component.leaves, interner,
          canonical_scratch_[tree_index]);
      component.lcas[tree_index] = result.lca_node;
      if (tree_index == 0) {
        reference_id = result.canonical_id;
      } else if (result.canonical_id != reference_id) {
        structures_match = false;
      }
    }

    if (trees_->size() > std::numeric_limits<uint32_t>::max()) {
      throw std::runtime_error("Too many input trees");
    }
    return {structures_match, static_cast<uint32_t>(trees_->size())};
  }

  std::optional<size_t> calculate_penalty() const {
    size_t penalty = 0;
    for (size_t tree_index = 0; tree_index < trees_->size(); ++tree_index) {
      if (!keep_running()) {
        return std::nullopt;
      }

      const Tree& tree = trees_->at(tree_index);
      StampedOwnerScratch& owner = owner_scratch_[tree_index];
      owner.begin(tree.nodes.size());

      for (const auto& item : component_order_) {
        const size_t component_id = item.first;
        const size_t record_index = item.second;
        const ComponentRecord& component = component_records_[record_index];
        const size_t lca = component.lcas[tree_index];

        for (size_t leaf_label : component.leaves) {
          size_t current = tree.find_leaf_idx(leaf_label);
          while (true) {
            if (owner.conflict_or_claim(current, component_id)) {
              ++penalty;
              break;
            }
            if (current == lca) {
              break;
            }
            current = *tree.nodes[current].parent;
          }
        }
      }
    }
    return penalty;
  }

  std::pair<size_t, size_t> interrupted_result() const {
    return {0, params_->interrupted_fitness_penalty};
  }
  void record_completed_result(const std::pair<size_t, size_t>&) const {}

  const std::vector<Tree>* trees_;
  size_t original_leaf_count_;
  const SolverParameters* params_;
  std::vector<size_t> leaf_iteration_order_;
  mutable std::vector<std::pair<size_t, size_t>> component_stack_;
  mutable std::vector<std::pair<size_t, size_t>> component_order_;
  mutable std::vector<ComponentRecord> component_records_;
  mutable std::vector<size_t> component_id_to_record_;
  mutable std::vector<uint64_t> component_id_stamp_;
  mutable uint64_t component_id_epoch_ = 0;
  mutable size_t active_component_count_ = 0;
  mutable std::vector<InducedCanonicalScratch> canonical_scratch_;
  mutable std::vector<StampedOwnerScratch> owner_scratch_;
  mutable std::unique_ptr<ExactComponentCompatibilityCache> component_cache_;
};
class EA {
 public:
  virtual ~EA() = default;
  virtual std::optional<
      std::tuple<size_t, std::vector<uint8_t>, std::pair<size_t, size_t>>>
  next() = 0;
};

class OnePlusOneEA : public EA {
 public:
  OnePlusOneEA(size_t bit_count, AgreementForestFitness fitness,
               size_t maximum_generations, uint64_t seed,
               const std::vector<uint8_t>& initial_solution,
               const SolverParameters& parameters)
      : fitness_fn_(std::move(fitness)),
        bit_count_(bit_count),
        max_generations_(maximum_generations),
        generation_(0),
        solution_(initial_solution),
        fitness_(fitness_fn_.calculate(solution_)),
        rng_(seed),
        stagnation_(0),
        parameters_(&parameters) {}

  std::optional<
      std::tuple<size_t, std::vector<uint8_t>, std::pair<size_t, size_t>>>
  next() override {
    if (generation_ > max_generations_ || !keep_running()) {
      return std::nullopt;
    }

    if (generation_ > 0) {
      std::vector<uint8_t> candidate = solution_;
      mutate(candidate);

      const std::pair<size_t, size_t> candidate_fitness =
          fitness_fn_.calculate(candidate);
      if (!keep_running()) {
        return std::nullopt;
      }

      update_current_solution(std::move(candidate), candidate_fitness);
    }

    auto result = std::make_tuple(generation_, solution_, fitness_);
    ++generation_;
    return result;
  }

 private:
  void mutate(std::vector<uint8_t>& candidate) {
    size_t primary_flips = 0;
    if (should_try_edge_swap()) {
      primary_flips = apply_edge_swap(candidate);
    }

    if (primary_flips == 0) {
      primary_flips = apply_uniform_mutation(candidate);
    }

    if (stagnation_ > parameters_->perturbation_stagnation_threshold) {
      apply_perturbation(candidate);
      stagnation_ = 0;
    }
  }

  bool should_try_edge_swap() {
    if (stagnation_ <= parameters_->swap_stagnation_threshold) {
      return false;
    }
    return std::bernoulli_distribution(parameters_->swap_probability)(rng_);
  }

  size_t apply_edge_swap(std::vector<uint8_t>& candidate) {
    std::vector<size_t> cut_edges;
    std::vector<size_t> uncut_edges;
    cut_edges.reserve(bit_count_);
    uncut_edges.reserve(bit_count_);

    for (size_t edge = 0; edge < bit_count_; ++edge) {
      if (candidate[edge] == 1) {
        cut_edges.push_back(edge);
      } else {
        uncut_edges.push_back(edge);
      }
    }

    if (cut_edges.empty() || uncut_edges.empty()) {
      return 0;
    }

    const size_t cut_index =
        std::uniform_int_distribution<size_t>(0, cut_edges.size() - 1)(rng_);
    const size_t uncut_index =
        std::uniform_int_distribution<size_t>(0, uncut_edges.size() - 1)(rng_);

    candidate[cut_edges[cut_index]] = 0;
    candidate[uncut_edges[uncut_index]] = 1;
    return 2;
  }

  size_t apply_uniform_mutation(std::vector<uint8_t>& candidate) {
    size_t flips = 0;
    double mutation_probability = 1.0;
    if (bit_count_ > 0) {
      mutation_probability = 1.0 / static_cast<double>(bit_count_);
    }

    for (size_t edge = 0; edge < bit_count_; ++edge) {
      if (std::bernoulli_distribution(mutation_probability)(rng_)) {
        candidate[edge] ^= 1;
        ++flips;
      }
    }

    if (flips == 0 && bit_count_ > 0) {
      const size_t forced_edge =
          std::uniform_int_distribution<size_t>(0, bit_count_ - 1)(rng_);
      candidate[forced_edge] ^= 1;
      flips = 1;
    }

    return flips;
  }

  void apply_perturbation(std::vector<uint8_t>& candidate) {
    for (size_t edge = 0; edge < bit_count_; ++edge) {
      if (std::bernoulli_distribution(parameters_->perturbation_probability)(
              rng_)) {
        candidate[edge] ^= 1;
      }
    }
  }

  void update_current_solution(std::vector<uint8_t> candidate,
                               std::pair<size_t, size_t> candidate_fitness) {
    const bool better_penalty = candidate_fitness.second < fitness_.second;
    const bool equal_penalty = candidate_fitness.second == fitness_.second;
    const bool no_more_components = candidate_fitness.first <= fitness_.first;

    if (better_penalty || (equal_penalty && no_more_components)) {
      if (better_penalty || candidate_fitness.first < fitness_.first) {
        stagnation_ = 0;
      } else {
        ++stagnation_;
      }
      solution_ = std::move(candidate);
      fitness_ = candidate_fitness;
    } else {
      ++stagnation_;
    }
  }

  AgreementForestFitness fitness_fn_;
  size_t bit_count_;
  size_t max_generations_;
  size_t generation_;
  std::vector<uint8_t> solution_;
  std::pair<size_t, size_t> fitness_;
  std::mt19937_64 rng_;
  size_t stagnation_;
  const SolverParameters* parameters_;
};

// Serialize the restriction of one original tree to one component.
std::string serialize_induced_component(const Tree& tree,
                                        std::vector<size_t> leaf_labels) {
  if (leaf_labels.empty()) {
    return "";
  }
  if (leaf_labels.size() == 1) {
    return std::to_string(leaf_labels.front());
  }

  std::vector<size_t> selected_nodes;
  selected_nodes.reserve(leaf_labels.size());
  for (size_t label : leaf_labels) {
    selected_nodes.push_back(tree.find_leaf_idx(label));
  }
  auto by_preorder = [&](size_t lhs, size_t rhs) {
    return tree.preorder_index[lhs] < tree.preorder_index[rhs];
  };
  std::sort(selected_nodes.begin(), selected_nodes.end(), by_preorder);
  selected_nodes.erase(
      std::unique(selected_nodes.begin(), selected_nodes.end()),
      selected_nodes.end());

  std::vector<size_t> virtual_nodes = selected_nodes;
  virtual_nodes.reserve(selected_nodes.size() * 2);
  for (size_t i = 1; i < selected_nodes.size(); ++i) {
    virtual_nodes.push_back(
        tree.get_lca(selected_nodes[i - 1], selected_nodes[i]));
  }
  std::sort(virtual_nodes.begin(), virtual_nodes.end(), by_preorder);
  virtual_nodes.erase(std::unique(virtual_nodes.begin(), virtual_nodes.end()),
                      virtual_nodes.end());

  std::vector<std::vector<size_t>> children(virtual_nodes.size());
  std::vector<size_t> stack;
  stack.reserve(virtual_nodes.size());
  for (size_t index = 0; index < virtual_nodes.size(); ++index) {
    const size_t node = virtual_nodes[index];
    while (!stack.empty() &&
           !tree.is_ancestor(virtual_nodes[stack.back()], node)) {
      stack.pop_back();
    }
    if (!stack.empty()) {
      children[stack.back()].push_back(index);
    }
    stack.push_back(index);
  }

  // The first node in preorder is the LCA of the full component.
  const size_t root_index = 0;
  enum class EventKind : uint8_t { Node, Comma, Close };
  struct Event {
    EventKind kind;
    size_t index = 0;
  };

  std::string result;
  result.reserve(leaf_labels.size() * 4);
  std::vector<Event> events = {{EventKind::Node, root_index}};
  while (!events.empty()) {
    const Event event = events.back();
    events.pop_back();
    if (event.kind == EventKind::Comma) {
      result.push_back(',');
      continue;
    }
    if (event.kind == EventKind::Close) {
      result.push_back(')');
      continue;
    }

    const size_t original_node = virtual_nodes[event.index];
    const Node& node = tree.nodes[original_node];
    if (node.label.has_value()) {
      result += std::to_string(*node.label);
      continue;
    }

    const auto& node_children = children[event.index];
    if (node_children.empty()) {
      throw std::runtime_error(
          "Virtual output tree contains an empty internal node");
    }
    if (node_children.size() == 1) {
      // Suppress unary nodes in the restricted output tree.
      events.push_back({EventKind::Node, node_children.front()});
      continue;
    }
    if (node_children.size() != 2) {
      throw std::runtime_error(
          "Virtual output tree is not binary after restriction");
    }

    result.push_back('(');
    events.push_back({EventKind::Close, 0});
    events.push_back({EventKind::Node, node_children[1]});
    events.push_back({EventKind::Comma, 0});
    events.push_back({EventKind::Node, node_children[0]});
  }
  return result;
}

std::vector<std::vector<size_t>> expand_components(
    const Tree& reduced_tree, const std::vector<uint8_t>& bits,
    const std::unordered_map<size_t, std::vector<size_t>>& meta_map,
    size_t original_leaf_count) {
  const std::vector<size_t> component_map =
      reduced_tree.get_leaf_components(bits);
  std::unordered_map<size_t, std::vector<size_t>> reduced_components;
  reduced_components.reserve(reduced_tree.leaf_to_idx.size());
  for (const auto& [label, node_index] : reduced_tree.leaf_to_idx) {
    (void)node_index;
    reduced_components[component_map[label]].push_back(label);
  }

  std::vector<std::vector<size_t>> expanded_components;
  expanded_components.reserve(reduced_components.size());
  std::vector<uint8_t> seen(original_leaf_count + 1, 0);
  for (const auto& [component_id, reduced_labels] : reduced_components) {
    (void)component_id;
    std::vector<size_t> expanded;
    for (size_t label : reduced_labels) {
      auto meta = meta_map.find(label);
      if (meta == meta_map.end()) {
        expanded.push_back(label);
      } else {
        expanded.insert(expanded.end(), meta->second.begin(),
                        meta->second.end());
      }
    }
    std::sort(expanded.begin(), expanded.end());
    expanded.erase(std::unique(expanded.begin(), expanded.end()),
                   expanded.end());
    for (size_t label : expanded) {
      if (label == 0 || label > original_leaf_count || seen[label]) {
        throw std::runtime_error(
            "Invalid or duplicate original label during solution lifting");
      }
      seen[label] = 1;
    }
    expanded_components.push_back(std::move(expanded));
  }

  for (size_t label = 1; label <= original_leaf_count; ++label) {
    if (!seen[label]) {
      throw std::runtime_error("Solution lifting omitted original leaf " +
                               std::to_string(label));
    }
  }
  std::sort(expanded_components.begin(), expanded_components.end(),
            [](const auto& lhs, const auto& rhs) {
              return lhs.front() < rhs.front();
            });
  return expanded_components;
}

std::string serialize_agreement_forest(
    const Tree& original_tree, const Tree& reduced_tree,
    const std::vector<uint8_t>& bits,
    const std::unordered_map<size_t, std::vector<size_t>>& meta_map,
    size_t original_leaf_count) {
  const auto components =
      expand_components(reduced_tree, bits, meta_map, original_leaf_count);
  std::string output;
  output.reserve(original_leaf_count * 6);
  for (const auto& component : components) {
    const std::string newick =
        serialize_induced_component(original_tree, component);
    if (!newick.empty()) {
      output += newick;
      output += ";\n";
    }
  }
  return output;
}

bool is_common_chain(const Tree& tree, const std::vector<size_t>& chain) {
  size_t k = chain.size();
  if (k < 2) {
    return false;
  }
  auto find_label = [&](size_t label) -> std::optional<size_t> {
    if (tree.leaf_to_idx.count(label)) {
      return tree.leaf_to_idx.at(label);
    }
    return std::nullopt;
  };
  auto u_opt = find_label(chain[k - 1]), v_opt = find_label(chain[k - 2]);
  if (!u_opt || !v_opt) {
    return false;
  }
  size_t u = *u_opt, v = *v_opt;
  if (tree.nodes[u].parent != tree.nodes[v].parent ||
      !tree.nodes[u].parent.has_value()) {
    return false;
  }
  size_t curr = tree.nodes[u].parent.value();
  for (int i = (int)k - 3; i >= 0; --i) {
    auto w_opt = find_label(chain[i]);
    if (!w_opt) {
      return false;
    }
    size_t w = *w_opt;
    if (tree.nodes[w].parent != tree.nodes[curr].parent ||
        !tree.nodes[w].parent.has_value()) {
      return false;
    }
    curr = tree.nodes[curr].parent.value();
  }
  return true;
}

// Greedily remove cuts while the penalty remains zero.
std::vector<uint8_t> apply_greedy_initialization(
    size_t n_bits, const AgreementForestFitness& fitness_fn,
    std::mt19937_64& rng) {
  std::vector<uint8_t> sol(n_bits, 1);
  std::vector<size_t> indices(n_bits);
  std::iota(indices.begin(), indices.end(), 0);
  std::shuffle(indices.begin(), indices.end(), rng);

  for (size_t i : indices) {
    if (!keep_running()) {
      break;
    }
    sol[i] = 0;
    if (fitness_fn.calculate(sol).second != 0) {
      sol[i] = 1;
    }
  }
  return sol;
}

// Reduce common chains across the two trees.
bool apply_chain_reduction(
    std::vector<Tree>& trees, size_t& n_leaves,
    std::unordered_map<size_t, std::vector<size_t>>& meta_map) {
  (void)n_leaves;
  if (trees.empty() || !keep_running()) {
    return false;
  }
  const Tree& t0 = trees[0];
  std::vector<std::vector<size_t>> chains;
  for (size_t i = 0; i < t0.nodes.size(); ++i) {
    if (!keep_running()) {
      return false;
    }
    if (t0.nodes[i].children.size() != 2) {
      continue;
    }
    std::vector<size_t> curr_chain;
    size_t curr = i;
    while (true) {
      const auto& n = t0.nodes[curr];
      if (n.children.size() != 2) {
        break;
      }
      const auto &c0 = t0.nodes[n.children[0]], c1 = t0.nodes[n.children[1]];
      if (c0.label.has_value() && !c1.label.has_value()) {
        curr_chain.push_back(c0.label.value());
        curr = n.children[1];
      } else if (c1.label.has_value() && !c0.label.has_value()) {
        curr_chain.push_back(c1.label.value());
        curr = n.children[0];
      } else if (c0.label.has_value() && c1.label.has_value()) {
        curr_chain.push_back(c0.label.value());
        curr_chain.push_back(c1.label.value());
        break;
      } else {
        break;
      }
    }
    if (curr_chain.size() > 3) {
      chains.push_back(curr_chain);
    }
  }
  for (const auto& chain : chains) {
    bool match_all = true;
    for (size_t i = 1; i < trees.size(); ++i) {
      if (!is_common_chain(trees[i], chain)) {
        match_all = false;
        break;
      }
    }
    if (match_all) {
      std::vector<size_t> to_remove(chain.begin() + 2, chain.end() - 1);
      size_t rep = chain[1];
      std::vector<Tree> new_trees;
      for (const auto& t : trees) {
        new_trees.push_back(t.prune_and_compact(to_remove));
      }
      trees = std::move(new_trees);
      std::vector<size_t> removed_origs;
      for (size_t r : to_remove) {
        if (meta_map.count(r)) {
          removed_origs.insert(removed_origs.end(), meta_map[r].begin(),
                               meta_map[r].end());
          meta_map.erase(r);
        } else {
          removed_origs.push_back(r);
        }
      }
      std::vector<size_t> origs_of_rep;
      if (meta_map.count(rep)) {
        origs_of_rep = std::move(meta_map[rep]);
      } else {
        origs_of_rep = {rep};
      }
      origs_of_rep.insert(origs_of_rep.end(), removed_origs.begin(),
                          removed_origs.end());
      meta_map[rep] = std::move(origs_of_rep);
      return true;
    }
  }
  return false;
}

// Collapse all maximal disjoint common subtrees in one batch.
size_t apply_subtree_reduction_batch(
    std::vector<Tree>& trees, size_t& next_label,
    std::unordered_map<size_t, std::vector<size_t>>& meta_map) {
  if (trees.empty() || !keep_running()) {
    return 0;
  }

  StructuralInterner interner;
  std::vector<std::vector<size_t>> canonical_ids;
  canonical_ids.reserve(trees.size());
  std::vector<std::unordered_map<size_t, size_t>> node_by_id(trees.size());

  for (size_t tree_index = 0; tree_index < trees.size(); ++tree_index) {
    if (!keep_running()) {
      return 0;
    }
    canonical_ids.push_back(trees[tree_index].full_canonical_ids(interner));
    auto& index = node_by_id[tree_index];
    index.reserve(trees[tree_index].nodes.size());
    for (size_t node_index = 0; node_index < trees[tree_index].nodes.size();
         ++node_index) {
      if (trees[tree_index].nodes[node_index].label.has_value()) {
        continue;
      }
      const size_t id = canonical_ids[tree_index][node_index];
      auto [iterator, inserted] = index.emplace(id, node_index);
      if (!inserted) {
        throw std::runtime_error(
            "Duplicate exact subtree identity within one phylogenetic tree");
      }
      (void)iterator;
    }
  }

  const Tree& first = trees.front();
  std::vector<uint8_t> common(first.nodes.size(), 0);
  for (size_t node_index = 0; node_index < first.nodes.size(); ++node_index) {
    const Node& node = first.nodes[node_index];
    if (node.label.has_value() || node_index == first.root) {
      continue;
    }
    const size_t id = canonical_ids.front()[node_index];
    bool present_everywhere = true;
    for (size_t tree_index = 1; tree_index < trees.size(); ++tree_index) {
      if (!node_by_id[tree_index].contains(id)) {
        present_everywhere = false;
        break;
      }
    }
    if (present_everywhere) {
      common[node_index] = 1;
    }
  }

  // Select highest common nodes because descendant selections are redundant.
  std::vector<size_t> selected_nodes;
  std::vector<size_t> stack;
  for (auto iterator = first.nodes[first.root].children.rbegin();
       iterator != first.nodes[first.root].children.rend(); ++iterator) {
    stack.push_back(*iterator);
  }
  while (!stack.empty()) {
    if (!keep_running()) {
      return 0;
    }
    const size_t node_index = stack.back();
    stack.pop_back();
    if (common[node_index]) {
      selected_nodes.push_back(node_index);
      continue;
    }
    for (auto iterator = first.nodes[node_index].children.rbegin();
         iterator != first.nodes[node_index].children.rend(); ++iterator) {
      stack.push_back(*iterator);
    }
  }
  if (selected_nodes.empty()) {
    return 0;
  }

  struct Selection {
    size_t canonical_id;
    size_t new_label;
    std::vector<size_t> current_labels;
    std::vector<size_t> original_labels;
  };
  std::vector<Selection> selections;
  selections.reserve(selected_nodes.size());

  const size_t first_new_label = next_label + 1;
  for (size_t selection_index = 0; selection_index < selected_nodes.size();
       ++selection_index) {
    const size_t node_index = selected_nodes[selection_index];
    const size_t id = canonical_ids.front()[node_index];
    const size_t new_label = first_new_label + selection_index;
    std::vector<size_t> current_labels =
        first.get_subtree_leaf_labels(node_index);
    std::vector<size_t> originals;
    for (size_t label : current_labels) {
      auto meta = meta_map.find(label);
      if (meta == meta_map.end()) {
        originals.push_back(label);
      } else {
        originals.insert(originals.end(), meta->second.begin(),
                         meta->second.end());
      }
    }
    std::sort(originals.begin(), originals.end());
    selections.push_back(
        {id, new_label, std::move(current_labels), std::move(originals)});
  }

  std::vector<Tree> reduced_trees;
  reduced_trees.reserve(trees.size());
  for (size_t tree_index = 0; tree_index < trees.size(); ++tree_index) {
    if (!keep_running()) {
      return 0;
    }
    auto nodes = trees[tree_index].nodes;
    for (const Selection& selection : selections) {
      const auto found = node_by_id[tree_index].find(selection.canonical_id);
      if (found == node_by_id[tree_index].end()) {
        throw std::runtime_error(
            "Common-subtree batch lost a corresponding node");
      }
      Node& node = nodes[found->second];
      node.label = selection.new_label;
      node.children.clear();
    }
    reduced_trees.push_back(
        Tree::compact(std::move(nodes), trees[tree_index].root));
  }
  // Update reduction metadata only after both rebuilt trees are ready.
  for (const Selection& selection : selections) {
    for (size_t label : selection.current_labels) {
      meta_map.erase(label);
    }
    meta_map.emplace(selection.new_label, selection.original_labels);
  }
  next_label += selections.size();
  trees = std::move(reduced_trees);
  return selections.size();
}

// CLI configuration; algorithmic choices are fixed for the submitted solver.
struct Args {
  std::string input_file;
  bool read_from_stdin = true;
  std::optional<size_t> max_gens;
  double timeout = 300.0;
  SolverParameters params;
};

std::string usage() {
  return "Usage:\n"
         "  solver [options] < input.nw > solution.txt\n"
         "  solver INSTANCE [options] > solution.txt\n"
         "\nSupported options:\n"
         "  -m, --max_gens N       Optional generation cap.\n"
         "  --timeout SECONDS      Soft time budget. Defaults to "
         "STRIDE_TIMEOUT or 300.\n"
         "  --seed N               Random seed. Default: 42.\n";
}

Args parse_args(int argc, char* argv[]) {
  Args args;
  bool timeout_was_set = false;
  std::vector<std::string> positional;

  auto require_value = [&](int& index,
                           const std::string& option) -> std::string {
    if (index + 1 >= argc) {
      throw std::runtime_error("Missing value for " + option);
    }
    ++index;
    return argv[index];
  };

  for (int i = 1; i < argc; ++i) {
    const std::string option = argv[i];
    if (option == "-m" || option == "--max_gens") {
      args.max_gens = std::stoull(require_value(i, option));
    } else if (option == "--timeout") {
      args.timeout = std::stod(require_value(i, option));
      timeout_was_set = true;
    } else if (option == "--seed") {
      args.params.seed = std::stoull(require_value(i, option));
    } else if (option == "--help" || option == "-h") {
      throw std::runtime_error(usage());
    } else if (!option.empty() && option[0] == '-') {
      throw std::runtime_error("Unknown option: " + option + "\n" + usage());
    } else {
      positional.push_back(option);
    }
  }

  if (!timeout_was_set) {
    if (const char* stride_timeout = std::getenv("STRIDE_TIMEOUT")) {
      args.timeout = std::stod(stride_timeout);
    }
  }

  if (positional.empty()) {
    args.read_from_stdin = true;
  } else if (positional.size() == 1) {
    args.read_from_stdin = false;
    args.input_file = positional[0];
  } else {
    throw std::runtime_error("Invalid positional arguments.\n" + usage());
  }

  if (args.timeout <= 0.0) {
    throw std::runtime_error("timeout must be positive");
  }
  return args;
}

std::pair<size_t, std::vector<Tree>> parse_instance_stream(
    std::istream& input, const std::string& source_name,
    size_t* announced_leaf_count) {
  std::string line;
  std::optional<size_t> leaf_count;
  std::optional<size_t> declared_tree_count;
  std::vector<std::string> newicks;
  while (std::getline(input, line)) {
    if (line.empty()) {
      continue;
    }
    if (line[0] == '#') {
      if (line.rfind("#p ", 0) == 0) {
        std::stringstream header(line.substr(3));
        size_t trees = 0;
        size_t leaves = 0;
        if (!(header >> trees >> leaves)) {
          throw std::runtime_error("Malformed #p line: " + line);
        }
        declared_tree_count = trees;
        leaf_count = leaves;
        if (announced_leaf_count != nullptr) {
          *announced_leaf_count = leaves;
        }
      }
      continue;
    }
    newicks.push_back(line);
  }

  if (!input.eof() && input.bad()) {
    throw std::runtime_error("Read error while reading instance from " +
                             source_name);
  }
  if (!leaf_count.has_value()) {
    throw std::runtime_error("No #p line found");
  }
  if (!declared_tree_count.has_value() || *declared_tree_count != 2) {
    throw std::runtime_error("Expected exactly two input trees");
  }
  if (newicks.size() != *declared_tree_count) {
    throw std::runtime_error(
        "Declared tree count does not match the number of Newick trees");
  }

  std::vector<Tree> trees;
  trees.reserve(*declared_tree_count);
  for (const std::string& newick : newicks) {
    trees.push_back(Tree::parse_newick(newick));
  }
  return {*leaf_count, std::move(trees)};
}

std::pair<size_t, std::vector<Tree>> parse_instance(
    const std::string& path, size_t* announced_leaf_count) {
  std::ifstream input(path);
  if (!input) {
    throw std::runtime_error("Cannot open input file: " + path);
  }
  return parse_instance_stream(input, path, announced_leaf_count);
}

void validate_instance(size_t leaf_count, const std::vector<Tree>& trees) {
  if (trees.size() != 2) {
    throw std::runtime_error("Expected exactly two trees");
  }
  for (size_t tree_index = 0; tree_index < trees.size(); ++tree_index) {
    const Tree& tree = trees[tree_index];
    if (tree.leaf_to_idx.size() != leaf_count) {
      throw std::runtime_error(
          "Tree " + std::to_string(tree_index + 1) +
          " does not contain the declared number of leaves");
    }
    for (size_t label = 1; label <= leaf_count; ++label) {
      if (!tree.leaf_to_idx.contains(label)) {
        throw std::runtime_error("Tree " + std::to_string(tree_index + 1) +
                                 " is missing leaf label " +
                                 std::to_string(label));
      }
    }
    for (size_t node_index = 0; node_index < tree.nodes.size(); ++node_index) {
      const Node& node = tree.nodes[node_index];
      if (node.label.has_value()) {
        if (!node.children.empty()) {
          throw std::runtime_error("Labelled node has children");
        }
      } else if (node.children.size() != 2) {
        throw std::runtime_error(
            "Every internal node must have exactly two children");
      }
      if (node_index == tree.root) {
        if (node.parent.has_value() || node.edge_idx.has_value()) {
          throw std::runtime_error("Root has stale parent or edge metadata");
        }
      } else if (!node.parent.has_value() || !node.edge_idx.has_value()) {
        throw std::runtime_error(
            "Non-root node is missing parent/edge metadata");
      }
    }
  }
}

std::string singleton_solution(size_t leaf_count) {
  std::string output;
  output.reserve(leaf_count * 4);
  for (size_t label = 1; label <= leaf_count; ++label) {
    output += std::to_string(label);
    output += ";\n";
  }
  return output;
}

class BestSolutionOutput {
 public:
  bool due() const {
    if (!initialized_) {
      return true;
    }
    return std::chrono::duration<double>(std::chrono::steady_clock::now() -
                                         last_store_)
               .count() >= interval_seconds_;
  }

  void store(std::string text) {
    text_ = std::move(text);
    initialized_ = true;
    last_store_ = std::chrono::steady_clock::now();
  }

  bool initialized() const { return initialized_; }

  void publish_final() const {
    if (!initialized_) {
      throw std::runtime_error("No valid solution is available");
    }
    std::cout << text_;
  }

 private:
  static constexpr double interval_seconds_ = 10.0;
  bool initialized_ = false;
  std::string text_;
  std::chrono::steady_clock::time_point last_store_{};
};

int main(int argc, char* argv[]) {
  std::signal(SIGINT, signal_handler);
  std::signal(SIGTERM, signal_handler);

  const auto wall_start = std::chrono::steady_clock::now();
  std::optional<Args> parsed_args;
  size_t announced_leaf_count = 0;
  bool valid_output_exists = false;

  try {
    Args args = parse_args(argc, argv);
    parsed_args = args;
    configure_internal_deadline(wall_start, args.timeout,
                                args.params.output_reserve_seconds);

    size_t leaf_count = 0;
    std::vector<Tree> trees;
    {
      if (args.read_from_stdin) {
        auto parsed = parse_instance_stream(std::cin, "standard input",
                                            &announced_leaf_count);
        leaf_count = parsed.first;
        trees = std::move(parsed.second);
      } else {
        auto parsed = parse_instance(args.input_file, &announced_leaf_count);
        leaf_count = parsed.first;
        trees = std::move(parsed.second);
      }
    }
    {
      validate_instance(leaf_count, trees);
    }

    BestSolutionOutput best_output;
    best_output.store(singleton_solution(leaf_count));
    valid_output_exists = true;

    // Use the first original tree to lift and serialize the output forest.
    const Tree original_output_tree = trees.front();
    const size_t original_leaf_count = leaf_count;
    size_t next_meta_label = leaf_count;
    std::unordered_map<size_t, std::vector<size_t>> meta_map;

    auto store_all_cuts = [&]() {
      const std::vector<uint8_t> all_cuts(trees.front().n_edges, 1);
      best_output.store(
          serialize_agreement_forest(original_output_tree, trees.front(),
                                     all_cuts, meta_map, original_leaf_count));
    };

    if (keep_running()) {
      while (keep_running()) {
        const size_t reduced =
            apply_subtree_reduction_batch(trees, next_meta_label, meta_map);
        if (reduced == 0) {
          break;
        }
        // Store the current best solution after each successful
        // subtree-reduction batch.
        store_all_cuts();
      }
    }

    if (keep_running()) {
      bool changed = false;
      while (keep_running() &&
             apply_chain_reduction(trees, next_meta_label, meta_map)) {
        changed = true;
        if (best_output.due()) {
          store_all_cuts();
        }
      }
      if (changed) {
        store_all_cuts();
      }
    }

    const size_t bit_count = trees.front().n_edges;
    std::vector<uint8_t> initial_solution(bit_count, 1);
    std::unique_ptr<EA> ea;

    if (keep_running()) {
      std::optional<AgreementForestFitness> fitness;
      {
        fitness.emplace(trees, original_leaf_count, args.params);
      }

      std::mt19937_64 initialization_rng(args.params.seed);
      initial_solution =
          apply_greedy_initialization(bit_count, *fitness, initialization_rng);

      // Store the greedy initial solution.
      best_output.store(serialize_agreement_forest(
          original_output_tree, trees.front(), initial_solution, meta_map,
          original_leaf_count));

      if (keep_running()) {
        const size_t max_generations =
            args.max_gens.value_or(args.params.default_max_generations);
        ea = std::make_unique<OnePlusOneEA>(bit_count, std::move(*fitness),
                                            max_generations, args.params.seed,
                                            initial_solution, args.params);
      }
    }

    std::pair<size_t, size_t> best_fitness = {args.params.initial_best_fitness,
                                              args.params.initial_best_fitness};
    std::vector<uint8_t> best_solution = initial_solution;
    bool best_output_dirty = false;

    if (ea) {
      while (keep_running()) {
        // Check the deadline once per generation.
        if (!check_internal_deadline_now()) {
          break;
        }
        auto result = ea->next();
        if (!result.has_value()) {
          break;
        }
        auto [generation, solution, fitness] = std::move(*result);
        (void)generation;
        if (fitness.second < best_fitness.second ||
            (fitness.second == best_fitness.second &&
             fitness.first < best_fitness.first)) {
          best_fitness = fitness;
          best_solution = std::move(solution);
          best_output_dirty = true;
        }

        if (best_output_dirty && best_output.due()) {
          best_output.store(serialize_agreement_forest(
              original_output_tree, trees.front(), best_solution, meta_map,
              original_leaf_count));
          best_output_dirty = false;
        }
      }
    }

    // Store the latest improved solution before writing the final answer.
    if (best_output_dirty) {
      best_output.store(serialize_agreement_forest(
          original_output_tree, trees.front(), best_solution, meta_map,
          original_leaf_count));
    }
    best_output.publish_final();
  } catch (const std::exception& error) {
    if (!keep_running() && parsed_args.has_value()) {
      // If interrupted after a valid solution is available, stop without
      // printing a partial fallback.
      if (valid_output_exists) {
        return 0;
      }
      if (announced_leaf_count > 0) {
        try {
          const std::string fallback = singleton_solution(announced_leaf_count);
          std::cout << fallback;
          return 0;
        } catch (...) {
        }
      }
    }
    std::cerr << "Error: " << error.what() << '\n';
    return 1;
  }
  return 0;
}