/*
 * Code example for CP264 Data Structures II
 * RBT insert and delete operations by iterative algorithms
 * HBF
 */
 
#include <stdio.h>
#include <stdlib.h>

#define FatalError( Str ) fprintf( stderr, "%s\n", Str ), exit( 1 )

typedef enum ColorType {
  red, black
} ColorType;

typedef struct node {
  int key;
  ColorType color;
  struct node *left;
  struct node *right;
} tnode;

tnode *search(tnode *rootp, int val);
tnode *insert(tnode *rootp, int val);
tnode *delete(tnode *rootp, int val);
void display_tree(tnode *root, int len);
void clean_tree(tnode **rootp);

tnode *rotate_left(tnode *y);
tnode *rotate_right(tnode *y);

typedef struct stack_node {
  tnode *tnp;
  struct stack_node *next;
} snode;
void push(snode **, tnode*);
tnode *pop(snode **);
void clean_stack(snode **);

int main(int argc, char *args[]) {
  int n = 8;
  if (argc > 1)
    n = atoi(args[1]);

  tnode *root = NULL, *p;

  int i;
  for (i = 0; i < n; i++) {
    root = insert(root, i);
  }
  printf("After RBT insertion:\n");
  display_tree(root, 0);

  for (i = 0; i < n; i++) {
    if (i % 2 == 0)
      root = delete(root, i);
  }

  printf("After RBT deletion:\n");
  display_tree(root, 0);

  for (i = 0; i < n; i++) {
    if ((p = search(root, i)) == NULL)
      printf("Not found %d\n", i);
    else
      printf("Found %d\n", i);
  }

  clean_tree(&root);
  return 0;
}

void display_tree(tnode *root, int prelen) {
  if (root != NULL) {
    int i;
    for (i = 0; i < prelen; i++)
      printf("%c", ' ');
    printf("%s", "|___");
    printf("%d %c\n", root->key, (root->color == 0) ? 'R' : 'B');
    display_tree(root->right, prelen + 4);
    display_tree(root->left, prelen + 4);
  }
}

void clean_tree(tnode **rootp) {
  if (*rootp) {
    clean_tree(&(*rootp)->left);
    clean_tree(&(*rootp)->right);
    free(*rootp);
  }
}

tnode *search(tnode *root, int val) {
  while (root) {
    if (val == root->key)
      return root;
    else if (val < root->key)
      root = root->left;
    else
      root = root->right;
  }
  return NULL;
}

tnode *rotate_right(tnode *y) {
  tnode *x = y->left;
  y->left = x->right;
  x->right = y;
  return x;
}

tnode *rotate_left(tnode *y) {
  tnode *x = y->right;
  y->right = x->left;
  x->left = y;
  return x;
}

