package lab7;

public class AVLTree {

    public TreeNode root;

    /**
     * Default constructor
     */
    public AVLTree() {
        root = null;
    }

    /**
     * Method to perform right rotation
     */
    private TreeNode rightRotate(TreeNode oldParent) {

        // YOUR CODE HERE// completed
        TreeNode y = oldParent.left;
        TreeNode y2 = y.right;

        y.right = oldParent;
        oldParent.left = y2;

        oldParent.height = Math.max(height(oldParent.left), height(oldParent.right)) + 1;
        y.height = Math.max(height(y.left), height(y.right)) + 1;
        return y;
    }

    /**
     * Method to perform left rotation
     */
    private TreeNode leftRotate(TreeNode oldParent) {

        // YOUR CODE HERE// completed
        TreeNode y = oldParent.right;
        TreeNode y2 = y.left;

        y.left = oldParent;
        oldParent.right = y2;
        oldParent.height = Math.max(height(oldParent.left), height(oldParent.right)) + 1;
        y.height = Math.max(height(y.left), height(y.right)) + 1;

        return y;
    }

    /**
     * Method to calculate the height of a tree
     */
    public int height() {
        return height(root);
    }

    // Recursive height
    private int height(TreeNode node) {

        // YOUR CODE HERE // completed
        if (node == null) {
            return 0;
        } else {
            int left = height(node.left);
            int right = height(node.right);

            if (left > right) {
                return left++;
            } else {
                return right++;
            }
        }

    }

    /**
     * Method to remove a key from the tree
     */
    public void remove(int key) {
        if (root == null) {
            System.out.println("Tree is empty.");
            return;
        }
        root = removeRecursive(root, key);
    }

// Recursive remove
    public TreeNode removeRecursive(TreeNode node, int key) {

        // YOUR CODE HERE
        if (node == null) {
            return node;
        }
        if (key < node.data) {
            node.left = removeRecursive(node.left, key);
        } else if (key > node.data) {
            node.right = removeRecursive(node.right, key);
        } else {
            if (node.left == null || node.right == null) {
                TreeNode t;
                if (node.left != null) {
                    t = node.left;
                } else {
                    t = node.right;
                }
                if (t == null) {
                    t = node;
                    node = null;
                } else {
                    node = t;

                }
                t = null;

            } else {
                TreeNode t = smallest(node.right);
                node.data = t.data;
                node.right = removeRecursive(node.right, t.data);

            }
        }
        if (node == null) {
            return node;
        }
        node.height = Math.max(height(node.left), height(node.right)) + 1;

        int balance = getBalance(node);

        if (balance > 1 && getBalance(node.left) >= 0) {
            return rightRotate(node);
        }
        if (balance > 1 && getBalance(node.left) < 0) {
            node.left = leftRotate(node.left);
            return rightRotate(node);
        }

        if (balance < -1 && getBalance(node.right) <= 0) {
            return leftRotate(node);
        }
        if (balance < -1 && getBalance(node.right) > 0) {
            node.right = rightRotate(node.right);
            return leftRotate(node);
        }

        return node;

    }

    /**
     * Method to add an item to the tree
     */
    public void add(int key) {
        root = addRecursive(root, key);
    }

    // Recursive add
    public TreeNode addRecursive(TreeNode node, int key) {

        // YOUR CODE HERE// complete
        if (node == null) {
            return new TreeNode(key);
        }
        if (key < node.data) {
            node.left = addRecursive(node.left, key);
        } else {
            node.right = addRecursive(node.right, key);
        }

        node.height = Math.max(height(node.left), height(node.right)) + 1;

        int balance = getBalance(node);
        if (balance > 1 && key < node.left.data) {
            return rightRotate(node);
        }

        if (balance < -1 && key > node.right.data) {
            return leftRotate(node);
        }

        if (balance > 1 && key > node.left.data) {
            node.left = leftRotate(node.left);
            return rightRotate(node);
        }

        if (balance < -1 && key < node.left.data) {
            node.right = leftRotate(node.right);
            return leftRotate(node);
        }

        return node;
    }

    /**
     * Method to calculate the balance factor of node N
     */
    int getBalance(TreeNode N) {

        // YOUR CODE HERE // completed
        if (N == null) {
            return 0;
        }
        return height(N.left) - height(N.right);
    }

    //------------------------------DO NOT CHANGE THE CODE BELOW------------------------------
    /**
     * A method to find the smallest item from the tree
     */
    private TreeNode smallest(TreeNode node) {
        if (node != null && node.left != null) {
            return smallest(node.left);
        }
        return node;
    }

    /**
     * Method to show the tree
     */
    @Override
    public String toString() {

        if (root == null) {
            return "Tree is empty";
        }

        return toString(root);
    }

    public String toString(TreeNode current) {
        String resutls = "";

        if (current != null) {
            String left = toString(current.left);
            String right = toString(current.right);
            if (left != "") {
                resutls += "(" + left + ")";
            }
            resutls += current.data;
            if (right != "") {
                resutls += "(" + right + ")";
            }
        }
        return resutls;
    }

}