// Copyright 2025 Oliver Theimer
#pragma once
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <optional>
#include <vector>

#include "mockturtle/networks/cover.hpp"
#include "mockturtle/utils/window_utils.hpp"
#include "mockturtle/views/color_view.hpp"

/**
 * @class block_t
 * @brief one subcircuit of a support-bounded partition
 *
 * Every gate is owned by exactly one block; a leaf is a primary input, a
 * constant, or a root of another block.
 */
struct block_t {
  /// input support of the block, at most `max_inputs` entries
  std::vector<uint64_t> leaves;
  /// gates owned by the block, disjoint between blocks, includes the roots
  std::vector<uint64_t> gates;
  /// gates that drive a primary output or a gate outside of this block
  std::vector<uint64_t> roots;
};

/**
 * @brief maps network nodes to the decision diagram variables of a single block
 *
 * Renumbers a block's leaves to the contiguous ids 2..k+1, because bbdd's
 * POS(in, cvo) is `in % cvo->size()`. BuDDy gets `id - 2`.
 */
struct leaf_map_t {
  /// node index -> diagram variable id in 2..k+1, 0 when not a leaf
  std::vector<uint32_t> local;
  /// diagram variable id - 2 -> node index
  std::vector<uint64_t> nodes;

  uint32_t operator[](uint64_t n) const { return local[n]; }
  size_t size() const { return nodes.size(); }
};

/**
 * @brief builds the leaf map of a block
 *
 * @param network_size number of nodes in the network the block came from
 * @param leaves leaves of the block, at most one diagram variable each
 * @return leaf map assigning the contiguous ids 2..k+1 to the leaves
 */
inline leaf_map_t make_leaf_map(size_t network_size,
                                std::vector<uint64_t> const &leaves) {
  leaf_map_t map;
  map.local.assign(network_size, 0);
  map.nodes.reserve(leaves.size());
  for (uint64_t const &l : leaves) {
    map.local[l] = static_cast<uint32_t>(map.nodes.size()) + 2;
    map.nodes.push_back(l);
  }
  return map;
}

