package collection.mutable;

import java.util.ArrayList;

import java.util.Comparator;

public class TreeMap<K,V> extends AbstractSortedMap<K,V> {
    protected BalanceableBinaryTree<K,V> tree = new BalanceableBinaryTree<>();

    public TreeMap() {
        super();
        tree.addRoot(null);
    }
    public TreeMap(Comparator<K> comp) {
        super(comp);
        tree.addRoot(null);
    }

    public int size() { return (tree.size() - 1) / 2; }

    /**
     * Returns the entry with smallest key (if any).
     * @return
     */
    public Entry<K,V> firstEntry() {
        if (tree.isEmpty()) return null;
        return treeMin(root()).getElement();
    }

    /**
     * Returns the entry with greatest key (if any).
     * @return
     */
    public Entry<K,V> lastEntry() {
        if (tree.isEmpty()) return null;
        return treeMax(root()).getElement();
    }

    /**
     * Utility used when inserting a new entry at a leaf ot the tree.
     * @param p
     * @param entry
     */
    private void expandExternal(Position<Entry<K,V>> p, Entry<K,V> entry) {
        tree.set(p, entry);
        tree.addLeft(p, null);
        tree.addRight(p, null);
    }

    /**
     * Returns the position in p's subtree having given key (or else the terminal left).
     * @param p
     * @param key
     * @return
     */
    private Position<Entry<K,V>> treeSearch(Position<Entry<K,V>> p, K key) {
        if (tree.isExternal(p)) return p;
        int comp = compare(key, p.getElement());
        if (comp == 0) return p;
        else if (comp < 0) return treeSearch(tree.left(p), key);
        else return treeSearch(tree.right(p), key);
    }

    protected Position<Entry<K,V>> root() { return tree.root(); }

    /**
     * Returns the value associated with the specified key (or else null).
     * @param key
     * @return
     * @throws IllegalArgumentException
     */
    public V get(K key) throws IllegalArgumentException {
        checkKey(key);
        Position<Entry<K,V>> p = treeSearch(root(), key);
        rebalanceAccess(p);
        if (tree.isExternal(p)) return null;
        return p.getElement().getValue();
    }

    /**
     * Associates the given value with the given key, returning any overridden value.
     * @param key
     * @param value
     * @return
     * @throws IllegalArgumentException
     */
    public V put(K key, V value) throws IllegalArgumentException {
        checkKey(key);
        Entry<K,V> newEntry = new MapEntry<>(key, value);
        Position<Entry<K,V>> p = treeSearch(root(), key);
        if (tree.isExternal(p)) {
            expandExternal(p, newEntry);
            rebalanceInsert(p);
            return null;
        } else {
            V old = p.getElement().getValue();
            tree.set(p, newEntry);
            rebalanceAccess(p);
            return old;
        }
    }

    /**
     * Removes the entry having key k (if any) and returns its associated value.
     * @param key
     * @return
     * @throws IllegalArgumentException
     */
    public V remove(K key) throws IllegalArgumentException {
        checkKey(key);
        Position<Entry<K,V>> p = treeSearch(root(), key);
        if (tree.isExternal(p)) {
            rebalanceAccess(p);
            return null;
        } else {
            V old = p.getElement().getValue();
            if (tree.isInternal(tree.left(p)) && tree.isInternal(tree.right(p))) {
                Position<Entry<K,V>> replacement = treeMax(tree.left(p));
                tree.set(p, replacement.getElement());
                p = replacement;
            }
            Position<Entry<K,V>> leaf = (tree.isExternal(tree.left(p)) ? tree.left(p) : tree.right(p));
            Position<Entry<K,V>> sib = tree.sibling(leaf);
            tree.remove(leaf);
            tree.remove(p);
            rebalanceDelete(sib);
            return old;
        }
    }

    /**
     * Returns the position with the maximum key in subtree rooted at Position p.
     * @param p
     * @return
     */
    protected Position<Entry<K,V>> treeMax(Position<Entry<K,V>> p) {
        Position<Entry<K,V>> walk = p;
        while (tree.isInternal(walk))
            walk = tree.right(walk);
        return tree.parent(walk);
    }

    /**
     * Returns the position with the minimum key in subtree rooted at Position p.
     * @param p
     * @return
     */
    protected Position<Entry<K,V>> treeMin(Position<Entry<K,V>> p) {
        Position<Entry<K,V>> walk = p;
        while (tree.isInternal(walk))
            walk = tree.left(walk);
        return tree.parent(walk);
    }


    /**
     * Returns the entry with greatest key less than or equal to given key (if any).
     * @param key
     * @return
     * @throws IllegalArgumentException
     */
    public Entry<K,V> floorEntry(K key) throws IllegalArgumentException {
        checkKey(key);
        Position<Entry<K,V>> p = treeSearch(root(), key);
        if (tree.isInternal(p)) return p.getElement();
        while (!tree.isRoot(p)) {
            if (p == tree.right(tree.parent(p))) return tree.parent(p).getElement();
            else p = tree.parent(p);
        }
        return null;
    }

    /**
     * Returns the entry with smallest key greater than or equal to given key (if any).
     * @param key
     * @return
     * @throws IllegalArgumentException
     */
    public Entry<K,V> ceilingEntry(K key) throws IllegalArgumentException {
        checkKey(key);
        Position<Entry<K,V>> p = treeSearch(root(), key);
        if (tree.isInternal(p)) return p.getElement();
        while (!tree.isExternal(p)) {
            if (p == tree.left(tree.parent(p))) return p.getElement();
            else p = tree.right(p);
        }
        return null;

    }

