/**
  * Copyright (c): Uwe Schmidt, FH Wedel
  *
  * You may study, modify and distribute this source code
  * FOR NON-COMMERCIAL PURPOSES ONLY.
  * This copyright message has to remain unchanged.
  *
  * Note that this document is provided 'as is',
  * WITHOUT WARRANTY of any kind either expressed or implied.
  */

package ds.persistent.map;

/** This is a translation of a Haskell implementation
    for red-black trees from
    http://matt.might.net/articles/red-black-delete/
    The Haskell source can be found under
    http://matt.might.net/articles/red-black-delete/code/RedBlackSet.hs

    This Java implementation is a PERSISTENT one and
    it supports deletion.
*/

import ds.interfaces.Map;

import ds.util.K;            // example class for keys
import ds.util.V;            // example class for values
import ds.util.KV;           // key value pair
import ds.util.Queue;        // used by iterators
import ds.util.Function2;
import ds.util.NullIterator;

import static ds.util.KV.mkPair;
import static ds.util.Boolean.toInt;
import static ds.util.Integer.max;
import static ds.util.Integer.min;
import static ds.util.Undef.undefined;

import java.util.Iterator;

import static java.lang.Integer.signum;

//----------------------------------------

public abstract
    class RedBlackTree
    implements Map {

    /* trc *
    static void msg(String s) {
        System.out.println(s);
    }

    static RedBlackTree trc(String s, RedBlackTree t) {
        msg(s + " " + t);
        return t;
    }
    /**/

    //----------------------------------------
    // the smart constructors

    // empty tree
    public static
	RedBlackTree empty() {
	return
	    EMPTY;
    }

    // singleton tree
    public static
	RedBlackTree singleton(K k, V v) {
         return
	     empty().insert(k, v);
    }

    // build search tree from arbitrary sequence (iterator)
    public static
        RedBlackTree fromIterator(Iterator<KV> elems) {

        RedBlackTree res = empty();
        while ( elems.hasNext() ) {
            KV p = elems.next();
            res = res.insert(p.fst, p.snd);
        }
        return
            res;
    }

    public RedBlackTree
	union(Map m2) {
	return
	    unionWith(const2, m2);
    }

    // general unionWith by iterating over m2
    public RedBlackTree
	unionWith(Function2<V,V,V> op,
		  Map m2) {
	Iterator<KV> i = m2.iterator();
	RedBlackTree res = this;

	while ( i.hasNext() ) {
	    KV kv2 = i.next();
	    K  k2  = kv2.fst;
	    V  v2  = kv2.snd;
	    V  v1  = res.lookup(k2);
	    V  vr  = v1 == null ? v2 : op.apply(v1, v2);
	    res = res.insert(k2, vr);
	}
	return
	    res;
    }

    public RedBlackTree
	difference(Map m2) {
	return
	    differenceWith(null2, m2);
    }

    // general differenceWith by iterating over m2
    public RedBlackTree
	differenceWith(Function2<V,V,V> op,
		       Map m2) {

	Iterator<KV> i = m2.iterator();
	RedBlackTree res = this;

	while ( i.hasNext() ) {
	    KV kv2 = i.next();
	    K  k2  = kv2.fst;
	    V  v2  = kv2.snd;
	    V  v1  = res.lookup(k2);

	    if ( v1 != null ) { // key found in res
		V  vr  = op.apply(v1, v2);
		if ( vr == null )
		    res = res.remove(k2);
		else
		    res = res.insert(k2, vr);
	    }
	}
	return
	    res;
    }

    public RedBlackTree copy() {
	return
	    this;
    }

    //----------------------------------------
    // special binary tree methods

    abstract public int depth();
    abstract public int minDepth();

    public boolean member(K k) {
        return
            lookup(k) != null;
    }

    //----------------------------------------
    // internal methods

    // getter
    Node         node()    { return nodeExpected(); }
    int          color()   { return nodeExpected(); }
    K            key()     { return nodeExpected(); }
    V            value()   { return nodeExpected(); }
    RedBlackTree left()    { return nodeExpected(); }
    RedBlackTree right()   { return nodeExpected(); }

    // tree manipulation
    public RedBlackTree insert(K k, V v) {
        assert k != null;

        return
            insert1(k, v).paintItBlack();
    }

    public RedBlackTree remove(K k) {
        assert k != null;

        return
            remove1(k).paintItBlack();
    }

    // 2. iterator for descending enumeration

    abstract public Iterator<KV> iteratorDesc();

    public String toString() {
        return
            iterator().toString();
    }

    public boolean inv() {
        return
            invOrdered()
            &&
            invRedWithBlackChildren()
            &&
            noOfBlackNodes() != -1
            &&
            isBlack();
    }


    //----------------------------------------
    // internal methods

    abstract RedBlackTree insert1(K k, V v);
    abstract RedBlackTree remove1(K k);

    abstract boolean invRedWithBlackChildren();
    abstract int     noOfBlackNodes();
    abstract boolean invOrdered();
    abstract boolean allLS(K k);
    abstract boolean allGT(K k);

    // helper for ***With functions
    private static <A> A nodeExpected() {
	return
	    undefined("Node expected");
    }

    private static Function2<V,V,V> const2
	= new Function2<V,V,V>() {
	public V apply(V x, V y) {
	    return y;
	}
    };

    private static Function2<V,V,V> null2
	= new Function2<V,V,V>() {
	public V apply(V x, V y) {
	    return null;
	}
    };

    //----------------------------------------
    // coloring

    static final int NEGATIVE_BLACK = -1;
    static final int RED            = 0;
    static final int BLACK          = 1;
    static final int DOUBLE_BLACK   = 2;

    static int blacker(int c) {
        assert c < DOUBLE_BLACK;
        return
            c + 1;
    }

    static int redder(int c) {
        assert c > NEGATIVE_BLACK;
        return
            c - 1;
    }

    boolean isRed()              { return color() == RED;            }
    boolean isBlack()            { return color() == BLACK;          }
    boolean isDoubleBlack()      { return color() == DOUBLE_BLACK;   }
    boolean isNegativeBlack()    { return color() == NEGATIVE_BLACK; }
    boolean isEmptyDoubleBlack() { return false;                     }

    abstract RedBlackTree paintItRed();
    abstract RedBlackTree paintItBlack();

    abstract RedBlackTree blacker();
    abstract RedBlackTree redder();

    //----------------------------------------
    // the empty tree implemented by applying the singleton design pattern

    // the "generic" empty list object (with a raw type)

    private static final RedBlackTree EMPTY
        = new Empty();

    // the singleton class for the empty list object

    private static
        class Empty extends RedBlackTree {

        public boolean isEmpty() {
	    return true;
	}

        public int size() {
	    return 0;
	}

        public int depth() {
	    return 0;
	}

        public int minDepth() {
	    return 0;
	}

        public V lookup(K k) {
	    return null;
	}

	public KV findMin() {
	    return nodeExpected();
	}

	public KV findMax() {
	    return nodeExpected();
	}

        public Iterator<KV> iterator() {
            ++cntIter;
            return
                new NullIterator<KV>();
        }
        public Iterator<KV> iteratorDesc() {
            return
                iterator();
        }

        /* *
           public String toString() {
           return
           ".";
           }
           /* */

        //----------------------------------------
        // not public stuff

        Empty() { }

        boolean invRedWithBlackChildren()  { return true; }
        int     noOfBlackNodes()           { return 1;    }
        boolean invOrdered()               { return true; }
        boolean allLS(K k)                 { return true; }
        boolean allGT(K k)                 { return true; }

        int color() {
	    return BLACK;
	}

        RedBlackTree insert1(K k, V v) {
            return
                new Node(k, v, RED, EMPTY, EMPTY);
        }
        RedBlackTree remove1(K k) {
            return
                this;
        }
        RedBlackTree paintItBlack() {
            return
                empty(); // works also for DoubleEmpty
        }
        RedBlackTree paintItRed() {
            return
                undefined("paintItRed on empty tree");
        }
        RedBlackTree blacker() {
            return
                doubleEmpty();
        }
        RedBlackTree redder() {
            return
		undefined("redder of empty tree");
        }

    }

    //----------------------------------------
    // double black empty tree
    // helper for deletion

    private static
        RedBlackTree doubleEmpty() {
        return
	    DOUBLE_EMPTY;
    }

    private static final RedBlackTree DOUBLE_EMPTY
        = new DoubleEmpty();

    // the singleton class for the double black empty tree

    private static final
        class DoubleEmpty extends Empty {

        /* *
           public String toString() {
           return
           "!";
           }
           /* */

        DoubleEmpty() { }

        int color() {
	    return DOUBLE_BLACK;
	}

        boolean isEmptyDoubleBlack() {
	    return true;
	}

        RedBlackTree blacker() {
            return
                undefined("blacker of double black");
        }
        RedBlackTree redder() {
            return
                empty();
        }
    }

    //----------------------------------------
    // the node class for none empty trees

    private static final
        class Node extends RedBlackTree {

        /* persistent */
        final K            k;
        final V            v;
        final int          c; // the color
        final RedBlackTree l;
        final RedBlackTree r;

        /* destructive *
	K            k;
	V            v;
	int          c;
	RedBlackTree l;
	RedBlackTree r;
	/**/

        Node(K k, V v, int c, RedBlackTree l, RedBlackTree r) {
            assert k != null;

            this.k = k;
            this.v = v;
            this.c = c;
            this.l = l;
            this.r = r;
            ++cntNode;
        }

	/* destructive *
	public RedBlackTree copy() {
	    return
		new Node(k, v, c, l, r);
	}
	/**/

        /* trc *
           String toString() {
           return
           "(" + l + " " + k + "," + v + "," + c2s(c) + " " + r + ")";
           }
           /**/

        private static String c2s(int c) {
            switch ( c ) {
            case NEGATIVE_BLACK: return "N";
            case RED:            return "R";
            case BLACK:          return "B";
            case DOUBLE_BLACK:   return "D";
            }
            return "";
        }

        //----------------------------------------
        // predicates and attributes

        public boolean isEmpty() {
            return false;
        }

        public int size() {
            return
                1 + l.size() + r.size();
        }

        public int depth() {
            return
                1 + max(l.depth(), r.depth());
        }

        public int minDepth() {
            return
                1 + min(l.minDepth(), r.minDepth());
        }

        public V lookup(K k) {
            assert k != null;

            switch (signum(k.compareTo(this.k))) {
            case -1:
                return
                    l.lookup(k);
            case 0:
                return
                    v;
            case +1:
                return
                    r.lookup(k);
            }
            // never executed, but required
            return
                null;
        }

        //----------------------------------------
        // getter

        Node         node()  { return this; }
        int          color() { return c; }
	K            key()   { return k; }
	V            value() { return v; }
	RedBlackTree left()  { return l; }
	RedBlackTree right() { return r; }

        public KV findMin() {
            if ( l.isEmpty() )
                return
                    mkPair(k, v);
            return
                l.findMin();
        }

        public KV findMax() {
            if ( r.isEmpty() )
                return
                    mkPair(k, v);
            return
                r.findMax();
        }

        public Iterator<KV> iterator() {
            ++cntIter;
            return
                new NodeIterator(this);
        }

        public Iterator<KV> iteratorDesc() {
            ++cntIter;
            return
                new NodeIteratorDescending(this);
        }

        //----------------------------------------
        // invariant helpers

        boolean invRedWithBlackChildren() {
            return
                ( ! isRed()
                  ||
                  ( l.isBlack() && r.isBlack() ) )
                &&
                l.invRedWithBlackChildren()
                &&
                r.invRedWithBlackChildren();
        }

	// -1 indicates # of black nodes differ
        int noOfBlackNodes() {
            int nl = l.noOfBlackNodes();
            int nr = r.noOfBlackNodes();
            return
                (nl == nr && nl != -1)
                ? nl + toInt(isBlack())
                : -1;
        }

        boolean invOrdered() {
            return
                l.allLS(k)
                &&
                r.allGT(k)
                &&
                l.invOrdered()
                &&
                r.invOrdered();
        }

        boolean allLS(K k) {
            return
                this.k.compareTo(k) < 0
                &&
                r.allLS(k)
                // &&         // redundant if only
                // l.allLS(k) // used by inv
                ;
        }
        boolean allGT(K k) {
            return
                this.k.compareTo(k) > 0
                &&
                l.allGT(k)
                // &&         // redundant if only
                // r.allGT(k) // used by inv
                ;
        }


        //----------------------------------------
        // getter

        //----------------------------------------
        // recoloring

        RedBlackTree paintItBlack() {
            return
                setColor(BLACK);
        }
        RedBlackTree paintItRed() {
            return
                setColor(RED);
        }
        RedBlackTree blacker() {
            return
                setColor(blacker(color()));
        }
        RedBlackTree redder() {
            return
                setColor(redder(color()));
        }

        //----------------------------------------
        // balancing

        Node balance(int c, RedBlackTree t1, K k, V v, RedBlackTree t2) {
            if ( c == BLACK )
                return
                    balanceBlack(t1, k, v, t2);

	    // DOUBLE_BLACK only used in remove1
            if ( c == DOUBLE_BLACK )
                return
                    balanceDoubleBlack(t1, k, v, t2);

	    // no balancing at red nodes
            return
                set(k, v, c, t1, t2);
        }

        // Okasaki's rules for insertion balancing

        Node balanceBlack(RedBlackTree l, K k, V v, RedBlackTree r) {
            // pattern matching with java is pain in the ass

            if ( l.isRed() ) {
                Node ln = l.node();
                if ( ln.l.isRed() ) {
                    Node lln = ln.l.node();
                    return
                        build3(RED, lln.l, lln.k, lln.v,
                                    lln.r,  ln.k,  ln.v,
                                     ln.r,     k,     v, r, ln, lln);
                }
                if ( ln.r.isRed() ) {
                    Node lrn = ln.r.node();
                    return
                        build3(RED,  ln.l,  ln.k,  ln.v,
                                    lrn.l, lrn.k, lrn.v,
                                    lrn.r,     k,     v, r, ln, lrn);
                }
            }
            if ( r.isRed() ) {
                Node rn = r.node();
                if ( rn.l.isRed() ) {
                    Node rln = rn.l.node();
                    return
                        build3(RED,
			           l,     k,     v,
			       rln.l, rln.k, rln.v,
			       rln.r,  rn.k,  rn.v,
			        rn.r,
			       rln, rn);
                }
                if ( rn.r.isRed() ) {
                    Node rrn = rn.r.node();
                    return
                        build3(RED,
			           l,     k,     v,
			        rn.l,  rn.k,  rn.v,
			       rrn.l, rrn.k, rrn.v,
			       rrn.r,
			       rrn, rn);
                }
            }
            return
                set(k, v, BLACK, l, r);
        }

        Node balanceDoubleBlack(RedBlackTree l, K k, V v, RedBlackTree r) {
            // same rules as in balanceBlack, double black is absorbed

            if ( l.isRed() ) {
                Node ln = l.node();
                if ( ln.l.isRed() ) {
                    Node lln = ln.l.node();
                    return
                        build3(BLACK,
			       lln.l, lln.k, lln.v,
			       lln.r,  ln.k,  ln.v,
			        ln.r,     k,     v,
			           r,
			       ln, lln);
                }
                if ( ln.r.isRed() ) {
                    Node lrn = ln.r.node();
                    return
                        build3(BLACK,
			        ln.l,  ln.k,  ln.v,
			       lrn.l, lrn.k, lrn.v,
			       lrn.r,     k,     v,
			           r,
			       ln, lrn);
                }
            }
            if ( r.isRed() ) {
                Node rn = r.node();
                if ( rn.l.isRed() ) {
                    Node rln = rn.l.node();
                    return
                        build3(BLACK,
			           l,     k,     v,
			       rln.l, rln.k, rln.v,
			       rln.r,  rn.k,  rn.v,
			        rn.r,
			       rln, rn);
                }
                if ( rn.r.isRed() ) {
                    Node rrn = rn.r.node();
                    return
                        build3(BLACK,
			           l,     k,     v,
			        rn.l,  rn.k,  rn.v,
			       rrn.l, rrn.k, rrn.v,
			       rrn.r,
			       rrn, rn);
                }
            }

            // extra rules for negative black none empty trees
            if ( r.isNegativeBlack() ) {
                Node rn = r.node();
                if ( rn.l.isBlack()
                     &&
                     rn.r.isBlack() ) {
                    Node rln = rn.l.node();
                    return
                        rln.set(rln.k, rln.v, BLACK,
                                set(k, v, BLACK, l, rln.l),
                                rn.balanceBlack(rln.r, rn.k, rn.v, rn.r.paintItRed()));
                }
            }
            if ( l.isNegativeBlack() ) {
                Node ln = l.node();
                if ( ln.l.isBlack()
                     &&
                     ln.r.isBlack() ) {
                    Node lrn = ln.r.node();
                    return
                        lrn.set(lrn.k, lrn.v, BLACK,
                                ln.balanceBlack(ln.l.paintItRed(), ln.k, ln.v, lrn.l),
                                set(k, v, BLACK, l, lrn.r));
                }
            }
            return
                set(k, v, DOUBLE_BLACK, l, r);
        }

        Node bubble(int c, RedBlackTree l, K k, V v, RedBlackTree r) {
            if ( l.isDoubleBlack()
                 ||
                 r.isDoubleBlack() ) {
                return
                    balance(blacker(c), l.redder(), k, v, r.redder());
            }
            return
                balance(c, l, k, v, r);
        }

        // build 3 nodes from a b c and d with labels x y z
        // when a destrutive set is used
        // the root is stored in this
        // the left node is stored in l
        // the right node in r

        Node build3(int rootColor,
		    RedBlackTree a, K xk, V xv,
		    RedBlackTree b, K yk, V yv,
		    RedBlackTree c, K zk, V zv,
		    RedBlackTree d,
		    Node l, Node r) {
            return
                set(yk, yv, rootColor,
                    l.set(xk, xv, BLACK, a, b),
                    r.set(zk, zv, BLACK, c, d));
        }

        //----------------------------------------
        // setter

        private Node set(K k, V v, int c,
			 RedBlackTree l,
			 RedBlackTree r) {
            /* persistent */
            if ( k == this.k && v == this.v
                 &&
                 c == this.c
                 &&
                 l == this.l && r == this.r )
                return
                    this;
            return
                new Node(k, v, c, l, r);

            /* destructive *
	    this.k = k;
	    this.v = v;
	    this.c = c;
	    this.l = l;
	    this.r = r;
	    return
		this;
	    /**/
        }
        private Node setV(V v) {
            /* persistent */
            return
                ( v == this.v )
                ? this
                : new Node(this.k, v, this.c, this.l, this.r);

            /* destructive *
	    this.v = v;
	    return
		this;
	    /**/
        }

        private Node setColor(int c) {
            /* persistent */
            return
                c == this.c
                ? this
                : new Node(this.k, this.v, c, this.l, this.r);

            /* destructive *
               this.c = c;
               return
               this;
               /**/
        }

        //----------------------------------------
        // tree manipulation

        RedBlackTree insert1(K k, V v) {
            switch ( signum(k.compareTo(this.k)) ) {
            case -1:
                return
                    balance(c, l.insert1(k, v), this.k, this.v, r);
            case 0:
                return
                    setV(v);
            case +1:
                return
                    balance(c, l, this.k, this.v, r.insert1(k, v));
            }
            // never executed, but required
            return
                null;
        }

        RedBlackTree remove1(K k) {
            switch ( signum(k.compareTo(this.k)) ) {
            case -1:
                return
                    bubble(c, l.remove1(k), this.k, this.v, r);
            case 0:
                return
                    removeRoot();
                //trc("removeRoot",removeRoot());
            case +1:
                return
                    bubble(c, l, this.k, this.v, r.remove1(k));
            }
            // never executed, but required
            return
                null;
        }

        RedBlackTree removeRoot() {
            if ( isRed()
                 &&
                 l.isEmpty()
                 &&
                 r.isEmpty()
                 )
                return
                    empty();
            if ( isBlack() ) {
                if ( l.isEmpty()
                     &&
                     r.isEmpty() )
                    return
                        doubleEmpty();
                if ( l.isRed()
                     &&
                     r.isEmpty() )
                    return
                        l.blacker();
                if ( r.isRed()
                     &&
                     l.isEmpty() )
                    return
                        r.blacker();
            }
            KV p = l.findMax();
            return
                bubble(c, l.node().removeMax(), p.fst, p.snd, r);
        }

        RedBlackTree removeMax() {
            if ( r.isEmpty() )
                return
                    removeRoot();
            return
                bubble(c, l, k, v, r.node().removeMax());
        }

        //----------------------------------------
        // iterators

        // ascending enumeration

        private static
            class NodeIterator
            extends ds.util.Iterator<KV> {

            Queue queue;

            NodeIterator(Node n) {
                queue = Queue.empty();
                add(n);
                ++cntIter;
            }

            void add(Node n) {
                queue = queue.cons(n);
                if ( ! n.l.isEmpty() )
                    add(n.l.node()); // downcast
            }

            void addSubNodes(Node n) {
                if ( ! n.r.isEmpty() )
                    add(n.r.node());
            }

            public boolean hasNext() {
                return
                    ! queue.isEmpty();
            }

            public KV next() {
                if ( queue.isEmpty() )
                    return
                        undefined("empty tree");

                Node n = (Node)queue.head();
                queue  = queue.tail();
                addSubNodes(n);
                return
                    mkPair(n.k, n.v);
            }
        }

        // descending enumeration

        private
            class NodeIteratorDescending
            extends NodeIterator {

            NodeIteratorDescending(Node n) {
                super(n);
            }

            void add(Node n) {
                queue = queue.cons(n);
                if ( ! n.r.isEmpty() )
                    add(n.r.node()); // downcast
            }

            void addSubNodes(Node n) {
                if ( ! n.l.isEmpty() )
                    add(n.l.node());
            }
        }
    }

    //----------------------------------------
    // profiling

    private static int cntNode = 0;
    private static int cntIter = 0;

    public static String stats() {
        return
            "stats for ds.persistent.map.RedBlackTree:\n" +
            "# new Node()               : " + cntNode + "\n" +
            "# new Iterator()           : " + cntIter + "\n";
    }

    public String objStats() {
        int s = size();
        int o = s;
        int f = 5 * s;
        int m = o + f;

        return
            "mem stats for ds.persistent.map.RedBlackTree object:\n" +
            "# elements (size)      : " + s + "\n" +
            "# objects              : " + o + "\n" +
            "# fields               : " + f + "\n" +
            "# mem words            : " + m + "\n";
    }
}