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)

    """
    edges_out, edges_in, result = yield [("RO", [a]), ("RI", [b]), ("CN", [])]
    python_result = set()
    options = set(edges_out) & set(edges_in)
    if options:
        all_destination_edges = yield [("RO", [i]) for i in options]
        for e in all_destination_edges:
            st, = yield [("RE", [edge]) for edge in e]
            s, t = st
            python_result.add(t)
    yield [("CE", [result, i]) for i in python_result]
    raise PrimitiveFinished(result)
    """

def reverseKeyLookup(a, b, **remainder):
    edges_out, edges_in = yield [("RO", [a]), ("RI", [b])]
    options = set(edges_out) & set(edges_in)
    if options:
        # Select one option randomly
        edge = options.pop()
        out_edges, = yield [("RO", [edge])]
        # Select one option randomly
        out_edge = out_edges.pop()
        e, = yield [("RE", [out_edge])]
        s, t = e
        raise PrimitiveFinished(t)
    else:
        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 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)

def string_substr(a, b, c, **remainder):
    a_val, b_val, c_val = yield [("RV", [a]),
                                 ("RV", [b]),
                                 ("RV", [c])]
    try:
        new_value = a_val[b_val:c_val]
    except:
        new_value = ""
    
    result, = yield [("CNV", [new_value])]
    raise PrimitiveFinished(result)