/*
 * Implementation for the dictionary with an AVL Tree
 *
 * Restrict your changes to the bottom of the file, and be sure your
 *   operations run in the specified time.
 *
 * This inherits all the functionality and helper methods from the
 *   BinarySearchTree, so you only need to implement a new insert method
 *   that handles rotations to keep the tree balanced.
 * Do not forget to update the heights of your TreeNodes when you insert!
 *
 * You can implement AVL deletion for a small bonus if you wish, but it
 *   is not required.  Please indicate clearly as a comment at the start
 *   the method if you completed the AVL delete operator.
 */
public class AVLTree extends BinarySearchTree {

/*******************************************************************/
/* Restrict your changes to BELOW this line to ensure you do not   */
/* lose any marks unneccessarily.                                  */
/*******************************************************************/

  /*
   * Create a new entry in the tree preserving the sorted binary property.
   *   Remeber to fix up the heights, and possibly perform a single/double
   *   rotation.
   *
   * You need to implement this in O(log n) where n is the number of nodes in
   *   the tree.
   */
	public void insert(int key, String data) {
		int iteratorKey = 0;
		TreeNode tmpRoot = null;
		TreeNode iterator = null;
		TreeNode treeNode = new TreeNode(key, data);

		if (root == null) {
			root = treeNode;

			// Increment number of nodes
			numNodes++;
		} else {
			// Traverse tree, starting at root till element found or added
			iterator = root;
			while (true) {
				// Keep track parent node
				tmpRoot = iterator;
				iteratorKey = iterator.getKey();

				// Key already exist, replace
				if (iteratorKey == key) {
					iterator.setData(data);
					break;
				} else if (key < iteratorKey) {
					// Go left
					iterator = iterator.getLeft();

					// We're at the end of left the branch
					if (iterator == null) {
						tmpRoot.setLeft(treeNode);

						// Increment number of nodes
						numNodes++;

						// We're done
						break;
					}
				} else if (key > iteratorKey) {
					// Go right
					iterator = iterator.getRight();

					// We're at the end of right branch
					if (iterator == null) {
						tmpRoot.setRight(treeNode);

						// Increment number of nodes
						numNodes++;

						// We're done
						break;
					}
				}
			}
			// Set parent
			treeNode.setParent(tmpRoot);

			// Check tree balance
			doRebalance(treeNode);
		}
	}

	/**
	 * Rebalance tree if needed.
	 * @param theNode Node just inserted
	 */
	private void doRebalance(TreeNode theNode) {
		// Keep track of height
		int heightCounter = 0;

		// Traverse tree from bottom till root, update the height for
		// each node, and check balance. If unbalanced do rotation as needed
		TreeNode iterator = theNode;
		TreeNode child = theNode;
		while (iterator != root) {
			child = iterator;
			iterator = iterator.getParent();
			heightCounter++;

			// New height
			if (iterator.getHeight() <= heightCounter) {
				iterator.setHeight(heightCounter);
			}

			// Check node added is balanced, if not, then do rotation(s) as needed
			if (!isBalanced(iterator)) {
				doRotation(iterator, child, theNode);

				// We're done updating tree
				return;
			}
		}
	}

	/**
	 * Do rotation using m, n, and o as pointer for nodes. m, n, and o are pointer names as in lecture notes.
	 * @param m TreeNode
	 * @param n TreeNode
	 * @param theNode TreeNode
	 */
	private void doRotation(TreeNode m, TreeNode n, TreeNode theNode) {
		TreeNode o = null;

		// Figure out what o is
		if (n.getKey() < theNode.getKey()) {
			o = n.getRight();
		} else {
			o = n.getLeft();
		}

		// Do rotations
		// Check for m, n, o key values to determine if nodes are
		// in straight line or zig-zag pattern
		if ((m.getKey() < n.getKey()) && (n.getKey() < o.getKey())) {
			doSingleLeftRotation(m, n);
		} else if ((m.getKey() < n.getKey()) && (n.getKey() > o.getKey())) {
			doDoubleLeftRotation(m, n, o);
		} else if ((m.getKey() > n.getKey()) && (n.getKey() > o.getKey())) {
			doSingleRightRotation(m, n);
		} else if ((m.getKey() > n.getKey()) && (n.getKey() < o.getKey())) {
			doDoubleRightRotation(m, n, o);
		}
	}

