#include "rb_tree.h"
#include <stdio.h>
#include <stdlib.h>

#define MIN_SIZE_BUF 10

enum { start_depth = -1 };
enum vert_border { empty, line };
struct scopes {
    unsigned int size;
    enum vert_border *mas;
};

void rb_tree_init(rb_tree *t) {
    t->nil = malloc(sizeof(node));
    node *nil = t->nil;
    nil->key = 0;
    nil->color = black;
    nil->p = NULL;
    nil->left = NULL;
    nil->right = NULL;
    t->root = nil;
}

static void left_rotate(rb_tree *t, node *cur_node)
{
    node *x = cur_node;
    node *y = x->right;
    x->right = y->left;     /* turn y's left subtee into x's right subtree  */
    if(y->left != t->nil)
        y->left->p = x;
    y->p = x->p;            /* link x's parent to y */
    if(x->p == t->nil)      /* if x was root  */
        t->root = y;
    else if(x == x->p->left)
        x->p->left = y;
    else
        x->p->right = y;
    y->left = x;            /* put x on y's left */
    x->p = y;               /* y is parent of x */
}

static void right_rotate(rb_tree *t, node *cur_node)
{
    node *y = cur_node;
    node *x = y->left;
    y->left = x->right;
    if(x->right != t->nil)
        x->right->p = y;
    x->p = y->p;
    if(y->p == t->nil)
        t->root = x;
    else if(y == y->p->left)
        y->p->left = x;
    else
        y->p->right = x;
    x->right = y;
    y->p = x;
}

static void rb_insert_fixup(rb_tree *t, node *cur_node)
{
    node *z = cur_node;
    node *y;
    while(z->p->color == red) {
        if(z->p == z->p->p->left) {/* parent of z is left child of its parent */
           y = z->p->p->right;     /* uncle of z */
           if(y->color == red) {
               z->p->color = black;
               y->color = black;
               z->p->p->color = red;
               z = z->p->p;
           } else {
               if(z == z->p->right) { /* z - right chid, y - black  */
                   z = z->p;
                   left_rotate(t, z);
               }
               /* z - left child, y - black */
               z->p->color = black;
               z->p->p->color = red;
               right_rotate(t, z->p->p);
           }
        } else {                 /* parent of z is right child of its parent */
            y = z->p->p->left;
            if(y->color == red) {
                z->p->color = black;
                y->color = black;
                z->p->p->color = red;
                z = z->p->p;
            } else {
                if(z == z->p->left) {
                    z = z->p;
                    right_rotate(t, z);
                }
                z->p->color = black;
                z->p->p->color = red;
                left_rotate(t, z->p->p);
            }
        }
    }
    t->root->color = black;
}

/* replacing one subtree with another subtree  */
static void rb_transplant(rb_tree *t, node *source, node *dest)
{
    if(source->p == t->nil)
        t->root = dest;
    else if(source == source->p->left)
        source->p->left = dest;
    else
        source->p->right = dest;
    dest->p = source->p;
}

static node* tree_minimum(rb_tree *t, node *cur_node)
{
    while(cur_node->left != t->nil)
        cur_node = cur_node->left;
    return cur_node;
}

static void rb_delete_fixup(rb_tree *t, node *x)
{
    node *w;
    while(x != t->root && x->color == black) {
        if(x == x->p->left) {       /* x is left child  */
            w = x->p->right;
            if(w->color == red) {   /* case 1 - x's sibling w is red */
                w->color = black;
                x->p->color = red;
                left_rotate(t, x->p);
                w = x->p->right;
            }
       /* case 2 - x's subling w is black, and both of w's children are black */
            if(w->left->color == black && w->right->color == black) {
                w->color = red;
                x = x->p;
            } else {
       /*case 3 - w's black, w's left child's red and w's right child's  black*/
                if(w->right->color == black) {
                    w->left->color = black;
                    w->color = red;
                    right_rotate(t, w);
                    w = x->p->right; /* new w */
                }
       /* case 4 - x's sibling w's black, and w's right child's red */
                w->color = x->p->color;
                x->p->color = black;
                w->right->color = black;
                left_rotate(t, x->p);
                x = t->root;        /* after that the loop's terminated  */
            }
        } else {                    /* x's right child  */
            w = x->p->left;
            if(w->color == red) {
                w->color = black;
                x->p->color = red;
                right_rotate(t, x->p);
                w = x->p->left;
            }
            if(w->right->color == black && w->left->color == black) {
                w->color = red;
                x = x->p;
            } else {
                if(w->left->color == black) {
                    w->right->color = black;
                    w->color = red;
                    left_rotate(t, w);
                    w = x->p->left;
                }
                w->color = x->p->color;
                x->p->color = black;
                w->left->color = black;
                right_rotate(t, x->p);
                x = t->root;
            }
        }
    }
    x->color = black;
}

