from typing import List, Dict
from Lib.Errors import MiniCInternalError
from Lib.Operands import Temporary, Operand, S, Register, Offset, DataLocation, GP_REGS
from Lib.Statement import Instruction
from Lib.Allocator import Allocator
from Lib.FunctionData import FunctionData
from Lib import RiscV
from Lib.Graphes import Graph  # For Graph coloring utility functions


class SmartAllocator(Allocator):

    _igraph: Graph  # interference graph

    def __init__(self, fdata: FunctionData, basename: str, liveness,
                 debug=False, debug_graphs=False):
        self._liveness = liveness
        self._basename: str = basename
        self._debug: bool = debug
        self._debug_graphs: bool = debug_graphs
        super().__init__(fdata)

    def replace(self, old_instr: Instruction) -> List[Instruction]:
        """
        Replace Temporary operands with the corresponding allocated
        physical register (Register) OR memory location (Offset).
        """
        before: List[Instruction] = []
        after: List[Instruction] = []
        subst: Dict[Operand, Operand] = {}
        
        old_args = old_instr.args()
        numreg = 1
        for arg in old_args:
            if(isinstance(arg, Temporary)): # Else we don't change anything
                dest = self._fdata._pool.get_alloced_loc(arg)
                if(isinstance(dest,Register)):
                    # We use the provided register
                    subst[arg] = dest
                elif(isinstance(dest,Offset)):
                    # We use a s register in which we will read|write
                    subst[arg] = S[numreg]
                    numreg += 1
                else:
                    raise MiniCInternalError("Unsupported dataLocation for smart replace")
        for arg in old_instr.used():
            if(isinstance(arg, Temporary)):
                # If the argument has an Offset substitution
                if(isinstance(self._fdata._pool.get_alloced_loc(arg),Offset)):
                    # We have to read it from memory
                    before.append(RiscV.ld(subst[arg],self._fdata._pool.get_alloced_loc(arg)))
        for arg in old_instr.defined():
            if(isinstance(arg, Temporary)):
                # If the argument has an Offset substitution
                if(isinstance(self._fdata._pool.get_alloced_loc(arg),Offset)):
                    # We have to write them after to memory
                    after.append(RiscV.sd(subst[arg],self._fdata._pool.get_alloced_loc(arg)))
        
        # And now return the new list!
        
        print(old_instr,"->",before,subst,after)
        try:
            new_instr = old_instr.substitute(subst)
        except Exception:
            # We have an instruction that doesn't need substitution
            return [old_instr]
        return before + [new_instr] + after
    
    def prepare(self) -> None:
        """
        Perform all preparatory steps related to smart register allocation:

        - Dataflow analysis to compute the liveness range of each
          temporary.
        - Interference graph construction.
        - Graph coloring.
        - Associating temporaries with actual locations.
        """
        # Liveness analysis
        self._liveness.run()
        # Interference graph
        self.build_interference_graph()
        if self._debug_graphs:
            print("Printing the interference graph")
            self._igraph.print_dot(self._basename + "interference.dot")
        # Smart Allocation via graph coloring
        self.smart_alloc()

    def build_interference_graph(self) -> None:
        """
        Build the interference graph (in self._igraph).
        Vertices of the graph are temporaries,
        and an edge exists between temporaries iff they are in conflict.
        """
        self._igraph: Graph = Graph()
        # Create a vertex for every temporary
        # There may be temporaries the code does not use anymore,
        # but it does not matter as they interfere with no one.
        for v in self._fdata._pool.get_all_temps():
            self._igraph.add_vertex(v)
        # Iterate over self._liveness._liveout (dictionary containing all
        # live out temporaries for each instruction), and for each conflict use
        # self._igraph.add_edge((t1, t2)) to add the corresponding edge.
        for st,varz in self._liveness._liveout.items():
            for x in varz:
                if(isinstance(x,Temporary)):
                    # First condition
                    for y in varz:
                        if x!=y and isinstance(y,Temporary):
                            self._igraph.add_edge((x,y))
                    # Second/third condition
                    for y in st.defined():
                        if x!=y and isinstance(y,Temporary): # x is alive after x definition
                            self._igraph.add_edge((x,y))

    def smart_alloc(self) -> None:
        """
        Allocates all temporaries via graph coloring.
        Prints the colored graph if self._debug_graphs is True.

        Precondition: the interference graph _igraph must have been built.
        """
        # Checking the interference graph has been built
        if not self._igraph:
            raise MiniCInternalError("Empty interference graph in the Smart Allocator")
        # Coloring of the interference graph
        coloringreg: Dict[Temporary, int] = self._igraph.color()
        if self._debug_graphs:
            print("coloring = " + str(coloringreg))
            self._igraph.print_dot(self._basename + "_colored.dot", coloringreg)
        # Temporary -> DataLocation (Register or Offset) dictionary,
        # specifying where a given Temporary should be allocated:
        alloc_dict: Dict[Temporary, DataLocation] = dict()
        # Use the coloring `coloringreg` to fill `alloc_dict`.
        # Our version is less than 5 lines of code.
        moreColors_alloc_dict: Dict[int, DataLocation] = dict()
        
        for var,col in coloringreg.items():
            if(col<len(GP_REGS)):
                alloc_dict[var] = GP_REGS[col]
            else:
                if(not col in moreColors_alloc_dict):
                    moreColors_alloc_dict[col] = self._fdata.fresh_offset()
                alloc_dict[var] = moreColors_alloc_dict[col]
        
        if self._debug:
            print("Allocation:")
            print(alloc_dict)
        self._fdata._pool.set_temp_allocation(alloc_dict)