#pragma once
#include <iostream>
using namespace std;

template<typename DT>
struct AVLNode {
    DT datum;
    struct AVLNode* left;
    struct AVLNode* right;
};

template<typename DT>
class AVLTree {
private:
    AVLNode<DT>* root;
    int height(AVLNode<DT>* node);
    int diff(AVLNode<DT>* node);
    AVLNode<DT>* rrRotation(AVLNode<DT>* parent);
    AVLNode<DT>* llRotation(AVLNode<DT>* parent);
    AVLNode<DT>* lrRotation(AVLNode<DT>* parent);
    AVLNode<DT>* rlRotation(AVLNode<DT>* parent);
    AVLNode<DT>* balance(AVLNode<DT>* node);
    AVLNode<DT>* remove(DT datum, AVLNode<DT>* node);
    AVLNode<DT>* insert(AVLNode<DT>* node, DT datum);
    void display(AVLNode<DT>* node, int level);
    void inOrder(AVLNode<DT>* node);
    void preOrder(AVLNode<DT>* node);
    void postOrder(AVLNode<DT>* node);
    AVLNode<DT>* minValueNode(AVLNode<DT>* node);
    AVLNode<DT>* maxValueNode(AVLNode<DT>* node);
public:
    AVLNode<DT>* insert(DT datum);
    AVLNode<DT>* remove(DT datum);
    void display();
    void inOrder();
    void preOrder();
    void postOrder();
    AVLTree();
};

template<typename DT>
int AVLTree<DT>::height(AVLNode<DT>* node) {
    int h = 0;
    if (node != NULL) {
        int lHeight = height(node->left);
        int rHeight = height(node->right);
        if (lHeight >= rHeight) h = lHeight + 1;
        else h = rHeight + 1;
    }
    return h;
}

template<typename DT>
int AVLTree<DT>::diff(AVLNode<DT>* node) {
    int lHeight = height(node->left);
    int rHeight = height(node->right);
    int bFactor = lHeight - rHeight;
    return bFactor;
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::rrRotation(AVLNode<DT>* parent) {
    AVLNode<DT>* temp;
    temp = parent->right;
    parent->right = temp->left;
    temp->left = parent;
    return temp;
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::llRotation(AVLNode<DT>* parent) {
    AVLNode<DT>* temp;
    temp = parent->left;
    parent->left = temp->right;
    temp->right = parent;
    return temp;
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::lrRotation(AVLNode<DT>* parent) {
    AVLNode<DT>* temp;
    temp = parent->left;
    parent->left = rrRotation(temp);
    return llRotation(parent);
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::rlRotation(AVLNode<DT>* parent) {
    AVLNode<DT>* temp;
    temp = parent->right;
    parent->right = llRotation(temp);
    return rrRotation(parent);
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::balance(AVLNode<DT>* node) {
    int balFactor = diff(node);
    if (balFactor > 1) {
        if (diff(node->left) > 0) node = llRotation(node);
        else node = lrRotation(node);
    }
    else if (balFactor < -1) {
        if (diff(node->right) > 0) node = rlRotation(node);
        else node = rrRotation(node);
    }
    return node;
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::remove(DT datum, AVLNode<DT>* node) {
    if (root == NULL) return root;
    if (datum < node->datum) node->left = remove(datum, node->left);
    else if (datum > node->datum) node->right = remove(datum, node->right);
    else {
        if (node->left == NULL || root->right == NULL) {
            AVLNode<DT>* temp = node->left ? node->left : node->right;
            if (temp == NULL) {
                temp = node;
                node = NULL;
            }
            else *node = *temp;
            delete(temp);
        }
        else {
            if (minValueNode(node->right) != NULL) {
                AVLNode<DT>* temp = minValueNode(node->right);
                node->datum = temp->datum;
                node->right = remove(temp->datum, node->right);
            }
            else {
                AVLNode<DT>* temp = maxValueNode(node->left);
                node->datum = temp->datum;
                node->left = remove(temp->datum, node->left);
            }


        }
    }

    if (node == NULL) return NULL;

    node = balance(node);
    return node;
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::insert(AVLNode<DT>* node, DT datum) {
    if (node == NULL) {
        node = new AVLNode<DT>;
        node->datum = datum;
        node->left = NULL;
        node->right = NULL;
        return node;
    }
    else if (datum < node->datum) {
        node->left = insert(node->left, datum);
        node = balance(node);
    }
    else if (datum >= node->datum)
    {
        node->right = insert(node->right, datum);
        node = balance(node);
    }
    return node;
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::insert(DT datum) {
    if (root == NULL) {
        root = new AVLNode<DT>;
        root->datum = datum;
        root->left = NULL;
        root->right = NULL;
        return root;
    }
    else return root = insert(root, datum);
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::remove(DT datum) {
    return root = remove(datum, root);
}

template<typename DT>
void AVLTree<DT>::display() {
    display(root, 1);
}

template<typename DT>
void AVLTree<DT>::inOrder() {
    cout << "Inorder: [";
    inOrder(root);
    cout << "]" << endl;
}

template<typename DT>
void AVLTree<DT>::preOrder() {
    cout << "Preorder: [";
    preOrder(root);
    cout << "]" << endl;
}

template<typename DT>
void AVLTree<DT>::postOrder() {
    cout << "Postorder: [";
    postOrder(root);
    cout << "]" << endl;
}

template<typename DT>
void AVLTree<DT>::display(AVLNode<DT>* node, int level) {
    if (node != NULL) {
        display(node->right, level + 1);
        cout << endl;
        if (node == root) cout << "Root -> ";
        for(int i = 0; i < level && node != root; i++) cout << "        ";
        cout << node->datum;
        display(node->left, level + 1);
    }
}

template<typename DT>
void AVLTree<DT>::inOrder(AVLNode<DT>* node) {
    if (node == NULL) return;
    inOrder(node->left);
    cout << node->datum << ", ";
    inOrder(node->right);
}

template<typename DT>
void AVLTree<DT>::preOrder(AVLNode<DT>* node) {
    if (node == NULL) return;
    cout << node->datum << ", ";
    preOrder(node->left);
    preOrder(node->right);
}

template<typename DT>
void AVLTree<DT>::postOrder(AVLNode<DT>* node) {
    if (node == NULL)
        return;
    postOrder(node->left);
    postOrder(node->right);
    cout << node->datum << ", ";
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::minValueNode(AVLNode<DT>* node) {
    if (node == NULL) return NULL;
    AVLNode<DT>* current = node;
    while (current->left != NULL) current = current->left;
    return current;
}

template<typename DT>
AVLNode<DT>* AVLTree<DT>::maxValueNode(AVLNode<DT>* node) {
    if (node == NULL) return NULL;
    AVLNode<DT>* current = node;
    while (current->right != NULL) current = current->right;
    return current;
}

template<typename DT>
AVLTree<DT>::AVLTree() {
    root = NULL;
}