from SoSe24.algodat.foundation import AlgoDatArray, AlgoDatValue, read_int_sequence
from SoSe24.lec04_trees.bin_tree import BinTree, BinTreeNode
from time import perf_counter as pfc

class AVLTreeNode(BinTreeNode):
    def __init__(self, value: AlgoDatValue):
        super().__init__(value)
        self.parent = None
        self.balance = 0

    def update_balance(self):
        left_height = self.left.height() if self.left else 0
        right_height = self.right.height() if self.right else 0
        self.balance = right_height - left_height

    def right_rotate(self) -> BinTreeNode:
        new_root = self.left
        new_root.parent = self.parent
        self.left = new_root.right
        if self.left:
            self.left.parent = self
        new_root.right = self
        self.parent = new_root
        if new_root.parent:
            if new_root.parent.left == self:
                new_root.parent.left = new_root
            else:
                new_root.parent.right = new_root
        self.update_balance()
        new_root.update_balance()
        return new_root

    def left_rotate(self) -> BinTreeNode:
        new_root = self.right
        new_root.parent = self.parent
        self.right = new_root.left
        if self.right:
            self.right.parent = self
        new_root.left = self
        self.parent = new_root
        if new_root.parent:
            if new_root.parent.left == self:
                new_root.parent.left = new_root
            else:
                new_root.parent.right = new_root
        self.update_balance()
        new_root.update_balance()
        return new_root

    def right_left_rotate(self) -> BinTreeNode:
        self.right = self.right.right_rotate()
        return self.left_rotate()

    def left_right_rotate(self) -> BinTreeNode:
        self.left = self.left.left_rotate()
        return self.right_rotate()

class AVLTree(BinTree):

    def new_node(self, value: AlgoDatValue):
        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: AlgoDatValue):
        node, parent = super().insert(value)
        node.parent = parent
        if parent:
            self.balance(parent)
        return node, parent

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



if __name__ == "__main__":
    z = read_int_sequence("../../seq0.txt")
    print(z, len(z))
    start = pfc()
    tree = AVLTree()
    for i in z:
        tree.insert(i)
    tree.walk()
    tree.tree_walk()
    tree.levelwalk()
    tree.graph_walk()
    tree.delete(AlgoDatValue(46))
    tree.delete(AlgoDatValue(48))
    #tree.graph_walk()
    print(f"Dauer: {pfc() - start:.4f}s")
    AlgoDatValue.summary()