// Copyright 2025 Oliver Theimer
#include <algorithm>
#include <cstdlib>
#include <filesystem>
#include <unistd.h>
#include <iostream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <cstdio>

#include <lorina/diagnostics.hpp>
#include <lorina/lorina.hpp>
#include <mockturtle/algorithms/cover_to_graph.hpp>
#include <mockturtle/mockturtle.hpp>

#include "bdd.h"
#include "mockturtle/io/write_blif.hpp"
#include "mockturtle/networks/cover.hpp"
#include "util/cover_to_bdd.hpp"
#include "util/partition.hpp"

/**
 * @brief initial size of the BuDDy node table for one block
 *
 * As block_table_size in convert_bbdd, but only an initial size: BuDDy grows
 * on demand.
 *
 * @param leaf_count number of variables of the block
 * @param table_size upper bound requested on the command line
 * @return initial node table size for this block
 */
static int block_node_table(size_t leaf_count, int table_size) {
  const uint32_t shift = static_cast<uint32_t>(std::min<size_t>(leaf_count, 17));
  uint64_t wanted = 8ull * (1ull << shift);
  if (wanted < 16384ull) {
    wanted = 16384ull;
  }
  if (wanted > static_cast<uint64_t>(table_size)) {
    wanted = static_cast<uint64_t>(table_size);
  }
  return static_cast<int>(wanted);
}

/**
 * @brief writes the diagram below one node as blif
 *
 * @param blif_file netlist that is being written
 * @param signal_map name of every variable of the block, indexed by variable
 * @param prefix namespace for the node ids of this block, empty for a single
 * block so that an unpartitioned run keeps writing the netlist it always did
 * @param node diagram node to write
 * @param emitted nodes already defined in this block
 * @param output whether this is a root, whose line carries the output name
 */
int write_blif_recursive(std::ofstream &blif_file,
                         std::vector<std::string> const &signal_map,
                         std::string const &prefix, const BDD &node,
                         std::unordered_set<BDD> &emitted, bool output) {
  if (node < 2) {
    return node;
  }
  // define each node once; roots still get their own line, which carries the
  // output name
  if (!output && !emitted.insert(node).second) {
    return node;
  }
  int low_out = write_blif_recursive(blif_file, signal_map, prefix,
                                     bdd_low(node), emitted, false);
  int high_out = write_blif_recursive(blif_file, signal_map, prefix,
                                      bdd_high(node), emitted, false);
  blif_file << ".names " << signal_map[bdd_var(node)] << " ";
  if (IS_TERM(low_out)) {
    if (low_out == 0) {
      blif_file << "$false ";
    } else {
      blif_file << "$true ";
    }
  } else {
    blif_file << prefix << std::to_string(low_out);
  }
  blif_file << " ";
  if (IS_TERM(high_out)) {
    if (high_out == 0) {
      blif_file << "$false ";
    } else {
      blif_file << "$true ";
    }
  } else {
    blif_file << prefix << std::to_string(high_out) << " ";
  }
  if (!output) {
    blif_file << prefix << std::to_string(node);
    blif_file << "\n0-1 1\n11- 1\n";
  }
  return node;
}

/**
 * @brief writes every diagram of one block into the open netlist
 *
 * @param blif_file netlist that is being written
 * @param signal_map name of every variable of the block, indexed by variable
 * @param prefix namespace for the node ids of this block
 * @param output_map root of every diagram of the block with its name
 */
void write_block_body(std::ofstream &blif_file,
                      std::vector<std::string> const &signal_map,
                      std::string const &prefix,
                      const std::vector<std::pair<bdd, std::string>> &output_map) {
  std::unordered_set<BDD> emitted;
  for (const std::pair<bdd, std::string> &out_map : output_map) {
    if (out_map.first == bdd_true()) {
      blif_file << ".names " << out_map.second << "\n1\n";
    } else if (out_map.first == bdd_false()) {
      blif_file << ".names " << out_map.second << "\n0\n";
    } else {
      write_blif_recursive(blif_file, signal_map, prefix, out_map.first.id(),
                           emitted, true);
      blif_file << out_map.second << "\n0-1 1\n11- 1\n";
    }
  }
}

