#include </home/jasper/uni/i-edge/application/plugins/cpp-dump/cpp-dump.hpp>
#include "kernel/log.h"
#include "kernel/rtlil.h"
#include "kernel/yosys.h"
#include "kernel/sigtools.h"
#include "kernel/celltypes.h"
#include "kernel/yosys_common.h"

#include <string>
#include <map>
#include <set>


USING_YOSYS_NAMESPACE
PRIVATE_NAMESPACE_BEGIN

//Set up pair to map signal to cell driver or trigger
typedef std::pair<RTLIL::Cell*, RTLIL::IdString> sig2driver_entry_t;
static SigSet<sig2driver_entry_t> sig2driver, sig2trigger;

static void populate_sig_driver_trigger_database(RTLIL::Design *design, RTLIL::Module *mod, SigMap *sigmap){
    log("Building database of all cells and their matching connections \n");

    CellTypes ct(design);
			
    for (auto cell : mod->cells()) {
        for (auto &conn_it : cell->connections()) {
            
            if (ct.cell_output(cell->type, conn_it.first) || !ct.cell_known(cell->type)) {
                RTLIL::SigSpec sig = conn_it.second;
                sigmap->apply(sig);
                //log("adding output driver %s:%s - %s:%s\n",log_id(cell),log_id(conn_it.first),(sig.is_wire() ? log_id(sig.as_wire()->name): "not wire"),log_id(sig[0].wire));
                sig2driver.insert(sig, sig2driver_entry_t(cell, conn_it.first));
            }
            if (ct.cell_input(cell->type, conn_it.first) || !ct.cell_known(cell->type)) {
                RTLIL::SigSpec sig = conn_it.second;
                sigmap->apply(sig);
                //log("adding input trigger %s\n",log_id(cell));
                sig2trigger.insert(sig, sig2driver_entry_t(cell, conn_it.first));
            }
        }
    }
}



static int find_longest_path(RTLIL::Design *design, SigMap *sigmap_ptr, RTLIL::Wire *current, RTLIL::Wire *wire_out, std::vector<RTLIL::SigSpec*>* visited, std::vector<RTLIL::SigSpec>* longest_path_cache_ptr, int length, int longest_length)
{
    SigMap &sigmap = *sigmap_ptr;
    std::vector<RTLIL::SigSpec> &longest_path_cache = *longest_path_cache_ptr;    

    CellTypes ct(design); //get celltype match to check for outputs
    //log("--------------------------------------------------------\n");
    //log("Checking path: %s -> %s\n",log_id(current->name),log_id(wire_out->name));

     // If we reached the destination, the path length is 0
    if (sigmap(current) == sigmap(wire_out)){
        //log("Found target!!!!!!!!!\n");
        //log("--------------------------------------------------------\n");
        if(length>longest_length){
            //log("new route is longer setting route\n");
            longest_path_cache.clear();
            for(auto sig: *visited){
                //log("%s->",(sig->is_wire() ? log_id(sig->as_wire()):log_id(sig->as_bit().wire)));
                longest_path_cache.push_back(*sig);
                longest_length = length;
            }
            //log("\n");
        }
        return longest_length;
    }

    RTLIL::SigSpec sig_wire = sigmap(current);

    // Mark this wire as visited
    visited->push_back(&sig_wire);

    //for(auto sig: *visited){
    //    log("%s->",(sig->is_wire() ? log_id(sig->as_wire()):log_id(sig->as_bit().wire)));
    //}
    //log("\n");
    
    std::set<sig2driver_entry_t> cellport_list; //create list to populate with the connected cells
    sig2trigger.find(sig_wire, cellport_list); //populate it with the found cells
    
    //log("before %s, after %s\n",log_id(current),(sig_wire.is_wire() ? log_id(sig_wire.as_wire()->name): "not wire"));
    // Explore all connections (edges) from the current wire
    for (auto cell: cellport_list){
        //log("wire %s on %s,\n",log_id(cell.second),log_id(cell.first));
        for (auto conn : cell.first->connections()) {
            // Avoid revisiting wires to prevent cycles 
            if ((conn.second.is_wire() || conn.second.is_bit()) && ct.cell_output(cell.first->type, conn.first)) {
                // Recursively calculate the longest path from next_wire to wire_out
                RTLIL::Wire* next_wire= (conn.second.is_wire() ? conn.second.as_wire(): conn.second.as_bit().wire); //check if it is a wire or only one bit of a wire
                
                //log("Going recursive on %s\n",log_id(next_wire));
                
                longest_length = find_longest_path(design, sigmap_ptr, next_wire, wire_out, visited, longest_path_cache_ptr, length+1,longest_length);
            }
        } 
    }

    // Unmark this wire (backtrack)
    visited->pop_back();

    // Cache the result for this wire
    //log("-----------------------returning %s---------------------\n",log_id(current->name));
    // Return the longest path found
    return longest_length;
    

}