	/**
	 * Do double left rotation.
	 * @param m TreeNode
	 * @param n TreeNode
	 * @param o TreeNode
	 */
	private void doDoubleLeftRotation(TreeNode m, TreeNode n, TreeNode o) {
		// Attach o's left child sub tree to m
		m.setRight(o.getLeft());
		if (o.getLeft() != null) {
			o.getLeft().setParent(m);
		}

		// Attach o's right child sub tree to n
		n.setLeft(o.getRight());
		if (o.getRight() != null) {
			o.getRight().setParent(n);
		}

		// Make o the parent of m and n
		o.setLeft(m);
		o.setRight(n);

		// We may have a new root
		if (m == root) {
			root = o;
			o.setParent(null);
		} else {
			// Reattach subtree
			if (m.getParent().getKey() < m.getKey()) {
				m.getParent().setRight(o);
			} else {
				m.getParent().setLeft(o);
			}

			// Set o's parent
			o.setParent(m.getParent());
		}

		// Make o the parent of m and n
		m.setParent(o);
		n.setParent(o);

		// Update heights
		// m decreases by 2, n by 1 and o increses by 2
		m.setHeight(m.getHeight() - 2);
		n.setHeight(n.getHeight() - 1);
		o.setHeight(o.getHeight() + 2);
	}

	/**
	 * Do double right rotation.
	 * @param m TreeNode
	 * @param n TreeNode
	 * @param o TreeNode
	 */
	private void doDoubleRightRotation(TreeNode m, TreeNode n, TreeNode o) {
		// Attach o's right child sub tree to m
		m.setLeft(o.getRight());
		if (o.getRight() != null) {
			o.getRight().setParent(m);
		}

		// Attach o's left child sub tree to n
		n.setRight(o.getLeft());
		if (o.getLeft() != null) {
			o.getLeft().setParent(n);
		}

		// Make o parent of m and n
		o.setLeft(n);
		o.setRight(m);

		// We may have a new root
		if (m == root) {
			root = o;
			o.setParent(null);
		} else {
			// Re-attach subtree
			if (m.getParent().getKey() > m.getKey()) {
				m.getParent().setLeft(o);
			} else {
				m.getParent().setRight(o);
			}

			// Set o's parent
			o.setParent(m.getParent());
		}

		// Make o the parent of m and n
		m.setParent(o);
		n.setParent(o);

		// Update heights
		// m decreases by 2, n by 1 and o increses by 2
		m.setHeight(m.getHeight() - 2);
		n.setHeight(n.getHeight() - 1);
		o.setHeight(o.getHeight() + 2);
	}

	/**
	 * Do single left rotation and update heights accordingly.
	 * @param m TreeNode
	 * @param n TreeNode
	 */
	private void doSingleLeftRotation(TreeNode m, TreeNode n) {
		// Attach n's left child sub tree to m
		m.setRight(n.getLeft());
		if (n.getLeft() != null) {
			n.getLeft().setParent(m);
		}

		// Attach m to n
		n.setLeft(m);

		// Reattach n subtree
		if (m == root) {
			// We have a new root
			root = n;
			n.setParent(null);
		} else {
			if (m.getParent().getKey() < m.getKey()) {
				m.getParent().setRight(n);
			} else {
				m.getParent().setLeft(n);
			}
			n.setParent(m.getParent());
		}

		// Update parent
		m.setParent(n);

		// Update heights
		// m decreases by 2 and n increases by 1
		m.setHeight(m.getHeight() - 2);
		n.setHeight(n.getHeight() + 1);
	}

	/**
	 * Do single right rotation and update heights accordingly.
	 * @param m TreeNode
	 * @param n TreeNode
	 */
	private void doSingleRightRotation(TreeNode m, TreeNode n) {
		// Attach n's right child sub tree to
		m.setLeft(n.getRight());
		if (n.getRight() != null) {
			n.getRight().setParent(m);
		}

		// Set n as parent of m
		n.setRight(m);

		// Reattach n sub tree
		if (m == root) {
			// New have a new root
			root = n;
			n.setParent(null);
		} else {
			if (m.getParent().getKey() < m.getKey()) {
				m.getParent().setRight(n);
			} else {
				m.getParent().setLeft(n);
			}
			n.setParent(m.getParent());
		}

		// Update parent
		m.setParent(n);

		// Update heights
		// m subtree decreases by 2 and n increases by 1
		m.setHeight(m.getHeight() - 2);
		n.setHeight(n.getHeight() + 1);
	}

	/**
	 * Get height of node.
	 * @param node TreeNode
	 * @return int
	 */
	private int getNodeHeight(TreeNode node) {
		if (node == null) {
			return -1;
		} else {
			return node.getHeight();
		}
	}

	/**
	 * Check if node is balanced. The balance is given by (1 + (RightSubTree's height)) - (1 + (LeftSubTree's height)).
	 * If difference is -1, 0 or 1 then node is balanced, else it's unbalanced.
	 * @param node TreeNode
	 * @return boolean
	 */
	private boolean isBalanced(TreeNode node) {
		int balance = (1 + getNodeHeight(node.getRight())) - (1 + getNodeHeight(node.getLeft()));
		switch (balance) {
			case - 1:
			case 0:
			case 1:
				return true;
		}
		return false;
	}

  /*
   * You do not need to implement (unless you are seeking the small bonus)
   */
	public boolean delete(int searchKey) {
		System.out.println("You are not required to implement deletion for an AVL tree");
		return false;
	}
}