// Written by Nic Roets. This code is in the public domain.
using namespace std;
#include "artree.h"

// The red - black rules :
// RB0. Operations does not change the order of the elements, i.e. if it was
//      sorted, it will stay sorted.
// RB1. Every leaf node (NULL) has the same number (N) of black ancestors.
// RB2. We choose the root node to be black. This choice does not affect the
//      validity of RB1, because it's an ancestor of all leaves, but it does
//      make it easier to satisfy RB3.
// RB3. A red node's parent must be black.

template <class t> static void Rotate (arnode<t> *x,
  int acwise, arnode<t> **root)
{ // Preserves RB0
  //printf (acwise ? "anticlockwise\n" : "clockwise\n");
  arnode<t> *child = acwise ? x->right : x->left,
    *gchild = acwise ? child->left : child->right;
  //if (child->left) printf ("%d\n", *(int*)&child->left->node);
  child->parent = x->parent;
  x->parent = child;
  *(acwise ? &x->right : &x->left) = gchild;
  *(acwise ? &child->left : &child->right) = x;
  if (gchild) gchild->parent = x;
  
  //printf ("%p\n", *(int*)&(*root)->node);
  if (child->parent) *(child->parent->left == x ?
    &child->parent->left : &child->parent->right) = child;
  else *root = child;
  
  x->s.branchsize -= child->s.branchsize;
  child->s.branchsize += x->s.branchsize;
  if (gchild) x->s.branchsize += gchild->s.branchsize;
  //printf ("%d\n", *(int*)&x->right->node);
  //printf ("%p\n", *(int*)&(*root)->node);
}

template <class t> int artree<t>::Next (arnode<t> **ptr)
{
  if (!*ptr || (*ptr)->right) {
    if (!*ptr && !root) return 0;
    for (*ptr = *ptr ? (*ptr)->right : root;
         (*ptr)->left; *ptr = (*ptr)->left) {}
    return 1;
  }
  for (;;) {
    arnode<t> *old = *ptr;
    *ptr = (*ptr)->parent;
    if (!*ptr) return 0;
    if ((*ptr)->left == old) return 1;
  }
}

template <class t> void artree<t>::InsertBefore (arnode<t> *me,
  arnode<t> *x)
{
  arnode<t> *leaf = me ? me->left : root, *gp;
  //printf (leaf ? "right\n" : me ? "left\n" : "root\n");
  if (leaf) {
    while (leaf->right) leaf = leaf->right;
    leaf->right = x;
    me = leaf;
  }
  else if (me) me->left = x; // 'me' has no children on it's left
  // else there's no root yet.
  
  x->left = NULL;
  x->right = NULL;
  x->parent = me;
  x->s.branchsize = 1;
  if (!root) {
    x->s.red = 0;
    root = x;
  }
  else {
    for (gp = x->parent; gp; gp = gp->parent) gp->s.branchsize++;
    x->s.red = 1;
    while (x->parent->s.red) {
    // Invarients : RB1 & RB2. 'x' is the only exception to RB3.
//      for (arnode<t> *q = NULL; Next (&q, *root); printf ("%d ", q->node)) {}
      gp = x->parent->parent; // RB2 => it must exit.    RB3 => it's black.
//      printf ("--%d\n", *(int*)&gp->node);
      me = gp->left == x->parent ? gp->right : gp->left;
      if (!me || !me->s.red) {
        int pIsLeft = gp->left == x->parent, xIsLeft = x->parent->left == x;
        // RB3 => sibling(x) is black.
        // RB3 => All x's children are black.
        // So if either of these 3 becomes gp's child, gp can become red
        // without violating RB3.
        if (xIsLeft != pIsLeft) Rotate (x->parent, pIsLeft, &root);
        Rotate (gp, !pIsLeft, &root);
//      for (arnode<t> *q = NULL; Next (&q, *root); printf ("%d ", q->node)) {}
//      printf ("dbug \n");
        gp->s.red = 1;
        gp->parent->s.red = 0;
        break;
      }
      me->s.red = 0;
      x->parent->s.red = 0;
      if (gp->parent == NULL) break; // N increases
      gp->s.red = 1;
      x = gp;
    } // while restoring RB3 after an insert
  } // If this is not the first insert.
}