// memoised: one visit per node, not per path
int get_bdd_height(BDD node, std::unordered_map<BDD, int> &heights) {
  if (node < 2) {
    return 0;
  }
  auto cached = heights.find(node);
  if (cached != heights.end()) {
    return cached->second;
  }
  int low_height = get_bdd_height(bdd_low(node), heights);
  int high_height = get_bdd_height(bdd_high(node), heights);
  int height = std::max(low_height, high_height) + 1;
  heights.emplace(node, height);
  return height;
}

int get_total_height(const std::vector<std::pair<bdd, std::string>> &output_map) {
  int max_height = -1;
  std::unordered_map<BDD, int> heights;
  for (const std::pair<bdd, std::string> &out_map : output_map) {
    int current_height = get_bdd_height(out_map.first.id(), heights);
    if (current_height > max_height) {
      max_height = current_height;
    }
  }
  return max_height;
}

/// shared node count of a block's diagrams, so that nodes several roots have in
/// common are counted once
static uint64_t block_node_count(
    const std::vector<std::pair<bdd, std::string>> &output_map) {
  // bdd.h maps bdd_anodecount to the C++ overload taking bdd, not raw ids
  std::vector<bdd> roots;
  roots.reserve(output_map.size());
  for (const std::pair<bdd, std::string> &out_map : output_map) {
    roots.push_back(out_map.first);
  }
  if (roots.empty()) {
    return 0;
  }
  return static_cast<uint64_t>(
      bdd_anodecount(roots.data(), static_cast<int>(roots.size())));
}

