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

def reverseKeyLookupMulti(a, b, **remainder):
    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]
    edges = yield [("CE", [result, result]) for value in values]
    yield [("CE", [edge, value[1]]) for edge, value in zip(edges, values)]

    raise PrimitiveFinished(result)

def reverseKeyLookup(a, b, **remainder):
    if a is None or b is None:
        result, = yield [("CNV", [""])]
        raise PrimitiveFinished(result)
    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", [""])]
    raise PrimitiveFinished(result)

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):
    keys, =         yield [("RDK", [b])]
    edges =         yield [("CE", [a, a]) for key in keys]
    _ =             yield [("CE", [edge, key]) for edge, key in zip(edges, keys)]
    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)

def integer_gt(a, b, **remainder):
    a_value, b_value =  yield [("RV", [a]), ("RV", [b])]
    result, = yield [("CNV", [a_value > b_value])]
    raise PrimitiveFinished(result)

def integer_neg(a, **remainder):
    a_value, =          yield [("RV", [a])]
    result, = yield [("CNV", [-a_value])]
    raise PrimitiveFinished(result)

def float_gt(a, b, **remainder):
    a_value, b_value =  yield [("RV", [a]), ("RV", [b])]
    result, = yield [("CNV", [a_value > b_value])]
    raise PrimitiveFinished(result)

def float_neg(a, **remainder):
    a_value, =          yield [("RV", [a])]
    result, = yield [("CNV", [-a_value])]
    raise PrimitiveFinished(result)

def value_neq(a, b, **remainder):
    a_value, b_value =  yield [("RV", [a]), ("RV", [b])]
    result, = yield [("CNV", [a_value != b_value])]
    raise PrimitiveFinished(result)

def element_neq(a, b, **remainder):
    result, = yield [("CNV", [a != b])]
    raise PrimitiveFinished(result)

def list_append(a, b, **remainder):
    a_outgoing, = yield [("RO", [a])]
    _ = yield [("CD", [a, len(a_outgoing), b])]
    raise PrimitiveFinished(a)

def list_read(a, b, **remainder):
    b_value, = yield [("RV", [b])]
    result, = yield [("RD", [a, b_value])]
    if result is None:
        raise Exception("List read out of bounds: %s" % b_value)
    raise PrimitiveFinished(result)

def list_len(a, **remainder):
    outgoings, = yield [("RO", [a])]
    result, = yield [("CNV", [len(outgoings)])]
    raise PrimitiveFinished(result)

def dict_add(a, b, c, **remainder):
    new_edge, = yield [("CE", [a, c])]
    yield [("CE", [new_edge, b])]
    raise PrimitiveFinished(a)

def dict_len(a, **remainder):
    outgoings, = yield [("RO", [a])]
    result, = yield [("CNV", [len(outgoings)])]
    raise PrimitiveFinished(result)

def set_add(a, b, **remainder):
    v, =        yield [("RV", [b])]
    is_in, =    yield [("RD", [a, v])]
    if not is_in:
        _, =    yield [("CD", [a, v, a])]
    raise PrimitiveFinished(a)

def set_add_node(a, b, **remainder):
    is_in, =    yield [("RDN", [a, b])]
    if not is_in:
        edge, = yield [("CE", [a, a])]
        _, =    yield [("CE", [edge, b])]
    raise PrimitiveFinished(a)

def set_pop(a, **remainder):
    outgoing, = yield [("RO", [a])]
    if outgoing:
        outgoing = outgoing[0]
        new_outgoing, = yield [("RO", [outgoing])]
        new_outgoing = new_outgoing[0]
        edge, _ = yield [("RE", [new_outgoing]), ("DE", [outgoing])]
        raise PrimitiveFinished(edge[1])
    else:
        raise Exception("POP from empty set")
        print("Pop from empty set!")
        raise PrimitiveFinished(remainder["root"])

def set_create(**remainder):
    result, = yield [("CN", [])]
    raise PrimitiveFinished(result)

def list_create(**remainder):
    result, = yield [("CN", [])]
    raise PrimitiveFinished(result)

def dict_create(**remainder):
    result, = yield [("CN", [])]
    raise PrimitiveFinished(result)

def create_tuple(a, b, **remainder):
    result, = yield [("CN", [])]
    _, _ =    yield [("CD", [result, 0, a]),
                     ("CD", [result, 1, b]),
                    ]
    raise PrimitiveFinished(result)

def set_overlap(a, b, **remainder):
    a_keys, b_keys, res = yield [("RDK", [a]), ("RDK", [b]), ("CN", [])]
    a_values = yield [("RV", [i]) for i in a_keys]
    b_values = yield [("RV", [i]) for i in b_keys]

    result = set(a_values) & set(b_values)
    yield [("CD", [res, value, res]) for value in result]

    raise PrimitiveFinished(res)

def get_superclasses(a, b, **remainder):
    model, name = a, b
    model_dict, tm_dict, name_value = yield [("RD", [a, "model"]), 
                                             ("RD", [a, "type_mapping"]),
                                             ("RV", [b])]

    worklist = set([name_value])
    found = set([])

    cache_value = {}

    while worklist:
        name = worklist.pop()
        if name in found:
            continue
        elem, = yield [("RD", [model_dict, name])]
        found.add(name)

        # Iterate over all outgoing links
        outgoing, = yield [("RO", [elem])]
        outgoing = set(outgoing)
        while (outgoing):
            link = outgoing.pop()

            # If the link is typed by "Inheritance", we add its destination
            link_name_node, = yield [("CALL_ARGS", [reverseKeyLookup, [model_dict, link]])]
            link_name, = yield [("RV", [link_name_node])]
            t_edge, = yield [("RD", [tm_dict, link_name])]
            t_edge, = yield [("RV", [t_edge])]
            if t_edge == "Inheritance":
                edge, = yield [("RE", [link])]
                src, dst = edge
                # Look up dst's name and add it
                if dst not in cache_value:
                    dst_name, = yield [("CALL_ARGS", [reverseKeyLookup, [model_dict, dst]])]
                    dst_name_value, = yield [("RV", [dst_name])]
                    cache_value[dst] = dst_name_value
                dst_name_value = cache_value[dst]
                worklist.add(dst_name_value)

    result, = yield [("CN", [])]
    yield [("CD", [result, i, result]) for i in found]

    raise PrimitiveFinished(result)

def list_pop_final(a, **remainder):
    lst, = yield [("RO", [a])]
    length = len(lst)
    result, result_edge = yield [("RD", [a, length - 1]),
                                 ("RDE", [a, length -1])]
    _, = yield [("DE", [result_edge])]
    raise PrimitiveFinished(result)

def instantiate_node(a, b, c, **remainder):
    node, dict_entry, typing, name = \
        yield [("CN", []),
               ("RD", [a, "model"]),
               ("RD", [a, "type_mapping"]),
               ("RV", [c]),
              ]

    if name == "":
        name = "__" + str(node)
        name_node, = yield [("CNV", [name])]
    else:
        name_node = c

    yield [("CD", [dict_entry, name, node])]
    yield [("CD", [typing, name, b])]

    raise PrimitiveFinished(name_node)