namespace partition_detail {

using cover_color_view =
    mockturtle::out_of_place_color_view<mockturtle::cover_network>;

/// number of times the cut is allowed to exceed the budget before giving up,
/// mirrors mockturtle::expand_towards_tfi
static constexpr uint32_t MAX_ITERATIONS = 5u;

/**
 * @brief zero cost expansion of a cut towards the transitive fanin
 *
 * mockturtle::expand0_towards_tfi, except that gates owned by another block
 * are terminal too, so a cone never duplicates foreign logic.
 *
 * @param ntk network wrapped in a color view
 * @param leaves cut that should be expanded, modified in place
 * @param is_terminal predicate that decides where the expansion has to stop
 * @return true if and only if every leaf is a terminal
 */
template <typename Ntk, typename TermFn>
bool expand0_bounded(Ntk const &ntk, std::vector<typename Ntk::node> &leaves,
                     TermFn &&is_terminal) {
  using node = typename Ntk::node;

  bool trivial_cut = true;
  bool changed = true;
  std::vector<node> new_leaves;

  while (changed) {
    trivial_cut = true;
    changed = false;

    for (auto it = std::begin(leaves); it != std::end(leaves);) {
      if (is_terminal(*it)) {
        ++it;
        continue;
      }
      trivial_cut = false;

      // count the fanins that are already part of the cut
      uint32_t count_fanin_inside = 0;
      std::optional<node> outside;
      ntk.foreach_fanin(*it, [&](auto const &fi) {
        node const f = ntk.get_node(fi);
        if (ntk.eval_color(
                f, [&ntk](auto c) { return c == ntk.current_color(); })) {
          ++count_fanin_inside;
        } else {
          outside = f;
        }
      });

      // expansion would grow the cut, leave this leaf alone
      if (count_fanin_inside + 1 < ntk.fanin_size(*it)) {
        ++it;
        continue;
      }

      if (outside) {
        if (ntk.eval_color(*outside, [&ntk](auto c) {
              return c != ntk.current_color();
            })) {
          new_leaves.push_back(*outside);
          ntk.paint(*outside);
        }
      }
      it = leaves.erase(it);
      changed = true;
    }

    std::copy(std::begin(new_leaves), std::end(new_leaves),
              std::back_inserter(leaves));
    new_leaves.clear();
  }

  return trivial_cut;
}

/**
 * @brief selects the next fanin that should be pulled into the cut
 *
 * select_next_fanin_to_expand_tfi with a custom terminal predicate: the most
 * referenced fanin wins, ties go to the higher fanout.
 *
 * @param ntk network wrapped in a color view
 * @param leaves current cut
 * @param is_terminal predicate that decides which leaves cannot be expanded
 * @return the fanin to expand, or nothing when the cut is terminal
 */
template <typename Ntk, typename TermFn>
std::optional<typename Ntk::node>
select_next_fanin(Ntk const &ntk, std::vector<typename Ntk::node> const &leaves,
                  TermFn &&is_terminal) {
  using node = typename Ntk::node;

  std::vector<std::pair<node, uint32_t>> candidates;
  for (auto const &l : leaves) {
    if (is_terminal(l)) {
      continue;
    }
    ntk.foreach_fanin(l, [&](auto const &fi) {
      node const f = ntk.get_node(fi);
      if (ntk.is_constant(f)) {
        return;
      }
      auto it = std::find_if(std::begin(candidates), std::end(candidates),
                             [&f](auto const &p) { return p.first == f; });
      if (it == std::end(candidates)) {
        candidates.emplace_back(f, 1u);
      } else {
        ++it->second;
      }
    });
  }

  if (candidates.empty()) {
    return std::nullopt;
  }

  std::pair<node, uint32_t> best = candidates.front();
  for (auto const &candidate : candidates) {
    if (candidate.second > best.second ||
        (candidate.second == best.second &&
         ntk.fanout_size(candidate.first) > ntk.fanout_size(best.first))) {
      best = candidate;
    }
  }
  return best.first;
}

/// how many merge partners one block examines per visit; caps the merge pass
/// on designs that start out with tens of thousands of blocks
static constexpr size_t MAX_MERGE_CANDIDATES = 64u;

/// how often the merge pass sweeps all blocks before giving up on a fixpoint
static constexpr uint32_t MAX_MERGE_PASSES = 4u;

/**
 * @brief leaf set that two blocks would have if they were merged
 *
 * Union of both leaf sets minus the nets either block drives, since those
 * become internal.
 *
 * @param a leaves of the first block, sorted
 * @param b leaves of the second block, sorted
 * @param owner gate -> block id + 1, 0 for anything that is not a gate
 * @param id_a id of the first block
 * @param id_b id of the second block
 * @return the merged leaf set, sorted
 */
inline std::vector<uint64_t> merge_leaf_sets(std::vector<uint64_t> const &a,
                                             std::vector<uint64_t> const &b,
                                             std::vector<uint32_t> const &owner,
                                             uint32_t id_a, uint32_t id_b) {
  std::vector<uint64_t> u;
  u.reserve(a.size() + b.size());
  std::set_union(std::begin(a), std::end(a), std::begin(b), std::end(b),
                 std::back_inserter(u));
  u.erase(std::remove_if(std::begin(u), std::end(u),
                         [&](uint64_t l) {
                           uint32_t const o = owner[l];
                           return o == id_a + 1 || o == id_b + 1;
                         }),
          std::end(u));
  return u;
}

/**
 * @brief merges blocks for as long as the budget allows
 *
 * Growth stops every cone on gates an earlier cone claimed, so no block ever
 * exceeds one output cone. This greedily merges pairs whose merged support
 * still fits `max_inputs`. Candidates are the block driving a leaf and the
 * blocks sharing a leaf, capped per visit.
 *
 * @param ntk network the blocks were derived from
 * @param cntk the same network in the color view, for collect_outputs
 * @param blocks partition to merge in place; roots are recomputed afterwards
 * @param owner gate -> block id + 1, updated in place and stale once the
 * surviving blocks are renumbered
 * @param refs scratch reference counts for collect_outputs, restored to 0
 * @param max_inputs maximum input support per block
 */
template <typename ColorNtk>
inline void merge_blocks(mockturtle::cover_network const &ntk,
                         ColorNtk const &cntk, std::vector<block_t> &blocks,
                         std::vector<uint32_t> &owner,
                         std::vector<uint32_t> &refs, uint32_t max_inputs) {
  using node = mockturtle::cover_network::node;

  uint32_t const count = static_cast<uint32_t>(blocks.size());
  if (count < 2) {
    return;
  }

  std::vector<char> alive(count, 1);

  // leaf -> blocks that have it as a leaf, so that two blocks over the same
  // support find each other without scanning the whole partition
  std::vector<std::vector<uint32_t>> leaf_blocks(ntk.size());
  for (uint32_t b = 0; b < count; b++) {
    for (node const &l : blocks[b].leaves) {
      leaf_blocks[l].push_back(b);
    }
  }

  // dedupes the candidate list without clearing a set on every visit
  std::vector<uint32_t> stamp(count, 0);
  uint32_t stamp_id = 0;
  std::vector<uint32_t> candidates;

  bool changed = true;
  for (uint32_t pass = 0; changed && pass < MAX_MERGE_PASSES; pass++) {
    changed = false;
    for (uint32_t b = 0; b < count; b++) {
      // saturate this block before moving on: every merge changes its leaf
      // set, which can open up a partner that did not fit a moment ago
      for (bool grown = true; alive[b] && grown;) {
        grown = false;

        ++stamp_id;
        stamp[b] = stamp_id;
        candidates.clear();
        for (node const &l : blocks[b].leaves) {
          if (owner[l] != 0 && stamp[owner[l] - 1] != stamp_id) {
            stamp[owner[l] - 1] = stamp_id;
            candidates.push_back(owner[l] - 1);
          }
          for (uint32_t const &o : leaf_blocks[l]) {
            if (stamp[o] == stamp_id) {
              continue;
            }
            stamp[o] = stamp_id;
            candidates.push_back(o);
            if (candidates.size() >= MAX_MERGE_CANDIDATES) {
              break;
            }
          }
          if (candidates.size() >= MAX_MERGE_CANDIDATES) {
            break;
          }
        }

        // smallest merged support first; first-fit gets stuck far earlier
        uint32_t best = UINT32_MAX;
        std::vector<node> best_leaves;
        for (uint32_t const &c : candidates) {
          if (c == b || !alive[c]) {
            continue;
          }
          std::vector<node> leaves =
              merge_leaf_sets(blocks[b].leaves, blocks[c].leaves, owner, b, c);
          if (leaves.size() > max_inputs) {
            continue;
          }
          if (best == UINT32_MAX || leaves.size() < best_leaves.size()) {
            best = c;
            best_leaves = std::move(leaves);
          }
        }

        if (best != UINT32_MAX) {
          uint32_t const c = best;
          std::vector<node> leaves = std::move(best_leaves);

          std::vector<node> gates;
          gates.reserve(blocks[b].gates.size() + blocks[c].gates.size());
          std::merge(std::begin(blocks[b].gates), std::end(blocks[b].gates),
                     std::begin(blocks[c].gates), std::end(blocks[c].gates),
                     std::back_inserter(gates));
          for (node const &g : blocks[c].gates) {
            owner[g] = b + 1;
          }

          blocks[b].leaves = std::move(leaves);
          blocks[b].gates = std::move(gates);
          blocks[c] = block_t{};
          alive[c] = 0;

          for (node const &l : blocks[b].leaves) {
            if (leaf_blocks[l].empty() || leaf_blocks[l].back() != b) {
              leaf_blocks[l].push_back(b);
            }
          }
          grown = true;
          changed = true;
        }
      }
    }
  }

  // drop the absorbed blocks and recompute the roots of the survivors: a
  // boundary net that a merge turned into an internal net is not a root any
  // more unless something outside the merged block still reads it
  std::vector<block_t> survivors;
  survivors.reserve(count);
  for (uint32_t b = 0; b < count; b++) {
    if (!alive[b]) {
      continue;
    }
    std::vector<node> roots = mockturtle::collect_outputs(
        cntk, blocks[b].leaves, blocks[b].gates, refs);
    std::stable_sort(std::begin(roots), std::end(roots));
    survivors.push_back(block_t{std::move(blocks[b].leaves),
                                std::move(blocks[b].gates), std::move(roots)});
  }
  blocks = std::move(survivors);
}

} // namespace partition_detail