void rb_tree_delete_node(rb_tree *t, node *cur_node)
{
    node *z = cur_node;
    node *y = z;
    node *x;
    enum color_type y_original_color = y->color;    /* y's color might change */
    if(z->left == t->nil) {/* if exists only right child (it's can be nil too)*/
        x = z->right;
        rb_transplant(t, z, z->right);
    } else if(z->right == t->nil) { /* or exists only left child  */
        x = z->left;
        rb_transplant(t, z, z->left);
    } else {
        y = tree_minimum(t, z->right); /* both children exist */
        y_original_color = y->color;
        x = y->right;
        if(y->p == z)       /* node y is right child of z  */
            x->p = y;
        else {
            rb_transplant(t, y, y->right); /* x subtree takes place of y  */
            y->right = z->right;           /* y takes place of z */
            y->right->p = y;
        }
        rb_transplant(t, z, y);
        y->left = z->left;
        y->left->p = y;
        y->color = z->color;
    }
    free(cur_node);
    if(y_original_color == black) /*if y with black color was moved or removed*/
        rb_delete_fixup(t, x);
}

void rb_tree_insert_node(rb_tree *t, node *cur_node)
{
    node *z = cur_node;
    node *y = t->nil;
    node *x = t->root;
    while(x != t->nil) {
        y = x; /* save parent of z  */
        x = z->key < x->key ? x->left : x->right;
    }
    z->p = y;
    if(y == t->nil)
        t->root = z;
    else if(z->key < y->key)
        y->left = z;
    else
        y->right = z;
    z->left = t->nil;
    z->right = t->nil;
    z->color = red;
    rb_insert_fixup(t, cur_node);
}

node* rb_tree_search(rb_tree *t, node *root, int num)
{
    if(root == t->nil || root->key == num)
        return root;
    if(num < root->key)
        return rb_tree_search(t, root->left, num);
    else
        return rb_tree_search(t, root->right, num);
}

node* rb_tree_create_node(rb_tree *t, int num)
{
    if(rb_tree_search(t, t->root, num) != t->nil)
        return NULL;
    node *new_node = malloc(sizeof(node));
    new_node->key = num;
    new_node->color = red;
    new_node->p = t->nil;
    new_node->left = t->nil;
    new_node->right = t->nil;
    return new_node;
}

void rb_tree_clear(rb_tree *t, node *root)
{
    if(root == t->nil)
        return;
    rb_tree_clear(t, root->left);
    rb_tree_clear(t, root->right);
    if(root->p == t->nil) {
        free(t->nil);
        t->nil = NULL;
    }
    free(root);
}

static void buf_copy(enum vert_border *source, enum vert_border *dest,
                     unsigned buf_size)
{
    int idx;
    for(idx = 0; idx < buf_size; ++idx)
        dest[idx] = source[idx];
}

static void allocate(struct scopes *cur_scopes)
{
    unsigned size = cur_scopes->size;
    cur_scopes->size *= 2;
    enum vert_border *new_mas =
        malloc(cur_scopes->size * sizeof(enum vert_border));
    buf_copy(cur_scopes->mas, new_mas, size);
    free(cur_scopes->mas);
    cur_scopes->mas = new_mas;
}

static void print(rb_tree *t, node *root, int depth, struct scopes *cur_scopes)
{
#if 0
    if(root == t->nil)
        return;
#endif
    int i;
    printf("\t");
#if 0
    if(depth >= cur_scopes->size)
        allocate(cur_scopes);
#endif
    for(i = 0; i <= depth;i++)
        if(i == depth)
            /* left |- or right |_  and ---   */
            printf("%s\u2014\u2014\u2014", cur_scopes->mas[depth] ? "\u0371" :
                "\u221F");
        else
            /* left | or right ' '  */
            printf("%s   ", cur_scopes->mas[i]? "\u23B8" : " ");
#if 1
    if(root == t->nil) {
        printf("\033[0;34;49m%s\n\033[0m", "nil");
        return;
    }
#endif
    if(root->color == red)
        printf("\033[0;31;49m%d\n\033[0m", root->key);
    else
        printf("\033[0;34;49m%d\n\033[0m", root->key);
    if(depth+1 >= cur_scopes->size)
        allocate(cur_scopes);
    /* vertical line if left branch  */
    cur_scopes->mas[depth+1] = line;
    print(t, root->left, depth+1, cur_scopes);
    cur_scopes->mas[depth+1] = empty;
    print(t, root->right, depth+1, cur_scopes);
}

void rb_tree_print(rb_tree *t, node *root)
{
    struct scopes cur_scopes;
    cur_scopes.mas = malloc(MIN_SIZE_BUF * sizeof(enum vert_border));
    cur_scopes.size = MIN_SIZE_BUF;
    print(t, root, start_depth, &cur_scopes);
    free(cur_scopes.mas);
}