public class AVLTree implements IAVLTree {

    private AVLNode root;

    /**
     * Node that might need restructuring after a remove, set in removeBST()
     */
    private AVLNode restructureNodeOnRemove = null;
    private int size = 0;


    /**
     * Default constructor. Initializes the AVL tree.
     */
    public AVLTree() {
        root = null;
    }

    /**
     * @returns the root of the AVLTree
     */
    public AVLNode getTreeRoot() {
        if (root == null) {
            return null;
        }
        AVLNode currNode = root;
        //iterating until the current AVL node has no parent node, so until the node is a root node
        while(currNode.parent != null) {
            currNode = currNode.parent;
        }
        //returning the root node of the tree
        return currNode;

    }


    /**
     * Retrieves tree height.
     *
     * @return -1 in case of empty tree, current tree height otherwise.
     */
    public int getTreeHeight() {
        //if there is no root, meaning tree is empty then height is -1
        if(root == null){
            return -1;
        }else{
            //otherwise the height the root is the height of the tree
            return root.height;
        }
    }

    /**
     * Yields number of key/value pairs in the tree.
     *
     * @return Number of key/value pairs.
     */
    public int getSize() {
        return size;
    }



    /**
     * Returns value of node with given key.
     *
     * @param key  Key to search.
     * @return 	  Corresponding value if key was found, -1 otherwise.
     */
    public float findByKey(int key) throws IllegalArgumentException {
        AVLNode current = root;
        while (current != null) {
            int cmp = Integer.compare(current.key, key);
            if (cmp == 0)
                return current.value;
            else if (cmp < 0)
                current = current.right;
            else
                current = current.left;
        }
        return -1f;
    }

    /**
     * Inserts a new node into AVL tree.
     *
     * @param key Key of the new node.
     * @param value Data of the new node. Must be >= 0.
     * @return   Nodes with the same key are not allowed. In this case false is returned. Otherwise true.
     */
    public boolean insertNode(int key, float value) throws IllegalArgumentException {
        if(value < 0  || findByKey(key) != -1f) {
            //returning false if node has a key that already exists or value is < 0
            return false;
        }
        AVLNode newNode = null;
        if (root == null){
                root = new AVLNode(key, value);
                updateHeights(root);
                size++;
                return true;
        }else {
            AVLNode current = root;
            while (true) {
                    int cmp = Integer.compare(current.key, key);
                    if (cmp == 0)
                        return false;
                    else if (cmp < 0) {
                        if (current.right != null)
                            current = current.right;
                        else {
                            newNode = new AVLNode(key, value);
                            size++;
                            setRight(current,newNode);
                            newNode.parent = current;
                            break;
                        }
                    } else {
                        if (current.left != null)
                            current = current.left;
                        else {
                            newNode = new AVLNode(key, value);
                            size++;
                            setLeft(current, newNode);
                            newNode.parent = current;

                            break;
                        }
                    }
                }
            }

        AVLNode n = newNode;
        updateHeights(n);

        while (n != null) {
            AVLNode x = n;
            if (x.parent != null && x.parent.parent != null && !isBalanced(x.parent.parent)) {
                AVLNode z = x.parent.parent;
                AVLNode y = x.parent;
                restructure(x, y, z);
                updateHeights(n);
            }else {
                n = n.parent;
            }
        }
        return true;
    }

    /**
     * Helper method responsible for restructuring the tree after the insert or remove operation
     * using rotation method
     *
     * @param x Node x
     * @param y Node y
     * @param z Node z
     */
    public void restructure(AVLNode x, AVLNode y, AVLNode z){
        NodeGroup g = getComponents(x,y,z);
        //Getting a NodeGroup object g that for x,y,z returns a,b,c node and subtrees t0,t1,t2 and t3
        g.b.left = g.a;
        g.b.right = g.c;

        g.a.left = g.t0;
        g.a.right = g.t1;

        g.c.left = g.t2;
        g.c.right = g.t3;

        g.a.parent = g.b;
        g.c.parent = g.b;

        if (z.parent != null) {
            if (z.parent.left == z) {
                z.parent.left = g.b;
            } else {
                z.parent.right = g.b;
            }
        } else {
            root = g.b;
        }
        g.b.parent = z.parent;

        updateHeights(g.a);
        updateHeights(g.b);
        updateHeights(g.c);
    }

    /**
     * Method to update the height of a node
     *
     * @param node AVLNode node
     */
    public void updateHeights(AVLNode node){
        if (node == null) {
            return;
        }

        //setting current the value of AVLNode parameter
        AVLNode current = node;
        int oldHeight = current.height;
        AVLNode leftNode = current.left;
        AVLNode rightNode = current.right;

        if (leftNode == null && rightNode == null) {
            current.height = 0;
        } else if (leftNode == null) {
            current.height = 1 + rightNode.height;
        } else if (rightNode == null) {
            current.height = 1 + leftNode.height;
        } else {
            current.height = 1 + Math.max(current.left.height, current.right.height);
        }

        //if current node has changed and its parent is not empty go up recursively and update the height
        if (current.height != oldHeight && current.parent != null) {
            updateHeights(current.parent);
        }
    }

