# Navrhněte efektivní algoritmus, která zjistí průměrnou výšku stromu definovanou jako `součet délek všech cest z kořene do listu / počet listů`. Délka cesty je počet hran.

from queue import Queue


class Node:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

    def __repr__(self):
        return f"Node({self.value})"


# BFS with saving depth into the nodes
def avg_depth(root):
    if root is None:
        return 0

    q = Queue()
    root.depth = 0
    q.put(root)

    leaf_count = 0
    leaf_distance = 0

    while not q.empty():
        node = q.get()

        if node.left is None and node.right is None:  # leaf
            leaf_count += 1
            leaf_distance += node.depth
        else:
            if node.left is not None:
                node.left.depth = node.depth + 1
                q.put(node.left)
            if node.right is not None:
                node.right.depth = node.depth + 1
                q.put(node.right)
    
    return leaf_distance / leaf_count
    

# BFS with layers
def avg_depth(root):
    if root is None:
        return 0

    q = Queue()
    q.put(root)

    current_layer_count = 1
    next_layer_count = 0
    current_depth = 0

    leaf_count = 0
    leaf_distance = 0

    while not q.empty():
        node = q.get()

        if node.left is None and node.right is None:  # leaf
            leaf_count += 1
            leaf_distance += current_depth
        else:
            if node.left is not None:
                q.put(node.left)
                next_layer_count += 1
            if node.right is not None:
                q.put(node.right)
                next_layer_count += 1
        
        current_layer_count -= 1
        if current_layer_count == 0:
            current_layer_count = next_layer_count
            next_layer_count = 0
            current_depth += 1
    
    return leaf_distance / leaf_count


# recursive solution
def avg_depth(root):
    if root is None:
        return 0

    sum_leaf_distances, leaf_count = avg_depth_recursive(root)
    
    return sum_leaf_distances / leaf_count

def avg_depth_recursive(root):
    if root.left is None and root.right is None:  # leaf
        return 0, 1

    leaf_distances_left = 0
    leaf_count_left = 0
    leaf_distances_right = 0
    leaf_count_right = 0

    if root.left is not None:
        leaf_distances_left, leaf_count_left = avg_depth_recursive(root.left)
    if root.right is not None:
        leaf_distances_right, leaf_count_right = avg_depth_recursive(root.right)

    leaf_distances_left += leaf_count_left  # increase the distance to each leaf by one edge
    leaf_distances_right += leaf_count_right  # increase the distance to each leaf by one edge

    return leaf_distances_left + leaf_distances_right, leaf_count_left + leaf_count_right


# for debugging purposes
def print_tree(root):
    if root is None:
        return

    print("(", end="")
    
    if root.left is not None:
        print_tree(root.left)

    print(root.value, end="")

    if root.right is not None:
        print_tree(root.right)

    print(")", end="")
    


if __name__ == "__main__":
    root = \
        Node(1, 
            Node(2, 
                Node(4,
                    None,
                    Node(8)
                ),
                Node(5,
                    Node(9),
                    None
                )
            ),
            Node(3, 
                Node(6),
                Node(7,
                    None,
                    Node(10, 
                        Node(11),
                        Node(12)
                    )
                )
            )
        )

    print_tree(root)
    print()

    print(avg_depth(root))