""" Red Black Tree implementation """
import logging


logger = logging.getLogger()


class Color:
    """ Node color """
    BLACK = 0
    RED = 1


class Node:
    """ Node """
    def __init__(self, val, item=None):
        self.val = val
        self.parent = None
        self.left = None
        self.right = None
        self.color = Color.RED
        self.items = [] if item is None else [item]

    def __repr__(self):
        """ repr """
        return f"<Node::{self.val}({len(self.items)})>"

    def get_node(self, val):
        """ get node """
        if self.val == val:
            return self
        elif val < self.val and self.left is not None:
            return self.left.get_node(val)
        elif val > self.val and self.right is not None:
            return self.right.get_node(val)

    def get_all(self):
        """ get all """
        nodes = []
        if self.left is not None:
            nodes.extend(self.left.get_all())
        if self.val != -1:
            nodes.append(self)
        if self.right is not None:
            nodes.extend(self.right.get_all())
        return nodes


class RBTree:
    """ Red Black Tree """
    def __init__(self):
        self.NULL = Node(-1)
        self.NULL.color = Color.BLACK
        self.NULL.left = None
        self.NULL.right = None
        self.root = self.NULL

    def insert_node(self, key, item):
        """ Insert node """
        node = Node(key, item)
        node.parent = None
        node.val = key
        node.left = self.NULL
        node.right = self.NULL
        node.color = Color.RED

        y = None
        x = self.root

        while x != self.NULL:  # Find position for new node
            y = x
            if node.val < x.val:
                x = x.left
            elif node.val > x.val:
                x = x.right
            elif node.val == x.val:
                x.items.append(item)
                return

        node.parent = y  # Set parent of Node as y
        if y is None:  # If parent i.e, is none then it is root node
            self.root = node
        elif node.val < y.val:  # Check if it is right Node or Left Node
                                # by checking the value
            y.left = node
        else:
            y.right = node

        if node.parent is None:  # Root node is always Black
            node.color = 0
            return

        if node.parent.parent is None:  # If parent of node is Root Node
            return

        self.fix_insert(node)

    def minimum(self, node):
        """ Find the min value """
        while node.left != self.NULL:
            node = node.left
        return node

    def rotate_left(self, x):
        """ Rotate left """
        y = x.right  # Y = Right child of x
        x.right = y.left  # Change right child of x to left child of y
        if y.left != self.NULL:
            y.left.parent = x

        y.parent = x.parent  # Change parent of y as parent of x
        if x.parent is None:  # If parent of x == None ie. root node
            self.root = y  # Set y as root
        elif x is x.parent.left:
            x.parent.left = y
        else:
            x.parent.right = y
        y.left = x
        x.parent = y

    def rotate_right(self, x):
        """ Rotate right """
        y = x.left  # Y = Left child of x
        x.left = y.right  # Change left child of x to right child of y
        if y.right != self.NULL:
            y.right.parent = x

        y.parent = x.parent  # Change parent of y as parent of x
        if x.parent is None:  # If x is root node
            self.root = y  # Set y as root
        elif x == x.parent.right:
            x.parent.right = y
        else:
            x.parent.left = y
        y.right = x
        x.parent = y

    def fix_insert(self, k):
        """ Fix insert """
        while k.parent.color == Color.RED:  # While parent is red
            if k.parent == k.parent.parent.right:  # if parent is right child of its parent
                u = k.parent.parent.left  # Left child of grandparent
                if u.color == Color.RED:  # if color of left child of grandparent i.e, uncle node is red
                    u.color = Color.BLACK  # Set both children of grandparent node as black
                    k.parent.color = Color.BLACK
                    k.parent.parent.color = Color.RED  # Set grandparent node as Red
                    k = k.parent.parent  # Repeat the algo with Parent node to check conflicts
                else:
                    if k == k.parent.left:  # If k is left child of it's parent
                        k = k.parent
                        self.rotate_right(k)  # Call for right rotation
                    k.parent.color = Color.BLACK
                    k.parent.parent.color = Color.RED
                    self.rotate_left(k.parent.parent)
            else:  # if parent is left child of its parent
                u = k.parent.parent.right  # Right child of grandparent
                if u.color == Color.RED:  # if color of right child of grandparent i.e, uncle node is red
                    u.color = Color.BLACK  # Set color of childs as black
                    k.parent.color = Color.BLACK
                    k.parent.parent.color = Color.RED  # set color of grandparent as Red
                    k = k.parent.parent  # Repeat algo on grandparent to remove conflicts
                else:
                    if k == k.parent.right:  # if k is right child of its parent
                        k = k.parent
                        self.rotate_left(k)  # Call left rotate on parent of k
                    k.parent.color = Color.BLACK
                    k.parent.parent.color = Color.RED
                    self.rotate_right(k.parent.parent)  # Call right rotate on grandparent
            if k == self.root:  # If k reaches root then break
                break
        self.root.color = Color.BLACK  # Set color of root as black

    def fix_delete(self, x):
        """ Fix Nodes after deletion """
        while x != self.root and x.color == Color.BLACK:  # Repeat until x reaches nodes and color of x is black
            if x == x.parent.left:  # If x is left child of its parent
                s = x.parent.right  # Sibling of x
                if s.color == Color.RED:  # if sibling is red
                    s.color = Color.BLACK  # Set its color to black
                    x.parent.color = Color.RED  # Make its parent red
                    self.rotate_left(x.parent)  # Call for left rotate on parent of x
                    s = x.parent.right
                # If both the child are black
                if s.left.color == Color.BLACK and s.right.color == Color.BLACK:
                    s.color = Color.RED  # Set color of s as red
                    x = x.parent
                else:
                    if s.right.color == Color.BLACK:  # If right child of s is black
                        s.left.color = Color.BLACK  # set left child of s as black
                        s.color = Color.RED  # set color of s as red
                        self.rotate_right(s)  # call right rotation on x
                        s = x.parent.right

                    s.color = x.parent.color
                    x.parent.color = Color.BLACK  # Set parent of x as black
                    s.right.color = Color.BLACK
                    self.rotate_left(x.parent)  # call left rotation on parent of x
                    x = self.root
            else:  # If x is right child of its parent
                s = x.parent.left  # Sibling of x
                if s.color == Color.RED:  # if sibling is red
                    s.color = Color.BLACK  # Set its color to black
                    x.parent.color = Color.RED  # Make its parent red
                    self.rotate_right(x.parent)  # Call for right rotate on parent of x
                    s = x.parent.left

                if s.right.color == Color.BLACK and s.right.color == Color.BLACK:
                    s.color = Color.RED
                    x = x.parent
                else:
                    if s.left.color == Color.BLACK:  # If left child of s is black
                        s.right.color = Color.BLACK  # set right child of s as black
                        s.color = Color.RED
                        self.rotate_left(s)  # call left rotation on x
                        s = x.parent.left

                    s.color = x.parent.color
                    x.parent.color = Color.BLACK
                    s.left.color = Color.BLACK
                    self.rotate_right(x.parent)
                    x = self.root
        x.color = Color.BLACK

    # Function to transplant nodes
    def __rb_transplant(self, u, v):
        if u.parent is None:
            self.root = v
        elif u == u.parent.left:
            u.parent.left = v
        else:
            u.parent.right = v
        v.parent = u.parent

    def delete_node_helper(self, node, key):
        """ Delete node """
        z = self.NULL
        while node != self.NULL:  # Search for the node having that value/ key and store it in 'z'
            if node.val == key:
                z = node

            if node.val <= key:
                node = node.right
            else:
                node = node.left

        if z == self.NULL:  # If Kwy is not present then deletion not possible so return
            print("Value not present in Tree !!")
            return

        y = z
        y_original_color = y.color  # Store the color of z- node
        if z.left == self.NULL:  # If left child of z is NULL
            x = z.right  # Assign right child of z to x
            self.__rb_transplant(z, z.right)  # Transplant Node to be deleted with x
        elif z.right == self.NULL:  # If right child of z is NULL
            x = z.left  # Assign left child of z to x
            self.__rb_transplant(z, z.left)   # Transplant Node to be deleted with x
        else:  # If z has both the child nodes
            y = self.minimum(z.right)  # Find minimum of the right sub tree
            y_original_color = y.color  # Store color of y
            x = y.right
            if y.parent == z:  # If y is child of z
                x.parent = y  # Set parent of x as y
            else:
                self.__rb_transplant(y, y.right)
                y.right = z.right
                y.right.parent = y

            self.__rb_transplant(z, y)
            y.left = z.left
            y.left.parent = y
            y.color = z.color
        if y_original_color == Color.BLACK:  # If color is black then fixing is needed
            self.fix_delete(x)

    # Deletion of node
    def delete_node(self, val):
        """ Delete node """
        self.delete_node_helper(self.root, val)  # Call for deletion

    def __printCall(self, node, indent, last):
        if node != self.NULL:
            print(indent, end=' ')
            if last:
                print("R----", end=' ')
                indent += "     "
            else:
                print("L----", end=' ')
                indent += "|    "

            s_color = "RED" if node.color == Color.RED else "BLACK"
            print(str(node.val) + "(" + s_color + "): ", node.items)
            self.__printCall(node.left, indent, False)
            self.__printCall(node.right, indent, True)

    def get_node(self, val):
        """ Get the Node by the value """
        return self.root.get_node(val)

    def get_all(self):
        """ Get all nodes """
        return self.root.get_all()

    def print_tree(self):
        """ print """
        self.__printCall(self.root, "", True)


if __name__ == "__main__":
    bst = RBTree()

    for i in range(1, 6):
        bst.insert_node(i, {1: 2})
    bst.insert_node(100, {2: 3})
    bst.insert_node(100, {3: 4})
    bst.insert_node(100, {4: 5})

    # bst.print_tree()
    print(bst.get_node(100))
    print(bst.get_node(120))
    # print("\nAfter deleting an element")
    # print(bst.get_all())
