#ifndef prototype1_src_l1_5_core_rb_tree_node_h
#define prototype1_src_l1_5_core_rb_tree_node_h

#include "../../l1/core/util.h"

typedef enum {
    RBTREE_BLACK = 0,
    RBTREE_RED = 1,
} RBTreeClr;

typedef struct RBTreeNode RBTreeNode;
struct RBTreeNode{
    RBTreeNode* parent;
    RBTreeNode* left;
    RBTreeNode* right;
    RBTreeClr color;
};

void RBTree_left_rotate(RBTreeNode** root, RBTreeNode* NIL, RBTreeNode* x){
    assert(x != NIL);
    RBTreeNode* y = x->right;
    assert(y != NIL);
    x->right = y->left;
    if (y->left != NIL)
        y->left->parent = x;

    y->parent = x->parent;
    if (x->parent == NIL)
        *root = y;
    else if (x == x->parent->left)
        x->parent->left = y;
    else
        x->parent->right = y;

    x->parent = y;
    y->left = x;
}

void RBTree_right_rotate(RBTreeNode** root, RBTreeNode* NIL, RBTreeNode* x){
    assert(x != NIL);
    RBTreeNode* y = x->left;
    assert(y != NIL);
    x->left = y->right;
    if (y->right != NIL)
        y->right->parent = x;

    y->parent = x->parent;
    if (x->parent == NIL)
        *root = y;
    else if (x->parent->right == x)
        x->parent->right = y;
    else
        x->parent->left = y;

    x->parent = y;
    y->right = x;
}

void RBTree_fix_after_insert(RBTreeNode** root, RBTreeNode* NIL, RBTreeNode* me){
    assert(me != NIL);
    while (true) {
        RBTreeNode* mom = me->parent;
        if (mom == NIL)
            break;
        if (mom->color == RBTREE_BLACK)
            return;
        RBTreeNode* grandma = mom->parent;
        assert(grandma != NIL);
        assert(grandma->color == RBTREE_BLACK);
        RBTreeNode* aunt = grandma->left == mom ? grandma->right : grandma->left;
        assert(aunt != mom);
        if (aunt->color == RBTREE_RED) {
            /* Easy case */
            mom->color = RBTREE_BLACK;
            aunt->color = RBTREE_BLACK;
            grandma->color = RBTREE_RED;
            me = grandma;
        } else if (grandma->left == mom) {
            /* Hard case: firstborn orientation */
            if (mom->right == me) {
                RBTree_left_rotate(root, NIL, mom);
                me->color = RBTREE_BLACK;
            } else {
                mom->color = RBTREE_BLACK;
            }
            RBTree_right_rotate(root, NIL, grandma);
            grandma->color = RBTREE_RED;
            return;
        } else {
            /* Hard case: benjamin orientation */
            if (mom->left == me) {
                RBTree_right_rotate(root, NIL, mom);
                me->color = RBTREE_BLACK;
            } else {
                mom->color = RBTREE_BLACK;
            }
            RBTree_left_rotate(root, NIL, grandma);
            grandma->color = RBTREE_RED;
            return;
        }
    }
    assert(*root == me);
    me->color = RBTREE_BLACK;
}


void RBTree_steal_neighbours(RBTreeNode** root, RBTreeNode* NIL, RBTreeNode* fr, RBTreeNode* to){
    if (fr->parent == NIL)
        *root = to;
    else if (fr->parent->left == fr)
        fr->parent->left = to;
    else
        fr->parent->right = to;
    fr->left->parent = to;
    fr->right->parent = to;
    *to = *fr;
}


RBTreeNode* RBTreeNode_minimum_in_subtree(RBTreeNode* x, RBTreeNode* NIL){
    assert(x != NIL);
    while (x->left != NIL)
        x = x->left;
    return x;
}

RBTreeNode* RBTreeNode_maximum_in_subtree(RBTreeNode* x, RBTreeNode* NIL){
    assert(x != NIL);
    while (x->right != NIL)
        x = x->right;
    return x;
}

/* Returns NULL instead of NIL */
RBTreeNode* RBTreeNode_find_next(RBTreeNode* x, RBTreeNode* NIL){
    assert(x != NIL);
    if (x->right != NIL)
        return RBTreeNode_minimum_in_subtree(x->right, NIL);
    while (true) {
        RBTreeNode* p = x->parent;
        if (p == NIL)
            return NULL;
        if (p->left == x)
            return p;
        x = p;
    }
}

/* Returns NULL instead of NIL */
RBTreeNode* RBTreeNode_find_prev(RBTreeNode* x, RBTreeNode* NIL){
    assert(x != NIL);
    if (x->left != NIL)
        return RBTreeNode_maximum_in_subtree(x->left, NIL);
    while (true) {
        RBTreeNode* p = x->parent;
        if (p == NIL)
            return NULL;
        if (p->right == x)
            return p;
        x = p;
    }
}

