from flowgraph import FlowGraph
from interferencegraph import InterferenceGraph


class RegisterAllocator:
    """
        Target independent register allocator.

        Algorithm is iterated .. (ITR) by Appel and George

        The process consists of the following steps:
        - build interference graph from the instruction list
        - (optional) coalesc registers to remove redundant moves
        - (optional) spill registers
        - select registers
    """
    def __init__(self):
        pass

    def useUnused(self, inslist):
        # Use unused temporaries at the end of the list
        defTemps = []
        useTemps = []
        for i in inslist:
            for d in iter(i.dst):
                defTemps.append(d)
            for s in iter(i.src):
                useTemps.append(s)
        defTemps = set(defTemps)
        useTemps = set(useTemps)
        unUsed = defTemps - useTemps
        assert not unUsed

    def build(self, f):
        """ 1. Construct interference graph from instruction list """
        self.useUnused(f.instructions)
        f.cfg = FlowGraph(f.instructions)
        f.ig = InterferenceGraph(f.cfg)
        self.pre_colored = set(n for n in f.ig.nodes if n.varname in f.tempMap.keys())
        print('pre-colored:', self.pre_colored)
        self.moves = [i for i in f.instructions if i.ismove]
        print('move instructions:', self.moves)
        self.nodestack = []

    def getMoveRelated(self, f):
        mr = set()
        for i in self.moves:
            mr.add(f.ig.getNode(i.src[0]))
            mr.add(f.ig.getNode(i.dst[0]))
        return mr

    def simplify(self, f):
        """ 2. Remove nodes from the graph """
        # Chaitin's algorithm:
        # remove all non move related nodes that have less than K neighbours:
        # Also do not remove pre-colored nodes.
        # move_related = self.getMoveRelated(f)
        #print('move related', move_related)
        # o_nodes = f.ig.nodes - (self.pre_colored | move_related)
        o_nodes = f.ig.nodes - self.pre_colored
        #print('o_nodes', o_nodes)
        while o_nodes:
            less_k = [n for n in o_nodes if n.Degree < self.K]
            if not less_k:
                raise NotImplementedError('Spilling not implemented')
            n = less_k[0]
            self.nodestack.append(n)
            f.ig.delNode(n)
            o_nodes.remove(n)

    def coalesc(self, f):
        # for mov in ins:
        # if mov.src is not connected to mov.dst
        # and the node being coalesced has less than K nodes
        # of significant degree, then the node can be coalesced.
        # This is briggs algorithm.
        for m in self.moves:
            src = f.ig.getNode(m.src[0])
            dst = f.ig.getNode(m.dst[0])
            assert not f.ig.hasEdge(src, dst)
            neighbours = src.Adj | dst.Adj
            # neighbours with significant degree:
            sd_nb = set(n for n in neighbours if n.Degree >= self.K)
            print('neighbours with significant degree', sd_nb)
            if len(sd_nb) < self.K:
                print('Coalesc:', m)
                if dst in self.pre_colored:
                    if src in self.pre_colored:
                        break
                v2 = src
                for i in f.instructions:
                    rf = lambda x: x
                    i.src = rf(i.src)
                    i.dst = rf(i.dst)
                f.instructions.remove(m)

    def select(self, f):
        """ Add nodes of less than K degree back to the graph to color it. """
        regMap = dict(f.tempMap) # start with pre-colored

        # Add nodes back to the graph: 
        while self.nodestack:
            n = self.nodestack.pop(-1)
            f.ig.addNode(n)
            takenregs = set(regMap[m.varname] for m in n.Adj)
            possibleregs = set(self.regs) - takenregs
            regMap[n.varname] = list(possibleregs)[0]
        # We have a register mapping!
        print(regMap)
        # Use allocated registers:
        lookup = lambda t: regMap[t]
        for i in f.instructions:
            i.src = tuple(map(lookup, i.src))
            i.dst = tuple(map(lookup, i.dst))

    def allocFrame(self, f):
        """ Do register allocation for a single stack frame. """
        self.regs = set(f.regs)
        self.K = len(self.regs)
        self.build(f)
        self.simplify(f)
        # TODO: implement spilling
        #self.coalesc(f) # Optional!
        self.select(f)