import logging
import math
import re
from typing import Dict

from TopRTLParser import TopRTLParser
from HLSProject import HLSProject

FIFO_TYPE_DEPTH_THRESHOLD = 1024


class Edge:
  def __init__(self, name:str):
    self.src : Vertex = None
    self.dst : Vertex = None
    self.width = -1
    self.depth = -1  # original depth as in the RTL generated by HLS
    self.addr_width = -1
    self.name = name
    self.pipeline_level = 0
    self.added_depth_for_rebalance = 0
    self.fifo_type = ""

    logging.debug(f'create edge {self.name} of width {self.width} and depth {self.depth}')

  def __hash__(self):
    return hash(self.name)

  def __eq__(self, other):
    return self.name == other.name

  def setDepth(self, depth):
    self.depth = depth
    self.addr_width = int(math.log2(self.depth)+1)

  def addDepth(self, depth_delta):
    self.setDepth(self.depth + depth_delta)

  def setWidth(self, width):
    self.width = width

  def getDefaultType(self) -> str:
    if self.depth >= FIFO_TYPE_DEPTH_THRESHOLD:
      return 'BRAM'
    else:
      return 'SRL'

  def getSRLFIFOArea(self, skip_balance_part: bool) -> Dict[str, int]:
    fifo_depth = self.depth
    if not skip_balance_part:
      fifo_depth += self.added_depth_for_rebalance

    lut_ram = self.width * ((fifo_depth - 1) / 16 + 1)
    lut_logic = 15 + 3 * int(math.log2(fifo_depth))
    ff = 7 + 3 * int(math.log2(fifo_depth))

    return {
      'BRAM': 0,
      'DSP': 0,
      'URAM': 0,
      'FF': ff,
      'LUT': lut_ram + lut_logic
    }

  def getBRAMFIFOArea(self, skip_balance_part: bool) -> Dict[str, int]:
    return {
      'BRAM': int(self.width/36 * 2),  # x2 because we are counting BRAM_18k
      'DSP': 0,
      'URAM': 0,
      'FF': self.width,  # approximation
      'LUT': 0  # based on HLS report
    }

  def getArea(self, skip_balance_part: bool = False) -> Dict[str, int]:
    if self.fifo_type == '':
      fifo_type = self.getDefaultType()
    else:
      fifo_type = self.fifo_type

    if fifo_type == 'BRAM':
      return self.getBRAMFIFOArea(skip_balance_part)
    elif fifo_type == 'SRL':
      return self.getSRLFIFOArea(skip_balance_part)
    else:
      assert False, fifo_type

  def getName(self):
    return self.name

class Vertex():
  def __init__(self, type:str, name : str):
    self.in_edges = [] # stores Edge objects
    self.out_edges = []
    self.in_edge_names = [] # stores Edge objects
    self.out_edge_names = []
    self.type = type
    self.name = name
    self.id = self.type + self.name
    self.area = {} # str_name -> count
    self.latency = 0

    logging.debug(f'create vertex {self.name} of type {self.type}')

  def __hash__(self):
    return hash(self.id)

  def __eq__(self, other):
    return self.id == other.id

  def getEdgeNames(self):
    return self.in_edge_names + self.out_edge_names

  def getEdges(self):
    return self.in_edges + self.out_edges
  
  def getInEdges(self):
    return self.in_edges
  
  def getOutEdges(self):
    return self.out_edges

  def getNeighborVertices(self):
    src_neighbors = [e.src for e in self.getInEdges()]
    dst_neighbors = [e.dst for e in self.getOutEdges()]
    return src_neighbors + dst_neighbors

  def getTotalWireWidth(self) -> int:
    """
    sum of the width of all edges of the vertex
    """
    return sum(e.width for e in self.getEdges())

  def getVertexAndInboundFIFOArea(self) -> Dict[str, int]:
    return {
      item: sum(e.getArea()[item] for e in self.in_edges) + val for item, val in self.area.items()
    }
  
  def getName(self) -> str:
    return self.name
  
  def getLatency(self) -> int:
    return self.latency
  
  def getArea(self) -> dict:
    return self.area