template <class t> void artree<t>::Remove (arnode<t> *old)
{
  // If this has 2 children, we need to find it's successor.
  arnode<t> *unfull = old->left && old->right ? old->right : old;
  arnode<t> *unfullChild, *p, *sibling;
  while (old->left && old->right && unfull->left) unfull = unfull->left;
  
  // Step 0 : Update branchsizes
  for (p = unfull->parent; p; p = p->parent) p->s.branchsize--;
  // Step 1 : Unlink unfull
  p = unfull->parent; // Remember it's parent
  unfullChild = unfull->left ? unfull->left : unfull->right; // And it's child
  
  int isLeft;
  if (p) {
    isLeft = p->left == unfull;
    *(isLeft ? &p->left : &p->right) = unfullChild; // Unhook p
  }
  else root = unfullChild; // Unhook root
  
  if (unfullChild) unfullChild->parent = p; // Unhook it's child
  
  if (unfull->s.red) p = NULL; // If removing a red one, then no fix
  
  if (unfull != old) {
  // Now unhook 'this' and rehook 'unfull' in it's place
    unfull->parent = old->parent;
    unfull->left = old->left;
    unfull->right = old->right;
    unfull->s = old->s;
    
    if (old->left) old->left->parent = unfull;
    if (old->right) old->right->parent = unfull;
    if (!old->parent) root = unfull;
    else *(old->parent->left == old ?
      &old->parent->left : &old->parent->right) = unfull;
    if (p == old) p = unfull;
  }
  
  while (p) {
  // Invarient : RB0, RB2. RB1 is true except that one of p's decendants have
  // one too few black ancestor. RB3 may be violated after unlinking 'unfull',
  // but in that case we will not enter this loop and restore RB3 below.
    sibling = isLeft ? p->right : p->left;
    // On the one side we removed a black node, so sibling must exsist.
    if (sibling->s.red) {
      // sibling must have children
      sibling->s.red = 0;
      p->s.red = 1; // Parent must have been black.
      Rotate (p, isLeft, &root);
      sibling = isLeft ? p->right : p->left;
      // new sibling must be internal
    }
    if (sibling->left && sibling->right &&
        (sibling->left->s.red || sibling->right->s.red)) {
      if (!(isLeft ? sibling->right : sibling->left)->s.red) {
        (isLeft ? sibling->left : sibling->right)->s.red = 0;
        // Now sibling has 2 black children
        sibling->s.red = 1;
        Rotate (sibling, 1 - isLeft, &root);
        sibling = isLeft ? p->right : p->left;
      }
      sibling->s.red = p->s.red; // Sibling was black.
      p->s.red = 0;
      (isLeft ? sibling->right : sibling->left)->s.red = 0;
      Rotate (p, isLeft, &root);
      break;
    }
    sibling->s.red = 1; // sibling has only black children, so is this is OK
    if (p->s.red) {
      p->s.red = 0;
      break;
    }
    isLeft = p->parent && p->parent->left == p;
    p = p->parent;
  } // While trying to restore the imbalance between p's children.
}

template <class t> int artree<t>::Search (arnode<t> **ptr, int diff)
{
  if (!*ptr) return 1;
  if (diff == 0) return 2;
  if (diff < 0) {
    if (!(*ptr)->left) return 1;
    *ptr = (*ptr)->left;
  }
  else {
    if (!(*ptr)->right) {
      Next (ptr);
      return 1;
    }
    *ptr = (*ptr)->right;
  }
  return 0;
}

template <class t> arnode<t> &artree<t>::operator[](int i)
{
  arnode<t> *x = root;
  while (i != (x->left ? x->left->s.branchsize : 0)) {
    if (i < (x->left ? x->left->s.branchsize : 0)) x = x->left;
    else {
      i -= x->left ? x->left->s.branchsize + 1 : 1;
      x = x->right;
    }
  }
  return *x;
}

/* // Unmaintained code for printing the tree.
int Print (arnode<int> *x, int l)
{
  int c = 0;
  if (l == 0) {
    if (x) printf (" %2d%c ", x->node, x->s.red ? 'r' : 'b');
    else printf ("     ");
    if (x && (x->left || x->right)) c++;
  }
  else {
    c += Print (x ? x->left : NULL, l - 1);
    c += Print (x ? x->right : NULL, l - 1);
  }
  return c;
}
    for (int l = 0;;l++) {
      printf ("\n%*c", 40 - (5 << l) / 2, ' ');
      if (!Print (root, l)) break;
    }
    printf ("\n");

*/