/**
 * @brief splits a network into subcircuits with a bounded input support
 *
 * Grows a cone back from every primary output while its support fits the
 * budget; a fanin that would exceed it becomes a boundary net and the root of
 * a new block. merge_blocks then joins blocks while the budget allows.
 * `n >= |PI|` short-circuits to the unpartitioned single block.
 *
 * @param ntk network that should be partitioned
 * @param max_inputs maximum input support per block, 0 disables partitioning
 * @return the blocks, in the order they were discovered
 */
inline std::vector<block_t> partition_cover(mockturtle::cover_network &ntk,
                                            uint32_t max_inputs) {
  using node = mockturtle::cover_network::node;

  const uint32_t size = ntk.size();
  std::vector<block_t> blocks;

  // n == 0 or n >= |PI|: one block over all inputs. Handled explicitly because
  // greedy pairwise merging can get stuck before reaching it.
  if (max_inputs == 0 || max_inputs >= ntk._storage->inputs.size()) {
    block_t all;
    all.leaves.assign(std::begin(ntk._storage->inputs),
                      std::end(ntk._storage->inputs));
    ntk.foreach_gate([&](auto const &n) { all.gates.push_back(n); });
    ntk.foreach_co([&](auto const &f) {
      node const d = ntk.get_node(f);
      if (std::find(std::begin(all.roots), std::end(all.roots), d) ==
          std::end(all.roots)) {
        all.roots.push_back(d);
      }
    });
    blocks.emplace_back(std::move(all));
    return blocks;
  }

  assert(max_inputs >= mockturtle::cover_network::min_fanin_size + 1 &&
         "budget too small to ever cover a single gate");

  // block id + 1 for owned gates, 0 while a gate is still unassigned
  std::vector<uint32_t> assigned(size, 0);
  // scratch reference counts for collect_outputs, always restored to 0
  std::vector<uint32_t> refs(size, 0);
  // guards against queueing the same root twice
  std::vector<char> queued(size, 0);

  partition_detail::cover_color_view cntk{ntk};

  auto is_terminal = [&](node n) {
    return ntk.is_constant(n) || ntk.is_ci(n) || assigned[n] != 0;
  };

  // seed with the primary output drivers, in a deterministic order so that
  // repeated runs on the same design produce the same partition
  std::vector<node> worklist;
  ntk.foreach_co([&](auto const &f) {
    node const d = ntk.get_node(f);
    if (!queued[d]) {
      queued[d] = 1;
      worklist.push_back(d);
    }
  });

  // the worklist grows while it is walked: every boundary net that is cut
  // becomes the root of a block of its own
  for (size_t w = 0; w < worklist.size(); ++w) {
    node const root = worklist[w];
    if (is_terminal(root)) {
      // a primary output driven straight by an input or a constant, or a
      // boundary net that a later cone happened to absorb
      continue;
    }

    // grow the cone from the root towards the inputs
    std::vector<node> leaves{root};
    cntk.new_color();
    cntk.paint(root);

    bool trivial = partition_detail::expand0_bounded(cntk, leaves, is_terminal);
    std::optional<std::vector<node>> best;
    if (leaves.size() <= max_inputs) {
      best = leaves;
    }

    uint32_t iterations = 0;
    while (!trivial && (leaves.size() <= max_inputs ||
                        iterations < partition_detail::MAX_ITERATIONS)) {
      auto next =
          partition_detail::select_next_fanin(cntk, leaves, is_terminal);
      if (!next) {
        break;
      }
      leaves.push_back(*next);
      cntk.paint(*next);
      trivial = partition_detail::expand0_bounded(cntk, leaves, is_terminal);

      iterations = leaves.size() > max_inputs ? iterations + 1 : 0;
      if (leaves.size() <= max_inputs &&
          (!best || best->size() <= leaves.size())) {
        best = leaves;
      }
    }

    if (best) {
      leaves = *best;
    }
    // the budget could not be met even by a single gate, fall back to a block
    // holding just the root so that the partition still covers the network
    if (leaves.size() > max_inputs) {
      leaves.clear();
      ntk.foreach_fanin(root, [&](auto const &fi) {
        node const f = ntk.get_node(fi);
        if (std::find(std::begin(leaves), std::end(leaves), f) ==
            std::end(leaves)) {
          leaves.push_back(f);
        }
      });
    }

    // everything strictly between the root and the leaves belongs to the block
    std::vector<node> gates = mockturtle::cover(cntk, root, leaves);
    if (gates.empty()) {
      continue;
    }

    const uint32_t id = static_cast<uint32_t>(blocks.size());
    for (node const &g : gates) {
      assert(assigned[g] == 0 && "gate claimed by two blocks");
      assigned[g] = id + 1;
    }

    // a gate whose fanout count is not fully explained by uses inside the
    // block drives either a primary output or a foreign block
    std::vector<node> roots =
        mockturtle::collect_outputs(cntk, leaves, gates, refs);

    // every leaf that is a gate is a boundary net and needs a block of its own
    for (node const &l : leaves) {
      if (!ntk.is_constant(l) && !ntk.is_ci(l) && !assigned[l] && !queued[l]) {
        queued[l] = 1;
        worklist.push_back(l);
      }
    }

    std::stable_sort(std::begin(leaves), std::end(leaves));
    std::stable_sort(std::begin(gates), std::end(gates));
    std::stable_sort(std::begin(roots), std::end(roots));

    blocks.push_back(
        block_t{std::move(leaves), std::move(gates), std::move(roots)});
  }

  partition_detail::merge_blocks(ntk, cntk, blocks, assigned, refs, max_inputs);

  return blocks;
}

