#!/usr/bin/env python
# coding: utf-8

import pyverilog
from   pyverilog.dataflow.dataflow_analyzer import VerilogDataflowAnalyzer

import graphviz
import networkx as nx
import os
import numpy as np

def toOrientedGraph(nodes):
    tmp = []
    for node in nodes.values():
        for ins in node.getIns():
            opts = {'style': 'solid'}
            
            if ins["inv"]:
                opts['style'] = 'dashed'
                
            tmp.append(
                (
                    ins["node"],
                    node.getName(),
                    opts
                )
            )
    return tmp

depthBuffer = {}

def getSubDelay(node, nodes, depth):
    global depthBuffer

    if str(node.getName()) in depthBuffer:
        return depthBuffer[str(node.getName())]

    maxLen = 0
    if depth >= 100:
        maxLen = np.inf
    else:
        for ins in node.getIns():
            subLen = getSubDelay(nodes[str(ins["node"])], nodes, depth + 1)
            if subLen > maxLen:
                maxLen = subLen

    depthBuffer[str(node.getName())] = maxLen + 1
    
    return maxLen + 1

def getDelay(nodes):
    global depthBuffer

    depthBuffer.clear() # clear deph buffer

    maxLen = 0
    for node in nodes.values():
        if node.isOutput():
            subLen = getSubDelay(node, nodes, 0)
            if subLen > maxLen:
                maxLen = subLen
    return maxLen - 2 # exclude input and output node

def countGates(nodes):
    count = 0
    for node in nodes.values():
        if node.tree is not None and node.tree.__class__.__name__ == "DFOperator":
            count += 1
    return count

class node:
    def __init__(self, name, types):
        self.name  = name
        self.types = types
        self.tree  = None
        self.ins  = []

    # ((a >> 0) & 0x01) & ((b >> 0) & 0x01)) | (((b >> 0) & 0x01) & (0x01)) | (((a >> 0) & 0x01) & (0x01))
    def addMajiIn(self, tree, pos1, pos2, pos3):
        a    = None
        aInv = False

        if not(str(tree) == "Or" and str(tree.nextnodes[0]) == "Or" and str(tree.nextnodes[1]) == "And" and str(tree.nextnodes[0].nextnodes[0]) == "And" and str(tree.nextnodes[0].nextnodes[1]) == "And"):
            return a, aInv        

        if tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].__class__.__name__ == "DFOperator": # is classic input
            if tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].operator == "Unot":
                if tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].nextnodes[0].__class__.__name__ == "DFIntConst": # is constant
                    a = str(tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].nextnodes[0])
                else:
                    a = tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].nextnodes[0].nextnodes[0].nextnodes[0].name
                aInv = True
            else:
                a = tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].nextnodes[0].nextnodes[0].name
        elif tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].__class__.__name__ == "DFIntConst": # is constant
            a = str(tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3])
        else:
            return None, False
        
        return a, aInv

    # ( ~x0 & n18 ) | ( ~x0 & n19 ) | ( n18 & n19 )
    def addMTMaji(self, tree, pos1, pos2, pos3):
        a    = None
        aInv = False

        if not(str(tree) == "Or" and str(tree.nextnodes[0]) == "Or" and str(tree.nextnodes[1]) == "And" and str(tree.nextnodes[0].nextnodes[0]) == "And" and str(tree.nextnodes[0].nextnodes[1]) == "And"):
            return a, aInv

        if tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].__class__.__name__ == "DFOperator": # is classic input
            if tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].operator == "Unot":
                a = str(tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3].nextnodes[0])
                aInv = True
            else:
                return None, False
        else:
            a = str(tree.nextnodes[pos1].nextnodes[pos2].nextnodes[pos3])
        
        return a, aInv

    def parseMaji(self, tree):
        if tree.__class__.__name__ == "DFOperator":
            a, aInv, b, bInv, c, cInv = None, None, None, None, None, None

            try:
                a, aInv = self.addMajiIn(tree, 0, 0, 0)
                b, bInv = self.addMajiIn(tree, 0, 0, 1)
                c, cInv = self.addMajiIn(tree, 0, 1, 1)
            except:
                a, aInv, b, bInv, c, cInv = None, None, None, None, None, None

            try:
                if a is None or b is None or c is None:
                    a, aInv = self.addMTMaji(tree, 0, 0, 0)
                    b, bInv = self.addMTMaji(tree, 0, 0, 1)
                    c, cInv = self.addMTMaji(tree, 0, 1, 1)
            except:
                a, aInv, b, bInv, c, cInv = None, None, None, None, None, None

            if a is None or b is None or c is None:
                self.addInFromTree(tree, False)
            else:
                self.addIn(a, aInv)
                self.addIn(b, bInv)
                self.addIn(c, cInv)
        else:
            self.addInFromTree(tree, False)

    def addInFromTree(self, tree, invert = False):
        if tree.__class__.__name__ == "DFOperator":
            for node in tree.nextnodes:
                self.addInFromTree(node, invert)
        elif tree.__class__.__name__ == "DFIntConst":
            self.addIn(str(tree), invert)
        else:
            self.addIn(tree.name, invert)

    def isInput(self):
        return "Input" in self.types
        
    def isOutput(self):
        return "Output" in self.types

    def getIns(self):
        return self.ins

    def getName(self):
        return self.name

    def addIn(self, name, invert = False):
        self.ins.append({
            "node": name,
            "inv":  invert
        })

