// Copyright 2025 Oliver Theimer
#pragma once
#include <climits>
#include <cstdint>
#include <string>
#include <vector>

#include "../bbdd/include/bbdd.hpp"
#include "../bbdd/include/bbdd_node.hpp"
#include "../bbdd/include/unique_table.hpp"
#include "mockturtle/traits.hpp"
#include "partition.hpp"
#include "util.hpp"

#define LEFT_CHILD(node) node.children[0].index
#define RIGHT_CHILD(node) node.children[1].index

#define CUBE(node) ntk._storage->data.covers[node.data[1].h1]

int nodes_visited;

using namespace mockturtle;

/**
 * @brief maps network nodes to the bbdd variables of a single block
 *
 * The bbdd chain variable ordering addresses its variables through
 * `POS(in, cvo)`, which is a plain `in % cvo->size()`, so the variable ids of
 * a block have to be the contiguous range `2 .. k+1` for that modulo to be a
 * bijection. Network node indices of a block are arbitrary, hence this
 * indirection: `local[n]` is the variable id of node `n`, or 0 when `n` is not
 * a leaf of the block.
 */
struct leaf_map_t {
  /// node index -> bbdd variable id in 2..k+1, 0 when not a leaf
  std::vector<uint32_t> local;
  /// bbdd 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 bbdd 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;
}

/**
 * @brief recursively creates a bbdd based on the given network
 *
 * The recursion stops at the leaves of the block, which are the primary inputs
 * for an unpartitioned run and additionally the boundary nets of the block
 * otherwise. Leaves become bbdd variables, everything above them is built as
 * bbdd nodes.
 *
 * @tparam Ntk template for which network type is used
 * @param table unique table in which the bbdd should be stored in
 * @param ntk input network data structure that should be converted
 * @param leaf_map maps the leaves of the block to their bbdd variable ids
 * @param node_i starting node from which onwards the bbdd should be created,
 * will be used for the recursion
 * @param output_completed statistics that are used in the recursion to properly
 * show the progress bar
 * @param show_progress whether the progress bar should be drawn
 * @return root node of the created bbdd
 */
template <typename Ntk>
static bbdd_node_t *create_bbdd(Unique_table *table, Ntk &ntk,
                                leaf_map_t const &leaf_map, const auto &node_i,
                                int output_completed, bool show_progress) {
  static_assert(has_is_ci_v<Ntk>, "Ntk does not implement the is_ci method");
  const auto &node = ntk._storage->nodes[node_i];
  bbdd_node_t *f, *g, *result;

  // constants are shared between every block and are always in the table
  if (node_i == 1 || node_i == 0) {
    return table->get_node_p(node_i);
  }

  // a leaf of the block: a primary input, or a boundary net that another
  // block drives. Either way it terminates the recursion and becomes a
  // variable of this block's bbdd.
  if (leaf_map[node_i]) {
    assert(node_i < pow(2, 31));
    return table->insert_node({{leaf_map[node_i], INT_MAX}, 0, 1});
  }

  bbdd_op_t op = Util::cube_to_bbdd_op(CUBE(node));
  if (ntk.visited(node_i)) {
    return table->get_node_p(ntk.visited(node_i));
  }
#ifdef DEBUG_COVER
  printf("[INFO] CREATE: node %d with op %s\n", node_i, bbdd_op_s[op].c_str());
#endif
  // a node that is neither a constant nor a leaf has to be a gate of the block
  assert((node.children.size() == 2 || node.children.size() == 1) &&
         "node is not a leaf of the block but has no usable fanin");

  // all the logic gates except buf and inv
  if (node.children.size() == 2) {
    node_index_t f_i = LEFT_CHILD(node), g_i = RIGHT_CHILD(node);
    if (leaf_map[f_i] && leaf_map[g_i]) {
      result = base_case_two_inputs(table, leaf_map[f_i], leaf_map[g_i], op);
#ifdef DEBUG_COVER
      printf("[INFO] CREATE: two input case %d: %d %s %d\n", result->index, f_i,
             bbdd_op_s[op].c_str(), g_i);
#endif
      return result;
    }
    f = create_bbdd(table, ntk, leaf_map, LEFT_CHILD(node), output_completed,
                    show_progress);
    g = create_bbdd(table, ntk, leaf_map, RIGHT_CHILD(node), output_completed,
                    show_progress);
    result = merge_bbdds(table, f, g, op);
#ifdef DEBUG_COVER
    std::cout << "[INFO] CREATE: merging " << result->index << ": " << f->index
              << " " << bbdd_op_s[op] << " " << g->index << "\n";
    dump_node(g);
#endif
    ntk.set_visited(node_i, result->index);
    if (show_progress) {
      Util::show_progress_bar(
          {{"Output", {output_completed, ntk._storage->outputs.size()}},
           {"Nodes", {++nodes_visited, ntk._storage->nodes.size()}}},
          50);
    }
    return result;
  } else {
    // buf and inverter nodes
    f = create_bbdd(table, ntk, leaf_map, LEFT_CHILD(node), output_completed,
                    show_progress);
    if (op == bbdd_inv) {
      result = negate_recursive(table, f);
      ntk.set_visited(node_i, result->index);
      if (show_progress) {
        Util::show_progress_bar(
            {{"Output", {output_completed, ntk._storage->outputs.size()}},
             {"Nodes", {++nodes_visited, ntk._storage->nodes.size()}}},
            50);
      }
      return result;
    }
    ntk.set_visited(node_i, f->index);
    if (show_progress) {
      Util::show_progress_bar(
          {{"Output", {output_completed, ntk._storage->outputs.size()}},
           {"Nodes", {++nodes_visited, ntk._storage->nodes.size()}}},
          50);
    }
    return f;
  }
}