/**
 * @brief merges blocks for as long as the merged diagram is actually smaller
 *
 * Like merge_blocks, but a merge is kept only if cost_fn(merged) is below the
 * sum of its parts, so the total node count never rises. Support is a poor
 * proxy for diagram size. `cap` bounds the support, and with it the cost of
 * a trial build; the smallest support is tried first.
 *
 * @param ntk network the blocks were derived from
 * @param blocks partition to merge in place; roots are recomputed as it goes
 * @param cap maximum input support a merged block may have
 * @param max_passes how often all blocks are swept before giving up
 * @param max_candidates how many partners one block examines per visit
 * @param cost_fn builds a block and returns its node count; must leave the
 * network as it found it
 */
template <typename CostFn>
inline void merge_blocks_by_gain(mockturtle::cover_network &ntk,
                                 std::vector<block_t> &blocks, uint32_t cap,
                                 uint32_t max_passes, uint32_t max_candidates,
                                 CostFn &&cost_fn) {
  using node = mockturtle::cover_network::node;

  uint32_t const count = static_cast<uint32_t>(blocks.size());
  if (count < 2 || cap == 0) {
    return;
  }

  partition_detail::cover_color_view cntk{ntk};
  std::vector<uint32_t> refs(ntk.size(), 0);
  std::vector<uint32_t> owner(ntk.size(), 0);
  for (uint32_t b = 0; b < count; b++) {
    for (node const &g : blocks[b].gates) {
      owner[g] = b + 1;
    }
  }

  std::vector<char> alive(count, 1);
  std::vector<uint64_t> cost(count, 0);
  for (uint32_t b = 0; b < count; b++) {
    cost[b] = cost_fn(blocks[b]);
  }

  std::vector<std::vector<uint32_t>> leaf_blocks(ntk.size());
  for (uint32_t b = 0; b < count; b++) {
    for (node const &l : blocks[b].leaves) {
      leaf_blocks[l].push_back(b);
    }
  }

  // dedupes the candidate list without clearing a set on every visit
  std::vector<uint32_t> stamp(count, 0);
  uint32_t stamp_id = 0;
  std::vector<std::pair<uint32_t, std::vector<uint64_t>>> candidates;

  bool changed = true;
  for (uint32_t pass = 0; changed && pass < max_passes; pass++) {
    changed = false;
    for (uint32_t b = 0; b < count; b++) {
      // saturate this block before moving on: a merge changes its leaf set,
      // which can open up a partner that did not fit a moment ago
      for (bool grown = true; alive[b] && grown;) {
        grown = false;

        ++stamp_id;
        stamp[b] = stamp_id;
        candidates.clear();

        auto consider = [&](uint32_t c) {
          if (c == b || !alive[c] || stamp[c] == stamp_id) {
            return;
          }
          stamp[c] = stamp_id;
          std::vector<uint64_t> leaves = partition_detail::merge_leaf_sets(
              blocks[b].leaves, blocks[c].leaves, owner, b, c);
          if (leaves.size() > cap) {
            return;
          }
          candidates.emplace_back(c, std::move(leaves));
        };

        for (node const &l : blocks[b].leaves) {
          if (owner[l] != 0) {
            consider(owner[l] - 1);
          }
          for (uint32_t const &o : leaf_blocks[l]) {
            consider(o);
          }
          if (candidates.size() >= max_candidates) {
            break;
          }
        }

        std::stable_sort(std::begin(candidates), std::end(candidates),
                         [](auto const &x, auto const &y) {
                           return x.second.size() < y.second.size();
                         });
        if (candidates.size() > max_candidates) {
          candidates.resize(max_candidates);
        }

        for (auto &candidate : candidates) {
          uint32_t const c = candidate.first;

          block_t merged;
          merged.leaves = std::move(candidate.second);
          merged.gates.reserve(blocks[b].gates.size() + blocks[c].gates.size());
          std::merge(std::begin(blocks[b].gates), std::end(blocks[b].gates),
                     std::begin(blocks[c].gates), std::end(blocks[c].gates),
                     std::back_inserter(merged.gates));
          merged.roots = mockturtle::collect_outputs(cntk, merged.leaves,
                                                     merged.gates, refs);
          std::stable_sort(std::begin(merged.roots), std::end(merged.roots));

          uint64_t const merged_cost = cost_fn(merged);
          if (merged_cost >= cost[b] + cost[c]) {
            continue;
          }

          for (node const &g : blocks[c].gates) {
            owner[g] = b + 1;
          }
          blocks[b] = std::move(merged);
          blocks[c] = block_t{};
          alive[c] = 0;
          cost[b] = merged_cost;

          for (node const &l : blocks[b].leaves) {
            if (leaf_blocks[l].empty() || leaf_blocks[l].back() != b) {
              leaf_blocks[l].push_back(b);
            }
          }
          grown = true;
          changed = true;
          break;
        }
      }
    }
  }

  std::vector<block_t> survivors;
  survivors.reserve(count);
  for (uint32_t b = 0; b < count; b++) {
    if (alive[b]) {
      survivors.push_back(std::move(blocks[b]));
    }
  }
  blocks = std::move(survivors);
}

