// Implements a red-black tree representation of an FMap<K,V>.
// 
// A red-black tree is a binary search tree that satisfies these
// red-black tree invariants:
// 
// No red node of the tree has a red child.
// For every path from root to a leaf, the number of black nodes is the same.
//
// A binary search tree is a binary tree that satisfies these
// binary search tree invariants:
// 
// For each non-empty node of the tree,
// 
//     the left subtree is a binary search tree
//     the right subtree is a binary search tree
//     every key within the left subtree is less than the key at this node
//     every key within the right subtree is greater than the key at this node
// 
// where "less than" and "greater than" are defined by the tree's
// Comparator, which must define a total order on the keys.

package fmap;

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Collections;                   // for its sort algorithm
import java.util.ArrayList;                     // for its iterator

abstract class FTree<K,V> extends FMapBase<K,V> {

    static <K,V> FMapBase<K,V> emptyTree (Comparator<? super K> c) {
        return new FTree.EmptyTree<K,V>(c);
    }

    // must define a total order on the keys
    // that's consistent with equals(Object)

    Comparator<? super K> c;

    // the color of this FTree

    char rb;

    // Initializes this FTree with the given Comparator and default color.

    FTree (Comparator<? super K> c) {
        this.rb = RED;
        this.c = c;
    }

    // Initializes this FTree with the given Color and Comparator.

    FTree (char rb, Comparator<? super K> c) {
        this.rb = rb;
        this.c = c;
    }

    public FMapBase<K,V> add (K k, V v) {
        FTree<K,V> result = insert (k, v).makeBlack();
        if (DEBUGGING) {
            result.checkInvariants();
        }
        return result;
    }

    // Colors are represented by characters.

    final static char RED = 'r';
    final static char BLACK = 'b';

    // returns true iff this FTree is red

    boolean isRed () { return rb == RED; }

    // returns true iff this FTree is red

    boolean isBlack () { return rb == BLACK; }

    // Returns a red-black tree that's equivalent (as a binary search tree)
    // to a non-empty tree whose root node associates the key k with value v,
    // whose left subtree is t1, and whose right subtree is t2.
    // 
    // That tree might not satisfy the red-black invariants, however.
    // The balance method detects situations in which the obvious result
    // would not satisfy the red-black invariants, and rotates the tree
    // as necessary to satisfy those invariants.
    // 
    // See Chris Okasaki. Red-black trees in a functional setting.
    // Journal of Functional Programming, 9(4), pages 474-477, July 1999.
    // 
    // In the notation of that paper, where B and R mean black and red,
    // and T is a creator of non-empty trees such that (T rb t1 k t2)
    // is a non-empty tree with color rb, left subtree t1, key k, and
    // right subtree t2:
    // 
    // balance B (T R (T R a x b) y c) z d = T R (T B a x b) y (T B c z d)
    // balance B (T R a x (T R b y c)) z d = T R (T B a x b) y (T B c z d)
    // balance B a x (T R (T R b y c) z d) = T R (T B a x b) y (T B c z d)
    // balance B a x (T R b y (T R c z d)) = T R (T B a x b) y (T B c z d)
    // 
    // otherwise
    // 
    // balance color a x b = T color a x b

