// Anmol Tiwari 
// Panel j
// PRN 1032233432

#include <iostream>
using namespace std;

class avl_node {
public:
    int data;
    avl_node *left;
    avl_node *right;

    avl_node(int value) {
        data = value;
        left = right = nullptr;
    }
};

class avlTree {
    avl_node *root;

public:
    avlTree() { root = nullptr; }

    int height(avl_node *);
    int diff(avl_node *);
    avl_node *rr_rotation(avl_node *);
    avl_node *ll_rotation(avl_node *);
    avl_node *lr_rotation(avl_node *);
    avl_node *rl_rotation(avl_node *);
    avl_node *balance(avl_node *);
    avl_node *insert(avl_node *, int);
    void display(avl_node *, int);
    void inorder(avl_node *);
    void preorder(avl_node *);
    void postorder(avl_node *);
   
    avl_node *getRoot() { return root; }
    void setRoot(avl_node *r) { root = r; }
};

int avlTree::height(avl_node *temp) {
    if (temp == nullptr)
        return 0;
    int l_height = height(temp->left);
    int r_height = height(temp->right);
    return max(l_height, r_height) + 1;
}

int avlTree::diff(avl_node *temp) {
    if (temp == nullptr)
        return 0;
    return height(temp->left) - height(temp->right);
}


avl_node *avlTree::rr_rotation(avl_node *parent) {
    avl_node *temp = parent->right;
    parent->right = temp->left;
    temp->left = parent;
    return temp;
}


avl_node *avlTree::ll_rotation(avl_node *parent) {
    avl_node *temp = parent->left;
    parent->left = temp->right;
    temp->right = parent;
    return temp;
}


avl_node *avlTree::lr_rotation(avl_node *parent) {
    parent->left = rr_rotation(parent->left);
    return ll_rotation(parent);
}


avl_node *avlTree::rl_rotation(avl_node *parent) {
    parent->right = ll_rotation(parent->right);
    return rr_rotation(parent);
}


avl_node *avlTree::balance(avl_node *temp) {
    int balance_factor = diff(temp);

    if (balance_factor > 1) {
        if (diff(temp->left) > 0)
            temp = ll_rotation(temp);
        else
            temp = lr_rotation(temp);
    } else if (balance_factor < -1) {
        if (diff(temp->right) < 0)
            temp = rr_rotation(temp);
        else
            temp = rl_rotation(temp);
    }
    return temp;
}


avl_node *avlTree::insert(avl_node *root, int value) {
    if (root == nullptr)
        return new avl_node(value);

    if (value < root->data)
        root->left = insert(root->left, value);
    else if (value > root->data)
        root->right = insert(root->right, value);
    else
        return root;

    return balance(root);
}

void avlTree::display(avl_node *node, int space) {
    if (node == nullptr)
        return;

    space += 5;

    display(node->right, space);

    cout << endl;
    for (int i = 5; i < space; i++)
        cout << " ";
    cout << node->data << endl;

    display(node->left, space);
}
    

void avlTree::inorder(avl_node *temp) {
    if (temp != nullptr) {
        inorder(temp->left);
        cout << temp->data;
        inorder(temp->right);
    }
}

void avlTree::preorder(avl_node *temp) {
    if (temp != nullptr) {
        cout << temp->data;
        preorder(temp->left);
        preorder(temp->right);
    }
}

void avlTree::postorder(avl_node *temp) {
    if (temp != nullptr) {
        postorder(temp->left);
        postorder(temp->right);
        cout << temp->data;
    }
}

int main() {
    avlTree avl;
    int choice, item;

    while (true) {
        cout << "\nMenu:\n";
        cout << "1. Insert Node\n";
        cout << "2. Display Tree Structure\n";
        cout << "3. Display Inorder Traversal\n";
        cout << "4. Display Preorder Traversal\n";
        cout << "5. Display Postorder Traversal\n";
        cout << "6. Exit\n";
        cout << "Enter your choice: ";
        cin >> choice;

        switch (choice) {
            case 1:
                cout << "Enter value to be inserted: ";
                cin >> item;
                avl.setRoot(avl.insert(avl.getRoot(), item));
                break;

            case 2:
                cout << "Tree Structure: \n";
                avl.display(avl.getRoot(), 0);
                cout << endl;
                break;

            case 3:
                cout << "Inorder traversal: ";
                avl.inorder(avl.getRoot());
                cout << endl;
                break;

            case 4:
                cout << "Preorder traversal: ";
                avl.preorder(avl.getRoot());
                cout << endl;
                break;

            case 5:
                cout << "Postorder traversal: ";
                avl.postorder(avl.getRoot());
                cout << endl;
                break;

            case 6:
                cout << "Exiting program.\n";
                return 0;

            default:
                cout << "Invalid choice! Please try again.\n";
        }
    }
}