/**
 * @brief checks the invariants a partition has to satisfy
 *
 * Meant to be called behind an assert or from a test: every gate is owned by
 * exactly one block, no block exceeds the budget, and every leaf is either a
 * primary input, a constant, or a root of some block.
 *
 * @param ntk network the blocks were derived from
 * @param blocks partition that should be checked
 * @param max_inputs budget the partition was built with, 0 skips the check
 * @return true if and only if the partition is well formed
 */
inline bool partition_is_valid(mockturtle::cover_network const &ntk,
                               std::vector<block_t> const &blocks,
                               uint32_t max_inputs) {
  using node = mockturtle::cover_network::node;

  std::vector<uint32_t> owner(ntk.size(), 0);
  std::vector<char> is_root(ntk.size(), 0);

  for (size_t i = 0; i < blocks.size(); i++) {
    if (max_inputs > 0 && blocks[i].leaves.size() > max_inputs) {
      std::cerr << "[ERROR] block " << i << " has " << blocks[i].leaves.size()
                << " leaves, budget is " << max_inputs << "\n";
      return false;
    }
    for (node const &g : blocks[i].gates) {
      if (owner[g] != 0) {
        std::cerr << "[ERROR] gate " << g << " owned by block " << owner[g] - 1
                  << " and block " << i << "\n";
        return false;
      }
      owner[g] = static_cast<uint32_t>(i) + 1;
    }
    for (node const &r : blocks[i].roots) {
      is_root[r] = 1;
    }
  }

  for (auto const &b : blocks) {
    for (node const &l : b.leaves) {
      if (ntk.is_constant(l) || ntk.is_ci(l)) {
        continue;
      }
      if (!is_root[l]) {
        std::cerr << "[ERROR] leaf " << l
                  << " is neither an input nor a root\n";
        return false;
      }
    }
  }

  bool complete = true;
  ntk.foreach_gate([&](auto const &n) {
    if (owner[n] == 0) {
      std::cerr << "[ERROR] gate " << n << " is not covered by any block\n";
      complete = false;
    }
  });
  return complete;
}