    FTree<K,V> balance (char rb, K k, V v, FTree<K,V> t1, FTree<K,V> t2) {
        if (rb == RED)
            return new Node<K,V> (rb, c, k, v, t1, t2);
        else if ((! (t1.isEmpty()))
                 && (t1.isRed())
                 && (! (t1.right().isEmpty()))
                 && (t1.right().isRed()))
            return new Node<K,V> (RED, c,
                                  t1.right().key(), t1.right().val(),
                                  new Node<K,V> (BLACK, c,
                                                 t1.key(), t1.val(),
                                                 t1.left(),
                                                 t1.right().left()),
                                  new Node<K,V> (BLACK, c,
                                                 k, v,
                                                 t1.right().right(),
                                                 t2));
        else if ((! (t2.isEmpty()))
                 && (t2.isRed())
                 && (! (t2.left().isEmpty()))
                 && (t2.left().isRed()))
            return new Node<K,V> (RED, c,
                                  t2.left().key(), t2.left().val(),
                                  new Node<K,V> (BLACK, c,
                                                 k, v,
                                                 t1,
                                                 t2.left().left()),
                                  new Node<K,V> (BLACK, c,
                                                 t2.key(), t2.val(),
                                                 t2.left().right(),
                                                 t2.right()));
        else if ((! (t1.isEmpty()))
                 && (t1.isRed())
                 && (! (t1.left().isEmpty()))
                 && (t1.left().isRed()))
            return new Node<K,V> (RED, c,
                                  t1.key(), t1.val(),
                                  t1.left().makeBlack(),
                                  new Node<K,V> (BLACK, c,
                                                 k, v,
                                                 t1.right(),
                                                 t2));
        else if ((! (t2.isEmpty()))
                 && (t2.isRed())
                 && (! (t2.right().isEmpty()))
                 && (t2.right().isRed()))
            return new Node<K,V> (RED, c,
                                  t2.key(), t2.val(),
                                  new Node<K,V> (BLACK, c,
                                                 k, v,
                                                 t1,
                                                 t2.left()),
                                  t2.right().makeBlack());
        else return new Node<K,V> (rb, c, k, v, t1, t2);
    }

    // returns the left subtree of a non-empty FTree

    FTree<K,V> left () {
        throw new UnsupportedOperationException ("left");
    }

    // returns the right subtree of a non-empty FTree

    FTree<K,V> right () {
        throw new UnsupportedOperationException ("right");
    }

    // returns the key at the root of a non-empty FTree

    K key () {
        throw new UnsupportedOperationException ("key");
    }

    // returns the value at the root of a non-empty FTree

    V val () {
        throw new UnsupportedOperationException ("val");
    }

    FTree<K,V> makeBlack () {
        throw new UnsupportedOperationException ("makeBlack");
    }

    // Returns the result of inserting the given key and value
    // into this FTree.

    FTree<K,V> insert (K k, V v) {
        throw new UnsupportedOperationException ("ins");
    }

    // Checks the red-black invariants for this FTree,
    // and returns the number of black nodes in each path to a leaf.

    int checkInvariants () {
        if (this.isEmpty()) {
            if (this.isRed())
                System.out.println ("Empty tree is red.");
            return 0;
        }
        int n1 = this.left().checkInvariants();
        int n2 = this.right().checkInvariants();
        if (n1 != n2)
            reportViolation (n1 + " != " + n2);
        if (this.isRed()) {
            if (this.left().isRed())
                reportViolation ("consecutive red (left)");
            if (this.right().isRed())
                reportViolation ("consecutive red (right)");
            return n1;
        }
        return 1 + n1;
    }

    private void reportViolation (String msg) {
        System.out.println ("VIOLATION: " + msg);
        System.out.println (toString());
        throw new RuntimeException ("Red-black invariant violated!");
    }

    // Instances of this class represent an empty red-black tree.

    private static class EmptyTree<K,V> extends FTree<K,V> {

        EmptyTree (Comparator<? super K> c) {
            super (BLACK, c);
        }

        // The purposes of the following methods are documented
        // in FMapBase.java or in the FTree abstract class above.

        public boolean isEmpty () { return true; }

        public int size () {
            return 0;
        }

        @Override
            FTree<K,V> insert (K k, V v) {
            return new Node<K,V> (RED, c, k, v, this, this);
        }

        public boolean containsKey (K x) {
            return false;
        }

        public V get (K x) {
            throw new RuntimeException ("get on empty tree");
        }

        boolean equalsLoop (FMap<K,V> m2) {
            return true;
        }

        ArrayList<K> allKeys (ArrayList<K> keys) {
            return keys;
        }
    }

    // Instances of this class represent a non-empty red-black tree.

    private static class Node<K,V> extends FTree<K,V> {

        private K k0;                 // the key at the root
        private V v0;                 // the value at the root
        private FTree<K,V> t1;        // the left subtree
        private FTree<K,V> t2;        // the right subtree

