/**
 * 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 BodyForceTreeMapItem 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 BodyForceTreeMapItem(key, value);
            size++;
            return null;
        }

        BodyForceTreeMapItem item = root;
        while (item != null) {
            if (item.key() == key) {
                Vector3 old = item.value();
                item.setValue(value);
                return old;
            } else if (item.key().mass() > key.mass()) {
                if (item.left() != null) {
                    item = item.left();
                } else {
                    item.setLeft(new BodyForceTreeMapItem(key, value));
                    size++;
                    break;
                }
            } else {
                if (item.right() != null) {
                    item = item.right();
                } else{
                    item.setRight(new BodyForceTreeMapItem(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) {
        BodyForceTreeMapItem 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) {
        BodyForceTreeMapItem 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(BodyForceTreeMapItem 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 toString(root);
    }
}

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

    public BodyForceTreeMapItem(Body key, Vector3 value) {
        this.key = key;
        this.value = value;
    }

    public Body key() {
        return this.key;
    }

    public void setValue(Vector3 value) {
        this.value = value;
    }

    public Vector3 value() {
        return this.value;
    }

    public BodyForceTreeMapItem left() {
        return this.left;
    }

    public BodyForceTreeMapItem right() {
        return this.right;
    }

    public BodyForceTreeMapItem parent() {
        return this.parent;
    }

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

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