    /**
     * Returns the entry with greatest key strictly less than given key (if any).
     * @param key
     * @return
     * @throws IllegalArgumentException
     */
    public Entry<K,V> lowerEntry(K key) throws IllegalArgumentException {
        checkKey(key);
        Position<Entry<K,V>> p = treeSearch(root(), key);
        if (tree.isInternal(p) && tree.isInternal(tree.left(p)))
            return treeMax(tree.left(p)).getElement();
        while (!tree.isRoot(p)) {
            if (p == tree.right(tree.parent(p))) return tree.parent(p).getElement();
            else p = tree.parent(p);
        }
        return null;
    }

    /**
     * Returns the entry with smallest key strictly greater than given key (if any).
     * @param key
     * @return
     * @throws IllegalArgumentException
     */
    public Entry<K,V> higherEntry(K key) throws IllegalArgumentException {
        checkKey(key);
        Position<Entry<K,V>> p = treeSearch(root(), key);
        if (tree.isInternal(p) && tree.isInternal(tree.right(p)))
            return treeMin(tree.right(p)).getElement();
        while (!tree.isExternal(p)) {
            if (p == tree.left(tree.parent(p))) return p.getElement();
            else p = tree.right(p);
        }
        return null;
    }

    /**
     * Returns an iterable collection of all key-value entries of the map.
     * @return
     */
    public Iterable<Entry<K,V>> entrySet() {
        ArrayList<Entry<K,V>> buffer = new ArrayList<>(size());
        for (Position<Entry<K,V>> p: tree.inorder())
            if (tree.isInternal(p)) buffer.add(p.getElement());
        return buffer;
    }

    /**
     * Returns an iterable of entries with keys in range [fromKey, toKey].
     * @param fromKey
     * @param toKey
     * @return
     */
    public Iterable<Entry<K,V>> subMap(K fromKey, K toKey) {
        ArrayList<Entry<K,V>> buffer = new ArrayList<>(size());
        if (compare(fromKey, toKey) < 0)
            subMapRecurse(fromKey, toKey, root(), buffer);
        return buffer;
    }

    private void subMapRecurse(K fromKey, K toKey, Position<Entry<K,V>> p, ArrayList<Entry<K,V>> buffer) {
        if (tree.isInternal(p))
            if (compare(p.getElement(), fromKey) < 0)
                subMapRecurse(fromKey, toKey, tree.right(p), buffer);
            else {
                subMapRecurse(fromKey, toKey, tree.left(p), buffer);
                if (compare(p.getElement(), toKey) < 0) {
                    buffer.add(p.getElement());
                    subMapRecurse(fromKey, toKey, tree.right(p), buffer);
                }
            }
    }

    protected boolean checkKey(K key) throws IllegalArgumentException {
        try {
            return (compare(key, key) == 0); // see if key can be compared to itself.
        } catch (ClassCastException e) {
            throw new IllegalArgumentException("Incompatible key");
        }
    }

    protected void rebalanceAccess(Position<Entry<K,V>> p) { }
    protected void rebalanceDelete(Position<Entry<K,V>> p) { }
    protected void rebalanceInsert(Position<Entry<K,V>> p) { }

    /**
     * A spefiablized version of LinkedBinaryTree with support for balancing.
     * @param <K>
     * @param <V>
     */
    protected static class BalanceableBinaryTree<K,V> extends LinkedBinaryTree<Entry<K,V>> {
        protected static class BSTNode<E> extends Node<E> {
            int aux = 0;
            BSTNode(E e, Node<E> parent, Node<E> leftChild, Node<E> rightChild) {
                super(e, parent, leftChild, rightChild);
            }
            public int getAux() { return aux; }
            public void setAux(int value) { aux = value; }
        }

        public int getAux(Position<Entry<K,V>> p) {
            return ((BSTNode<Entry<K,V>>)p).getAux();
        }
        public void setAux(Position<Entry<K,V>> p, int value) {
            ((BSTNode<Entry<K,V>>)p).setAux(value);
        }

        // Override with returning a BSTNode
        protected Node<Entry<K,V>> createNode(Entry<K,V> e, Node<Entry<K,V>> parent, Node<Entry<K,V>> left, Node<Entry<K,V>> right) {
            return new BSTNode<Entry<K,V>>(e, parent, left, right);
        }

        /**
         * Relink a parent node with its oriented child node.
         * @param parent
         * @param child
         * @param makeLeftChild
         */
        private void relink(Node<Entry<K,V>> parent, Node<Entry<K,V>> child, boolean makeLeftChild) {
            child.setParent(parent);
            if (makeLeftChild)
                parent.setLeft(child);
            else
                parent.setRight(child);
        }

        /**
         * Rotates Position p above its parent.
         * @param p
         */
        public void rotate(Position<Entry<K,V>> p) {
            Node<Entry<K,V>> x = validate(p);
            Node<Entry<K,V>> y = x.getParent();
            Node<Entry<K,V>> z = y.getParent();
            if (z == null) {
                root = x;
                x.setParent(null);
            } else
                relink(z, x, y == z.getLeft());
            if (x == y.getLeft()) {
                relink(y, x.getRight(), true);
                relink(x, y, false);
            } else {
                relink(y, x.getLeft(), false);
                relink(x, y, true);
            }
        }

        /**
         * Performs a trinode restructuring of Position x with its parent/grandparent.
         * @param x
         * @return
         */
        public Position<Entry<K,V>> restructure(Position<Entry<K,V>> x) {
            Position<Entry<K,V>> y = parent(x);
            Position<Entry<K,V>> z = parent(y);
            if ((x == right(y)) == (y == right(z))) {
                // Single rotation
                rotate(y);
                return y;
            } else {
                // Double rotation
                rotate(x);
                rotate(x);
                return x;
            }
        }
    }
}