class DataflowGraph:
  def __init__(self, hls_prj_manager : HLSProject, top_rtl_parser : TopRTLParser):
    self.hls_prj_manager = hls_prj_manager
    self.top_rtl_parser = top_rtl_parser

    self.vertices = {}  # name -> Vertex
    self.edges = {}     # name -> Edge
    logging.info([_.name for _ in self.top_rtl_parser.traverseVertexInAST()])
    for v_node in self.top_rtl_parser.traverseVertexInAST():
      if v_node.name in self.hls_prj_manager.getSubInst() or re.search("mul_\d+[a-z]+_\d+[a-z]+", v_node.name) : # or re.search("data_split", v_node.name)
        logging.warning("Excluding module {} - name {}".format(v_node.module, v_node.name))
        continue
      logging.info("v_node.module: {}, v_node.name: {}".format(v_node.module, v_node.name))
      self.__initVertices(v_node)

    for e_node in self.top_rtl_parser.traverseEdgeInAST():
      self.__initEdges(e_node)  #! set width and depth of FIFO

    self.__linkEdgeAndVertex()
    
    self.__checker()

    self.v_type_2_int = {} # map each vertex type     to an integer
    self.v_name_2_int = {} # map each vertex instance to an integer
    self.int_2_v_type = {}
    self.int_2_v_name = {}
    self.__initVTypeToInt()
    self.__initVInstToInt()
    
    # for v_node in self.top_rtl_parser.traverseVertexInAST():
    #   logging.debug("In: {}".format(self.vertices[v_node.name].in_edge_names))
    #   logging.debug("Out: {}".format(self.vertices[v_node.name].out_edge_names))
    
  # assign an integer to v type
  def __initVTypeToInt(self):
    id = 1
    for v_node in self.vertices.values():
      if v_node.type not in self.v_type_2_int:
        self.v_type_2_int[v_node.type] = id
        self.int_2_v_type[id] = v_node.type
        logging.debug(f'{v_node.type} : {id}')
        id += 1

  # assign an integer to v name
  def __initVInstToInt(self):
    id = 1
    for v_node in self.vertices.values():
      if v_node.name not in self.v_name_2_int:
        self.v_name_2_int[v_node.name] = id
        self.int_2_v_name[id] = v_node.name
        logging.debug(f'{v_node.name} : {id}')
        id += 1
      else:
        assert False, f'Found two modules of the same name {v_node.name}'

  # note that we need to merge repetitve edges
  def getIntegerGraph(self):
    int_e_list = [ (self.v_name_2_int[e.src.name], self.v_name_2_int[e.dst.name]) \
                  for e in self.edges.values() ]
    int_e_list = list(set(int_e_list))

    int_v_labels = [ [self.v_name_2_int[v.name], self.v_type_2_int[v.type] ] \
                  for v in self.vertices.values()]
    return int_e_list, int_v_labels

  def getIntEdgeToName(self):
    int_edge2name = { (self.v_name_2_int[e.src.name], self.v_name_2_int[e.dst.name]) : e.name \
                  for e in self.edges.values() }
    return int_edge2name

  def getIntIdToVType(self):
    return self.int_2_v_type

  def getIntIdToVName(self):
    return self.int_2_v_name

  def __checker(self):
    v_name_list = [v.type + v.name for v in self.getAllVertices()]
    e_name_list = [e.name for e in self.getAllEdges()]
    assert len(v_name_list) == len(set(v_name_list)), 'Find repeated modules'
    assert len(e_name_list) == len(set(e_name_list))

  def __initVertices(self, v_node):
    # print(v_node.__dir__())
    v = Vertex(v_node.module, v_node.name)
    # get area
    # ! DLF: deal with BRAM between sub-functions
    if re.search("temp_[A-Za-z_]+\d+", v_node.name) or re.search("data_split", v_node.name):
      # ! previously v_node.module, an counter example is like 
      '''
      top_nondf_kernel_2mm_A_V #( // -> v_node.module
          .DataWidth( 512 ),
          .AddressRange( 1024 ),
          .AddressWidth( 10 ))
      temp_xout0_V_U(       // -> v_node.name
          .clk(ap_clk),
          .reset(ap_rst_n_inv),
          .address0(temp_xout0_V_address0),
          .ce0(temp_xout0_V_ce0),
          .we0(temp_xout0_V_we0),
          .d0(grp_kernel3_x0_fu_92_C_d0),
          .q0(temp_xout0_V_q0)
      );
      ''' # ! v_node.module sometimes doesn't include that pattern, so v_node.name is more precise.
      v.area = {'BRAM': 0, 'DSP':0, 'FF': 0, 'LUT': 0, 'URAM':0}
    else:
      v.area = self.hls_prj_manager.getAreaFromModuleType(v.type)
    
    if not re.search("control_s_axi", v_node.module) and not re.search("gmem_", v_node.module) and not re.search("temp_[A-Za-z_]+\d+", v_node.name) and not re.search("data_split", v_node.name):  # ! same here, previously v_node.module
      v.latency = self.hls_prj_manager.getLatencyFromModuleType(v.type)
    
    v.in_edge_names = self.top_rtl_parser.getInFIFOsOfModuleInst(v.name)
    v.out_edge_names = self.top_rtl_parser.getOutFIFOsOfModuleInst(v.name)
    # if v_node.name.find("C_c1") != -1:
    logging.info(" in  edges: {}".format(v.in_edge_names))
    logging.info(" out edges: {}".format(v.out_edge_names))
    # if v_node.name.find("_nondf_") != -1:
    #   print(v.area)
    #   print(v.latency)
    #   print(v.in_edge_names)
    #   print(v.out_edge_names)
    #   exit()

    self.vertices[v_node.name] = v
    logging.info("self.vertices.keys() = {}".format(self.vertices.keys()))

  def __initEdges(self, e_node):

    e = Edge(e_node.name)

    # extract width
    e.setWidth(self.top_rtl_parser.getFIFOWidthFromFIFOType(e_node.module))
    e.setDepth(self.top_rtl_parser.getFIFODepthFromFIFOType(e_node.module))

    self.edges[e_node.name] = e

  def __linkEdgeAndVertex(self):
    for v in self.vertices.values():
      for fifo_in_name in v.in_edge_names:
        fifo_in = self.edges[fifo_in_name]
        fifo_in.dst = v
        v.in_edges.append(fifo_in)
        logging.debug("Adding dst {} to fifo {}".format(v.name, fifo_in_name))
      for fifo_out_name in v.out_edge_names:
        fifo_out = self.edges[fifo_out_name]
        fifo_out.src = v
        v.out_edges.append(fifo_out)
        logging.debug("Adding src {} to fifo {}".format(v.name, fifo_out_name))

  def printVertices(self):
    for v in self.vertices.values():
      logging.debug(f'{v.name}: latency={v.latency} area={v.area}')
      for e in v.in_edges:
        logging.debug(f'  <- {e.name}')
      for e in v.out_edges:
        logging.debug(f'  -> {e.name}')

  def printEdges(self):
    for e in self.edges.values():
      logging.debug(f'{e.name}: {e.src.name} -> {e.dst.name}')

  def getAllVertices(self):
    return self.vertices.values()

  def getAllEdges(self):
    return self.edges.values()

  def getNameToVertexMap(self):
    return self.vertices

  def getNameToEdgeMap(self):
    return self.edges

  def getVertex(self, v_name) -> Vertex:
    return self.vertices[v_name]