package org.xlj;

/**
 * red black tree
 * @param <K> key
 * @param <V> value
 */
public class RBTree<K, V> {

    Comparetor<K> comparetor;
    Node<K, V> root;

    public Comparetor<K> getComparetor() {
        return comparetor;
    }

    public void setComparetor(Comparetor<K> comparetor) {
        this.comparetor = comparetor;
    }

    public Node<K, V> getRoot() {
        return root;
    }

    public void setRoot(Node<K, V> root) {
        this.root = root;
    }

    /**
     * get node by key from this rb-tree
     * @param key
     * @return
     * @throws Exception
     */
    public Node get(K key) throws Exception {
        if(this.comparetor == null)
        {
            throw new Exception("this comparetor of the tree is null.");
        }

        Node<K, V> current = root;
        while (current != null)
        {
            //in right
            if(this.comparetor.compare(key, current.getKey()) > 0)
            {
                current = current.getRightChild();
            }
            else if(this.comparetor.compare(key, current.getKey()) < 0) //in left
            {
                current = current.getLeftChild();
            }
            else // equals 0
            {
                break;
            }
        }

        return current;
    }

    /**
     * remove a node from this tree and keep rb-tree's feature
     * @param key
     * @throws Exception
     */
    public void remove(K key) throws Exception {
        Node node = this.get(key);
        if(node == null)
        {
            throw new Exception("does not find this key: " + key);
        }

        Object[] nodes = this.removeNode(node);

        //removed node color is black
        if((Node.NodeColors)nodes[2] == Node.NodeColors.BLACK
                && this.getRoot() != null)
        {
            this.removeFixup((Node) nodes[0], (Node)nodes[1]);
        }
    }

    /**
     * remove
     * @param node
     * @return x at index 1, x's parent at index 2, node origin color at 3
     */
    private Object[] removeNode(Node node)
    {
        Node.NodeColors originColor = node.color;
        //node's  successor(y)
        Node successor = null;
        // x
        Node x = null;
        // if x == null, then have to use parentOfx
        Node parentOfx = null;

        //case 1: left child is null and right not null
        if(node.getLeftChild() == null && node.getRightChild() !=null)
        {
            successor = node.getRightChild();
            x = successor;
            parentOfx = node.parent;
        }

        //case 2: right child is null and left not null
        if(node.getRightChild() == null && node.getLeftChild() != null)
        {
            successor = node.getLeftChild();
            x = successor;
            parentOfx = node.parent;
        }

        //case 3 - 4: all is not null
        if(node.getLeftChild() != null && node.getRightChild() != null)
        {
            //find successor
            successor = node.getRightChild();
            while(successor.getLeftChild() != null)
            {
                successor = successor.getLeftChild();
            }

            //x = successor.right
            x = successor.getRightChild();


            //case 3: if node is parent of the successor
            if(node == successor.getParent())
            {
                //x's parent does not changed
                parentOfx = successor;

                //node.left transplant to left of successor
                this.transPlantToLeft(node.getLeftChild(), successor);

                successor.setColor(originColor);
            }
            else //case 4: if node is not parent of the successor
            {
                //need change x's parent to origin successor's parent
                parentOfx = successor.getParent();

                this.replaceAndTransPlant(successor, x);
                this.transPlantToLeft(node.getLeftChild(), successor);

                successor.setColor(originColor);
            }
        }

        //to replace
        this.replaceAndTransPlant(node, successor);

        if(node == this.getRoot())
        {
            this.setRoot(successor);
        }

//        if(successor != null)
//        {
////            successor.setColor(originColor);
//
//            // set root
//            if(successor.getParent() == null)
//            {
//                this.setRoot(successor);
//            }
//        }

        return new Object[]{x, parentOfx, originColor};
    }

