/**
 * A map that associates a Body with a Vector3 (typically this is the force exerted on the body).
 * The number of key-value pairs is not limited.
 */
public class BodyForceTreeMap {
    private int size = 0;
    private Item root = null;

    /**
     * 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(Body key, Vector3 value) {
        if (root == null) {
            root = new Item(key, value);
            size++;
            return null;
        }

        Item item = root;
        while (item != null) {
            if (item.key == 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(Body key) {
        Item item = root;
        while (item != null) {
            if (item.key == 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.
     */
    public boolean containsKey(Body key) {
        Item item = root;
        while (item != null) {
            if (item.key == key) {
                return true;
            } else if (item.key.mass() > key.mass()) {
                item = item.left;
            } else {
                item = item.right;
            }
        }
        return false;
    }

    public int size() {
        return this.size;
    }

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

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

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

        public Item(Body 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;
        }
    }
}