from utils.memory_array import MemoryArray
from vorlesung.L05_binaere_baeume.avl_tree_node import AVLTreeNode
from vorlesung.L05_binaere_baeume.bin_tree import BinaryTree
import logging

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

    def insert(self, value):
        insert_generator = self.insert_stepwise(value)
        node, parent = None, None
        while True:
            try:
                node, parent = next(insert_generator)
            except StopIteration:
                break
        return node, parent

    def insert_stepwise(self, value):
        node, parent = super().insert(value)
        yield None, None
        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 graph_filename(self):
        return "AVLTree"

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]
    values = MemoryArray.create_array_from_file("data/seq2.txt")

    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)