def getNodes(file, top = "top"):
    nodes = {}

    nodes['0'] = node('0', {"Input"})
    nodes['1'] = node('1', {"Input"})


    analyzer = VerilogDataflowAnalyzer([file], top)
    analyzer.generate()

    # First create nodes
    terms = analyzer.getTerms()
    for tk, tv in sorted(terms.items(), key=lambda x: str(x[0])):
        nodes[str(tv.name)] = node(tv.name, tv.termtype)

    # Create connections between nodes
    binddict = analyzer.getBinddict()
    for bk, bv in sorted(binddict.items(), key=lambda x: str(x[0])):
        for bvi in bv:
            nodes[str(bvi.dest)].parseMaji(bvi.tree)
            nodes[str(bvi.dest)].tree = bvi.tree

    return nodes

def plotNodes(nodes, name):
    # create a directed multi-graph
    G = nx.MultiDiGraph()
    G.add_edges_from(toOrientedGraph(nodes))
    # plot the graph
    a = nx.nx_agraph.to_agraph(G)
    src = graphviz.Source(a.to_string())
    src.render(name)

def generateAg(name):
    return """from io import StringIO
import os, sys

DIR_PATH = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(DIR_PATH, '..'))

from ariths_gen.core.arithmetic_circuits import GeneralCircuit
from ariths_gen.multi_bit_circuits.adders import UnsignedRippleCarryAdder, UnsignedCarryLookaheadAdder, UnsignedBrentKungAdder, UnsignedCarryIncrementAdder, UnsignedCarrySkipAdder, UnsignedCarrySelectAdder, UnsignedConditionalSumAdder, UnsignedHanCarlsonAdder, UnsignedKnowlesAdder, UnsignedKoggeStoneAdder, UnsignedLadnerFischerAdder, UnsignedPGRippleCarryAdder, UnsignedSklanskyAdder
from ariths_gen.multi_bit_circuits.multipliers import UnsignedArrayMultiplier, UnsignedCarrySaveMultiplier, UnsignedDaddaMultiplier, UnsignedWallaceMultiplier
from ariths_gen.wire_components import Bus

from ariths_gen.pdk import *
import os

if __name__ == "__main__":
    uut = """ + name + """
    maji = uut(Bus("a", 8), Bus("b", 8), prefix="", name="top")

    maji.get_v_code_flat(open("mig.v", "w"))
"""

def generateMT(name, outputA = True):
    if outputA:
        return """#include <lorina/aiger.hpp>
#include <mockturtle/mockturtle.hpp>

int main() {
  mockturtle::mig_network mig;

  std::vector<mockturtle::signal<mockturtle::mig_network>> a;
  std::vector<mockturtle::signal<mockturtle::mig_network>> b;
  auto cin = mig.get_constant(false);

  for (unsigned i = 0; i < 8; i++) {
    a.push_back(mig.create_pi());
    b.push_back(mig.create_pi());
  }
  
  mockturtle::""" + name + """(mig, a, b, cin);

  auto out = a;

  for (unsigned i = 0; i < out.size(); i++) {
    mig.create_po(out[i]);
  }
  
  // output
  mockturtle::write_verilog(mig, "mig.v");

  return 0;
}
"""
    else:
        return """#include <lorina/aiger.hpp>
#include <mockturtle/mockturtle.hpp>

int main() {
  mockturtle::mig_network mig;

  std::vector<mockturtle::signal<mockturtle::mig_network>> a;
  std::vector<mockturtle::signal<mockturtle::mig_network>> b;

  for (unsigned i = 0; i < 8; i++) {
    a.push_back(mig.create_pi());
    b.push_back(mig.create_pi());
  }
  
  auto out = mockturtle::""" + name + """(mig, a, b);

  for (unsigned i = 0; i < out.size(); i++) {
    mig.create_po(out[i]);
  }
  
  // output
  mockturtle::write_verilog(mig, "mig.v");

  return 0;
}
"""