    /**
     * fix to rb-tree feature
     * @param node
     * @param parent
     */
    private void removeFixup(Node node, Node parent)
    {
        // node is null or black and not root of this tree
        while ((node == null || node.getColor() == Node.NodeColors.BLACK) &&
                node != this.getRoot())
        {
            Node brother = this.getBrother(node, parent);

            // left of parent
            if(node == parent.getLeftChild())
            {
                // case 2: brother is red
                if(brother != null
                        && brother.getColor() == Node.NodeColors.RED)
                {
                    parent.setColor(Node.NodeColors.RED);
                    brother.setColor(Node.NodeColors.BLACK);
                    this.leftRotate(this, parent);
                    brother = this.getBrother(node, parent);

                    assert brother == null || brother.getColor() == Node.NodeColors.BLACK;
                    assert parent == null || parent.getColor() == Node.NodeColors.RED;
                }

                // case 3 - 5: brother is black.
                // case 3: brother's children are all black
                if(brother != null
                        && (brother.getLeftChild() == null
                        || brother.getLeftChild().getColor() == Node.NodeColors.BLACK)
                        &&
                        (brother.getRightChild() == null
                                || brother.getRightChild().getColor() == Node.NodeColors.BLACK))
                {
                    brother.setColor(Node.NodeColors.RED);
//                    parent.setColor(Node.NodeColors.BLACK);
                    node = parent;
                    parent = node.getParent();
                    brother = this.getBrother(node, parent);
                }
                else // case 4 - 5:
                {
                    // case 4: brother left is red, and right is black
                    if(brother != null
                            && (brother.getRightChild() == null
                            || brother.getRightChild().getColor() == Node.NodeColors.BLACK))
                    {
                        brother.setColor(Node.NodeColors.RED);
                        brother.getLeftChild().setColor(Node.NodeColors.BLACK);
                        this.rightRotate(this, brother);
                        brother = this.getBrother(node, parent);
                    }

                    // case 5: right is red
                    if(brother != null)
                    {
                        brother.setColor(parent.getColor());
                    }
                    parent.setColor(Node.NodeColors.BLACK);
                    if(brother != null)
                    {
                        brother.getRightChild().setColor(Node.NodeColors.RED);
                    }
                    this.leftRotate(this, parent);

                    // redeay to exit;
                    node = this.getRoot();
                }
            }
            else// right of parent
            {
                // case 2: brother is red
                if(brother.getColor() == Node.NodeColors.RED)
                {
                    brother.setColor(Node.NodeColors.BLACK);
                    parent.setColor(Node.NodeColors.RED);
                    this.rightRotate(this, brother);
                    brother = this.getBrother(node, parent);

                    assert brother == null || brother.getColor() == Node.NodeColors.BLACK;
                    assert parent == null || parent.getColor() == Node.NodeColors.RED;
                }

                // case 3 - 5: brother is black
                // case 3: both brother's children are blcak
                if((brother.getLeftChild() == null
                        || brother.getLeftChild().getColor() == Node.NodeColors.BLACK)
                        &&
                        (brother.getRightChild() == null
                                || brother.getRightChild().getColor() == Node.NodeColors.BLACK))
                {
                    brother.setColor(Node.NodeColors.RED);
//                    parent.setColor(Node.NodeColors.BLACK);
                    node = parent;
                    parent = node.getParent();
                    brother = this.getBrother(node, parent);
                }
                else // case 4 - 5:
                {
                    // case 4: brother left is black, right is red
                    if (brother.getLeftChild() == null
                            || brother.getLeftChild().getColor() == Node.NodeColors.BLACK)
                    {
                        brother.setColor(Node.NodeColors.RED);
                        brother.getRightChild().setColor(Node.NodeColors.BLACK);
                        this.leftRotate(this, brother);
                        brother = this.getBrother(node, parent);
                    }

                    // case 5: left is red
                    brother.setColor(parent.getColor());
                    parent.setColor(Node.NodeColors.BLACK);
                    brother.getLeftChild().setColor(Node.NodeColors.RED);
                    this.rightRotate(this, parent);

                    // redeay to exit;
                    node = this.getRoot();
                }
            }
        }

        //case 1: node color to back
        if(node != null)
        {
            node.setColor(Node.NodeColors.BLACK);
        }
    }
    
    /**
     * replace oring use replace
     */
    @SuppressWarnings("unchecked")
    public void replaceAndTransPlant(Node origin, Node replace)
    {
        //origin's parent
        Node parent = origin.getParent();
        if(replace != null)
        {
            replace.setParent(parent);
        }

        //parent is null to return
        if(parent != null)
        {
            //if origin is left node
            if(parent.getLeftChild() == origin)
            {
                parent.setLeftChild(replace);
            }
            else //else is right
            {
                parent.setRightChild(replace);
            }
        }

    }
    
    /**
     * transplant node to left of parent
     * @param node
     * @param parent
     */
    @SuppressWarnings("unchecked")
    private void transPlantToLeft(Node node, Node parent)
    {
        if(parent != null)
        {
            parent.setLeftChild(node);
        }

        if(node != null)
        {
            node.setParent(parent);
        }
    }

