#include <filesystem>
#include <iostream>
#include <lorina/verilog.hpp>
#include <mockturtle/algorithms/cleanup.hpp>
#include <mockturtle/io/verilog_reader.hpp>
#include <mockturtle/properties/migcost.hpp>
#include <mockturtle/utils/stopwatch.hpp>
#include <string>

#include "lorina/blif.hpp"
#include "mockturtle/algorithms/cover_to_graph.hpp"
#include "mockturtle/networks/cover.hpp"
#include "mockturtle/networks/muxig.hpp"

#include "mockturtle/algorithms/sim_resub.hpp"
#include "mockturtle/io/blif_reader.hpp"
#include "mockturtle/io/write_blif.hpp"

#define SEL_INDEX 0
#define IN0_INDEX 2
#define IN1_INDEX 1

#define SEL(ntk, node) ntk._storage->nodes[node].children[SEL_INDEX]
#define IN0(ntk, node) ntk._storage->nodes[node].children[IN0_INDEX]
#define IN1(ntk, node) ntk._storage->nodes[node].children[IN1_INDEX]

int main(int argc, char *argv[]) {
  // Default values
  int max_pis = 10; // 8;
  int max_divisors = 150;
  int max_inserts = 20; // 2;

  bool output_path_given = false;
  std::string output_dir;

  int opt;
  while ((opt = getopt(argc, argv, "o:p:d:i:")) != -1) {
    switch (opt) {
    case 'o':
      output_path_given = true;
      output_dir = optarg;
      break;
    case 'p':
      max_pis = std::atoi(optarg);
      break;
    case 'd':
      max_divisors = std::atoi(optarg);
      break;
    case 'i':
      max_inserts = std::atoi(optarg);
      break;
    default:
      std::cerr << "Usage: " << argv[0]
                << " [-p max_pi] [-d max_divisors] [-i max_inserts] [-o "
                   "output_file for performance statistics]"
                   "<positional_arg>\n";
      return 1;
    }
  }

  std::string benchmark;
  if (optind < argc) {
    benchmark = argv[optind];
  } else {
    std::cerr << "Error: Missing required positional argument.\n";
    std::cerr
        << "Usage: " << argv[0]
        << " [-p max_pi] [-d max_divisors] [-i max_inserts] <benchmark_file>\n";
    return 1;
  }

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

  std::filesystem::path order_output_dir_path;
  std::filesystem::path base_name_path;
  std::filesystem::path order_file_path;
  std::ofstream order_file;

  mockturtle::muxig_network muxig_original =
      mockturtle::convert_cover_to_graph<mockturtle::muxig_network>(cover);

  mockturtle::resubstitution_params ps;
  mockturtle::resubstitution_stats st;

  ps.max_inserts = max_inserts;
  ps.max_pis = max_pis;
  ps.max_divisors = max_divisors;
  ps.progress = false;
  ps.verbose = false;

  int size_before = muxig_original.num_gates();

  mockturtle::sim_resubstitution(muxig_original, ps, &st);

  muxig_original = mockturtle::cleanup_dangling(muxig_original);

  int size_after = muxig_original.num_gates();
  /*mockturtle::muxig_network muxig_sel = muxig_original.clone();
  mockturtle::muxig_network muxig_and = muxig_original.clone();
  mockturtle::muxig_network muxig_or = muxig_original.clone();
  mockturtle::muxig_network muxig_all = muxig_original.clone();*/
  mockturtle::muxig_network muxig = muxig_original.clone(); // jaspers
  /*printf("before: %d, after: %d, sel: %d, and: %d, or: %d, all: %d\n",
         size_before, muxig_original.num_gates(), muxig_sel.num_gates(),
         muxig_and.num_gates(), muxig_or.num_gates(), muxig_all.num_gates());*/

  // Info from mockturtle library
  /* mux(a b c) = a b + a' c */
  muxig.foreach_node([&](auto node) {
    if (muxig.is_mux(node) && (muxig.fanin_size(node) == 3)) {
      // if a and b are inverted push inverter to output
      if (muxig.is_complemented(IN0(muxig, node)) &&
          muxig.is_complemented(IN1(muxig, node))) {
        muxig.substitute_node(
            node, muxig.create_not(muxig.create_gate(
                      SEL(muxig, node), muxig.create_not(IN1(muxig, node)),
                      muxig.create_not(IN0(muxig, node)))));
      }

      if (IN0(muxig, node).index == IN1(muxig, node).index){
        assert(muxig.is_complemented(IN0(muxig, node)) || muxig.is_complemented(IN1(muxig, node)));
        printf("xor possibly\n");
      }
    }
  });

  mockturtle::fanout_view fnntk(muxig);
  muxig.foreach_node([&](auto node) {
    if (muxig.is_mux(node) && (muxig.fanin_size(node) == 3)) {
      uint32_t inv = muxig.fanout_size(node);
      std::vector<std::pair<int32_t, int32_t>> fanout_stack;
      fnntk.foreach_fanout(node, [&](auto fan_out) {
        int i = 0;
        muxig.foreach_fanin(fan_out, [&](auto signal) {
          if (muxig.get_node(signal) == node) {
            if (signal.complement) {
              inv--;
            }
            fanout_stack.emplace_back(fan_out, i);
          }
          i++;
        });
      });

      int output_count = 0;

      muxig.foreach_po([&](auto po) {
        if (muxig.get_node(po) == node && muxig.is_complemented(po)) {
          inv--;
          fanout_stack.emplace_back(output_count, -1);
        }
        output_count++;
      });

      // if all fanouts are inverted
      if (inv == 0 && muxig.fanout_size(node) != 0) {
        if (IN1(muxig, node) == 0 && IN0(muxig, node) != 0 &&
            !muxig.is_complemented(IN1(muxig, node))) {
          // 0 is represented as 0 and 1 as !0
          // switching not(not a and b) to (a or not b)
          // this gets rid of all inverters at the output
          std::swap(SEL(muxig, node), IN0(muxig, node));
          std::swap(IN1(muxig, node), IN0(muxig, node));
          IN0(muxig, node) = muxig.get_constant(true);
          for (const auto &[node, child] : fanout_stack) {
            if (child >= 0) {
              muxig._storage->nodes[node].children[child] =
                  muxig.create_not(muxig._storage->nodes[node].children[child]);
            } else {
              muxig._storage->outputs[node] =
                  muxig.create_not(muxig._storage->outputs[node]);
            }
          }
        } else if (IN0(muxig, node) == 0 && IN1(muxig, node) != 0 &&
                   muxig.is_complemented(IN0(muxig, node))) {
          // switching not(not a or b) to (a and not b)
          std::swap(SEL(muxig, node), IN1(muxig, node));
          std::swap(IN1(muxig, node), IN0(muxig, node));
          IN1(muxig, node) = muxig.get_constant(false);
          for (const auto &[node, child] : fanout_stack) {
            if (child >= 0) {
              muxig._storage->nodes[node].children[child] =
                  muxig.create_not(muxig._storage->nodes[node].children[child]);
            } else {
              muxig._storage->outputs[node] =
                  muxig.create_not(muxig._storage->outputs[node]);
            }
          }
        }
      }
    }
  });

  muxig.foreach_node([&](auto node) {
    if (muxig.is_mux(node) && (muxig.fanin_size(node) == 3)) {
      if (muxig.is_complemented(SEL(muxig, node))) {
        std::swap(IN0(muxig, node), IN1(muxig, node));
        SEL(muxig, node) = muxig.create_not(SEL(muxig, node));
        assert(!muxig.is_complemented(SEL(muxig, node)));
      }
      // and not
      if (IN0(muxig, node) == 0 && !muxig.is_complemented(IN0(muxig, node))) {
        if (muxig.is_complemented(IN1(muxig, node))) {
          // (sel, IN1, IN0)
          std::swap(IN1(muxig, node), SEL(muxig, node));
          std::swap(IN0(muxig, node), IN1(muxig, node));
          SEL(muxig, node) = muxig.create_not(SEL(muxig, node));
        }
      }
      // or not
      if (IN1(muxig, node) == 1) {
        if (muxig.is_complemented(IN0(muxig, node))) {
          std::swap(IN0(muxig, node), SEL(muxig, node));
          std::swap(IN0(muxig, node), IN1(muxig, node));
          SEL(muxig, node) = muxig.create_not(SEL(muxig, node));
        }
      }
    }
  });
  if (output_path_given) {
    order_output_dir_path = output_dir;
    base_name_path = base_name;
    base_name_path.replace_extension(".csv");
    order_file_path = order_output_dir_path / base_name_path;
    printf("Writing stats to %s\n", order_file_path.c_str());
    order_file.open(order_file_path);
    order_file << "before;after;inv_before;inv_all\n";
    order_file << size_before << ";" << size_after << ";"
               << mockturtle::num_inverters(muxig_original) << ";"
               << mockturtle::num_inverters(muxig) << "\n";
  }
  printf("inverter before: %d after: %d\n",
         mockturtle::num_inverters(muxig_original),
         mockturtle::num_inverters(muxig));
  /*printf("removed by sel: %d\n", mockturtle::num_inverters(muxig_original) -
                                     mockturtle::num_inverters(muxig_sel));
  printf("removed by and not: %d\n", mockturtle::num_inverters(muxig_original) -
                                         mockturtle::num_inverters(muxig_and));
  printf("removed by or not: %d\n", mockturtle::num_inverters(muxig_original) -
                                        mockturtle::num_inverters(muxig_or));*/
  mockturtle::write_blif(muxig, "temp/" + base_name + "_muxig_temp.blif");

  return EXIT_SUCCESS;
}