/*
 *  yosys -- Yosys Open SYnthesis Suite
 *
 *  Copyright (C) 2024  Silimate Inc.
 *
 *  Permission to use, copy, modify, and/or distribute this software for any
 *  purpose with or without fee is hereby granted, provided that the above
 *  copyright notice and this permission notice appear in all copies.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 *  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 *  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 *  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 *  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 *  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 *  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 */

#include "kernel/yosys.h"
#include "kernel/sigtools.h"
#include "kernel/ff.h"
#include "kernel/satgen.h"
#include <fstream>

USING_YOSYS_NAMESPACE
PRIVATE_NAMESPACE_BEGIN

// Maximum depth for BFS exploration of input cone
static const int MAX_INPUT_DEPTH = 10;

// If true, exclude Q (feedback) bits from enable input set
static const bool EXCLUDE_Q_FROM_ENABLE = true;

struct SatClockgateWorker
{
	Module *module;
	SigMap sigmap;
	
	// Maps output signal bits to their driver cells
	dict<SigBit, Cell*> sig_to_driver;
	
	// Q bits to exclude from enable input set (set per-FF)
	pool<SigBit> q_bits;

	// SAT solver and generator - created once per module
	ezSatPtr ez;
	SatGen satgen;
	
	SatClockgateWorker(Module *module) : module(module), sigmap(module), ez(), satgen(ez.get(), &sigmap)
	{
		// Build driver map: for each signal bit, find which cell drives it
		for (auto cell : module->cells()) {
			for (auto &conn : cell->connections()) {
				if (cell->output(conn.first)) {
					for (auto bit : sigmap(conn.second))
						sig_to_driver[bit] = cell;
				}
			}
		}

		// Import all cells once - circuit constraints are permanent
		for (auto cell : module->cells())
			satgen.importCell(cell);
	}

	// Set Q bits to exclude for current FF
	void set_excluded_q_bits(SigSpec sig_q)
	{
		q_bits.clear();
		if (EXCLUDE_Q_FROM_ENABLE) {
			for (auto bit : sigmap(sig_q))
				q_bits.insert(bit);
		}
	}

	// Get the set of input signals feeding into a given signal (one level back)
	pool<SigBit> get_input_signals(SigBit bit)
	{
		pool<SigBit> inputs;
		bit = sigmap(bit);
		
		if (!sig_to_driver.count(bit))
			return inputs; // Primary input or constant
		
		Cell *driver = sig_to_driver[bit];
		for (auto &conn : driver->connections()) {
			if (driver->input(conn.first)) {
				for (auto input_bit : sigmap(conn.second)) {
					if (input_bit.wire != nullptr && !q_bits.count(input_bit))
						inputs.insert(input_bit);
				}
			}
		}
		return inputs;
	}

	// Check if OR(input_set) is exactly equivalent to (D != Q)
	// Returns true if COI ↔ (D≠Q) for all circuit states (exact clock gating)
	// TODO: Consider pruning the expressions list — expressions accumulate across
	//       calls (OR, XOR, NE per query). For large designs with many FFs, this
	//       could cause memory growth. Options: solver checkpoints, fresh solver
	//       per FF with COI-only cell import, or periodic expression cleanup.
	bool input_set_is_enable(const pool<SigBit> &input_set, SigSpec sig_d, SigSpec sig_q)
	{
		if (input_set.empty())
			return false;

		// Import D and Q (uses cached literals if already imported)
		std::vector<int> d_vec = satgen.importSigSpec(sig_d);
		std::vector<int> q_vec = satgen.importSigSpec(sig_q);
		
		// Build COI = OR(input_set)
		std::vector<int> input_vars;
		for (auto bit : input_set)
			input_vars.push_back(satgen.importSigSpec(SigSpec(bit))[0]);
		int coi = ez->expression(ezSAT::OpOr, input_vars);
		
		// Build D != Q (single bit: is any bit different?)
		int d_ne_q = ez->vec_ne(d_vec, q_vec);
		
		// Query: COI XOR (D≠Q) — want this UNSAT (meaning COI ↔ D≠Q)
		// Use solve() with assumption instead of permanent assume()
		int query = ez->XOR(coi, d_ne_q);
		
		std::vector<int> assumptions = {query};
		std::vector<int> dummy_model_exprs;
		std::vector<bool> dummy_model_vals;
		
		// If UNSAT: COI is exactly when D≠Q → perfect enable
		return !ez->solve(dummy_model_exprs, dummy_model_vals, assumptions);
	}