struct Critical_path : public Pass {
    Critical_path() : Pass("critical_path", "will take the current design and calculate the longest path it can find. Afterwards will return a argument to print this in a show") { }

    void execute(std::vector<std::string> args, RTLIL::Design *design) override
    {
        sig2driver.clear();
        sig2trigger.clear();
        
        for (auto mod : design->modules())
        {
            //Check if we are in the main module
            if(GetSize(mod->cells())>0){
                // Instantiate SigMap for the current module
                SigMap sigmap(mod);

                
                
                //print module name
                log(" %s (%d wires, %d cells)\n", log_id(mod), GetSize(mod->wires()), GetSize(mod->cells()));
                populate_sig_driver_trigger_database(design, mod,&sigmap);

                std::vector<RTLIL::SigSpec> abst_longest_path;
                int abst_path = 0;

                for (auto wire_in : mod->wires())
                {
                    if(wire_in->port_input){ //Is input

                        for(auto wire_out: mod->wires()){

                            if(wire_out->port_output){ //is output
                                std::vector<RTLIL::SigSpec*> visited;
                                std::vector<RTLIL::SigSpec> longest_path_cache;
                                visited.clear();
                                longest_path_cache.clear();
                                //log("************************************** starting %s->%s ******************************\n",log_id(wire_in),log_id(wire_out));
                                int longest_path = find_longest_path(design,&sigmap ,wire_in, wire_out,&visited,&longest_path_cache,0,0);
                                
                                if(longest_path>abst_path){
                                    abst_path = longest_path;
                                    abst_longest_path.clear();
                                    for(auto sig_pasted: longest_path_cache){
                                        abst_longest_path.push_back(sig_pasted);
                                    }
                                    abst_longest_path.push_back(sigmap(wire_out));
                                }

                                //for(auto sig_pasted: longest_path_cache){
                                //    log("%s->",(sig_pasted.is_wire() ? log_id(sig_pasted.as_wire()) : log_id(sig_pasted.as_bit().wire)));
                                //}
                                //log("%s -| longest path %s->%s returned is: %d \n", log_id(wire_out),log_id(wire_in),log_id(wire_out),longest_path);
                            }
                        }
                    }
                }

                //Longest path
                std::string show_cmd = "show";

                log("Longest path found in total is %d\n",abst_path);
                for(auto sig_pasted: abst_longest_path){
                    show_cmd.append(" -color red w:");
                    show_cmd.append(sig_pasted.is_wire() ? log_id(sig_pasted.as_wire()) : log_id(sig_pasted.as_bit().wire));
                    log("%s,",(sig_pasted.is_wire() ? log_id(sig_pasted.as_wire()) : log_id(sig_pasted.as_bit().wire)));
                }
                log("\n Created following show: \n");
                
                log("%s\n", show_cmd.c_str());

                //calling the show
                Pass::call(design, show_cmd);

            }
        }

    }
} Critical_path;

PRIVATE_NAMESPACE_END