from modelverse_kernel.primitives import PrimitiveFinished
import modelverse_jit.runtime as jit_runtime 

def reverseKeyLookupMulti(a, b, **remainder):
    edges, = yield [("RO", [a])]
    b_val, = yield [("RV", [b])]
    expanded_edges = yield [("RE", [i]) for i in edges]
    values = yield [("RV", [i[1]]) for i in expanded_edges]
    result, = yield [("CN", [])]

    for i, edge in enumerate(values):
        if b_val == edge:
            # Found our edge: edges[i]
            outgoing, = yield [("RO", [edges[i]])]
            value, = yield [("RE", [outgoing[0]])]
            yield [("CE", [result, value[1]])]

    raise PrimitiveFinished(result)

def reverseKeyLookup(a, b, **remainder):
    # Guess, which might work
    guess, = yield [("RD", [a, "__%s" % b])]
    if guess == b:
        result, = yield [("CNV", ["__%s" % b])]
        raise PrimitiveFinished(result)

    # We failed: do the slow way
    edges, = yield [("RO", [a])]
    expanded_edges = yield [("RE", [i]) for i in edges]
    for i, edge in enumerate(expanded_edges):
        if b == edge[1]:
            # Found our edge: edges[i]
            outgoing, = yield [("RO", [edges[i]])]
            result, = yield [("RE", [outgoing[0]])]
            raise PrimitiveFinished(result[1])

    result, = yield [("CNV", ["(unknown: %s)" % b])]
    raise PrimitiveFinished(result)

def set_copy(a, **remainder):
    b, =         yield [("CN", [])]
    links, =     yield [("RO", [a])]
    exp_links = yield [("RE", [i]) for i in links]
    _ =         yield [("CE", [b, i[1]]) for i in exp_links]
    raise PrimitiveFinished(b)

def check_symbols(a, b, c, **remainder):
    symbols = {}
    function_name, = yield [("RV", [b])]
    symbols[function_name] = False
    object_links, = yield [("RO", [c])]
    set_elements = yield [("RE", [i]) for i in object_links]
    set_elements = [i[1] for i in set_elements]
    set_values = yield [("RV", [i]) for i in set_elements]
    set_elements = yield [("RD", [a, i]) for i in set_values]
    symbols_set = yield [("RD", [i, "symbols"]) for i in set_elements]
    all_keys = yield [("RDK", [i]) for i in symbols_set]
    for i, s in zip(all_keys, symbols_set):
        # For each object we have found
        keys = yield [("RV", [j]) for j in i]
        values = yield [("RD", [s, j]) for j in keys]
        values = yield [("RV", [j]) for j in values]
        for key, value in zip(keys, values):
            k = key
            v = value
            if v and symbols.get(k, False):
                result, = yield [("CNV", ["ERROR: multiple definition of symbol " + str(key)])]
                raise PrimitiveFinished(result)
            elif v and not symbols.get(k, False):
                symbols[k] = True
            elif not v and k not in symbols:
                symbols[k] = False

    for i, j in symbols.items():
        if i == "input" or i == "output":
            continue
        if not j:
            result, = yield [("CNV", ["ERROR: undefined symbol " + str(i)])]
            raise PrimitiveFinished(result)

    result, = yield [("CNV", ["OK"])]
    raise PrimitiveFinished(result)

def construct_const(**remainder):
    v, = yield [("CNV", [{"value": "constant"}])]

    # Get input: keep trying until we get something
    inp, = yield [("CALL_KWARGS", [jit_runtime.get_input, remainder])]

    yield [("CD", [v, "node", inp])]

    raise PrimitiveFinished(v)

def instantiated_name(a, b, **remainder):
    name_value, = yield [("RV", [b])]
    if name_value == "":
        b, = yield [("CNV", ["__" + str(a)])]
    raise PrimitiveFinished(b)

def set_merge(a, b, **remainder):
    outputs, = yield [("RO", [b])]
    values = yield [("RE", [i]) for i in outputs]
    yield [("CE", [a, i[1]]) for i in values]
    raise PrimitiveFinished(a)

def has_value(a, **remainder):
    v, = yield [("RV", [a])]
    if v is None:
        result, = yield [("CNV", [False])]
    else:
        result, = yield [("CNV", [True])]
    raise PrimitiveFinished(result)

def make_reverse_dictionary(a, **remainder):
    reverse, = yield [("CN", [])]
    key_nodes, = yield [("RDK", [a])]
    values = yield [("RDN", [a, i]) for i in key_nodes]
    yield [("CD", [reverse, str(v), k]) for k, v in zip(key_nodes, values)]
    raise PrimitiveFinished(reverse)

def dict_eq(a, b, **remainder):
    key_nodes, = yield [("RDK", [a])]
    key_values = yield [("RV", [i]) for i in key_nodes]
    values = yield [("RD", [a, i]) for i in key_values]
    values = yield [("RV", [i]) for i in values]
    a_dict = dict(zip(key_values, values))

    key_nodes, = yield [("RDK", [b])]
    key_values = yield [("RV", [i]) for i in key_nodes]
    values = yield [("RD", [b, i]) for i in key_values]
    values = yield [("RV", [i]) for i in values]
    b_dict = dict(zip(key_values, values))

    result, = yield [("CNV", [a_dict == b_dict])]
    raise PrimitiveFinished(result)