import codedraw.CodeDraw;

import java.util.NoSuchElementException;

/**
 * A map that associates an object of 'Massive' with a Vector3. The number of key-value pairs
 * is not limited.
 */
public class MassiveForceTreeMap implements MassiveSet {
    private int size = 0;
    private Item root;

    /**
     * Adds a new key-value association to this map. If the key already exists in this map,
     * the value is replaced and the old value is returned. Otherwise 'null' is returned.
     * Precondition: key != null.
     */
    public Vector3 put(Massive key, Vector3 value) {
        if (root == null) {
            root = new Item(key, value);
            size++;
            return null;
        }

        Item item = root;
        while (item != null) {
            if (item.key.equals(key)) {
                Vector3 old = item.value;
                item.value = value;
                return old;
            } else if (item.key.mass() > key.mass()) {
                if (item.left != null) {
                    item = item.left;
                } else {
                    item.setLeft(new Item(key, value));
                    size++;
                    break;
                }
            } else {
                if (item.right != null) {
                    item = item.right;
                } else{
                    item.setRight(new Item(key, value));
                    size++;
                    break;
                }
            }
        }

        return null;
    }

    /**
     * Returns the value associated with the specified key, i.e. the method returns the force vector
     * associated with the specified key. Returns 'null' if the key is not contained in this map.
     * Precondition: key != null.
     */
    public Vector3 get(Massive key) {
        Item item = root;
        while (item != null) {
            if (item.key.equals(key)) {
                return item.value;
            } else if (item.key.mass() > key.mass()) {
                item = item.left;
            } else {
                item = item.right;
            }
        }
        return null;
    }

    /**
     * Returns 'true' if this map contains a mapping for the specified key.
     * Precondition: key != null
     */
    public boolean containsKey(Massive key) {
        Item item = root;
        while (item != null) {
            if (item.key.equals(key)) {
                return true;
            } else if (item.key.mass() > key.mass()) {
                item = item.left;
            } else {
                item = item.right;
            }
        }
        return false;
    }

    private String toString(Item item) {
        StringBuilder s = new StringBuilder();
        if (item == null) {
            return s.toString();
        }
        s.append(this.toString(item.right));
        s.append(String.format("{%s: %s}\n", item.key, item.value));
        s.append(this.toString(item.left));
        return s.toString();
    }

    /**
     * Returns a readable representation of this map, in which key-value pairs are ordered
     * descending according to 'key.getMass()'.
     */
    public String toString() {
        return (root != null) ? toString(root) : "";
    }

    /**
     * Returns a `MassiveSet` view of the keys contained in this tree map. Changing the
     * elements of the returned `MassiveSet` object also affects the keys in this tree map.
     */
    public MassiveSet getKeys() {
        return this;
    }

    @Override
    public void draw(CodeDraw cd) {

    }

    @Override
    public MassiveIterator iterator() {
        return new MassiveIterator() {
            private Item last = null;
            private Item next = root.getLeftLeaf();

            @Override
            public Massive next() {
                if (!hasNext()) throw new NoSuchElementException();

                last = next;
                Item newNext = (next.right != null) ? next.right.getLeftLeaf() : next.parent;

                while (newNext != null && newNext.right == next) {
                    next = newNext;
                    newNext = newNext.parent;
                }

                next = newNext;
                return last.key;
            }

            @Override
            public boolean hasNext() {
                return next != null;
            }

            @Override
            public void remove() {
                if (last == null) throw new IllegalStateException();
                removeItem(last);
                last = null;
            }
        };
    }

    @Override
    public boolean contains(Massive element) {
        return containsKey(element);
    }

    @Override
    public void remove(Massive element) {
        Item item = root;
        while (item != null) {
            if (item.key.equals(element)) {
                removeItem(item);
                return;
            } else if (item.key.mass() > element.mass()) {
                item = item.left;
            } else {
                item = item.right;
            }
        }
        throw new NoSuchElementException();
    }

    private void removeItem(Item item) {
        size--;
        Item newP = null;

        if (item.left != null) {
            newP = item.left.getRightLeaf();
            if (newP != item.left) {
                newP.parent.setRight(newP.left);
                newP.setLeft(item.left);
            }
            newP.setRight(item.right);
        } else if (item.right != null) {
            newP = item.right.getLeftLeaf();
            if (newP != item.right) {
                newP.parent.setLeft(newP.right);
                newP.setRight(item.right);
            }
            newP.setLeft(item.left);
        }

        if (newP == null) {
            root = null;
            return;
        }

        if (item.parent != null) {
            if (item.parent.left == item) {
                item.parent.setLeft(newP);
            } else {
                item.parent.setRight(newP);
            }
        } else {
            root = newP;
            newP.parent = null;
        }
    }

    @Override
    public void clear() {
        size = 0;
        root = null;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public MassiveLinkedList toList() {
        MassiveLinkedList list = new MassiveLinkedList();
        for (Massive m : this) {
            list.addLast(m);
        }
        return list;
    }

    private static class Item {
        private final Massive key;
        private Vector3 value;
        private Item parent;
        private Item left;
        private Item right;

        public Item(Massive key, Vector3 value) {
            this.key = key;
            this.value = value;
        }

        public void setLeft(Item left) {
            this.left = left;
            if (left != null) left.parent = this;
        }

        public void setRight(Item right) {
            this.right = right;
            if (right != null) right.parent = this;
        }

        public Item getLeftLeaf() {
            Item cur = this;
            while (cur.left != null) cur = cur.left;
            return cur;
        }

        public Item getRightLeaf() {
            Item cur = this;
            while (cur.right != null) cur = cur.right;
            return cur;
        }
    }
}
