// 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
 *
 * A block is a cone of logic whose input support is limited to a configurable
 * number of signals. Blocks are disjoint with respect to their gates, so every
 * gate of the network is owned by exactly one block, and they are connected to
 * each other through boundary nets: a leaf of one block is either 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;
};

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
 *
 * Mirrors mockturtle::expand0_towards_tfi, with one difference: a node that is
 * already owned by another block terminates the expansion just like a primary
 * input does. Without that the cone would grow through foreign blocks and
 * duplicate their logic.
 *
 * Only replaces a leaf by its fanins when doing so does not increase the size
 * of the cut, hence "zero cost". Expects all leaves to be painted in the
 * current color and paints every node it pulls into the cut.
 *
 * @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
 *
 * Mirrors mockturtle::detail::select_next_fanin_to_expand_tfi with a caller
 * supplied terminal predicate. Picks the fanin that is referenced most often
 * by the current cut, because expanding it is the most likely to be paid back
 * by a later zero cost merge, and breaks ties towards the higher fanout node.
 *
 * @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;
}

} // namespace partition_detail

/**
 * @brief splits a network into subcircuits with a bounded input support
 *
 * Grows a cone backwards from every primary output for as long as its input
 * support stays within the budget. A fanin that would push the support over
 * the budget is cut instead: it becomes a boundary net, which is a leaf of the
 * current block and the root of a block of its own. Repeating this until no
 * boundary net is left covers every gate of the network exactly once.
 *
 * With a budget of at least the number of primary inputs the result is one
 * block per primary output cone whose leaves are the primary inputs, which is
 * the behaviour of the unpartitioned flow.
 *
 * @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;

  // no partitioning: a single block holding the whole network
  if (max_inputs == 0) {
    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)});
  }

  return blocks;
}

/**
 * @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;
}