package trees;

import java.util.*;
import java.io.*;

/**
 * RBTree rules:
 * (1) Every node is either red or black
 * (2) The root is always black
 * (3) A red parent cannot have a red child
 * (4) All descendants of a node have the same number of black nodes
 * (5) All leaf nodes are black nodes
*/
public class RedBlackTree {
    static class Node extends BaseNode {
        boolean isRed;
        Node parent;

        /**
         * Returning true means we have a double red.
         */
        boolean add(Node n) {
            if(n.c == c) {
                return false; // nothing to do
            } else if(n.c < c) {
                if(left == null) {
                    left = n;
                    n.parent = this;
                    return n.isRed && isRed;
                } else {
                    if(((Node)left).add(n))
                        fixDoubleRed();
                }
            } else if(c < n.c) {
                if(right == null) {
                    right = n;
                    n.parent = this;
                    return n.isRed && isRed;
                } else {
                    if(((Node)right).add(n))
                        fixDoubleRed();
                }
            }
            return false;
        }

        public boolean add(char c) {
            Node n = new Node();
            n.c = c;
            n.isRed = true;
            return add(n);
        }

        public void toString(StringBuilder sb) {
            if(left != null) left.toString(sb);
            sb.append(String.format("(%c:%d:%s%s) ",c,height(), isRed ? "R" : " ", parent != null ? " " : "H"));
            if(right != null) right.toString(sb);
        }
        public String toString() {
            StringBuilder sb = new StringBuilder();
            toString(sb);
            return sb.toString();
        }

        void fixDoubleRed() {
            boolean found = false;
            Node n1=null, n2=null, n3=null;
            BaseNode t1=null, t2=null, t3=null, t4=null;
            Node gp = parent;
            if(left != null && ((Node)left).isRed) {
                if(left.left != null && ((Node)left.left).isRed) {
                    found = true;
                    //
                    //          n3
                    //         /  \
                    //       n2    t4
                    //      /  \
                    //    n1    t3
                    //   /  \
                    // t1    t2
                    //
                    n1 = (Node)left.left;
                    n2 = (Node)left;
                    n3 = this;
                    t1 = n1.left; t2 = n1.right; t3 = n2.right; t4 = n3.right;
                } else if(left.right != null && ((Node)left.right).isRed) {
                    found = true;
                    //
                    //         n3
                    //        /  \
                    //       /    t4
                    //      /   
                    //    n1    
                    //   /  \
                    // t1    n2
                    //      /  \
                    //     t2   t3
                    //
                    n1 = (Node)left;
                    n2 = (Node)left.right;
                    n3 = this;
                    t1 = n1.left; t2 = n2.left; t3 = n2.right; t4 = n3.right;
                }
            }
            if(!found && right != null && ((Node)right).isRed) {
                if(right.right != null && ((Node)right.right).isRed) {
                    found = true;
                    //
                    //     n1
                    //    /  \
                    //  t1    n2
                    //       /  \
                    //     t2    n3
                    //          /  \
                    //        t3    t4
                    //
                    n1 = this;
                    n2 = (Node)right;
                    n3 = (Node)right.right;
                    t1 = n1.left; t2 = n2.left; t3 = n3.left; t4 = n3.right;
                } else if(right.left != null && ((Node)right.left).isRed) {
                    found = true;
                    //
                    //    n1
                    //   /  \
                    // t1    \
                    //        \
                    //         n3
                    //        /  \
                    //      n2    t4
                    //     /  \
                    //    t2   t3
                    //
                    n1 = this;
                    n3 = (Node)right;
                    n2 = (Node)right.left;
                    t1 = n1.left; t2 = n2.left; t3 = n2.right; t4 = n3.right;
                } else {
                    assert false : toString();
                }
            }
            assert found : toString();
            //
            //        n2
            //       /  \
            //      /    \
            //    n1      n3
            //   /  \    /  \
            // t1    t2 t3   t4
            //
            n3.left = t3; n3.right = t4;
            n1.left = t1; n1.right = t2;
            n2.left = n1; n2.right = n3;
            n1.parent = n2;
            n3.parent = n2;
            n2.parent = gp;
            if(gp != null) {
                if(gp.right == this) gp.right = n2;
                else if(gp.left == this) gp.left = n2;
                else assert false;
            }
            n1.isRed = false; n3.isRed = false;
            n2.isRed = gp != null;
            if(t1 != null) ((Node)t1).parent = n1;
            if(t2 != null) ((Node)t2).parent = n1;
            if(t3 != null) ((Node)t3).parent = n3;
            if(t4 != null) ((Node)t4).parent = n3;

            // Checking...
            if(t1 != null) assert t1.c < n1.c;
            if(t2 != null) assert n1.c < t2.c;
            if(t2 != null) assert t2.c < n2.c;
            if(t3 != null) assert n2.c < t3.c;
            if(t3 != null) assert t3.c < n3.c;
            if(t4 != null) assert n3.c < t4.c;
            assert n1.c < n2.c;
            assert n2.c < n3.c;

            if(n2.isRed && n2.parent != null && n2.parent.isRed) {
                if(n2.parent.parent != null) {
                    n2.parent.parent.fixDoubleRed();
                } else {
                    // red is at the top of the tree
                    assert false;
                }
            }
        }
    
        public void colorCheck() {
            if(isRed) {
                if(left != null) assert !((Node)left).isRed;
                if(right != null) assert !((Node)right).isRed;
            }
            if(left != null) ((Node)left).colorCheck();
            if(right != null) ((Node)right).colorCheck();
        }
    }

    Node head;

    public void add(char c) {
        if(head == null) {
            head = new Node();
            head.c = c;
        } else {
            if(head.add(c))
                head.fixDoubleRed();
            // Head might get moved by fix
            if(head.parent != null)
                head = head.parent;
        }
    }

    public int height() {
        if(head == null) return 0;
        else return head.height();
    }

    public int nodeCount() {
        if(head == null) return 0;
        else return head.nodeCount();
    }

    public String toString() {
        if(head == null) return "";
        else return head.toString();
    }

    public static RedBlackTree test(String str) {
        Set<Character> sc = new HashSet<>();
        RedBlackTree r = new RedBlackTree();
        int d = 0;
        for(int i = 0;i < str.length();i++) {
            char c = str.charAt(i);
            r.add(c);
            sc.add(c);
            r.colorCheck();
            int d1 = r.height();
            System.out.println("adding "+c+" -> "+r+": "+r.nodeCount()+"/"+d1);
            assert sc.size() == r.nodeCount();
            assert d1 <= d + 1;
            d = d1;
        }
        return r;
    }

    public void colorCheck() {
        if(head != null) head.colorCheck();
    }

    public void visit(BinaryTreeVisitor v) {
        if(head != null) head.visit(v);
    }

    public static void main(String[] args) throws IOException {
        test("fishes");
        test("Avacado");
        test("abcdefghijklmnop");
        RedBlackTree bt = test("the_quick_brown_fox_jumped_over_the_lazy_dog");
        try(final var pw = new PrintWriter("viz.dot")) {
            pw.println("digraph G {");
            bt.visit(new BinaryTreeVisitor() {
                @Override
                public void visit(BaseNode n) {
                    if(n.left != null)
                        pw.printf("%s -> %s;%n", n.c, n.left.c);
                    if(n.right != null)
                        pw.printf("%s -> %s;%n", n.c, n.right.c);
                }
            });
            pw.println("}");
        }
    }
}