/*
Write a program in ‘C’ language for the creation of a Red Black tree. Also, implement insertion and deletion operations.
*/

#include 
#include 

// Enum for node colors
typedef enum { RED, BLACK } Color;

// Red-Black Tree node structure
typedef struct Node {
    int data;
    Color color;
    struct Node *left, *right, *parent;
} Node;

// Global root node
Node *root = NULL;

// Function to create a new node
Node* createNode(int data) {
    Node *newNode = (Node*)malloc(sizeof(Node));
    newNode->data = data;
    newNode->color = RED; // New nodes are always red initially
    newNode->left = newNode->right = newNode->parent = NULL;
    return newNode;
}

// Function to perform left rotation
void leftRotate(Node *x) {
    Node *y = x->right;
    x->right = y->left;
    
    if (y->left != NULL)
        y->left->parent = x;
    
    y->parent = x->parent;
    
    if (x->parent == NULL)
        root = y;
    else if (x == x->parent->left)
        x->parent->left = y;
    else
        x->parent->right = y;
    
    y->left = x;
    x->parent = y;
}

// Function to perform right rotation
void rightRotate(Node *y) {
    Node *x = y->left;
    y->left = x->right;
    
    if (x->right != NULL)
        x->right->parent = y;
    
    x->parent = y->parent;
    
    if (y->parent == NULL)
        root = x;
    else if (y == y->parent->left)
        y->parent->left = x;
    else
        y->parent->right = x;
    
    x->right = y;
    y->parent = x;
}

// Function to fix violations after insertion
void fixInsert(Node *k) {
    Node *uncle;
    
    while (k->parent != NULL && k->parent->color == RED) {
        if (k->parent == k->parent->parent->left) {
            uncle = k->parent->parent->right;
            
            // Case 1: Uncle is red
            if (uncle != NULL && uncle->color == RED) {
                k->parent->color = BLACK;
                uncle->color = BLACK;
                k->parent->parent->color = RED;
                k = k->parent->parent;
            } else {
                // Case 2: k is right child
                if (k == k->parent->right) {
                    k = k->parent;
                    leftRotate(k);
                }
                // Case 3: k is left child
                k->parent->color = BLACK;
                k->parent->parent->color = RED;
                rightRotate(k->parent->parent);
            }
        } else {
            uncle = k->parent->parent->left;
            
            // Case 1: Uncle is red
            if (uncle != NULL && uncle->color == RED) {
                k->parent->color = BLACK;
                uncle->color = BLACK;
                k->parent->parent->color = RED;
                k = k->parent->parent;
            } else {
                // Case 2: k is left child
                if (k == k->parent->left) {
                    k = k->parent;
                    rightRotate(k);
                }
                // Case 3: k is right child
                k->parent->color = BLACK;
                k->parent->parent->color = RED;
                leftRotate(k->parent->parent);
            }
        }
        
        if (k == root) break;
    }
    
    root->color = BLACK;
}

// Insertion function
void insert(int data) {
    Node *newNode = createNode(data);
    Node *y = NULL;
    Node *x = root;
    
    // Standard BST insertion
    while (x != NULL) {
        y = x;
        if (newNode->data < x->data)
            x = x->left;
        else
            x = x->right;
    }
    
    newNode->parent = y;
    
    if (y == NULL)
        root = newNode;
    else if (newNode->data < y->data)
        y->left = newNode;
    else
        y->right = newNode;
    
    // Fix Red-Black Tree violations
    fixInsert(newNode);
}

// Function to find the minimum value node
Node* minValueNode(Node *node) {
    Node *current = node;
    while (current->left != NULL)
        current = current->left;
    return current;
}