/* me may be NIL at the beginning. ->parent field of NIL is meaningful */
void RBTree_fix_after_delete(RBTreeNode** root, RBTreeNode* NIL, RBTreeNode* me){
    assert((*root)->parent == NIL);
    while (me != *root && me->color == RBTREE_BLACK) {
        RBTreeNode* mom = me->parent;
        assert(mom != NIL);
        if (me == mom->left) { /* We are on the left */
            RBTreeNode* sister = mom->right;
            assert(sister != NIL);
            if (sister->color == RBTREE_RED) { /* Case 1 */
                mom->color = RBTREE_RED;
                sister->color = RBTREE_BLACK;
                RBTree_left_rotate(root, NIL, mom);
                /* Reassignation required */
                sister = mom->right;
                assert(sister != NIL);
            }
            /* Cases 2,3,4 */
            RBTreeNode* nephew_firstborn = sister->left;
            RBTreeNode* nephew_benjamin = sister->right;
            if (nephew_firstborn->color == RBTREE_BLACK && nephew_benjamin->color == RBTREE_BLACK) {
                /* Case 2 */
                sister->color = RBTREE_RED;
                me = mom;
                continue;
            }
            /* Cases 3,4 */
            if (nephew_benjamin->color == RBTREE_BLACK) {
                /* Case 3 */
                nephew_firstborn->color = RBTREE_BLACK;
                sister->color = RBTREE_RED;
                RBTree_right_rotate(root, NIL, sister);
                /* Reassignation required */
                nephew_benjamin = sister;
                sister = nephew_firstborn;
                assert(sister != NIL);
                /*nephew_firstborn = sister->left;*/
            }
            /* Case 4 */
            sister->color = mom->color;
            mom->color = RBTREE_BLACK;
            nephew_benjamin->color = RBTREE_BLACK;
            RBTree_left_rotate(root, NIL, mom);
            me = *root;
        } else { /* We are on the right */
            assert(me == mom->right);
            RBTreeNode* sister = mom->left;
            assert(sister != NIL);
            if (sister->color == RBTREE_RED) { /* Case 1 */
                mom->color = RBTREE_RED;
                sister->color = RBTREE_BLACK;
                RBTree_right_rotate(root, NIL, mom);
                /* Reassignation required */
                sister = mom->left;
                assert(sister != NIL);
            }
            /* Cases 2,3,4 (every instance of red-black tree has an itchy substring in source code containing 2,3,4) */
            RBTreeNode* nephew_firstborn = sister->left;
            RBTreeNode* nephew_benjamin = sister->right;
            if (nephew_firstborn->color == RBTREE_BLACK && nephew_benjamin->color == RBTREE_BLACK) {
                /* Case 2 */
                sister->color = RBTREE_RED;
                me = mom;
                continue;
            }
            /* Cases 3,4 */
            if (nephew_firstborn->color == RBTREE_BLACK) {
                /* Case 3 */
                nephew_benjamin->color = RBTREE_BLACK;
                sister->color = RBTREE_RED;
                RBTree_left_rotate(root, NIL, sister);
                /* Reassignation required */
                nephew_firstborn = sister;
                sister = nephew_benjamin;
                assert(sister != NIL);
                /*nephew_benjamin = sister->right;*/
            }
            /* Case 4 */
            sister->color = mom->color;
            mom->color = RBTREE_BLACK;
            nephew_firstborn->color = RBTREE_BLACK;
            RBTree_right_rotate(root, NIL, mom);
            me = *root;
        }
    }
    me->color = RBTREE_BLACK;
}

/* Assumes that z->key and z->value were already dropped properly. Does not free z node */
void RBTree_erase_empty_by_iter(RBTreeNode** root, RBTreeNode* NIL, RBTreeNode* z) {
    assert(z != NULL);
    assert(z != NIL);
    RBTreeNode* y = (z->left == NIL || z->right == NIL) ? z : RBTreeNode_minimum_in_subtree(z->right, NIL);
    RBTreeNode* x = y->left == NIL ? y->right : y->left;
    assert(x != y && x != z);
    RBTreeNode* x_adopter = y->parent;
    x->parent = x_adopter;
    if (x_adopter == NIL)
        *root = x;
    else if (x_adopter->left == y)
        x_adopter->left = x;
    else
        x_adopter->right = x;
    RBTreeClr y_org_clr = y->color;
    if (z != y) {
        RBTree_steal_neighbours(root, NIL, z, y);
        // if (x_adopter == z)
        //     x_adopter = y;
    }
    // x->parent = x_adopter;
    if (y_org_clr == RBTREE_BLACK)
        RBTree_fix_after_delete(root, NIL, x);
}

#endif