tnode *insert(tnode *root, int val) {
  snode *top = NULL; //stack to hold the top-down path for back tracking
  tnode *x = root, *p = NULL, *gp = NULL, *ggp = NULL, *uncle = NULL; //current node, parent, grand parrent, and grand grand parent
  int pci = 0, gpci = 0, ggpci = 0; //to present parent child relaiton on the top-down path, 0:left child; 1:right child
  while (x != NULL) {
    if (val == x->key) {
      clean_stack(&top);
      return root;
    } else {
      p = x;
      push(&top, x);
      if (val < x->key)
        x = x->left;
      else
        x = x->right;
    }
  }

  x = malloc(sizeof(tnode));
  if (x == NULL) {
    FatalError("out of space");
  } else {
    x->key = val;
    x->color = red;
    x->left = x->right = NULL;

    if (p == NULL) {
      x->color = black;
      root = x;
    } else {
      if (val < p->key) {
        p->left = x;
      } else {
        p->right = x;
      }

//rebalance
      repeat:    //this is     
      p = pop(&top);

      if (p && p->color == red) {
        pci = (x == p->left) ? 0 : 1;

        gp = pop(&top);
        uncle = NULL;

        if (gp)
          if (gp->left == p) {
            gpci = 0;
            uncle = gp->right;
          } else {
            gpci = 1;
            uncle = gp->left;
          }

        if (uncle && uncle->color == red) {
          p->color = black;
          uncle->color = black;

          if (gp != root) {
            gp->color = red;
            x = gp;
            goto repeat;
            //progate to upper grand parent of the top-down path
          }
        } else {
          ggp = pop(&top);

          if (ggp)
            ggpci = (ggp->left == gp) ? 0 : 1;

          if (gpci == 0) {
            if (pci == 0) {
              gp->color = red;
              p->color = black;
            } else {
              gp->color = red;
              x->color = black;
              gp->left = rotate_left(p);
            }

            if (ggp) {
              if (ggpci == 0)
                ggp->left = rotate_right(gp);
              else
                ggp->right = rotate_right(gp);
            } else {
              root = rotate_right(gp);
            }
          }

          else {
            if (pci == 1) {
              gp->color = red;
              p->color = black;
            } else {
              gp->color = red;
              x->color = black;
              gp->left = rotate_right(p);
            }

            if (ggp) {
              if (ggpci == 0)
                ggp->left = rotate_left(gp);
              else
                ggp->right = rotate_left(gp);
            } else {
              root = rotate_left(gp);
            }
          }
        }
      }
    }
  }
  clean_stack(&top);
  return root;
}

