// Copyright 2025 Oliver Theimer
#pragma once

#include <unordered_map>

#include "../bbdd/include/bbdd_node.hpp"
#include "bdd.h"
#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]

using namespace mockturtle;

/// node -> bdd root; refcounting keeps the roots alive across autoreorder
using bdd_memo_t = std::unordered_map<uint64_t, bdd>;

/**
 * @brief recursively creates a bdd based on the given network
 *
 * Recursion stops at leaf_map leaves (PIs, plus boundary nets when
 * partitioned).
 *
 * @tparam Ntk which network type is used as an input to the function
 * @param ntk input network which should be converted to bdd
 * @param leaf_map maps the leaves of the block to their bdd variable ids
 * @param node_i current network node which is currently converted is used for
 * the recursion
 * @param memo maps an already converted network node to its bdd root
 * @return return root bdd node of the created binary decision diagram
 */
template <typename Ntk>
static bdd create_bdd(Ntk &ntk, leaf_map_t const &leaf_map, const auto &node_i,
                      bdd_memo_t &memo) {
  static_assert(has_is_ci_v<Ntk>, "Ntk does not implement the is_ci method");
  const auto &node = ntk._storage->nodes[node_i];
  bdd f, g, result;

  if (node_i == 1) {
#ifdef DEBUG_BDD
    printf("[INFO] cover to bdd: constant 1 node\n");
#endif
    return bdd_true();
  }
  if (node_i == 0) {
#ifdef DEBUG_BDD
    printf("[INFO] cover to bdd: constant 0 node\n");
#endif
    return bdd_false();
  }

  // block leaf (PI or boundary net) -> variable. Must precede the fanin
  // dispatch: boundary nets have children.
  if (leaf_map[node_i]) {
    assert(node_i < pow(2, 31));
    assert(leaf_map[node_i] - 2 < leaf_map.size());
#ifdef DEBUG_BDD
    printf("[INFO] cover to bdd: base case leaf: %d\n", node_i);
#endif
    return bdd_ithvar(leaf_map[node_i] - 2);
  }

  // memoised: reconvergent fanout is otherwise exponential
  auto cached = memo.find(node_i);
  if (cached != memo.end()) {
    return cached->second;
  }

  assert(node.children.size() == 2 || node.children.size() == 1);
  bbdd_op_t op = Util::cube_to_bbdd_op(CUBE(node));

  // all the logic gates except buf and inv
  if (node.children.size() == 2) {
    f = create_bdd(ntk, leaf_map, LEFT_CHILD(node), memo);
#ifdef DEBUG_BDD
    printf("[INFO] cover to bdd: return from f\n");
#endif
    g = create_bdd(ntk, leaf_map, RIGHT_CHILD(node), memo);
#ifdef DEBUG_BDD
    printf("[INFO] cover to bdd: merge %d %s %d\n", f, bbdd_op_s[op].c_str(),
           g);
#endif
    switch (op) {
    case bbdd_and:
      result = bdd_apply(f, g, bddop_and);
      break;
    case bbdd_xor:
      result = bdd_apply(f, g, bddop_xor);
      break;
    case bbdd_xnor:
      result = bdd_not(bdd_apply(f, g, bddop_xor));
      break;
    case bbdd_or:
      result = bdd_apply(f, g, bddop_or);
      break;
    case bbdd_nand:
      result = bdd_apply(f, g, bddop_nand);
      break;
    case bbdd_nor:
      result = bdd_apply(f, g, bddop_nor);
      break;
    default:
      printf("%s\n", bbdd_op_s[op].c_str());
      assert(false && "not implemented yet");
      break;
    }
    memo.emplace(node_i, result);
    return result;
  } else {
    // buf and inverter nodes
    f = create_bdd(ntk, leaf_map, LEFT_CHILD(node), memo);
    if (op == bbdd_inv) {
#ifdef DEBUG_BDD
      printf("[INFO] cover to bdd: negate\n");
#endif
      result = bdd_not(f);
      memo.emplace(node_i, result);
      return result;
    }
#ifdef DEBUG_BDD
    printf("[INFO] cover to bdd: buffer\n");
#endif
    memo.emplace(node_i, f);
    return f;
  }
}

/**
 * @brief converts one block of a partition into a bdd forest
 *
 * One bdd per root, appended to output_map. The caller reinitialises BuDDy per
 * block with at least one variable per leaf.
 *
 * @tparam Ntk template of network type that should be used
 * @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 bdd variable ids
 * @param root_names name of each root of the block, in the same order
 * @param output_map collects the root of every diagram with its name
 */
template <typename Ntk>
void block_to_bdd(Ntk &ntk, block_t const &block, leaf_map_t const &leaf_map,
                  std::vector<std::string> const &root_names,
                  std::vector<std::pair<bdd, std::string>> &output_map) {
  assert(block.roots.size() == root_names.size());
  // shared so common cones are built once
  bdd_memo_t memo;
  for (size_t i = 0; i < block.roots.size(); i++) {
    output_map.emplace_back(create_bdd(ntk, leaf_map, block.roots[i], memo),
                            root_names[i]);
  }
}

/**
 * @brief main interface function that handles the conversion to a bdd from a given input network type
 *
 * @tparam Ntk network type that is used as an input to the conversion
 * @param ntk network that should be converted to a bdd
 * @param output_map outputmap that stores the pointers to the root nodes of each output in the network
 */
template <typename Ntk>
void cover_to_bdd(Ntk &ntk,
                  std::vector<std::pair<bdd, std::string>> &output_map) {

  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;

  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)));

  // shared so common cones are built once
  bdd_memo_t memo;
  ntk.foreach_co([&ntk, &index, &output_map, &memo,
                  &leaf_map](const auto &node_i) {
    const auto &node = ntk.get_node(node_i);
    bdd output = create_bdd(ntk, leaf_map, node, memo);
    output_map.emplace_back(output, ntk.get_signal_name(index));
#ifdef DEBUG_BDD
    bdd_printtable(output);
#endif
    index++;
  });
}