package AVL;

import java.util.Iterator;

public class AVLTree<T> {

    private AVLNode<T> root;
    private int size;

    public AVLTree() {
        this.root = null;
        this.size = 0;
    }

    public void addElement(T element) {
        if (element == null)
            return;
        else
            root = addElementHelper(root, element);

    }

    private AVLNode<T> addElementHelper(AVLNode<T> newRoot, T element) {
        if (root == null) {
            size++;
            return new AVLNode<T>(element);
        }
        Comparable<T> compareElement = (Comparable<T>) element;
        if (compareElement.compareTo(newRoot.getElement()) < 0) {
            newRoot.setLeftLeaf(addElementHelper(newRoot.getLeftLeaf(), element));
            if (height(newRoot.getLeftLeaf()) - height(newRoot.getRightLeaf()) == 2) {
                newRoot = leftRotation(newRoot);
            } else {
                newRoot = rightLeftRotation(newRoot);
            }
        } else if (compareElement.compareTo(newRoot.getElement()) > 0) {
            newRoot.setRightLeaf(addElementHelper(newRoot.getRightLeaf(),
                    element));
            if (height(newRoot.getRightLeaf()) - height(newRoot.getLeftLeaf()) == 2) {
                if (height(newRoot.getRightLeaf())
                        - height(newRoot.getLeftLeaf()) == 2) {
                    newRoot = rightRotation(newRoot);
                } else {
                    newRoot = leftRightRotation(newRoot);
                }
            }
        }
        newRoot.setHeight(Math.max(height(newRoot.getLeftLeaf()),
                height(newRoot.getRightLeaf())) + 1);
        return newRoot;
    }

    public T find(T element) throws EmptyCollectionException,
            ElementNotFoundException {
        if (isEmpty()) {
            throw new EmptyCollectionException("AVLTree");
        }
        if (element == null) {
            return null;
        }
        AVLNode<T> resultsNode = findHelper(root, element);
        if (resultsNode == null) {
            throw new ElementNotFoundException("AVLTree");
        } else {
            return resultsNode.getElement();
        }
    }

    private AVLNode<T> findHelper(AVLNode<T> newRoot, T element) {
        if (newRoot == null)
            return null;
        @SuppressWarnings("unchecked")
        Comparable<T> comparableElement = (Comparable<T>) element;
        if (comparableElement.compareTo(newRoot.getElement()) < 0)
            return findHelper(newRoot.getLeftLeaf(), element);
        if (comparableElement.compareTo(newRoot.getElement()) > 0)
            return findHelper(newRoot.getRightLeaf(), element);
        else
            return newRoot;
    }

    public T removeElement(T targetElement) throws EmptyCollectionException {
        if (isEmpty())
            throw new EmptyCollectionException("AVL Tree");
        root = removeElementHelper(root, targetElement);
        return targetElement;
    }

    private AVLNode<T> removeElementHelper(AVLNode<T> newRoot, T element)
            throws ElementNotFoundException {
        if (newRoot == null) {
            throw new ElementNotFoundException("AVL Tree");
        }
        Comparable<T> comparableElement = (Comparable<T>) element;
        if (comparableElement.compareTo(newRoot.getElement()) < 0) {
            newRoot.setLeftLeaf(removeElementHelper(newRoot.getLeftLeaf(),
                    element));
            if (height(newRoot.getRightLeaf()) - height(newRoot.getLeftLeaf()) == 2) {
                if (height(newRoot.getRightLeaf().getRightLeaf()) >= height(newRoot
                        .getRightLeaf().getLeftLeaf()))
                    newRoot = rightRotation(newRoot);
                else
                    newRoot = leftRightRotation(newRoot);
            }
        } else if (comparableElement.compareTo(newRoot.getElement()) > 0) {
            newRoot.setRightLeaf(removeElementHelper(newRoot.getRightLeaf(),
                    element));
            if (height(newRoot.getLeftLeaf()) - height(newRoot.getRightLeaf()) == 2) {
                if (height(newRoot.getLeftLeaf().getLeftLeaf()) >= height(newRoot
                        .getLeftLeaf().getRightLeaf()))
                    newRoot = leftRotation(newRoot);
                else
                    newRoot = rightLeftRotation(newRoot);
            }
        } else {
            return getNewNode(newRoot);
        }
        newRoot.setHeight(Math.max(height(newRoot.getLeftLeaf()),
                height(newRoot.getRightLeaf())) + 1);
        return newRoot;
    }

    private AVLNode<T> getNewNode(AVLNode<T> node) {
        AVLNode<T> newNode;

        if (node.getLeftLeaf() != null && node.getRightLeaf() != null) {

            newNode = findMinHelper(node.getRightLeaf());
            node.setRightLeaf(removeElementHelper(node.getRightLeaf(),
                    newNode.getElement()));
            newNode.setLeftLeaf(node.getLeftLeaf());
            newNode.setRightLeaf(node.getRightLeaf());

            if (height(newNode.getLeftLeaf()) - height(newNode.getRightLeaf()) == 2) {

                if (height(newNode.getLeftLeaf().getLeftLeaf()) >= height(newNode
                        .getLeftLeaf().getRightLeaf()))
                    newNode = leftRotation(newNode);

                else
                    newNode = rightLeftRotation(newNode);
            }

            newNode.setHeight(Math.max(height(newNode.getLeftLeaf()),
                    height(newNode.getRightLeaf())) + 1);

        }

        else {
            if (node.getLeftLeaf() != null)
                newNode = node.getLeftLeaf();
            else
                newNode = node.getRightLeaf();

            size--;
        }

        node.setLeftLeaf(null);
        node.setRightLeaf(null);

        return newNode;
    }

    private AVLNode<T> findMinHelper(AVLNode<T> newNode) {
        if (newNode.getLeftLeaf() == null)
            return newNode;
        else
            return findMinHelper(newNode.getLeftLeaf());
    }

    private AVLNode<T> leftRotation(AVLNode<T> node) {

        AVLNode<T> newRoot = node.getLeftLeaf();
        node.setLeftLeaf(newRoot.getRightLeaf());
        newRoot.setRightLeaf(node);
        node.setHeight(Math.max(height(node.getLeftLeaf()),
                height(node.getRightLeaf())) + 1);
        newRoot.setHeight(Math.max(height(newRoot.getLeftLeaf()),
                node.getHeight()) + 1);
        return newRoot;
    }

    private AVLNode<T> rightRotation(AVLNode<T> node) {

        AVLNode<T> newRoot = node.getRightLeaf();
        node.setRightLeaf(newRoot.getLeftLeaf());
        newRoot.setLeftLeaf(node);
        node.setHeight(Math.max(height(node.getLeftLeaf()),
                height(node.getRightLeaf())) + 1);
        newRoot.setHeight(Math.max(node.getHeight(),
                height(newRoot.getRightLeaf())) + 1);
        return newRoot;

    }

    private AVLNode<T> rightLeftRotation(AVLNode<T> node) {

        node.setLeftLeaf(rightRotation(node.getLeftLeaf()));
        return leftRotation(node);
    }

    private AVLNode<T> leftRightRotation(AVLNode<T> node) {

        node.setRightLeaf(leftRotation(node.getRightLeaf()));
        return rightRotation(node);
    }

    public int height(AVLNode<T> newNode) {
        if (newNode == null)
            return -1;
        else
            return newNode.getHeight();

    }

    public boolean isEmpty() {
        if (size() == 0)
            return true;
        else
            return false;
    }

    public int size() {
        return size;
    }

}