    /**
     * Method to check if a node is balanced or not
     *
     * @param node AVLNode node
     * @return true if balanced
     */
    public boolean isBalanced(AVLNode node){
        int diff;
        if(node.left == null){
            //if left child of the node has is empty then set height to -1
            int lHeight = -1;
            diff = node.right.height - lHeight;
        } else if (node.right == null) {
            //if right child of the node has is empty then set height to -1
            int rHeight = -1;
            diff = node.left.height - rHeight;
        } else if (node.right == null && node.left == null) {
            //if no child node than height difference is 0
            diff = 0;
        } else{
            diff = node.left.height - node.right.height;
        }
        //if difference <=1 then true is returned, otherwise false
        return Math.abs(diff) <= 1;
	}

    /**
     * Removes node with given key.
     *
     * @param key Key of node to remove.
     * @return true if node was found and deleted.
     */
    public boolean removeNode(int key) throws IllegalArgumentException {
        AVLNode parent = null;
        AVLNode current = root;
        AVLNode newSubRoot = null;

        while (current != null) {
            int cmp = Integer.compare(current.key, key);
            if (cmp == 0) {
                //String prev = toString(parent);
                if (parent == null) {
                    root = removeBST(current);
                    if (root != null)
                        root.parent = null;
                } else if (parent.left == current) {
                    newSubRoot = removeBST(current);
                    setLeft(parent, newSubRoot);
                } else if (parent.right == current) {
                    newSubRoot = removeBST(current);
                    setRight(parent, newSubRoot);
                } else throw new IllegalArgumentException();

                // restructureNodeOnRemove is the node from which the search for the first unbalanced node is started
                if (restructureNodeOnRemove != null) {
                    //TODO update heights
                    updateHeights(restructureNodeOnRemove);
                    while (restructureNodeOnRemove != null) {
                        //TODO restructure until tree is balanced
                        if (!isBalanced(restructureNodeOnRemove)) {
                            AVLNode z = restructureNodeOnRemove;
                            AVLNode y = (z.left != null && (z.right == null || z.left.height > z.right.height)) ? z.left : z.right;
                            AVLNode x = (y.left != null && (y.right == null || y.left.height > y.right.height)) ? y.left : y.right;
                            restructure(x, y, z);}
                            size--;
                            restructureNodeOnRemove = restructureNodeOnRemove.parent;
                    }
                }

                return true;
            } else {
                parent = current;
                if (cmp > 0) {
                    current = current.left;
                } else {
                    current = current.right;
                }
            }

        }
        return false;
    }

    //auxiliary functions

    private AVLNode removeBST(AVLNode oldSubRoot) {
        AVLNode newSubRoot;
        if (oldSubRoot.left == null && oldSubRoot.right == null) {
            // case 1: no children, trivial
            newSubRoot = null;
            restructureNodeOnRemove = oldSubRoot.parent;
        } else if (oldSubRoot.left == null) {
            // case 2a: only (alone) right child
            newSubRoot = oldSubRoot.right;
            restructureNodeOnRemove = newSubRoot;
        } else if (oldSubRoot.right == null) {
            // case 2b: only (alone) left child
            newSubRoot = oldSubRoot.left;
            restructureNodeOnRemove = newSubRoot;
        } else if (oldSubRoot.left.right == null) {
            // case 3a: 2 children, left has no right child
            newSubRoot = oldSubRoot.left;
            setRight(newSubRoot, oldSubRoot.right);
            restructureNodeOnRemove = newSubRoot; //.right
        } else if (oldSubRoot.right.left == null) {
            // case 3b: 2 children, right has no left child
            newSubRoot = oldSubRoot.right;
            setLeft(newSubRoot, oldSubRoot.left);
            restructureNodeOnRemove = newSubRoot; //.left
        } else {// case 4: 2 children, both have inner side children
            for (newSubRoot = oldSubRoot.left; newSubRoot.right != null; newSubRoot = newSubRoot.right)
                ;
            AVLNode predecessorP = newSubRoot.parent;
            setRight(predecessorP, newSubRoot.left);
            setRight(newSubRoot, oldSubRoot.right);
            setLeft(newSubRoot, oldSubRoot.left);
            restructureNodeOnRemove = predecessorP;
        }
        return newSubRoot;
    }

    private void setLeft(AVLNode parent, AVLNode child) {
        parent.left = child;
        if (child != null)
            child.parent = parent;
    }

    private void setRight(AVLNode parent, AVLNode child) {
        parent.right = child;
        if (child != null)
            child.parent = parent;
    }

    private NodeGroup getComponents(AVLNode x, AVLNode y, AVLNode z) {
        NodeGroup g = new NodeGroup();
        if (z.left == y) { // left version
            g.c = z; // in every left case
            if (y.left == x) { // single rot
                g.a = x;
                g.b = y;
                g.t1 = g.a.right;
            } else {// double rot
                g.a = y;
                g.b = x;
                g.t1 = g.b.left;
            }
            g.t2 = g.b.right; // in every left case
        } else { // right version
            g.a = z; // in every right case
            if (y.right == x) { // single rot
                g.b = y;
                g.c = x;
                g.t2 = g.c.left;
            } else {// double rot
                g.b = x;
                g.c = y;
                g.t2 = g.b.right;
            }
            g.t1 = g.b.left; // in every right case
        }
        g.t0 = g.a.left; // in every case
        g.t3 = g.c.right; // in every case
        return g;
    }

}