# https://mockturtle.readthedocs.io/en/latest/generators/arithmetic.html#addition-and-subtraction
circs = [
    # ADDRS
    {
        "name":  "UnsignedRippleCarryAdder",
        "ag":    "UnsignedRippleCarryAdder",
        "mt":    "carry_ripple_adder_inplace",
        "mtAOut": True
    },
    {
        "name":  "UnsignedCarryLookaheadAdder",
        "ag":    "UnsignedCarryLookaheadAdder",
        "mt":    "carry_lookahead_adder_inplace",
        "mtAOut": True
    },
    {
        "name":  "UnsignedBrentKungAdder",
        "ag":    "UnsignedBrentKungAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedCarryIncrementAdder",
        "ag":    "UnsignedCarryIncrementAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedCarrySkipAdder",
        "ag":    "UnsignedCarrySkipAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedCarrySelectAdder",
        "ag":    "UnsignedCarrySelectAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedConditionalSumAdder",
        "ag":    "UnsignedConditionalSumAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedHanCarlsonAdder",
        "ag":    "UnsignedHanCarlsonAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedKnowlesAdder",
        "ag":    "UnsignedKnowlesAdder",
        "mt":    None
    },
        {
        "name":  "UnsignedKoggeStoneAdder",
        "ag":    "UnsignedKoggeStoneAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedLadnerFischerAdder",
        "ag":    "UnsignedLadnerFischerAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedPGRippleCarryAdder",
        "ag":    "UnsignedPGRippleCarryAdder",
        "mt":    None
    },
    {
        "name":  "UnsignedSklanskyAdder",
        "ag":    "UnsignedSklanskyAdder",
        "mt":    None
    },
    # MULS
    {
        "name": "UnsignedArrayMultiplier",
        "ag": "UnsignedArrayMultiplier",
        "mt": "carry_ripple_multiplier",
        "mtAOut": False
    },
    {
        "name": "UnsignedCarrySaveMultiplier",
        "ag":   "UnsignedCarrySaveMultiplier",
        "mt":   None
    },
    {
        "name": "UnsignedDaddaMultiplier",
        "ag":   "UnsignedDaddaMultiplier",
        "mt":   None
    },
    {
        "name": "UnsignedWallaceMultiplier",
        "ag":   "UnsignedWallaceMultiplier",
        "mt":   None
    }
]

#
# Begin main
#

MT_PATH  = "./mockturtle/"
AGM_PATH = "./ariths-gen-mig/"
AG_PATH  = "./ariths-gen/"
PWD      = os.getcwd()

table = open(f"{PWD}/compare.csv", "w")

print("type;MT (GATES);MT (DELAY);ariths-gen->AIG->MT (GATES);ariths-gen->AIG->MT (DELAY);ariths-gen->AIG->MT+RESUBSTITUTION (GATES);ariths-gen->AIG->MT+RESUBSTITUTION (DELAY);ariths-gen-mig (GATES);ariths-gen-mig (DELAY);ariths-gen-mi->AIG->MT+RESUBSTITUTION (GATES);ariths-gen-mi->AIG->MT+RESUBSTITUTION (DELAY)", file=table)

