from treelib import Node, Tree
import random
import time
import os
import multiprocessing


letter_morse = {'A': '.-', 'B': '-...',
                'C': '-.-.', 'D': '-..', 'E': '.',
                'F': '..-.', 'G': '--.', 'H': '....',
                'I': '..', 'J': '.---', 'K': '-.-',
                'L': '.-..', 'M': '--', 'N': '-.',
                'O': '---', 'P': '.--.', 'Q': '--.-',
                'R': '.-.', 'S': '...', 'T': '-',
                'U': '..-', 'V': '...-', 'W': '.--',
                'X': '-..-', 'Y': '-.--', 'Z': '--..'}
morse_letter = dict()
for k, v in letter_morse.items():
    morse_letter[v] = k


def build_tree(tree, morse_left, parent_nid, original, nid):
    for k in tree.keys():
        print(k, end = " ")
    print()
    # We only want to keep doing this if we have too
    #
    # Added optimization. REALLY only add a new leaf if we absolutely have to
    #
    # current_leaf = tree.get_node(parent_nid)
    done = False
    while not done:
        try:
            current_leaf = tree[str(parent_nid)]
            done = True
        except:
            print("Stuck, can't find {}".format(parent_nid))
            time.sleep(.5)
    current_branch_length = 1
    while not current_leaf.is_root():
        current_leaf = tree[str(current_leaf.bpointer)]
        current_branch_length += 1
    #
    #  Added the end to the if statement
    #
    # This is a CRITICAL piece of this code.
    # If there is nothing left to evaluate, then dont.
    # but also, if the current branch is already big enough, only evaluate ONE more.
    if morse_left != "" and current_branch_length < len(original):
        # loop through the morse string that's left, matching as we go
        for i in range(len(morse_left) + 1):
            # for every letter in the alphabet as morse, see if it matches
            for j in morse_letter.keys():
                if j == morse_left[:i]:
                    # tree.create_node(morse_letter[j], nid.value, parent=parent_nid)
                    temp = Node(morse_letter[j], nid.value)
                    temp.bpointer = parent_nid
                    tree[str(nid.value)] = temp
                    parent = nid.value
                    nid.value += 1
                    left_to_match = morse_left[i:]
                    print("Created {}, nid:{} parent:{}".format(morse_letter[j], nid.value, parent_nid))
                    build_tree(tree, left_to_match, parent, original, nid)

def build_roots(morse, original, nid):
    # nid is the global node identifier. It must be unique
    # trees list to return
    manager = multiprocessing.Manager()
    trees = manager.list()
    # loop through the morse string, matching as we go
    processes = []
    for i in range(len(morse)):
        # for every letter in the alphabet as morse, see if it matches
        for l in morse_letter.keys():
            # If it does match, make a tree and add it to the list of possible roots
            if l == morse[:i]:
                left_to_match = morse[i:]
                # tree = Tree()
                # tree.create_node(morse_letter[l], nid.value)
                tree = manager.dict()
                tree[str(nid.value)] = Node(morse_letter[l], nid.value)
                print("Created {}, nid:{}".format(morse_letter[l], nid.value))

                parent = nid.value
                trees.append(tree)
                nid.value += 1
                p = multiprocessing.Process(target=build_tree,
                                            args=(tree, left_to_match, parent, original, nid))
                processes.append(p)
    cpus = multiprocessing.cpu_count()
    running = []
    i = 1
    while len(processes) > 0:
        while len(running) < cpus:
            if len(processes) > 0:
                current_p = processes.pop(0)
                running.append(current_p)
                print("appending")
                i += 1
            if len(processes) == 0:
                break
        for r in running:
            print("starting")
            r.start()
        time.sleep(2)
        for r in running:
            r.join()
            print("Removing")
            running.remove(r)
    return trees


def build_test_string(n=10):
    return "".join([chr(97 + random.randint(0, 25)) for e in range(n)])


def get_answer(list_of_trees, original):
    # Make a list for all possibilities
    all_possibilities = []

    # We need to loop through all of the roots. The roots will be saved as trees.
    # Its not fast, but it's very complete.
    for tree in list_of_trees:
        # tree.show()
        # Paths to leaves is perfect, allows us to get all paths from root to each leaf
        tree_path_lists = tree.paths_to_leaves()
        for path_list in tree_path_lists:

            if len(path_list) != len(original):
                #
                # Added optimization, don't words we dont have to.
                #
                continue
            for i in range(len(path_list)):
                # we need to get the actual tag of the node, that's the letter.
                path_list[i] = tree.get_node(path_list[i]).tag
            # Add what we just made to the list of possibilities
            possible = "".join(path_list)
            if len(possible) == len(original):
                all_possibilities.append(possible)
    # p = 1
    # for word in all_possibilities:
    #     print(word, end=" ")
    #     if p % 5 == 0:
    #         print()
    #     p += 1
    # print()
    return len(all_possibilities)


def get_morse(s):
    """
    get the morse code
    :param s: string of letters
    :return: morse code
    """
    global letter_morse
    to_ret = ""
    for c in s:
        to_ret += letter_morse[c.upper()]
    return to_ret


if __name__ == "__main__":
    start = time.time()
    nid = multiprocessing.Value('i', 0)

    # input = build_test_string(n=10)
    input = "eta"
    morse = get_morse(input)
    # print("\n"*5)
    print(input, morse, len(morse))
    actual_trees = build_roots(morse, input, nid)
    # print(actual_trees)
    trees = []
    for tree in actual_trees:
        t = Tree()
        for k,v in tree.items():
            tag = v.tag
            nid = v.identifier
            parent = v.bpointer
            # print("Received {}, nid:{} parent:{}".format(tag,nid,parent))
            if parent is None:
                t.create_node(tag,nid)
            else:
                t.create_node(tag,nid,parent=parent)
        trees.append(t)
    print(input, get_answer(trees, input))
    print(time.time()-start)