	// Recursively determine the enable input set via BFS expansion
	// Seeds initial input set from sig_d, Q bits filtered via get_input_signals
	bool determine_enable_recursive(pool<SigBit> &input_set, SigSpec sig_d, SigSpec sig_q, int depth)
	{
		if (depth > MAX_INPUT_DEPTH) {
			log_debug("  Max depth reached, giving up\n");
			return false;
		}

		// Seed initial input set from sig_d on first call
		if (depth == 1 && input_set.empty()) {
			for (auto bit : sigmap(sig_d)) {
				if (bit.wire != nullptr) {
					for (auto input_bit : get_input_signals(bit))
						input_set.insert(input_bit);
				}
			}

			if (input_set.empty()) {
				log_debug("  No inputs to D (besides Q)\n");
				return false;
			}
			log_debug("  Initial input set has %zu signals\n", input_set.size());
		}

		// Check if current input set works as enable
		if (input_set_is_enable(input_set, sig_d, sig_q)) {
			log_debug("  Found enable at depth %d with %zu signals\n", depth, input_set.size());
			return true;
		}

		// Expand input set via BFS (one more level)
		pool<SigBit> new_inputs;
		for (auto bit : input_set) {
			for (auto input_bit : get_input_signals(bit)) {
				if (!input_set.count(input_bit))
					new_inputs.insert(input_bit);
			}
		}

		if (new_inputs.empty()) {
			log_debug("  No more inputs to explore at depth %d\n", depth);
			return false;
		}

		// Add new inputs and recurse
		for (auto bit : new_inputs)
			input_set.insert(bit);

		return determine_enable_recursive(input_set, sig_d, sig_q, depth + 1);
	}

	// Create CE logic based on the input set and modify the FF
	// The enable condition is: when all input_set bits are 0, D == Q (hold)
	// So CE should be: OR of all input_set bits (active high: CE=1 means update)
	void create_ce_logic(const pool<SigBit> &input_set, FfData &ff)
	{
		if (input_set.empty())
			return;

		log("  Creating CE from %zu input signals\n", input_set.size());

		// Build CE as OR of all input signals
		// CE = 1 when any input is 1 (meaning: update the register)
		// CE = 0 when all inputs are 0 (meaning: hold, since D == Q)
		SigSpec ce_inputs;
		for (auto bit : input_set)
			ce_inputs.append(bit);

		SigBit ce_signal;
		if (GetSize(ce_inputs) == 1) {
			ce_signal = ce_inputs[0];
		} else {
			// Create OR gate: CE = |input_set
			Wire *ce_wire = module->addWire(NEW_ID);
			module->addReduceOr(NEW_ID, ce_inputs, ce_wire);
			ce_signal = ce_wire;
		}

		// Set the CE on the FF
		ff.has_ce = true;
		ff.sig_ce = ce_signal;
		ff.pol_ce = true;  // Active high

		log("  CE signal: %s\n", log_signal(ce_signal));
	}

