# -*- coding: utf-8 -*-
"""
Created on Fri Apr 10 11:40:16 2020


Balance_factor = h(left_child) - h(right_child)
h(node) = 1 + max(h(node.leftchild),h(node.right_child))

add left_child, currNode.bf +=1
add right_child, currNode.bf -=1

@author: Durant
"""

import turtle
t = turtle.Turtle()
t.penup()
t.speed('fast')#'fastest' 'fast' 'normal' 'slow' 'slowest'
t.shape('turtle')#arrow", "turtle", "circle", "square", "triangle", "classic"
class drawTree():
    def __init__(self,tree):
        self.radius = 32
        self.level = 0
        self.branch_length = 4*self.radius
        self.node_memo = {}
        self.draw_binary_tree(tree,(0,300),self.level)
        
    def draw_node(self,val,payload,balance_factor,start_pos):
        # draw the node
        (x,y) = start_pos
        length = len(str(val)+str(payload))+1
        t.setx(x);t.sety(y)
        t.pendown()
        t.setheading(0)
        t.fillcolor('#FFF0F5')
        t.begin_fill()
        t.circle(self.radius)
        t.end_fill()
        t.circle(self.radius)
        t.penup();t.goto(x-length/2*8,y+self.radius-8);t.pendown()
        t.write(str(val)+':'+str(payload), font=("Arial", 8, "normal"))
        t.penup();t.goto(x,y+self.radius-20);t.pendown()
        t.write(str(balance_factor), font=("Arial", 8, "normal"))
        
        t.penup()
        t.setx(x)
        t.sety(y)

    def draw_branch(self,start_pos,orient,level):
        #draw the branch
        (x,y) = start_pos
        t.setx(x);t.sety(y)
        t.pendown()
        t.setheading(180)
        left_angle=15+10*level# This should be edited in case you get badplot
        right_angle = 180-left_angle
        if (right_angle - left_angle)<=30:
            left_angle = 75
            right_angle = 105
        if orient == 'left':
            t.left(left_angle)
            t.forward(self.branch_length-level*30)
        elif orient == 'right':
            t.left(right_angle)
            t.forward(self.branch_length-level*30)
        t.penup()
        return t.pos()


    
    def draw_binary_tree(self,p,start_pos,level):
        (x,y) = start_pos
        self.node_memo[p] = (x,y+self.radius)
        if p!=None:
            level+=1
            print("level:",level)
            self.draw_node(p.val,p.payload,p.balance_factor,(x,y))
            if p.left_child!=None:
                (x1,y1) = self.draw_branch((x,y),'left',level)
                self.draw_binary_tree(p.left_child,(x1,y1-2*self.radius),level)
            
            if p.right_child!=None:
                (x2,y2)=self.draw_branch((x,y),'right',level)
                self.draw_binary_tree(p.right_child,(x2,y2-2*self.radius),level)
    
    def get_node_coord(self,node)->tuple:
        #return the coord of the node circle centre
        return self.node_memo[node]
    
    def draw_line(self,coord1, coord2,pensize=5,pen_color='red'): #from 1 to 2
        (x,y) = coord1  
        t.penup()
        t.pensize(pensize)
        t.pencolor(pen_color)
        t.setx(x);t.sety(y)
        t.pendown()
        t.goto(coord2[0],coord2[1])
        
class AVLNode:
    def __init__(self,val,payload,balance_factor=0,left=None,right=None,parent=None):
        self.val = val
        self.payload = payload
        self.balance_factor = balance_factor
        self.left_child = left
        self.right_child = right
        self.parent = parent
    
    def isLeftChild(self):
        return self.parent and self.parent.left_child==self
    
    def isRightChild(self):
        return self.parent and self.parent.right_child==self
    
    def isRoot(self):
        return not self.parent
    
    def isLeaf(self):
        return not (self.left_child!=None or self.right_child!=None)
    
    def hasAtLeastOneChild(self):
        return self.left_child!=None or self.right_child!=None
    
    def hasBothChildren(self):
        return self.left_child!=None and self.right_child!=None
    
    def destroy_leaf(self):
        if not self.parent: self = None 
        elif self.parent.left_child == self:
            self.parent.left_child = None
        else:
            self.parent.right_child = None
            
    def __iter__(self):
        #implement the iteration mechanism utilising inorder transverse
        if self:
            if self.left_child!=None:
                for item in self.left_child:
                    yield item
            yield self
            if self.right_child!=None:
                for item in self.right_child:
                    yield item

    def setOneNode(self,val,payload,lc,rc):
        self.val = val
        self.payload = payload
        self.right_child = rc
        self.left_child = lc
        if self.left_child!=None:
            self.left_child.parent = self
        if self.right_child!=None:
            self.right_child.parent = self
        
        
