from structures.StsToken import StsToken


def count_node_children_and_live_nodes(node):
    nodes_count = 1
    live_nodes_count = 0
    if not node.expired:
        live_nodes_count += 1
    for child in node.children:
        node_result = count_node_children_and_live_nodes(child)
        nodes_count += node_result[0]
        live_nodes_count += node_result[1]
    return nodes_count, live_nodes_count


def print_tree(current_node, indent="", last='updown'):
    # This function taken from https://stackoverflow.com/users/4768304/clemtoy
    nb_children = lambda node: sum(nb_children(child) for child in node.children) + 1
    size_branch = {child: nb_children(child) for child in current_node.children}

    """ Creation of balanced lists for "up" branch and "down" branch. """
    up = sorted(current_node.children, key=lambda node: nb_children(node))
    down = []
    while up and sum(size_branch[node] for node in down) < sum(size_branch[node] for node in up):
        down.append(up.pop())

    """ Printing of "up" branch. """
    for child in up:
        next_last = 'up' if up.index(child) is 0 else ''
        next_indent = '{0}{1}{2}'.format(indent, ' ' if 'up' in last else '│', " " * len(current_node.token))
        print_tree(child, indent=next_indent, last=next_last)

    """ Printing of current node. """
    if last == 'up': start_shape = '┌'
    elif last == 'down': start_shape = '└'
    elif last == 'updown': start_shape = ' '
    else: start_shape = '├'

    if up: end_shape = '┤'
    elif down: end_shape = '┐'
    else: end_shape = ''

    print('{0}{1}{2}{3}'.format(indent, start_shape, current_node.token, end_shape))

    """ Printing of "down" branch. """
    for child in down:
        next_last = 'down' if down.index(child) is len(down) - 1 else ''
        next_indent = '{0}{1}{2}'.format(indent, ' ' if 'down' in last else '│', " " * len(current_node.token))
        print_tree(child, indent=next_indent, last=next_last)

def print_root_access_key_sts_tree(root_tokens):
    for root_token in root_tokens:
        print("Print the tokens generated by the root token {}:".format(root_token))
        temp_root_sts_token = StsToken()
        temp_root_sts_token.token = root_token
        for child_token in root_tokens[root_token]:
            temp_root_sts_token.children.append(child_token)
        print_tree(temp_root_sts_token)
from structures.StsToken import StsToken


def count_node_children_and_live_nodes(node):
    nodes_count = 1
    live_nodes_count = 0
    if not node.expired:
        live_nodes_count += 1
    for child in node.children:
        node_result = count_node_children_and_live_nodes(child)
        nodes_count += node_result[0]
        live_nodes_count += node_result[1]
    return nodes_count, live_nodes_count


def print_tree(current_node, indent="", last='updown'):
    # This function taken from https://stackoverflow.com/users/4768304/clemtoy
    nb_children = lambda node: sum(nb_children(child) for child in node.children) + 1
    size_branch = {child: nb_children(child) for child in current_node.children}

    """ Creation of balanced lists for "up" branch and "down" branch. """
    up = sorted(current_node.children, key=lambda node: nb_children(node))
    down = []
    while up and sum(size_branch[node] for node in down) < sum(size_branch[node] for node in up):
        down.append(up.pop())

    """ Printing of "up" branch. """
    for child in up:
        next_last = 'up' if up.index(child) is 0 else ''
        next_indent = '{0}{1}{2}'.format(indent, ' ' if 'up' in last else '│', " " * len(current_node.token))
        print_tree(child, indent=next_indent, last=next_last)

    """ Printing of current node. """
    if last == 'up': start_shape = '┌'
    elif last == 'down': start_shape = '└'
    elif last == 'updown': start_shape = ' '
    else: start_shape = '├'

    if up: end_shape = '┤'
    elif down: end_shape = '┐'
    else: end_shape = ''

    print('{0}{1}{2}{3}'.format(indent, start_shape, current_node.token, end_shape))

    """ Printing of "down" branch. """
    for child in down:
        next_last = 'down' if down.index(child) is len(down) - 1 else ''
        next_indent = '{0}{1}{2}'.format(indent, ' ' if 'down' in last else '│', " " * len(current_node.token))
        print_tree(child, indent=next_indent, last=next_last)

def print_root_access_key_sts_tree(root_tokens):
    for root_token in root_tokens:
        print("Print the tokens generated by the root token {}:".format(root_token))
        temp_root_sts_token = StsToken()
        temp_root_sts_token.token = root_token
        for child_token in root_tokens[root_token]:
            temp_root_sts_token.children.append(child_token)
        print_tree(temp_root_sts_token)
# This function taken from https://stackoverflow.com/users/4768304/clemtoy Creation of balanced lists for "up" branch and "down" branch. Printing of "up" branch. Printing of current node. Printing of "down" branch.