        Node (Comparator<? super K> c,
              K k0, V v0, FTree<K,V> t1, FTree<K,V> t2) {
            super (c);
            this.k0 = k0;
            this.v0 = v0;
            this.t1 = t1;
            this.t2 = t2;
            if (SIZE_IS_FAST)
                this.theSize = 1 + t1.size() + t2.size();
        }

        Node (char rb, Comparator<? super K> c,
              K k0, V v0, FTree<K,V> t1, FTree<K,V> t2) {
            super (rb, c);
            this.k0 = k0;
            this.v0 = v0;
            this.t1 = t1;
            this.t2 = t2;
            if (SIZE_IS_FAST)
                this.theSize = 1 + t1.size() + t2.size();
        }

        // The purposes of the following methods are documented
        // in FMapBase.java or in the FTree abstract class above.

        public boolean isEmpty () { return false; }

        public int size () {
            if (SIZE_IS_FAST)
                return theSize;
            return 1 + t1.size() + t2.size();
        }

        FTree<K,V> left () { return t1; }

        FTree<K,V> right () { return t2; }

        K key () { return k0; }

        V val () { return v0; }

        FTree<K,V> makeBlack () {
            return new Node<K,V> (BLACK, c, k0, v0, t1, t2);
        }

        @Override
        FTree<K,V> insert (K k, V v) {
            if (k.equals(k0))
                return new Node<K,V> (rb, c, k0, v, t1, t2);
            int direction = c.compare (k, k0);
            if (direction < 0)
                return balance (rb, k0, v0, t1.insert (k, v), t2);
            else if (direction > 0)
                return balance (rb, k0, v0, t1, t2.insert (k, v));
            else
                throw new RuntimeException ("bad Comparator?");
        }

        public boolean containsKey (K x) {
            if (x.equals(k0))
                return true;
            else if (c.compare (x, k0) < 0)
                return t1.containsKey (x);
            else if (c.compare (x, k0) > 0)
                return t2.containsKey (x);
            else throw new RuntimeException ("bad Comparator?");
        }

        public V get (K x) {
            if (x.equals(k0))
                return v0;
            else if (c.compare (x, k0) == 0) {
                System.out.println ("Ooops (get)");  // FIXME
                return v0;
            }
            else if (c.compare (x, k0) < 0)
                return t1.get (x);
            else if (c.compare (x, k0) > 0)
                return t2.get (x);
            else throw new RuntimeException ("bad Comparator?");
        }

        // FIXME: temporary hack for debugging

        V get (K x, int h) {
            checkInvariants();
            System.out.println (x + ", " + k0 + ", " + c.compare(x, k0));
            // return get (k);
            if (x.equals(k0))
                return v0;
            else if (c.compare (x, k0) == 0) {
                System.out.println ("Ooops (get)");  // FIXME
                return v0;
            }
            else if (c.compare (x, k0) < 0)
                return t1.get (x, h + 1);
            else if (c.compare (x, k0) > 0)
                return t2.get (x, h + 1);
            else throw new RuntimeException ("bad Comparator?");
        }

        boolean equalsLoop (FMap<K,V> m2) {
            return t1.equalsLoop (m2)
                && m2.containsKey (k0)
                && m2.get (k0).equals (v0)
                && t2.equalsLoop (m2);
        }

        ArrayList<K> allKeys (ArrayList<K> keys) {
            t1.allKeys (keys);
            keys.add (k0);
            t2.allKeys (keys);
            return keys;
        }

        // FIXME: just for debugging

        @Override
        public String toString () {
            if (DEBUGGING) {
                checkInvariants();
                return toString(0);
            }
            else return super.toString();
        }

        // Returns a String representation of this FTree for debugging,
        // indented n spaces.

        private String toString (int n) {
            String result = "";
            for (int i = 0; i < n; i = i + 1)
                result = result + " ";
            result = result + rb;
            result = result + " " + k0;
            result = result + "\n";
            if (t1.isEmpty())
                result = result + "\n";
            else {
                Node<K,V> left = (Node<K,V>) t1;
                result = result + left.toString (n + 1);
            }
            if (t2.isEmpty())
                result = result + "\n";
            else {
                Node<K,V> right = (Node<K,V>) t2;
                result = result + right.toString (n + 1);
            }
            return result;
        }
    }
}