class AVLTree:
    def __init__(self,IsRebalance=False):
        self.root = None
        self.size = 0
        self.isrebalance = IsRebalance
    
    def __len__(self):
        return self.size
    
    def __iter__(self):
        #enpower the usage like 'for i in BinarySearchTree()'
        return self.root.__iter__()
    
    def initBST(self,data:dict):
        for key in data.keys():
            self[key] = data[key]
    def destroy_tree(self):
        self.root = None
        
    def put (self, val, payload):
        if self.root:
            self._put(val, payload,self.root)
        else:
            self.root = AVLNode(val,payload)
        self.size += 1
#_________________________________ADD_CORECODE_________________________________

    def renew_balance_factor(self,currNode:AVLNode):
        if currNode.balance_factor>1 or currNode.balance_factor<-1:
            if self.isrebalance: self.rebalance(currNode)
            return 
        if currNode.isLeftChild(): 
            currNode.parent.balance_factor+=1
        elif currNode.isRightChild():
            currNode.parent.balance_factor-=1
        if currNode.parent!=None and currNode.parent.balance_factor != 0 :         
            self.renew_balance_factor(currNode.parent)
            
    def LR_doubRot(self,currNode:AVLNode):
        '''
           16        16
          /          /
         3    ->    7    ->   7
          \        /         / \
           7      3         3   16
        '''
        Rotleftchild = currNode.left_child
        p = Rotleftchild.right_child
        currNode.left_child = p
        m = p.left_child
        p.left_child = Rotleftchild
        Rotleftchild.right_child = m
        p.parent = currNode
        Rotleftchild.parent = p
        if m: m.parent = Rotleftchild
        
        k = p.right_child
        p.right_child = currNode
        currNode.left_child = k
        if k: k.parent = currNode
        if not currNode.parent:
            self.root = p
            p.parent = None
        else:
            if currNode.isLeftChild():
                currNode.parent.left_child = p
            elif currNode.isRightChild():
                currNode.parent.right_child = p
            p.parent = currNode.parent
        currNode.parent = p
        currNode.balance_factor = self.get_node_bf(currNode)
        Rotleftchild.balance_factor = self.get_node_bf(Rotleftchild) 
        p.balance_factor = self.get_node_bf(p) 
        
    def RL_doubRot(self,currNode:AVLNode):
        '''
       16       16   
        \         \  
         26  ->    18    -> 18
        /           \       / \
       18            26    16  26   
        '''
        #step1 : commit a line
        RotRightchild = currNode.right_child
        p = RotRightchild.left_child
        currNode.right_child = p
        m = p.right_child
        p.right_child = RotRightchild
        RotRightchild.left_child = m
        p.parent = currNode
        RotRightchild.parent = p
        if m: m.parent = RotRightchild
        # step2: L Rotate
        k = p.left_child
        p.left_child = currNode
        currNode.right_child = k
        if k: k.parent = currNode
        if not currNode.parent:
            self.root = p
            p.parent = None
        else:
            if currNode.isLeftChild():
                currNode.parent.left_child = p
            elif currNode.isRightChild():
                currNode.parent.right_child = p
            p.parent = currNode.parent
        currNode.parent = p
        currNode.balance_factor = self.get_node_bf(currNode)
        RotRightchild.balance_factor = self.get_node_bf(RotRightchild) 
        p.balance_factor = self.get_node_bf(p)
        
    def R_Rot(self,currNode:AVLNode):
        '''
             16
            /
           11   ->   11
          /         /  \
         9         9   16
        '''
        Rotleftchild = currNode.left_child
        p = Rotleftchild.right_child
        if not currNode.parent:
            self.root = Rotleftchild
        elif currNode.isLeftChild():
            currNode.parent.left_child = Rotleftchild
        elif currNode.isRightChild():
            currNode.parent.right_child = Rotleftchild
        Rotleftchild.right_child = currNode
        Rotleftchild.parent = currNode.parent
        currNode.parent = Rotleftchild
        currNode.left_child = p
        if p!=None: p.parent = currNode
        #update balance factor
        currNode.balance_factor = currNode.balance_factor -1 +\
                                    min(-Rotleftchild.balance_factor,0)
                                    
        Rotleftchild.balance_factor = Rotleftchild.balance_factor - 1 +\
                                    min(currNode.balance_factor,0)
    
    def L_Rot(self,currNode:AVLNode):
        '''
          11      
           \
            16   ->     16
             \         /  \
              26      11  26
        pay attention to case in which RotRight_child has left_child
        '''
        RotRightchild = currNode.right_child
        p = RotRightchild.left_child
        if not currNode.parent:
            self.root = RotRightchild
        elif currNode.isLeftChild():
            currNode.parent.left_child = RotRightchild
        elif currNode.isRightChild():
            currNode.parent.right_child = RotRightchild
        RotRightchild.left_child = currNode
        RotRightchild.parent = currNode.parent
        currNode.parent = RotRightchild
        currNode.right_child = p
        if p!=None: p.parent = currNode
        #updating balance factor
        currNode.balance_factor = currNode.balance_factor + 1 -\
                                min(0,RotRightchild.balance_factor)
        
        RotRightchild.balance_factor = RotRightchild.balance_factor + 1 +\
                                max(currNode.balance_factor,0)
    

        
    def rebalance(self,currNode:AVLNode):
        if currNode.balance_factor>1:# left-heavy sub tree
            if not currNode.right_child and currNode.left_child.right_child:
                self.LR_doubRot(currNode)
                return 
            else:
                self.R_Rot(currNode)
        elif currNode.balance_factor<-1:# right-heavy sub tree
            if not currNode.left_child and currNode.right_child.left_child:
                self.RL_doubRot(currNode)
            else:
                self.L_Rot(currNode)
            
        
    def _put(self,val,payload,currentNode):
        
        if val < currentNode.val:
            if currentNode.left_child==None:
                left_new = AVLNode(val,payload,parent=currentNode)
                currentNode.left_child = left_new
                
                self.renew_balance_factor(left_new)
            else:
                self._put(val,payload,currentNode.left_child)
                
        elif val > currentNode.val:
            if currentNode.right_child==None:
                right_new = AVLNode(val,payload,parent=currentNode)
                currentNode.right_child = right_new
                
                self.renew_balance_factor(right_new)
            else:
                self._put(val,payload,currentNode.right_child)
        else:
            currentNode.payload = payload
