import graph

class InterferenceGraphNode(graph.Node):
    def __init__(self, g, varname):
        super().__init__(g)
        self.varname = varname

    def __repr__(self):
        return '{}'.format(self.varname)


class InterferenceGraph(graph.Graph):
    def __init__(self, flowgraph):
        """ 
            Create a new interference graph from a flowgraph 
        """
        super().__init__()
        # Mapping from node to live set
        self.liveMap = {}
        self.tempMap = {}
        # Calculate liveness in CFG:
        ###
        # Liveness:
        #  in[n] = use[n] UNION (out[n] - def[n])
        #  out[n] = for s in n.succ in union in[s]
        ###
        for n in flowgraph.nodes:
            n.live_in = set()
            n.live_out = set()
        # Dataflow fixed point iteration:
        change = True
        while change:
            change = False
            for n in flowgraph.nodes:
                _in = n.live_in
                _out = n.live_out
                n.live_in = n.uses | (n.live_out - n.defs)
                if n.succ:
                    n.live_out = set.union(*(s.live_in for s in n.succ))
                else:
                    n.live_out = set()
                change = change or (_in != n.live_in) or (_out != n.live_out)
        # Construct interference graph:
        for n in flowgraph.nodes:
            for tmp in n.live_out:
                n1 = self.getNode(tmp)
                for tmp2 in (n.live_out - {tmp}):
                    n2 = self.getNode(tmp2)
                    self.addEdge(n1, n2)

    def getNode(self, tmp):
        if tmp not in self.tempMap:
            self.tempMap[tmp] = self.newNode(tmp)
        return self.tempMap[tmp]

    def newNode(self, varname):
        n = InterferenceGraphNode(self, varname)
        self.nodes.append(n)
        return n


class RegisterAllocator:
    """ Target independent register allocator """
    def registerAllocate(self, ig, regs, precolored):
        """ 
            Color the given interference graph with the provided registers 
        """
        regMap = dict(precolored)
        regs = set(regs)
        K = len(regs)
        allvars = list(n for n in ig.nodes if n.varname not in regMap)

        # Chaitin's algorithm:
        # remove all nodes that have less than K neighbours:
        worklist = []
        while allvars:
            less_k = [n for n in allvars if len(n.Adj - set(worklist)) < K]
            if not less_k:
                raise NotImplementedError('Spilling not implemented')
            n = less_k[0]
            worklist.append(n)
            allvars.remove(n)

        # Add nodes back to the graph: 
        while worklist:
            n = worklist.pop(-1)
            adj = n.Adj - set(worklist)
            possibleregs = set(regs) - set(regMap[n.varname] for n in adj)
            regMap[n.varname] = list(possibleregs)[0]
        # We have a register mapping!
        # print(regMap)
        # TODO: implement spilling
        # TODO: implement coalescing
        return regMap