    /**
     * transplant node to right of parent
     * @param node
     * @param parent
     */
    @SuppressWarnings("unchecked")
    private void transPlantToRight(Node node, Node parent)
    {
        if(parent != null)
        {
            parent.setRightChild(node);
        }

        if(node != null)
        {
            node.setParent(parent);
        }
    }

    /**
     * 獲取兄弟節點
     * @param node
     * @param parent
     * @return
     */
    private Node getBrother(Node node, Node parent)
    {
        if(parent == null)
        {
            return null;
        }

        return node == parent.getLeftChild()? parent.getRightChild(): parent.getLeftChild();
    }

    /**
     * insert
     * @param node
     */
    public void insert(Node<K, V> node) throws Exception {
        if(this.comparetor == null)
        {
            throw new Exception("this comparetor of the tree is null.");
        }

        if(this.get(node.key) != null)
        {
            throw new Exception("has exists this key: " + node.getKey());
        }

        if(this.getRoot() == null)
        {
            this.setRoot(node);
        }

        Node<K, V> currentNode = this.getRoot();
        int nCompareResult = 0;

        while (currentNode != null)
        {
            nCompareResult = 0;
            nCompareResult = this.comparetor.compare(currentNode.getKey(), node.getKey());

            if(nCompareResult > 0)
            {
                if(currentNode.getLeftChild() != null)
                {
                    currentNode = currentNode.getLeftChild();
                }
                else
                {
                    currentNode.setLeftChild(node);
                    node.setParent(currentNode);
                    break;
                }
            }

            if(nCompareResult < 0)
            {
                if(currentNode.getRightChild() != null)
                {
                    currentNode = currentNode.getRightChild();
                }
                else
                {
                    currentNode.setRightChild(node);
                    node.setParent(currentNode);
                    break;
                }
            }

            //has already exist
            if (nCompareResult == 0)
            {
                break;
            }
        }

        this.insertFixup(node);

        this.getRoot().setColor(Node.NodeColors.BLACK);
    }

    /**
     * insert fix of the tree
     * @param addedNode be inserted node
     */
    private void insertFixup(Node addedNode)
    {
        //directly return if addedNode is root
        if(addedNode == this.getRoot()
                || addedNode.getParent().getColor() == Node.NodeColors.BLACK)
        {
            return;
        }

        while (addedNode.getParent() != null
                && addedNode.getParent().getParent() != null
                && addedNode.getParent().getColor() == Node.NodeColors.RED ) {
            //parent
            Node parent = addedNode.getParent();
            //grandParent
            Node grandParent = parent.getParent();
            //uncle
            Node uncle = grandParent == null ?
                    null :
                    grandParent.getLeftChild() == parent ?
                            grandParent.getRightChild() : grandParent.getLeftChild();

            // case 1: brother of parent and uncle color is red
            if (uncle != null
                    && uncle.getColor() == Node.NodeColors.RED) {

                parent.setColor(Node.NodeColors.BLACK);
                uncle.setColor(Node.NodeColors.BLACK);
                grandParent.setColor(Node.NodeColors.RED);

                addedNode = grandParent;
                continue;
            }

            //parent is left of the grandparent
            if (parent == grandParent.getLeftChild()) {

                // case 2: red of the uncle color, addedNode is right node of the parent
                if (addedNode == parent.getRightChild()) {

                    addedNode = parent;
                    this.leftRotate(this, addedNode);

                    //parent
                    parent = addedNode.getParent();
                    if (parent == null) {
                        break;
                    }
                    //parent of parent
                    grandParent = parent.getParent();
                }

                // case 3: red of the uncle color, addedNode is left node of the parent
                parent.setColor(Node.NodeColors.BLACK);
                if (grandParent == null) {
                    break;
                }
                grandParent.setColor(Node.NodeColors.RED);
                this.rightRotate(this, grandParent);
            }
            else //parent == parent.right
            {
                if(addedNode == parent.getLeftChild())
                {
                    addedNode = parent;
                    this.rightRotate(this, addedNode);

                    //parent
                    parent = addedNode.getParent();
                    if (parent == null) {
                        break;
                    }
                    //parent of parent
                    grandParent = parent.getParent();
                }

                //case 3:
                parent.setColor(Node.NodeColors.BLACK);
                if(grandParent == null)
                {
                    break;
                }
                grandParent.setColor(Node.NodeColors.RED);
                this.leftRotate(this, grandParent);
            }
        }
    }

