/**
 *****************************************************************************
 * file name : AVLTree.cpp
 * author    : Hung Ngo
 * description : implement selected functions from the AVLTree interface
 *****************************************************************************
 */
#include <iostream>
#include <stdexcept>
using namespace std;

/**
 * -----------------------------------------------------------------------------
 *  remove node with given key, return true if found and removed
 *  false otherwise
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
bool AVLTree<Key, Value>::remove(Key key) {
    cout << "TBD\n"; // perhaps a homework problem
    return false;
}

/**
 * -----------------------------------------------------------------------------
 *  return a node with a given key, NULL if not found
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
typename AVLTree<Key, Value>::AVLNode* 
AVLTree<Key, Value>::search(
        typename AVLTree<Key, Value>::AVLNode* root, 
        Key key
        ) 
{
    AVLNode* node = root;
    while (node != NULL && node->key != key) {
        if (key < node->key) node = node->left;
        else node = node->right;
    }
    return node;
}

/**
 * -----------------------------------------------------------------------------
 *  insert a new key, value pair, return 
 *    true if the insertion was successful
 *    false if the key is found already
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
bool AVLTree<Key, Value>::insert(Key key, Value value) {
    AVLNode* p   = NULL;
    AVLNode* cur = root;
    while (cur != NULL) {
        p = cur;
        if (key < cur->key) 
            cur = cur->left;
        else if (key > cur->key)
            cur = cur->right;
        else
            return false; // key found, no insertion
    }

    // insert new node at a leaf position
    AVLNode* node = new AVLNode(key, value);
    node->parent = p;
    if (p == NULL) // empty tree to start with
        root = node; 
    else if (node->key < p->key)
        p->left = node;
    else
        p->right = node;

    // readjust balance of all nodes up to the root if necessary
    rebalance(node);
    return true;
}

/**
 * -----------------------------------------------------------------------------
 *  left rotate around node c
 *                       c             b
 *                      / \           / \
 *                     C   b  -->    c   B
 *                        / \       / \
 *                       A   B     C   A
 * adjust parent pointers accordingly
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
void AVLTree<Key, Value>::left_rotate(AVLNode*& node) {
    if (node == NULL || node->right == NULL) return;

    AVLNode* c = node;
    AVLNode* b = c->right;
    AVLNode* p = c->parent;

    // first, adjust all parent pointers
    b->parent = p;
    c->parent = b;
    if (b->left != NULL) b->left->parent = c;

    // make sure c's parent points to b now
    if (p != NULL) {    
        if (p->right == c) p->right = b;
        else p->left = b;
    }

    // finally, adjust downward pointers
    c->right = b->left;
    b->left  = c;

    node = b;                // new local root
    if (root == c) root = b; // new root if necessary
}


/**
 * -----------------------------------------------------------------------------
 *  right rotate around node c
 *                       c             b
 *                      / \           /  \
 *                     b   C  -->    A    c
 *                    / \                / \
 *                   A   B              B   C
 * adjust parent pointers accordingly
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
void AVLTree<Key, Value>::right_rotate(AVLNode*& node) {
    if (node == NULL || node->left == NULL) return;

    AVLNode* c = node;
    AVLNode* b = c->left;
    AVLNode* p = c->parent;

    // first, adjust all parent pointers
    b->parent = p;
    c->parent = b;
    if (b->right != NULL) b->right->parent = c;

    // next, make sure c's parent points to b
    if (p != NULL) {    // make sure c's parent points to b now
        if (p->right == c) p->right = b;
        else p->left = b;
    }

    // finally, adjust the downward pointers
    c->left  = b->right;
    b->right = c;

    node = b;                // new local root
    if (root == c) root = b; // new root if necessary
}

/**
 * -----------------------------------------------------------------------------
 *  node points to the root of a sub-tree which just had a height increase
 *  we assume the invariance that node's parent is not unbalanced, which
 *  certainly holds for the first node that got inserted
 *  the 'balance' field of the parent may not be correct and we have to
 *  adjust that too
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
void AVLTree<Key, Value>::rebalance(AVLNode* node) {
    AVLNode* p = node->parent;
    if (p == NULL) return;
    
    // first, recompute 'balance' of the parent; node got a heigh increase
    if (p->left == node) 
        p->balance++;
    else 
        p->balance--;

    // if there's no grandparent or if the parent is balanced then we're done
    AVLNode* gp = p->parent; // the grand parent
    if (gp == NULL || p->balance == AVLNode::BALANCED) return;

    // if we get here then the parent p just got a height increase
    // next, see if the grand parent is unbalanced
    if (node == p->left) {
        if (p == gp->left) {
            if (gp->balance == AVLNode::LEFT_HEAVY) { 
                // this is the LL case
                //        gp(+2)          p (0)
                //       /   \           /  \
                //      p(+1) B  -->   node  gp (0)
                //     / \                  / \
                //   node A                A   B
                p->balance = gp->balance = AVLNode::BALANCED;
                right_rotate(gp);
                return;
            }
        } else { // p == gp->right
            if (gp->balance == AVLNode::RIGHT_HEAVY) { // the RL case
                // this is the RL case
                //        gp(-2)               node(0)
                //       /   \                 /   \
                //      A    p(+1)    -->     gp(x) p(y)
                //           / \              /\   / \
                //         node D            A  B C   D
                //         / \
                //        B   C
                //  computing the new balance is a little trickier, depending on
                //  with of B & C is heavier
                switch (node->balance) {
                case AVLNode::LEFT_HEAVY:
                    p->balance  = AVLNode::RIGHT_HEAVY;
                    gp->balance = AVLNode::BALANCED;
                    break;
                case AVLNode::BALANCED: // only happens if B & C are NULL
                    p->balance  = AVLNode::BALANCED;
                    gp->balance = AVLNode::BALANCED;;
                    break;
                case AVLNode::RIGHT_HEAVY:
                    p->balance  = AVLNode::BALANCED;;
                    gp->balance = AVLNode::LEFT_HEAVY;
                    break;
                }
                node->balance = AVLNode::BALANCED;
                right_rotate(p);
                left_rotate(gp);
                return;
            }
        }
    } else { // node == p->right
        if (p == gp->right) {
            if (gp->balance == AVLNode::RIGHT_HEAVY) {
                // this is the RR case
                //        gp(-2)              p(0)
                //       /   \               /  \
                //      A   p(-1)    -->  gp(0) node
                //           / \          / \   / \
                //          B  node      A   B
                p->balance = gp->balance = AVLNode::BALANCED;
                left_rotate(gp);
                return;
            } 
        } else { // p == gp->left
            if (gp->balance == AVLNode::LEFT_HEAVY) {
                // this is the LR case
                //        gp(+2)          node(0)
                //       /   \            /   \
                //     p(-1)  D   -->   p(x)   gp(y)
                //     /  \             /\     / \
                //    A    node        A  B   C   D
                //         / \
                //        B   C
                //  computing the new balance is a little trickier, depending on
                //  with of B & C is heavier
                switch (node->balance) {
                case AVLNode::LEFT_HEAVY:
                    p->balance  = AVLNode::BALANCED;
                    gp->balance = AVLNode::RIGHT_HEAVY;
                    break;
                case AVLNode::BALANCED: // only happens if B & C are NULL
                    p->balance  = AVLNode::BALANCED;
                    gp->balance = AVLNode::BALANCED;;
                    break;
                case AVLNode::RIGHT_HEAVY:
                    p->balance  = AVLNode::LEFT_HEAVY;
                    gp->balance = AVLNode::BALANCED;;
                    break;
                }
                node->balance = AVLNode::BALANCED;
                left_rotate(p);
                right_rotate(gp);
                return;
            } 
        }
    }

    rebalance(p);
}

/**
 * -----------------------------------------------------------------------------
 *  print all keys inorder, kind of bad design here because we assume "keys"
 *  can be cout; but this is great for testing;
 *  a better design would need inorder to return iterators and the client uses
 *  the (inorder) iterators to print;
 *  this is sufficient for our illustrative purposes
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
void AVLTree<Key, Value>::inorder_print(AVLNode* node) {
    if (node != NULL) {
        inorder_print(node->left);
        cout << node->to_string() << " ";
        inorder_print(node->right);
    }
}

/**
 * -----------------------------------------------------------------------------
 *  postorder print all nodes
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
void AVLTree<Key, Value>::postorder_print(AVLNode* node) {
    if (node != NULL) {
        postorder_print(node->left);
        postorder_print(node->right);
        cout << node->to_string() << " ";
    }
}

/**
 * -----------------------------------------------------------------------------
 *  preorder print all nodes
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
void AVLTree<Key, Value>::preorder_print(AVLNode* node) {
    if (node != NULL) {
        cout << node->to_string() << " ";
        preorder_print(node->left);
        preorder_print(node->right);
    }
}

/**
 * -----------------------------------------------------------------------------
 *  clear all nodes, release memory
 * -----------------------------------------------------------------------------
 */
template <typename Key, typename Value>
void AVLTree<Key, Value>::clear(AVLNode*& node) {
    if (node != NULL) {
        clear(node->left);
        clear(node->right);
        delete node;
        node = NULL;
    }
}