tnode *delete(tnode *root, int val) {
  snode *top = NULL; //stack to hold the top-down path for back tracking
  tnode *x = root, *p = NULL, *z = NULL; //current node, parent, grand parrent, and grand grand parent
  int pci = 0; //to present parent child relaiton on the top-down path, 0:left child; 1:right child

  //top-down to find the value node
  while (x && x->key != val) {
    push(&top, x);
    if (val < x->key)
      x = x->left;
    else
      x = x->right;
  }

  if (x == NULL) {  //not found
    return root;
  } else {  //found

    if (x->left != NULL && x->right != NULL) {
      //find the smallest node right subtree
      p = x;
      push(&top, x);
      x = x->right;

      while (x->left) {
        push(&top, x);
        x = x->left;
      }
      p->key = x->key; //replace the deleted node by the right smallest
    }

    p = pop(&top);

    //delete node and set up the re-balance node z if necessary
    if (x->left == NULL && x->right == NULL) {
      if (p == NULL) {
        root = NULL;
      } else {
        pci = (p->left == x) ? 0 : 1;
        if (pci == 0) {  // x is left child of p
          p->left = NULL;  // set to NULL as x will be deleted

          if (p->right || x->color == red) { // remove x does not affect rb property
            z = NULL;
          } else if (p->color == red) { //p->right == NULL, change p->color to balck keep the property
            p->color = black;
            z = NULL;
          } else { // remove x won't keep rb property, need to do rebalancing at parent of p
            z = pop(&top);
            pci = (z->left == p) ? 0 : 1;
          }
        } else { //similarly for pci == 1
          p->right = NULL;

          if (p->left || x->color == red) {
            z = NULL;
          } else if (p->color == red) {
            p->color = black;
            z = NULL;
          } else {
            z = pop(&top);
            pci = (z->left == p) ? 0 : 1;
          }
        }
      }
    }

    else if (x->left == NULL && x->right != NULL) {
      if (p == NULL) {
        root = x->right;
        root->color = black;
        z = NULL;
      } else {
        pci = (p->left == x) ? 0 : 1;

        if (x->color == black && x->right->color == black) {
          z = p;
        } else {
          x->right->color = black;
          z = NULL;
        }

        if (pci == 0)
          p->left = x->right;
        else
          p->right = x->right;
      }
    }

    else if (x->left != NULL && x->right == NULL) {
      if (p == NULL) {
        root = x->left;
        root->color = black;
        z = NULL;
      } else {
        pci = (p->left == x) ? 0 : 1;
        if (x->color == black && x->left->color == black) {
          z = p;
        } else {
          x->left->color = black;
          z = NULL;
        }

        if (pci == 0)
          p->left = x->left;
        else
          p->right = x->left;
      }
    }

    free(x);

    //do rebalance if z != NULL        
    while (z) { //z is not balanced, i.e. the number of black nodes on left paths and right paths are different

      if (pci == 0) { //left path has number of black nodes 1 less that that of right paths 

        if (z->color == red && z->right && z->right->color == black) // case 1.  
            {

          if ((z->right->left == NULL
              || z->right->left->color == black)
              && (z->right->right == NULL
                  || z->right->right->color == black)) {
            z->color = black;
            z->right->color = red;
          }

          else if ((z->right->left && z->right->left->color == red)
              && (z->right->right && z->right->right->color == red)) {
            z->color = black;
            z->right->color = red;
            z->right->right->color = black;

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_left(z);
              else
                p->right = rotate_left(z);
            } else {
              root = rotate_left(z);
            }

          }

          else if ((z->right->left && z->right->left->color == black)
              && (z->right->right && z->right->right->color == red)) {

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_left(z);
              else
                p->right = rotate_left(z);
            } else {
              root = rotate_left(z);
            }

          }

          else if ((z->right->left && z->right->left->color == red)
              && (z->right->right
                  && z->right->right->color == black)) {

            z->right->color = red;
            z->right->left->color == black;
            z->right = rotate_right(z->right);

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_left(z);
              else
                p->right = rotate_left(z);
            } else {
              root = rotate_left(z);
            }

          }

          z = NULL;

        }

        else if (z->color == black && z->right && z->right->color == red
            && (z->right->left == NULL
                || z->right->left->color == black)
            && (z->right->right == NULL
                || z->right->right->color == black)) // case 2.
            {

          z->right->color = black;

          p = pop(&top);
          if (p) {
            if (p->left == z) {
              p->left = rotate_left(z);
              push(&top, p);
              push(&top, p->left);

            } else {
              p->right = rotate_left(z);
              push(&top, p);
              push(&top, p->right);
            }
          } else {
            root = rotate_left(z);
            push(&top, root);
          }

          if (z->right
              && (z->right->left == NULL
                  || z->right->left->color == black)
              && (z->right->right == NULL
                  || z->right->right->color == black)) {
            z->right->color = red;
            z = NULL;
          } else {
            //z = z; // reset the balanceing node
          }
        }

        else if (z->color == black && z->right
            && z->right->color == black) //case 3
                {

          if ((z->right->left == NULL
              || z->right->left->color == black)
              && (z->right->right == NULL
                  || z->right->right->color == black)) {

            z->right->color = red;
            p = pop(&top);
            pci = (p->left == z) ? 0 : 1;
            z = p;
          }

          else if ((z->right->left == NULL
              || z->right->left->color == black)
              && (z->right->right && z->right->right->color == red)) {
            z->right->right->color = black;

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_left(z);
              else
                p->right = rotate_left(z);
            } else {
              root = rotate_left(z);
            }
            z = NULL;
          }

          else if (z->right->left && z->right->left->color == red
              && (z->right->right == NULL
                  || z->right->right->color == black)) {
            z->right->left->color = black;

            z->right = rotate_right(z->right);

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_left(z);
              else
                p->right = rotate_left(z);
            } else {
              root = rotate_left(z);
            }
            z = NULL;
          }
        }

      }

      else { // pci == 1, right path has number of black nodes 1 less that that of left paths 

        if (z->color == red && z->left && z->left->color == black) //case 1
            {
          if ((z->left->left == NULL || z->left->left->color == black)
              && (z->left->right == NULL
                  || z->left->right->color == black)) {
            z->color = black;
            z->left->color = red;
          }

          else if ((z->left->left && z->left->left->color == red)
              && (z->left->right && z->left->right->color == red)) {
            z->color = black;
            z->left->color = red;
            z->left->left->color = black;

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_right(z);
              else
                p->right = rotate_right(z);
            } else {
              root = rotate_right(z);
            }

          }

          else if ((z->left->left && z->left->left->color == red)
              && (z->left->right && z->left->right->color == black)) {

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_right(z);
              else
                p->right = rotate_right(z);
            } else {
              root = rotate_right(z);
            }
          }

          else if ((z->left->left && z->left->left->color == black)
              && (z->left->right && z->left->right->color == red)) {

            z->left->color = red;
            z->left->right->color == black;

            z->left = rotate_left(z->left);

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_right(z);
              else
                p->right = rotate_right(z);
            } else {
              root = rotate_right(z);
            }
          }

          z = NULL;

        }

        else if (z->color == black && z->left && z->left->color == red //case 2
        && (z->left->left == NULL || z->left->left->color == black)
            && (z->left->right == NULL
                || z->left->right->color == black)) {
          z->left->color = black;
          p = pop(&top);
          if (p) {
            if (p->left == z) {
              p->left = rotate_right(z);
              push(&top, p);
              push(&top, p->left);
            } else {
              p->right = rotate_right(z);
              push(&top, p);
              push(&top, p->right);
            }
          } else {
            root = rotate_right(z);
            push(&top, root);
          }

          if (z->left->right
              && (z->left->right == NULL
                  || z->left->right->color == black)
              && (z->left->left == NULL
                  || z->left->left->color == black)) {
            z->left->color = red;
            z = NULL;
          } else {
            //z = z
          }
        }

        else if (z->color == black && z->left
            && z->left->color == black)  //case 3.
                {
          if ((z->left->left == NULL || z->left->left->color == black)
              && (z->left->right == NULL
                  || z->left->right->color == black)) {
            z->left->color = red;
            p = pop(&top);
            pci = (p->left == z) ? 0 : 1;
            z = p;
          }

          else if (z->left->left && z->left->left->color == red
              && (z->left->right == NULL
                  || z->left->right->color == black)) {
            z->left->left->color = black;
            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_right(z);
              else
                p->left = rotate_right(z);
            } else {
              root = rotate_right(z);
            }
            z = NULL;

          }

          else if (z->left->left && z->left->left->color == black
              && z->left->right && z->left->right->color == red) {
            z->left->right->color = black;
            z->left->color = black;
            z->left = rotate_left(z->left);

            p = pop(&top);
            if (p) {
              if (p->left == z)
                p->left = rotate_right(z);
              else
                p->left = rotate_right(z);
            } else {
              root = rotate_left(z);
            }
            z = NULL;
          }
        }

        else {
          //z = NULL;
        }

      }

    } //end of balance

  }

  clean_stack(&top);
  return root;
}

void push(snode **topp, tnode *tp) {
  snode *ptr = (snode*) malloc(sizeof(snode));
  ptr->tnp = tp;
  ptr->next = NULL;
  if (*topp == NULL) {
    *topp = ptr;
  } else {
    ptr->next = *topp;
    *topp = ptr;
  }
}

tnode *pop(snode **topp) {
  if (*topp == NULL) {
    return NULL;
  } else {
    snode *np = *topp;
    *topp = np->next;
    tnode *tnp = np->tnp;
    free(np);
    return tnp;
  }
}

void clean_stack(snode **topp) {
  snode *top = *topp, *temp;
  while (top) {
    temp = top;
    top = top->next;
    free(temp);
  }
  *topp = NULL;
}

/*
After RBT insertion:
|___3 B
    |___5 R
        |___6 B
            |___7 R
        |___4 B
    |___1 R
        |___2 B
        |___0 B
After RBT deletion:
|___3 B
    |___5 R
        |___7 B
    |___1 B
Not found 0
Found 1
Not found 2
Found 3
Not found 4
Found 5
Not found 6
Found 7
 */