#______________________________________________________________________________
    def __setitem__(self,k,v):
        self.put(k,v)
    
    def get(self,val):
        if self.root:
            return self._get(val,self.root)
        else:
            return None
    
    def _get(self,val,currentNode):
        if not currentNode:
            return None
        elif currentNode.val == val:
            return currentNode.payload
        elif currentNode.val>val:
            return self._get(val,currentNode.left_child)
        else:
            return self._get(val,currentNode.right_child)
        
    def getnode(self,val,currentNode):
        if not currentNode:
            return None
        elif currentNode.val == val:
            return currentNode
        elif currentNode.val>val:
            return self.getnode(val,currentNode.left_child)
        else:
            return self.getnode(val,currentNode.right_child)
        
    def __getitem__(self,k):
        return self.get(k)
    
    def __contains__(self,k):
        if self.get(k):
            return True
        else:
            return False
    def printTree(self):
        for node in self:
            print(node.val,',',node.payload,',', self.get_node_bf(node))
        print('')
#___________________________DEL CORE CODE______________________________________
        def get_node_bf(self,currNode:AVLNode)->int:
            # obtain current node's bf
            max_l = max_r =0
            path_level_l = []
            path_level_r = []
            def get_level(node,depth,path_level):
                if not node: return 0
                if node.isLeaf(): path_level.append(depth);return  
                depth+=1
                if node.left_child:
                    get_level(node.left_child,depth,path_level)
                if node.right_child:
                    get_level(node.right_child,depth,path_level)
            
            if currNode.isLeaf():
                return (max_l-max_r)
            else:
                get_level(currNode.left_child,1,path_level_l)
                get_level(currNode.right_child,1,path_level_r)
                if not path_level_l: max_l = 0
                else:max_l = max(path_level_l)
                if not path_level_r: max_r = 0
                else: max_r = max(path_level_r)
                return (max_l - max_r)
        
    def find_successor(self,currNode:AVLNode):
        if currNode.left_child:
#            print(currNode.left_child.val)
            return self.find_successor(currNode.left_child)
        else:
            return currNode
    
    def update_del_bf(self,currNode:AVLNode):
        if not currNode:return
        currNode.balance_factor = self.get_node_bf(currNode)
        if currNode.balance_factor>1 or currNode.balance_factor<-1:
            if self.isrebalance: self.rebalance(currNode)
            self.update_del_bf(currNode.parent)
            return 
        
        if currNode.parent!=None:#update all parent bf         
            self.update_del_bf(currNode.parent)
            
    def __delitem__(self,val):
        if not self.root: return None
        target_node = self.getnode(val,self.root)
        if not target_node: print("Not Found!");return None
        
        if target_node.isLeaf():
            if target_node ==self.root: self.destroy_tree()
            else:target_node.destroy_leaf() # Do not silly directly write: target_node = None
            self.update_del_bf(target_node.parent)
        elif target_node.left_child and not target_node.right_child:
            if not target_node.parent:
                self.root = target_node.left_child
                target_node.left_child.parent = None

            else:
                target_node.left_child.parent = target_node.parent
                if target_node.isLeftChild(): target_node.parent.left_child = target_node.left_child
                elif target_node.isRightChild(): target_node.parent.right_child = target_node.left_child

            self.update_del_bf(target_node.parent)
            target_node = None
        elif target_node.right_child and not target_node.left_child:
            if not target_node.parent:
                self.root = target_node.right_child
                target_node.right_child.parent = None
                
            else:
                target_node.right_child.parent = target_node.parent
                if target_node.isRightChild():target_node.parent.right_child = target_node.right_child
                elif target_node.isLeftChild():target_node.parent.left_child = target_node.right_child
                
            self.update_del_bf(target_node.parent)
            target_node = None
        else:
            successor = self.find_successor(target_node.right_child)
            target_node.val = successor.val
            target_node.payload = successor.payload
            if successor.isLeaf():
                successor.destroy_leaf()
            else:
                successor.right_child.parent = successor.parent
                successor.parent.left_child = successor.right_child
            self.update_del_bf(successor.parent)
            successor= None
#______________________________________________________________________________

#data = {16:'A',3:'B',7:'C',11:'D',9:'E',26:'F',18:'G',14:'H',15:'I'}
data = {120:'K',150:'A',260:'B',180:'C',160:'D',170:'E',130:'F',140:'G',155:'H',142:\
        'I',157:'J'}
#data = {15:'A',10:'B',20:'C',17:'D',26:'E'}
bst = AVLTree(IsRebalance=True)
bst.initBST(data)
#bst.printTree()
#print(bst.get_node_bf(bst.getnode(160,bst.root)))
#bst[65]='Edison'
9#bst[65]='Fisher'
#bst.printTree()
#print(bst[29])
#print(23 in bst)
#del bst[150]
#del bst[140]
#del bst[155]
#del bst[157]
#del bst[160]
#del bst[130]
t = drawTree(bst.root)
turtle.done()