from vorlesung.L05_binaere_baeume.avl_tree_node import AVLTreeNode
from vorlesung.L05_binaere_baeume.bin_tree import BinaryTree


class AVLTree(BinaryTree):

    def __init__(self):
        super().__init__()

    def new_node(self, value):
        return AVLTreeNode(value)


    def balance(self, node: AVLTreeNode):
        node.update_balance()
        if node.balance == -2:
            if node.left.balance <= 0:
                node = node.right_rotate()
            else:
                node = node.left_right_rotate()
        elif node.balance == 2:
            if node.right.balance >= 0:
                node = node.left_rotate()
            else:
                node = node.right_left_rotate()
        if node.parent:
            self.balance(node.parent)
        else:
            self.root = node
        # self.check_circle(self.root)

    def insert(self, value):
        node, parent = super().insert(value)
        node.parent = parent
        if parent:
            self.balance(parent)
        return node, parent

    def delete(self, value):
        node, parent = super().delete(value)
        if node:
            node.parent = parent
        if parent:
            self.balance(parent)

    def check_circle(self, node: AVLTreeNode, check_nodes = []):
        if node is None:
            return
        if id(node) in check_nodes:
            print("Circle detected")
            return
        self.check_circle(node.left, check_nodes + [id(node)])
        self.check_circle(node.right, check_nodes + [id(node)])


if __name__ == "__main__":

    def print_node(node, indent=0, level=0):
        print((indent * 3) * " ", node.value)

    tree = AVLTree()
    values = [5, 3, 7, 2, 4, 6, 5, 8]

    for value in values:
        tree.insert(value)


    print("In-order traversal:")
    tree.in_order_traversal(print_node)
    print("\nLevel-order traversal:")
    tree.level_order_traversal(print_node)
    print("\nTree structure traversal:")
    tree.tree_structure_traversal(print_node)
    print("\nGraph traversal:")
    tree.graph_traversal()

    tree.insert(9)
    tree.graph_traversal()

    print("\nDeleting 5:")
    tree.delete(5)

    print("In-order traversal after deletion:")
    tree.in_order_traversal(print_node)
    print("\nLevel-order traversal after deletion:")
    tree.level_order_traversal(print_node)
    print("\nTree structure traversal after deletion:")
    tree.tree_structure_traversal(print_node)