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

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

    # Keep results in a local Python set, as we want to bundle as many requests as possible
    todo = set()
    for i, edge in enumerate(values):
        if b_val == edge:
            todo.add(i)

    outgoings = yield [("RO", [edges[i]]) for i in todo]
    values = yield [("RE", [outgoing[0]]) for outgoing in outgoings]
    yield [("CE", [result, value[1]]) for value in values]

    #PROFILE print("[COMPILED]reverseKeyLookupMulti : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(result)

def reverseKeyLookup(a, b, **remainder):
    #PROFILE start = time.time()
    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])]
        result = e[1]
    else:
        result, = yield [("CNV", ["(unknown: %s)" % b])]
    #PROFILE print("[COMPILED]reverseKeyLookup : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(result)

def set_copy(a, **remainder):
    #PROFILE start = time.time()
    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]
    #PROFILE print("[COMPILED]set_copy : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(b)

def instantiated_name(a, b, **remainder):
    #PROFILE start = time.time()
    name_value, = yield [("RV", [b])]
    if name_value == "":
        b, = yield [("CNV", ["__" + str(a)])]
    #PROFILE print("[COMPILED]instantiated_name : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(b)

def set_merge(a, b, **remainder):
    #PROFILE start = time.time()
    outputs, = yield [("RO", [b])]
    values = yield [("RE", [i]) for i in outputs]
    yield [("CE", [a, i[1]]) for i in values]
    #PROFILE print("[COMPILED]set_merge : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(a)

def has_value(a, **remainder):
    #PROFILE start = time.time()
    v, = yield [("RV", [a])]
    if v is None:
        result, = yield [("CNV", [False])]
    else:
        result, = yield [("CNV", [True])]
    #PROFILE print("[COMPILED]has_value : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(result)

def make_reverse_dictionary(a, **remainder):
    #PROFILE start = time.time()
    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)]
    #PROFILE print("[COMPILED]make_reverse_dictionary : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(reverse)

def dict_eq(a, b, **remainder):
    #PROFILE start = time.time()
    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])]
    #PROFILE print("[COMPILED]dict_eq : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(result)

def string_substr(a, b, c, **remainder):
    #PROFILE start = time.time()
    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])]
    #PROFILE print("[COMPILED]string_substr : %s : %s" % (time.time() - start, time.time() - start))
    raise PrimitiveFinished(result)