/**
 * @brief converts one block of a partition into a bbdd forest
 *
 * Builds one bbdd per root of the block over the block's leaves and registers
 * it in the table under the matching name. The caller owns the table: it has
 * to be initialised with a chain variable ordering over the leaf ids of this
 * block, and it has to be freed and reinitialised between blocks.
 *
 * The visited flags of the network are used to memoize nodes to table indices.
 * They therefore have to be cleared by the caller between two blocks, because
 * indices of a previous table do not mean anything in the current one.
 *
 * @tparam Ntk template of network type that should be used
 * @param table unique table in which the data structure should be stored
 * @param ntk network that holds the block
 * @param block block that should be converted
 * @param leaf_map maps the leaves of the block to their bbdd variable ids
 * @param root_names name of each root of the block, in the same order
 * @param use_height if true the height of the bbdd is used as the target
 * function for the sifting algorithm, otherwise the number of nodes is used
 * @param sifting_repetitions how often sifting should be repeated at the end
 * @param sift_height_limit skip sifting when the forest is taller than this,
 * 0 disables the limit
 * @param show_progress whether the per-node progress bar should be drawn
 */
template <typename Ntk>
void block_to_bbdd(Unique_table *table, Ntk &ntk, block_t const &block,
                   leaf_map_t const &leaf_map,
                   std::vector<std::string> const &root_names, bool use_height,
                   int sifting_repetitions, uint32_t sift_height_limit = 30,
                   bool show_progress = false) {
  assert(block.roots.size() == root_names.size());

  for (size_t i = 0; i < block.roots.size(); i++) {
    bbdd_node_t *output = add_ref(create_bbdd(table, ntk, leaf_map,
                                              block.roots[i], static_cast<int>(i),
                                              show_progress));
    table->add_output(root_names[i], output);
  }

  for (int i = 0; i < sifting_repetitions; i++) {
    if (sift_height_limit == 0 || table->get_total_height() < sift_height_limit) {
      if (use_height) {
        sift(table, &Unique_table::get_total_height, table);
      } else {
        sift(table, &Unique_table::get_total_number_nodes, table);
      }
    }
  }
}

/**
 * @brief main conversion function that should be used as an interface, it
 * interact with the recursive function, instantiates the progress bar and
 * registers output in the unique table
 *
 * Converts the whole network as a single block, which is the unpartitioned
 * behaviour: every primary input is a bbdd variable and every primary output
 * gets its own bbdd.
 *
 * @tparam Ntk template of network type that should be used
 * @param table unique table in which the data structure should be stores
 * @param ntk data strucutre that stores the input network type which should be
 * converted
 * @param use_height if true the height of the bbdd is used as the target
 * function for the sifting algorithm, otherwise the number of nodes i used
 * @param sifting_repetitions indicates how often the sifting algorithm should be repeated at the end of the conversion
 */
template <typename Ntk>
void cover_to_bbdd(Unique_table *table, Ntk &ntk, bool use_height,
                   int sifting_repetitions) {

  static_assert(has_foreach_co_v<Ntk>,
                "Ntk does not implement the foreach_pi method");
  static_assert(has_get_node_v<Ntk>,
                "Ntk does not implement the get_node method");
  static_assert(has_set_visited_v<Ntk>,
                "Ntk does not implement the set_visited method");
  static_assert(has_visited_v<Ntk>,
                "Ntk does not implement the set_visited method");
  int index = ntk._storage->inputs.size() + 2;
  int output_completed = 0;
  nodes_visited = 0;

  leaf_map_t leaf_map = make_leaf_map(
      ntk._storage->nodes.size(),
      std::vector<uint64_t>(std::begin(ntk._storage->inputs),
                            std::end(ntk._storage->inputs)));

  Util::show_progress_bar(
      {{"Output", {output_completed, ntk._storage->outputs.size()}},
       {"Nodes", {++nodes_visited, ntk._storage->nodes.size()}}},
      50);
  ntk.foreach_co([&ntk, &table, &index, &output_completed, &leaf_map](
                     const auto &node_i) {
    const auto &node = ntk.get_node(node_i);
#ifdef DEBUG_COVER
    auto &n = ntk._storage->nodes[node_i];
    std::cout << "output with node_i: " << node_i << " and n->index: " << " : "
              << n.children.size() << "\n";
#endif
    bbdd_node_t *output =
        add_ref(create_bbdd(table, ntk, leaf_map, node, output_completed, true));
    Util::show_progress_bar(
        {{"Output", {++output_completed, ntk._storage->outputs.size()}},
         {"Nodes", {++nodes_visited, ntk._storage->nodes.size()}}},
        50);
    table->add_output(ntk.get_signal_name(index), output);
    index++;
  });
  for (int i = 0; i < sifting_repetitions; i++) {
    if (table->get_total_height() < 30) {
      if (use_height) {
        sift(table, &Unique_table::get_total_height, table);
      } else {
        sift(table, &Unique_table::get_total_number_nodes, table);
      }
    }
  }
}