for circ in circs:
    
    MTDirectGates = "-"
    MTDirectDelay = "-"

    ## MockTurtle direct generator

    if circ["mt"] is not None:
	
        with open(f"{MT_PATH}/examples/generate.cpp", "w") as Vfile:
            print(generateMT(circ["mt"], circ["mtAOut"]), file=Vfile)

        os.system(f"cd {MT_PATH}/build/examples && make && ./generate && cp mig.v {PWD}/mig.v")

        # generate blif for verification
        os.system(f"yosys formal.ys")

        os.system(f"mv formal.blif formalMT.blif && mv mig.v {circ['name']}_MT.v")

        nodes = getNodes(f"{circ['name']}_MT.v")

        MTDirectGates = countGates(nodes)
        MTDirectDelay = getDelay(nodes)

        plotNodes(nodes, f"{circ['name']}_MT")

    ## ArithsGen

    with open(f"{AG_PATH}/generate.py", "w") as Vfile:
        print(generateAg(circ["ag"]), file=Vfile)

    os.system(f"cd {AG_PATH} && python3 ./generate.py && cp mig.v {PWD}/mig.v")

    os.system(f"sed -i -e 's/0x00/0/g' mig.v && sed -i -e 's/0x01/1/g' mig.v") # 0x00 and 0x01 to 0 and 1

    # generate blif for verification
    os.system(f"yosys formal.ys")

    os.system(f"yosys toAig.ys") # yosys VERILOG to AIG

    os.system(f"./{MT_PATH}/build/examples/passAig") # aig.aig to mig.v

    os.system(f"mv mig.v {circ['name']}_AG.v")

    os.system(f"./{MT_PATH}/build/examples/resubstitution") # aig.aig to mig.v (resubstitution)

    os.system(f"mv mig.v {circ['name']}_AG_RE.v")

    nodes = getNodes(f"{circ['name']}_AG.v")

    print(f"Working on {circ['name']}_AG")

    AGDirectGates = countGates(nodes)
    AGDirectDelay = getDelay(nodes)

    plotNodes(nodes, f"{circ['name']}_AG")

    nodes = getNodes(f"{circ['name']}_AG_RE.v")

    print(f"Working on {circ['name']}_AG_RE")

    AGREDirectGates = countGates(nodes)
    AGREDirectDelay = getDelay(nodes)

    plotNodes(nodes, f"{circ['name']}_AG_RE")

    ## ArithsGenMig
    
    with open(f"{AGM_PATH}/generate.py", "w") as Vfile:
        print(generateAg(circ["ag"]), file=Vfile)

    os.system(f"cd {AGM_PATH} && python3 ./generate.py && cp mig.v {PWD}/{circ['name']}_AGM.v")

    os.system(f"sed -i -e 's/0x00/0/g' {circ['name']}_AGM.v && sed -i -e 's/0x01/1/g' {circ['name']}_AGM.v") # 0x00 and 0x01 to 0 and 1

    nodes = getNodes(f"{circ['name']}_AGM.v")

    print(f"Working on {circ['name']}_AGM")

    AGMDirectGates = countGates(nodes)
    AGMDirectDelay = getDelay(nodes)

    plotNodes(nodes, f"{circ['name']}_AGM")

    os.system(f"cp {circ['name']}_AGM.v mig.v")

    os.system(f"yosys toAig.ys") # yosys VERILOG to AIG

    os.system(f"./{MT_PATH}/build/examples/resubstitution") # aig.aig to mig.v (resubstitution)

    os.system(f"mv mig.v {circ['name']}_AGM_RE.v")

    nodes = getNodes(f"{circ['name']}_AGM_RE.v")

    print(f"Working on {circ['name']}_AGM_RE")

    AGMREDirectGates = countGates(nodes)
    AGMREDirectDelay = getDelay(nodes)

    plotNodes(nodes, f"{circ['name']}_AGM_RE")

    ec = os.system(f"bash yosys_equiv_check.sh -v {circ['name']}_AGM.v    -b formal.blif   -m sat")
    if ec != 0:
        print("Formal verification fail")
        exit(1)

    print(f"{circ['name']};{MTDirectGates};{MTDirectDelay};{AGDirectGates};{AGDirectDelay};{AGREDirectGates};{AGREDirectDelay};{AGMDirectGates};{AGMDirectDelay};{AGMREDirectGates};{AGMREDirectDelay}", file=table)

os.system("zip data.zip *.pdf *.v")
os.system("mkdir stash")
os.system("mv toAig.ys stash && mv test.py stash && mv compare.csv stash && mv setup.sh stash && mv data.zip stash && mv yosys_equiv_check.sh stash && mv formal.ys stash")
os.system("rm -f *")
os.system("mv stash/* . && rm -rf stash")
table.close()