int main(int argc, char *argv[]) {
  // -o: where the netlist and stats go
  std::string out_dir = "temp";
  // maximum input support per block, 0 disables partitioning
  uint32_t max_inputs = 0;
  // -g: support cap for gain-driven merging, 0 = off
  uint32_t gain_cap = 0;
  int table_size = 300000;
  int cache_size = 10000;
  int opt;
  while ((opt = getopt(argc, argv, "o:n:g:t:c:")) != -1) {
    switch (opt) {
    case 'o':
      out_dir = optarg;
      break;
    case 'n': {
      char *endptr = nullptr;
      long parsed = std::strtol(optarg, &endptr, 10);
      if (*endptr != '\0' || parsed < 0) {
        std::cerr << "[ERROR] Invalid integer for option -n: " << optarg
                  << std::endl;
        return EXIT_FAILURE;
      }
      if (parsed != 0 && parsed < 2) {
        std::cerr << "[ERROR] -n must be 0 (disabled) or at least 2\n";
        return EXIT_FAILURE;
      }
      max_inputs = static_cast<uint32_t>(parsed);
      break;
    }
    case 'g': {
      char *endptr = nullptr;
      long parsed = std::strtol(optarg, &endptr, 10);
      if (*endptr != '\0' || parsed < 0) {
        std::cerr << "[ERROR] Invalid integer for option -g: " << optarg
                  << std::endl;
        return EXIT_FAILURE;
      }
      gain_cap = static_cast<uint32_t>(parsed);
      break;
    }
    case 't': {
      char *endptr = nullptr;
      table_size = std::strtol(optarg, &endptr, 10);
      if (*endptr != '\0' || table_size <= 0) {
        std::cerr << "[ERROR] Invalid integer for option -t: " << optarg
                  << std::endl;
        return EXIT_FAILURE;
      }
      break;
    }
    case 'c': {
      char *endptr = nullptr;
      cache_size = std::strtol(optarg, &endptr, 10);
      if (*endptr != '\0' || cache_size <= 0) {
        std::cerr << "[ERROR] Invalid integer for option -c: " << optarg
                  << std::endl;
        return EXIT_FAILURE;
      }
      break;
    }
    default:
      std::cerr << "Usage: " << argv[0]
                << " [-n max_inputs] [-g gain_cap] [-t table_size] "
                   "[-c cache_size] [-o output_dir] <blif-file>"
                << std::endl;
      return EXIT_FAILURE;
    }
  }

  if (optind >= argc) {
    std::cout << "[ERROR] input file in blif format is missing\n";
    std::cerr << "[INFO] Usage: " << argv[0]
              << " [-n max_inputs] [-o output_dir] <blif-file>" << std::endl;
    return EXIT_FAILURE;
  }

  std::string benchmark(argv[optind]);
  if (!std::filesystem::exists(benchmark)) {
    std::cout << "[ERROR] File does not exist\n";
    return EXIT_FAILURE;
  }

  std::filesystem::path path(benchmark);
  std::string base_name = path.stem().string();

  mockturtle::cover_network cover;

  if (lorina::read_blif(benchmark, mockturtle::blif_reader(cover)) !=
      lorina::return_code::success) {
    std::cout << "[ERROR] While Lorina tried to read in file\n";
    return EXIT_FAILURE;
  }

  printf("[INFO] cover with %zu nodes\n", cover._storage->nodes.size());

  //////////////////////////////////////////////////////////////////////////
  // partition the network into blocks with a bounded input support
  //////////////////////////////////////////////////////////////////////////
  std::vector<block_t> blocks = partition_cover(cover, max_inputs);
  assert(partition_is_valid(cover, blocks, max_inputs));

  //////////////////////////////////////////////////////////////////////////
  // merge further, but only where the merged diagram is measurably smaller
  //////////////////////////////////////////////////////////////////////////
  if (gain_cap > max_inputs && max_inputs != 0 && blocks.size() > 1) {
    size_t const before = blocks.size();
    // trial build, same as the writer
    auto block_cost = [&](block_t const &b) -> uint64_t {
      leaf_map_t leaf_map =
          make_leaf_map(cover._storage->nodes.size(), b.leaves);
      std::vector<std::string> root_names(b.roots.size(), "t");
      const size_t vars = std::max<size_t>(b.leaves.size(), 1);

      bdd_init(block_node_table(b.leaves.size(), table_size), cache_size);
      bdd_setvarnum(static_cast<int>(vars));
      uint64_t nodes = 0;
      {
        std::vector<std::pair<bdd, std::string>> outs;
        block_to_bdd(cover, b, leaf_map, root_names, outs);
        nodes = block_node_count(outs);
        // the roots have to be released before the instance they live in
      }
      bdd_done();
      return nodes;
    };

    merge_blocks_by_gain(cover, blocks, gain_cap, /*max_passes=*/3u,
                         /*max_candidates=*/8u, block_cost);
    assert(partition_is_valid(cover, blocks, gain_cap));
    printf("[INFO] gain merge (cap %u): %zu block(s) -> %zu\n", gain_cap,
           before, blocks.size());
  }

  size_t max_support = 0, total_support = 0, max_gates = 0;
  for (block_t const &b : blocks) {
    max_support = std::max(max_support, b.leaves.size());
    total_support += b.leaves.size();
    max_gates = std::max(max_gates, b.gates.size());
  }
  printf("[INFO] %zu block(s), support max %zu avg %.1f, largest block %zu "
         "gates\n",
         blocks.size(), max_support,
         blocks.empty() ? 0.0
                        : static_cast<double>(total_support) / blocks.size(),
         max_gates);

  //////////////////////////////////////////////////////////////////////////
  // give every signal that appears in the output netlist a name
  //////////////////////////////////////////////////////////////////////////
  std::vector<std::string> canonical(cover._storage->nodes.size());
  canonical[0] = "$false";
  canonical[1] = "$true";
  for (uint64_t const &pi : cover._storage->inputs) {
    canonical[pi] = cover.get_signal_name(pi);
  }

  // primary outputs in declaration order, together with their driver
  std::vector<std::pair<uint64_t, std::string>> po_list;
  {
    uint32_t index = cover._storage->inputs.size() + 2;
    cover.foreach_co([&](auto const &f) {
      po_list.emplace_back(cover.get_node(f), cover.get_signal_name(index));
      index++;
    });
  }
  // a root that drives a primary output is named after it, so that the common
  // case needs no extra buffer
  std::vector<std::string> po_first(cover._storage->nodes.size());
  for (auto const &po : po_list) {
    if (po_first[po.first].empty()) {
      po_first[po.first] = po.second;
    }
  }

  // every root has to be named before any block is written, because a block
  // may use a root of a later block as one of its leaves
  for (block_t const &b : blocks) {
    for (uint64_t const &r : b.roots) {
      if (canonical[r].empty()) {
        canonical[r] =
            !po_first[r].empty() ? po_first[r] : "part_w" + std::to_string(r);
      }
    }
  }

  //////////////////////////////////////////////////////////////////////////
  // write the netlist: one header, then the diagrams of every block
  //////////////////////////////////////////////////////////////////////////
  std::string out_file = out_dir + "/" + base_name + "_bdd.blif";
  std::ofstream blif(out_file);
  if (!blif) {
    std::cerr << "[ERROR] cannot open " << out_file << " for writing\n";
    return EXIT_FAILURE;
  }
  blif << ".model " << cover.get_module_name() << "\n";
  blif << ".inputs";
  for (uint64_t const &pi : cover._storage->inputs) {
    blif << " " << canonical[pi];
  }
  blif << "\n.outputs";
  for (auto const &po : po_list) {
    blif << " " << po.second;
  }
  blif << "\n.names $true\n1\n";
  blif << ".names $false\n";

  uint64_t total_nodes = 0;
  int max_height = 0;
  int peak_table = 0;
  for (size_t id = 0; id < blocks.size(); id++) {
    block_t const &b = blocks[id];

    leaf_map_t leaf_map =
        make_leaf_map(cover._storage->nodes.size(), b.leaves);

    // names of the block's variables, indexed by BuDDy variable number, which
    // is the leaf map's id minus 2
    std::vector<std::string> signal_map;
    signal_map.reserve(b.leaves.size());
    for (uint64_t const &l : b.leaves) {
      signal_map.push_back(canonical[l]);
    }
    std::vector<std::string> root_names;
    root_names.reserve(b.roots.size());
    for (uint64_t const &r : b.roots) {
      root_names.push_back(canonical[r]);
    }

    // a fresh instance per block: a BuDDy variable number means a different
    // signal in every block, so blocks must not share a variable ordering
    const size_t vars = std::max<size_t>(b.leaves.size(), 1);
    bdd_init(block_node_table(b.leaves.size(), table_size), cache_size);
    bdd_setvarnum(static_cast<int>(vars));
    bdd_autoreorder(BDD_REORDER_SIFT);
    if (blocks.size() == 1) {
      printf("[INFO] init with %zu variables\n", b.leaves.size());
    }

    {
      std::vector<std::pair<bdd, std::string>> output_map;
      block_to_bdd(cover, b, leaf_map, root_names, output_map);
      // bdd_reorder_siftite is the same but repeats until no more progress is
      // done
      bdd_reorder(BDD_REORDER_SIFT);

      const int height = get_total_height(output_map);
      if (blocks.size() == 1) {
        printf("height after: %d\n", height);
      }
      max_height = std::max(max_height, height);
      total_nodes += block_node_count(output_map);
      peak_table = std::max(peak_table, bdd_getnodenum());

      // a single block owns the whole module, so its node indices are already
      // unique and the netlist stays identical to an unpartitioned run
      write_block_body(blif, signal_map,
                       blocks.size() == 1 ? "" : "b" + std::to_string(id) + "_",
                       output_map);
      // the roots have to be released before the instance they live in
    }
    bdd_done();
  }

  // outputs that are not named after their driver need an explicit buffer:
  // a primary input or constant wired straight through, or a second primary
  // output on the same driver
  for (auto const &po : po_list) {
    if (canonical[po.first] != po.second) {
      blif << ".names " << canonical[po.first] << " " << po.second << "\n1 1\n";
    }
  }
  blif << ".end\n";
  blif.close();

  printf("[INFO] BDD created: %lu nodes, max height %d, peak table use %d\n",
         total_nodes, max_height, peak_table);
  std::cout << "[INFO] Wrote bdd into " << out_file << "\n";

  // machine readable summary for the regression runner
  std::ofstream stats(out_dir + "/" + base_name + "_bdd_stats.csv");
  if (stats) {
    stats << "blocks;max_support;avg_support;max_block_gates;bdd_nodes;"
             "max_height;peak_table;seed_n;gain_cap\n";
    stats << blocks.size() << ";" << max_support << ";"
          << (blocks.empty()
                  ? 0.0
                  : static_cast<double>(total_support) / blocks.size())
          << ";" << max_gates << ";" << total_nodes << ";" << max_height << ";"
          << peak_table << ";" << max_inputs << ";" << gain_cap << "\n";
  }

  return EXIT_SUCCESS;
}