	// Process a single FF to find and insert CE
	bool process_ff(Cell *cell)
	{
		FfData ff(nullptr, cell);

		// Skip if already has CE, or doesn't have clock/data
		if (ff.has_ce) {
			log_debug("  Skipping %s: already has CE\n", log_id(cell));
			return false;
		}
		if (!ff.has_clk) {
			log_debug("  Skipping %s: no clock\n", log_id(cell));
			return false;
		}
		if (GetSize(ff.sig_d) == 0 || GetSize(ff.sig_q) == 0) {
			log_debug("  Skipping %s: no D or Q\n", log_id(cell));
			return false;
		}

		log("Processing FF: %s\n", log_id(cell));

		// Set Q bits to exclude from enable candidates
		set_excluded_q_bits(ff.sig_q);

		// Find enable via recursive BFS + SAT validation
		pool<SigBit> input_set;
		if (determine_enable_recursive(input_set, ff.sig_d, ff.sig_q, 1)) {
			create_ce_logic(input_set, ff);
			
			// Emit the modified FF
			ff.emit();
			return true;
		}

		log_debug("  Could not find enable for %s\n", log_id(cell));
		return false;
	}
};

void dump_flipflops_to_file(RTLIL::Design *design, const std::string &filename)
{
	std::ofstream outfile(filename);
	if (!outfile.is_open()) {
		log_error("Cannot open file %s for writing\n", filename.c_str());
		return;
	}

	for (auto module : design->selected_modules()) {
		outfile << "Module: " << log_id(module) << "\n";
		log("Module: %s\n", log_id(module));

		for (auto cell : module->cells()) {
			if (cell->is_builtin_ff()) {
				outfile << "  FF: " << log_id(cell) << " (type: " << log_id(cell->type) << ")\n";
				log("  FF: %s (type: %s)\n", log_id(cell), log_id(cell->type));
			}
		}
		outfile << "\n";
	}

	outfile.close();
	log("Wrote flip-flop list to %s\n", filename.c_str());
}

struct SatClockgatePass : public Pass {
	SatClockgatePass() : Pass("sat_clockgate", "SAT-based inferred clock gating") { }

	void help() override
	{
		log("\n");
		log("    sat_clockgate [options] [selection]\n");
		log("\n");
		log("This command performs SAT-based inferred clock gating insertion.\n");
		log("It analyzes flip-flops without explicit clock enables and uses SAT\n");
		log("to find input conditions under which D == Q (register holds value).\n");
		log("These conditions become the inferred clock enable.\n");
		log("\n");
		log("    -threshold <n>\n");
		log("        minimum number of FFs that must share an enable for clock gating\n");
		log("        to be inserted (default: 1)\n");
		log("\n");
	}

	void execute(std::vector<std::string> args, RTLIL::Design *design) override
	{
		log_header(design, "Executing SAT_CLOCKGATE pass (SAT-based inferred clock gating).\n");

		int threshold = 1;

		size_t argidx;
		for (argidx = 1; argidx < args.size(); argidx++) {
			if (args[argidx] == "-threshold" && argidx+1 < args.size()) {
				threshold = std::stoi(args[++argidx]);
				continue;
			}
			break;
		}
		extra_args(args, argidx, design);

		log("Using threshold: %d\n", threshold);

		// Dump all flip-flops to file (debug)
		dump_flipflops_to_file(design, "flip_flops.txt");

		int total_converted = 0;

		for (auto module : design->selected_modules()) {
			log("Processing module %s...\n", log_id(module));

			SatClockgateWorker worker(module);

			// Collect FFs to process (can't modify while iterating)
			std::vector<Cell*> ffs_to_process;
			for (auto cell : module->cells()) {
				if (cell->is_builtin_ff())
					ffs_to_process.push_back(cell);
			}

			int converted = 0;
			for (auto cell : ffs_to_process) {
				if (worker.process_ff(cell))
					converted++;
			}

			if (converted > 0) {
				log("Converted %d FFs in module %s\n", converted, log_id(module));
				total_converted += converted;
			}
		}

		log("Total FFs with inferred CE: %d\n", total_converted);

		// TODO: Call clockgate pass to convert CEs to ICG cells
		// if (total_converted >= threshold) {
		//     Pass::call(design, "clockgate");
		// }
	}
} SatClockgatePass;

PRIVATE_NAMESPACE_END