    /**
     * node left rotate
     */
    private void leftRotate(RBTree tree, Node node)
    {
        if(node.rightChild == null)
        {
            return;
        }

        Node rightChild = node.rightChild;

        //change parent child reference
        if(node.parent != null)
        {
            //if is left node
            if(node.parent.leftChild == node)
            {
                node.parent.leftChild = rightChild;
            }
            else //if is right node
            {
                node.parent.rightChild = rightChild;
            }
        }
        else //if node is root of the tree
        {
            tree.root = rightChild;
            tree.getRoot().setParent(null);
        }

        //change node and node.right child's parent
        rightChild.parent = node.parent;

        //set parent of the node
        node.parent = rightChild;

        node.rightChild = rightChild.leftChild;

        if(node.leftChild != null) {
            node.leftChild.parent = node;
        }

        rightChild.leftChild = node;
    }

    /**
     * node right rotate
     */
    private void rightRotate(RBTree tree, Node node)
    {
        //can not to rotate
        if(node.leftChild == null)
        {
            return;
        }

        Node leftChild = node.leftChild;

        //change parent of this nodes
        if(node.parent != null)
        {
            if(node == node.parent.leftChild)
            {
                node.parent.leftChild = leftChild;
            }
            else
            {
                node.parent.rightChild = leftChild;
            }
        }
        else //if node is root of the tree
        {
            tree.root = node.leftChild;
            tree.getRoot().setParent(null);
        }

        //change parent of the node.left
        leftChild.parent = node.parent;

        node.parent = leftChild;

        node.leftChild = leftChild.rightChild;
        if(node.leftChild != null) {
            node.leftChild.parent = node;
        }

        leftChild.rightChild = node;

    }

    /**
     * rb-tree's node
     * @param <K>
     * @param <V>
     */
    public static class Node<K,V>
    {
        Node<K, V> parent;
        Node<K, V> leftChild;
        Node<K, V> rightChild;

        K key;
        V value;

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }

        NodeColors color = NodeColors.RED;

        public Node<K, V> getParent() {
            return parent;
        }

        public void setParent(Node<K, V> parent) {
            this.parent = parent;
        }

        public Node<K, V> getLeftChild() {
            return leftChild;
        }

        public void setLeftChild(Node<K, V> leftChild) {
            this.leftChild = leftChild;
        }

        public Node<K, V> getRightChild() {
            return rightChild;
        }

        public void setRightChild(Node<K, V> rightChild) {
            this.rightChild = rightChild;
        }

        public NodeColors getColor() {
            return color;
        }

        public void setColor(NodeColors color) {
            this.color = color;
        }

        /**
         * rb-trees's child node colors
         */
        public enum NodeColors
        {
            RED,
            BLACK
        }

        @Override
        public String toString() {
            return "Node{" +
                    "key=" + key +
                    '}';
        }
    }

    /**
     * comparetor
     * @param <T>
     */
    public interface Comparetor<T>
    {
        /**
         * 比較 t1 和 t2
         * @param t1
         * @param t2
         * @return t1 > t2 return 1
         * t1 == t2 return 0
         * t1 < t2 return -1
         */
        int compare(T t1, T t2);
    }
}

package org.xlj;

import org.junit.Test;

public class RBTreeTest {

    @Test
    public void insertTest()
    {

        RBTree<Integer, Integer> tree = new RBTree<>();
        tree.setComparetor(new RBTree.Comparetor<Integer>() {
            @Override
            public int compare(Integer t1, Integer t2) {
                return t1 > t2 ? 1 : t1 == t2 ? 0 : - 1;
            }
        });

        try {
//            int[] values={8, 12, 31, 19, 40, 38};
            int[] values={40, 38, 31, 12, 19, 8};

            for (int i = 0; i < values.length; i++)
            {
                RBTree.Node<Integer, Integer> node = new RBTree.Node<>();
                node.setKey(values[i]);
                node.setValue(values[i]);

                tree.insert(node);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        this.printTree(tree.root, 0);


        while (tree.getRoot() != null) {
            System.out.println();
            System.out.println();

            try {
//                tree.remove(19);
//                tree.remove(12);
                tree.remove(tree.getRoot().getKey());
            } catch (Exception e) {
                e.printStackTrace();
            }

            this.printTree(tree.root, 0);
        }
    }

    public void printTree(RBTree.Node node, int index)
    {
        if(node == null)
        {
            return;
        }

        System.out.println(node.getKey() + " " + node.getColor() + " " + index++);
        this.printTree(node.getLeftChild(), index);
        this.printTree(node.getRightChild(), index);
        index--;
    }

}