import random, itertools

from helpers import *


global_cache = {}

def randomfacet_analytic(uso, n, cache=global_cache):
    """
    Return the expected runtime of RandomFacet on the USO generated by uso.

    Start positions of the algorithm are all elements of the vertices iterator. This allow easy
    splitting into separate tasks.
    Due to numerical stability concerns, the algorithm actually gives 2^n times the sum of runtimes
    over all random choices. To get the expectation of a single execution, divide by N! * 2^n
    """

    # as soon as this function leaves the namespace, the associated cache is garbage collected.
    @memoize()
    def fingerprint_subcube(cube):
        return uso.fingerprint_subcube(cube)

    def cache_key(vertex, cube):
        return global2local(vertex, cube), fingerprint_subcube(cube)

    @memoize(cache, cache_key=cache_key)
    def RF(vertex, cube):
        total_runtime = 0
        sink = uso.find_sink(cube)
        for i in range(n):
            if cube[i] != "*":
                continue

            # in the first branch, its necessary for the runtime to be summed twice in order
            # to keep everything in integers. The other branch does the same to keep results
            # compatible
            if vertex is None:
                # vertex u.a.r.
                good_cube = cube[:i] + sink[i] + cube[i+1:]
                bad_cube = cube[:i] + ("1" if sink[i] == "0" else "0") + cube[i+1:]
                w = uso.find_sink(bad_cube)
                w = w[:i] + good_cube[i] + w[i+1:]
                total_runtime += RF(None, good_cube) + RF(None, bad_cube) + RF(w, good_cube)
            else:
                # recurse on the side where reference vertex is
                subcube = cube[:i] + vertex[i] + cube[i+1:]
                total_runtime += 2 * RF(vertex, subcube)
                if sink[i] != subcube[i]:
                    # sink not found in first recursive call, replace vertex with w ^ {i}
                    w = uso.find_sink(subcube)
                    w = w[:i] + ("1" if w[i] == "0" else "0") + w[i+1:]
                    subcube = cube[:i] + w[i] + cube[i+1:]
                    total_runtime += 2 * RF(w, subcube)

        res = max(1, total_runtime)
        return res

    return RF(None, "*"*n)


def randomfacet_analytic_old(uso, n, vertices=None, cube=None, cache=global_cache):
    
    @memoize()
    def find_sink(subcube):
        return uso.find_sink(subcube)
    
    @memoize(cache)
    def RF(vertex, cube):
        total_runtime = 0
        sink = find_sink(cube)
        for i in range(n):
            if cube[i] != "*":
                continue

            # recurse on the side where reference vertex is
            subcube = cube[:i] + vertex[i] + cube[i+1:]
            runtime = RF(vertex, subcube)
            total_runtime += runtime

            if sink[i] != subcube[i]:
                # sink not found in first recursive call, replace vertex with w ^ {i}
                w = find_sink(subcube)
                w = w[:i] + ("1" if w[i] == "0" else "0") + w[i+1:]
                subcube = cube[:i] + w[i] + cube[i+1:]
                runtime = RF(w, subcube)
                total_runtime += runtime

        return max(1, total_runtime)

    if cube is None:
        cube = "*" * n

    if vertices is None:
        # iterate over whole set of vertices
        vertices = generate_vertices(n)

    return sum(RF(vertex, cube) for vertex in vertices)


def randomfacet_sample(uso, N, cube=None, vertex=None):
    """
    Return the runtime of 1 execution of RandomFacet on the USO generated by uso.
    """
    
    def RF(vertex, cube):
        if vertex == cube:
            return 1, vertex

        sink = uso.find_sink(cube)

        choices = [i for (i, dim) in enumerate(cube) if dim == "*"]
        i = random.choice(choices)
        #i = choices[0]

        # recurse on the side where the reference vertex is
        subcube = cube[:i] + vertex[i] + cube[i+1:]
        runtime, w = RF(vertex, subcube)
        total_runtime = runtime

        if sink[i] != w[i]:
            # sink not found in first recursive call, replace vertex with w ^ {i}
            w = w[:i] + ("1" if w[i] == "0" else "0") + w[i+1:]
            subcube = cube[:i] + w[i] + cube[i+1:]
            runtime, _ = RF(w, subcube)
            total_runtime += runtime

        return total_runtime, sink

    if cube is None:
        cube = "*" * N

    if vertex is None:
        # choose random vertex
        v = random.randint(0, 2**N-1)
        vertex = int_to_vertex(v, N)

    return RF(vertex, cube)[0]


