import sys
from pypy.objspace.flow.model import Variable, Constant, SpaceOperation
from pypy.objspace.flow.model import Block, Link, checkgraph, traverse, flatten
from pypy.tool.unionfind import UnionFind
from pypy.rpython import lltype
from pypy.translator.simplify import transform_dead_op_vars
from pypy.translator.translator import Translator
from pypy.translator import backendoptimization
from pypy.annotation import model as annmodel



def fn(x, y):
    total = 0
    for i in range(10):
        total += i
    return total


import matfunc

def measure_median_execution_cost(graph):
    linktargets = [graph.startblock]
    linkmap = {}
    for node in flatten(graph):
        if isinstance(node, Link):
            linkmap[node] = len(linktargets)
            linktargets.append(node.target)
    matrix = []
    vector = []
    for i, target in zip(range(len(linktargets)), linktargets):
        vector.append(len(target.operations))
        row = [0.0] * len(linktargets)
        row[i] = 1.0
        if target.exits:
            f = 1.0 / len(target.exits)
            for nextlink in target.exits:
                row[linkmap[nextlink]] -= f
        matrix.append(row)
    M = matfunc.Mat(matrix)
    V = matfunc.Vec(vector)
    # we must solve: M * (vector x1...xn) = V
    try:
        Solution = M._solve(V)
    except (OverflowError, ValueError):
        return sys.maxint
    else:
        return Solution[0]

def static_instruction_count(graph):
    count = 0
    for node in flatten(graph):
        if isinstance(node, Block):
            count += len(node.operations)
    return count

def inlining_heuristic(graph):
    # XXX ponderation factors?
    return (0.819487132 * measure_median_execution_cost(graph) +
            static_instruction_count(graph))


def static_callers(translator):
    result = []
    def build_call_graph(node):
        if isinstance(node, Block):
            for op in node.operations:
                if (op.opname == "direct_call" and
                    isinstance(op.args[0], Constant)):
                    funcobj = op.args[0].value._obj
                    graph = getattr(funcobj, 'graph', None)
                    if graph is not None:
                        result.append((parentgraph, graph))
    for parentgraph in translator.flowgraphs.itervalues():
        traverse(build_call_graph, parentgraph)
    return result


def auto_inlining(translator, threshold=20):
    from heapq import heappop, heapreplace
    callers = {}     # {graph: {graphs-that-call-it}}
    callees = {}     # {graph: {graphs-that-it-calls}}
    for graph1, graph2 in static_callers(translator):
        callers.setdefault(graph2, {})[graph1] = True
        callees.setdefault(graph1, {})[graph2] = True
    fiboheap = [(0.0, graph) for graph in callers]
    valid_weight = {}

    while fiboheap:
        weight, graph = fiboheap[0]
        if not valid_weight.get(graph):
            weight = inlining_heuristic(graph)
            print '  + cost %7.2f %50s' % (weight, graph.name)
            heapreplace(fiboheap, (weight, graph))
            valid_weight[graph] = True
            continue

        if weight >= threshold:
            break   # finished

        heappop(fiboheap)
        print 'Inlining %7.2f %50s' % (weight, graph.name)
        for parentgraph in callers[graph]:
            if parentgraph == graph:
                continue
            print '\t\t-> in %s' % parentgraph.name
            try:
                if backendoptimization.inline_function(translator, graph,
                                                       parentgraph):
                    valid_weight[parentgraph] = False
                    for graph2 in callees.get(graph, {}):
                        callees[parentgraph][graph2] = True
                        callers[graph2][parentgraph] = True
            except NotImplementedError:
                pass


def tget():
    t = Translator(fn)
    a = t.annotate([int]*2)
    t.specialize()
    t.backend_optimizations()
    t.checkgraphs()
    return t


def test():
    t = tget()
    #t.view()
    print repr(inlining_heuristic(t.getflowgraph()))
    auto_inlining(t)
    t.backend_optimizations()
    t.checkgraphs()
    t.view()

if __name__ == '__main__':
    test()