// Function to fix violations after deletion
void fixDelete(Node *x) {
    Node *sibling;
    
    while (x != root && (x == NULL || x->color == BLACK)) {
        if (x == x->parent->left) {
            sibling = x->parent->right;
            
            // Case 1: Sibling is red
            if (sibling->color == RED) {
                sibling->color = BLACK;
                x->parent->color = RED;
                leftRotate(x->parent);
                sibling = x->parent->right;
            }
            
            // Case 2: Both sibling's children are black
            if ((sibling->left == NULL || sibling->left->color == BLACK) &&
                (sibling->right == NULL || sibling->right->color == BLACK)) {
                sibling->color = RED;
                x = x->parent;
            } else {
                // Case 3: Sibling's right child is black
                if (sibling->right == NULL || sibling->right->color == BLACK) {
                    if (sibling->left != NULL)
                        sibling->left->color = BLACK;
                    sibling->color = RED;
                    rightRotate(sibling);
                    sibling = x->parent->right;
                }
                
                // Case 4: Sibling's right child is red
                sibling->color = x->parent->color;
                x->parent->color = BLACK;
                if (sibling->right != NULL)
                    sibling->right->color = BLACK;
                leftRotate(x->parent);
                x = root;
            }
        } else {
            sibling = x->parent->left;
            
            // Case 1: Sibling is red
            if (sibling->color == RED) {
                sibling->color = BLACK;
                x->parent->color = RED;
                rightRotate(x->parent);
                sibling = x->parent->left;
            }
            
            // Case 2: Both sibling's children are black
            if ((sibling->left == NULL || sibling->left->color == BLACK) &&
                (sibling->right == NULL || sibling->right->color == BLACK)) {
                sibling->color = RED;
                x = x->parent;
            } else {
                // Case 3: Sibling's left child is black
                if (sibling->left == NULL || sibling->left->color == BLACK) {
                    if (sibling->right != NULL)
                        sibling->right->color = BLACK;
                    sibling->color = RED;
                    leftRotate(sibling);
                    sibling = x->parent->left;
                }
                
                // Case 4: Sibling's left child is red
                sibling->color = x->parent->color;
                x->parent->color = BLACK;
                if (sibling->left != NULL)
                    sibling->left->color = BLACK;
                rightRotate(x->parent);
                x = root;
            }
        }
    }
    
    if (x != NULL)
        x->color = BLACK;
}

// Transplant function for deletion
void transplant(Node *u, Node *v) {
    if (u->parent == NULL)
        root = v;
    else if (u == u->parent->left)
        u->parent->left = v;
    else
        u->parent->right = v;
    
    if (v != NULL)
        v->parent = u->parent;
}

// Deletion function
void deleteNode(int data) {
    Node *z = root;
    Node *x, *y;
    
    // Find the node to delete
    while (z != NULL) {
        if (data == z->data)
            break;
        else if (data < z->data)
            z = z->left;
        else
            z = z->right;
    }
    
    if (z == NULL) {
        printf("Node with data %d not found!\n", data);
        return;
    }
    
    y = z;
    Color yOriginalColor = y->color;
    
    if (z->left == NULL) {
        x = z->right;
        transplant(z, z->right);
    } else if (z->right == NULL) {
        x = z->left;
        transplant(z, z->left);
    } else {
        y = minValueNode(z->right);
        yOriginalColor = y->color;
        x = y->right;
        
        if (y->parent == z) {
            if (x != NULL)
                x->parent = y;
        } else {
            transplant(y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }
        
        transplant(z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }
    
    free(z);
    
    if (yOriginalColor == BLACK && x != NULL)
        fixDelete(x);
}

// Inorder traversal
void inorder(Node *node) {
    if (node == NULL) return;
    inorder(node->left);
    printf("%d(%s) ", node->data, node->color == RED ? "RED" : "BLACK");
    inorder(node->right);
}

// Preorder traversal
void preorder(Node *node) {
    if (node == NULL) return;
    printf("%d(%s) ", node->data, node->color == RED ? "RED" : "BLACK");
    preorder(node->left);
    preorder(node->right);
}

// Function to display the tree
void displayTree() {
    printf("Inorder traversal: ");
    inorder(root);
    printf("\n");
    
    printf("Preorder traversal: ");
    preorder(root);
    printf("\n\n");
}

// Function to search for a node
Node* search(int data) {
    Node *current = root;
    while (current != NULL) {
        if (data == current->data)
            return current;
        else if (data < current->data)
            current = current->left;
        else
            current = current->right;
    }
    return NULL;
}

int main() {
    int choice, data;
    Node *result;
    
    printf("Red-Black Tree Implementation\n");
    printf("=============================\n");
    
    while (1) {
        printf("\nMenu:\n");
        printf("1. Insert\n");
        printf("2. Delete\n");
        printf("3. Search\n");
        printf("4. Display Tree\n");
        printf("5. Exit\n");
        printf("Enter your choice: ");
        scanf("%d", &choice);
        
        switch (choice) {
            case 1:
                printf("Enter data to insert: ");
                scanf("%d", &data);
                insert(data);
                printf("Node %d inserted successfully.\n", data);
                displayTree();
                break;
                
            case 2:
                printf("Enter data to delete: ");
                scanf("%d", &data);
                deleteNode(data);
                printf("Node %d deleted successfully.\n", data);
                displayTree();
                break;
                
            case 3:
                printf("Enter data to search: ");
                scanf("%d", &data);
                result = search(data);
                if (result != NULL)
                    printf("Node %d found in the tree.\n", data);
                else
                    printf("Node %d not found in the tree.\n", data);
                break;
                
            case 4:
                displayTree();
                break;
                
            case 5:
                printf("Exiting program...\n");
                exit(0);
                
            default:
                printf("Invalid choice! Please try again.\n");
        }
    }
    
    return 0;
}