/**
 * @file rbtree.h
 * @brief Red/Black Tree
 */

#include <pj/types.h>

PJ_BEGIN_DECL

/**
 * @defgroup PJ_RBTREE Red/Black Balanced Tree
 * @ingroup PJ_DS
 * @brief
 * Red/Black tree is the variant of balanced tree, where the search, insert, 
 * and delete operation is \b guaranteed to take at most \a O( lg(n) ).
 * @{
 */
/**
 * Color type for Red-Black tree.
 */
typedef enum pj_rbcolor_t
{
    PJ_RBCOLOR_BLACK,
    PJ_RBCOLOR_RED
} pj_rbcolor_t;

/**
 * The type of the node of the R/B Tree.
 */
typedef struct pj_rbtree_node
{
    /** Pointers to the node's parent, and left and right siblings. */
    struct pj_rbtree_node *parent, *left, *right;

    /** Key associated with the node. */
    const void *key;

    /** User data associated with the node. */
    void *user_data;

    /** The R/B Tree node color. */
    pj_rbcolor_t color;

} pj_rbtree_node;


/**
 * The type of function use to compare key value of tree node.
 * @return
 *  0 if the keys are equal
 * <0 if key1 is lower than key2
 * >0 if key1 is greater than key2.
 */
typedef int pj_rbtree_comp(const void *key1, const void *key2);


/**
 * Declaration of a red-black tree. All elements in the tree must have UNIQUE
 * key.
 * A red black tree always maintains the balance of the tree, so that the
 * tree height will not be greater than lg(N). Insert, search, and delete
 * operation will take lg(N) on the worst case. But for insert and delete,
 * there is additional time needed to maintain the balance of the tree.
 */
typedef struct pj_rbtree
{
    pj_rbtree_node null_node;   /**< Constant to indicate NULL node.    */
    pj_rbtree_node *null;       /**< Constant to indicate NULL node.    */
    pj_rbtree_node *root;       /**< Root tree node.                    */
    unsigned size;              /**< Number of elements in the tree.    */
    pj_rbtree_comp *comp;       /**< Key comparison function.           */
} pj_rbtree;


/**
 * Guidance on how much memory required for each of the node.
 */
#define PJ_RBTREE_NODE_SIZE      (sizeof(pj_rbtree_node))


/**
 * Guidance on memory required for the tree.
 */
#define PJ_RBTREE_SIZE           (sizeof(pj_rbtree))


/**
 * Initialize the tree.
 * @param tree the tree to be initialized.
 * @param comp key comparison function to be used for this tree.
 */
void pj_rbtree_init( pj_rbtree *tree, pj_rbtree_comp *comp);

/**
 * Get the first element in the tree.
 * The first element always has the least value for the key, according to
 * the comparison function.
 * @param tree the tree.
 * @return the tree node, or NULL if the tree has no element.
 */
pj_rbtree_node* pj_rbtree_first( pj_rbtree *tree );

/**
 * Get the last element in the tree.
 * The last element always has the greatest key value, according to the
 * comparison function defined for the tree.
 * @param tree the tree.
 * @return the tree node, or NULL if the tree has no element.
 */
pj_rbtree_node* pj_rbtree_last( pj_rbtree *tree );

/**
 * Get the successive element for the specified node.
 * The successive element is an element with greater key value.
 * @param tree the tree.
 * @param node the node.
 * @return the successive node, or NULL if the node has no successor.
 */
pj_rbtree_node* pj_rbtree_next( pj_rbtree *tree,
                   pj_rbtree_node *node );

/**
 * The the previous node for the specified node.
 * The previous node is an element with less key value.
 * @param tree the tree.
 * @param node the node.
 * @return the previous node, or NULL if the node has no previous node.
 */
pj_rbtree_node* pj_rbtree_prev( pj_rbtree *tree,
                   pj_rbtree_node *node );

/**
 * Insert a new node. 
 * The node will be inserted at sorted location. The key of the node must 
 * be UNIQUE, i.e. it hasn't existed in the tree.
 * @param tree the tree.
 * @param node the node to be inserted.
 * @return zero on success, or -1 if the key already exist.
 */
int pj_rbtree_insert( pj_rbtree *tree,
                 pj_rbtree_node *node );

/**
 * Find a node which has the specified key.
 * @param tree the tree.
 * @param key the key to search.
 * @return the tree node with the specified key, or NULL if the key can not
 *         be found.
 */
pj_rbtree_node* pj_rbtree_find( pj_rbtree *tree,
                   const void *key );

/**
 * Erase a node from the tree.
 * @param tree the tree.
 * @param node the node to be erased.
 * @return the tree node itself.
 */
pj_rbtree_node* pj_rbtree_erase( pj_rbtree *tree,
                    pj_rbtree_node *node );

/**
 * Get the maximum tree height from the specified node.
 * @param tree the tree.
 * @param node the node, or NULL to get the root of the tree.
 * @return the maximum height, which should be at most lg(N)
 */
unsigned pj_rbtree_max_height( pj_rbtree *tree,
                  pj_rbtree_node *node );

/**
 * Get the minumum tree height from the specified node.
 * @param tree the tree.
 * @param node the node, or NULL to get the root of the tree.
 * @return the height
 */
unsigned pj_rbtree_min_height( pj_rbtree *tree,
                  pj_rbtree_node *node );

static void left_rotate( pj_rbtree *tree, pj_rbtree_node *node )
{
    pj_rbtree_node *rnode, *parent;

    PJ_CHECK_STACK();

    rnode = node->right;
    if (rnode == tree->null)
        return;

    node->right = rnode->left;
    if (rnode->left != tree->null)
        rnode->left->parent = node;
    parent = node->parent;
    rnode->parent = parent;
    if (parent != tree->null) {
        if (parent->left == node)
     parent->left = rnode;
        else
     parent->right = rnode;
    } else {
        tree->root = rnode;
    }
    rnode->left = node;
    node->parent = rnode;
}

static void right_rotate( pj_rbtree *tree, pj_rbtree_node *node )
{
    pj_rbtree_node *lnode, *parent;

    PJ_CHECK_STACK();

    lnode = node->left;
    if (lnode == tree->null)
        return;

    node->left = lnode->right;
    if (lnode->right != tree->null)
  lnode->right->parent = node;
    parent = node->parent;
    lnode->parent = parent;

    if (parent != tree->null) {
        if (parent->left == node)
      parent->left = lnode;
  else
      parent->right = lnode;
    } else {
        tree->root = lnode;
    }
    lnode->right = node;
    node->parent = lnode;
}

static void insert_fixup( pj_rbtree *tree, pj_rbtree_node *node )
{
    pj_rbtree_node *temp, *parent;

    PJ_CHECK_STACK();

    while (node != tree->root && node->parent->color == PJ_RBCOLOR_RED) {
        parent = node->parent;
        if (parent == parent->parent->left) {
      temp = parent->parent->right;
      if (temp->color == PJ_RBCOLOR_RED) {
          temp->color = PJ_RBCOLOR_BLACK;
          node = parent;
          node->color = PJ_RBCOLOR_BLACK;
          node = node->parent;
          node->color = PJ_RBCOLOR_RED;
      } else {
          if (node == parent->right) {
         node = parent;
         left_rotate(tree, node);
          }
          temp = node->parent;
          temp->color = PJ_RBCOLOR_BLACK;
          temp = temp->parent;
          temp->color = PJ_RBCOLOR_RED;
          right_rotate( tree, temp);
      }
        } else {
      temp = parent->parent->left;
      if (temp->color == PJ_RBCOLOR_RED) {
          temp->color = PJ_RBCOLOR_BLACK;
          node = parent;
          node->color = PJ_RBCOLOR_BLACK;
          node = node->parent;
          node->color = PJ_RBCOLOR_RED;
      } else {
          if (node == parent->left) {
          node = parent;
          right_rotate(tree, node);
          }
          temp = node->parent;
          temp->color = PJ_RBCOLOR_BLACK;
          temp = temp->parent;
          temp->color = PJ_RBCOLOR_RED;
          left_rotate(tree, temp);
     }
        }
    }
  
    tree->root->color = PJ_RBCOLOR_BLACK;
}


static void delete_fixup( pj_rbtree *tree, pj_rbtree_node *node )
{
    pj_rbtree_node *temp;

    PJ_CHECK_STACK();

    while (node != tree->root && node->color == PJ_RBCOLOR_BLACK) {
        if (node->parent->left == node) {
      temp = node->parent->right;
      if (temp->color == PJ_RBCOLOR_RED) {
          temp->color = PJ_RBCOLOR_BLACK;
          node->parent->color = PJ_RBCOLOR_RED;
          left_rotate(tree, node->parent);
          temp = node->parent->right;
      }
      if (temp->left->color == PJ_RBCOLOR_BLACK &&
          temp->right->color == PJ_RBCOLOR_BLACK)
      {
          temp->color = PJ_RBCOLOR_RED;
          node = node->parent;
      } else {
          if (temp->right->color == PJ_RBCOLOR_BLACK) {
          temp->left->color = PJ_RBCOLOR_BLACK;
          temp->color = PJ_RBCOLOR_RED;
          right_rotate( tree, temp);
          temp = node->parent->right;
          }
          temp->color = node->parent->color;
          temp->right->color = PJ_RBCOLOR_BLACK;
          node->parent->color = PJ_RBCOLOR_BLACK;
          left_rotate(tree, node->parent);
          node = tree->root;
      }
        } else {
      temp = node->parent->left;
      if (temp->color == PJ_RBCOLOR_RED) {
          temp->color = PJ_RBCOLOR_BLACK;
          node->parent->color = PJ_RBCOLOR_RED;
          right_rotate( tree, node->parent);
          temp = node->parent->left;
      }
      if (temp->right->color == PJ_RBCOLOR_BLACK &&
      temp->left->color == PJ_RBCOLOR_BLACK)
      {
          temp->color = PJ_RBCOLOR_RED;
          node = node->parent;
      } else {
          if (temp->left->color == PJ_RBCOLOR_BLACK) {
          temp->right->color = PJ_RBCOLOR_BLACK;
          temp->color = PJ_RBCOLOR_RED;
          left_rotate( tree, temp);
          temp = node->parent->left;
          }
          temp->color = node->parent->color;
          node->parent->color = PJ_RBCOLOR_BLACK;
          temp->left->color = PJ_RBCOLOR_BLACK;
          right_rotate(tree, node->parent);
          node = tree->root;
      }
        }
    }
  
    node->color = PJ_RBCOLOR_BLACK;
}


void pj_rbtree_init( pj_rbtree *tree, pj_rbtree_comp *comp )
{
    PJ_CHECK_STACK();

    tree->null = tree->root = &tree->null_node;
    tree->null->key = NULL;
    tree->null->user_data = NULL;
    tree->size = 0;
    tree->null->left = tree->null->right = tree->null->parent = tree->null;
    tree->null->color = PJ_RBCOLOR_BLACK;
    tree->comp = comp;
}

pj_rbtree_node* pj_rbtree_first( pj_rbtree *tree )
{
    register pj_rbtree_node *node = tree->root;
    register pj_rbtree_node *null = tree->null;

    PJ_CHECK_STACK();

    while (node->left != null)
  node = node->left;
    return node != null ? node : NULL;
}

pj_rbtree_node* pj_rbtree_last( pj_rbtree *tree )
{
    register pj_rbtree_node *node = tree->root;
    register pj_rbtree_node *null = tree->null;

    PJ_CHECK_STACK();

    while (node->right != null)
  node = node->right;
    return node != null ? node : NULL;
}

pj_rbtree_node* pj_rbtree_next( pj_rbtree *tree,
                  register pj_rbtree_node *node )
{
    register pj_rbtree_node *null = tree->null;

    PJ_CHECK_STACK();

    if (node->right != null) {
  for (node=node->right; node->left!=null; node = node->left)
      /* void */;
    } else {
        register pj_rbtree_node *temp = node->parent;
        while (temp!=null && temp->right==node) {
      node = temp;
      temp = temp->parent;
  }
  node = temp;
    }
    return node != null ? node : NULL;
}

pj_rbtree_node* pj_rbtree_prev( pj_rbtree *tree,
                  register pj_rbtree_node *node )
{
    register pj_rbtree_node *null = tree->null;

    PJ_CHECK_STACK();

    if (node->left != null) {
        for (node=node->left; node->right!=null; node=node->right)
     /* void */;
    } else {
        register pj_rbtree_node *temp = node->parent;
        while (temp!=null && temp->left==node) {
      node = temp;
      temp = temp->parent;
        }
        node = temp;
    }
    return node != null ? node : NULL;
}

int pj_rbtree_insert( pj_rbtree *tree,
                pj_rbtree_node *element )
{
    int rv = 0;
    pj_rbtree_node *node, *parent = tree->null,
         *null = tree->null;
    pj_rbtree_comp *comp = tree->comp;
  
    PJ_CHECK_STACK();

    node = tree->root;   
    while (node != null) {
        rv = (*comp)(element->key, node->key);
        if (rv == 0) {
      /* found match, i.e. entry with equal key already exist */
      return -1;
  }
  parent = node;
        node = rv < 0 ? node->left : node->right;
    }

    element->color = PJ_RBCOLOR_RED;
    element->left = element->right = null;

    node = element;
    if (parent != null) {
        node->parent = parent;
        if (rv < 0)
     parent->left = node;
        else
     parent->right = node;
        insert_fixup( tree, node);
    } else {
        tree->root = node;
        node->parent = null;
        node->color = PJ_RBCOLOR_BLACK;
    }
  
    ++tree->size;
    return 0;
}


pj_rbtree_node* pj_rbtree_find( pj_rbtree *tree,
                  const void *key )
{
    int rv;
    pj_rbtree_node *node = tree->root;
    pj_rbtree_node *null = tree->null;
    pj_rbtree_comp *comp = tree->comp;

    while (node != null) {
        rv = (*comp)(key, node->key);
        if (rv == 0)
      return node;
        node = rv < 0 ? node->left : node->right;
    }
    return node != null ? node : NULL;
}

pj_rbtree_node* pj_rbtree_erase( pj_rbtree *tree,
                   pj_rbtree_node *node )
{
    pj_rbtree_node *succ;
    pj_rbtree_node *null = tree->null;
    pj_rbtree_node *child;
    pj_rbtree_node *parent;

    PJ_CHECK_STACK();

    if (node->left == null || node->right == null) {
        succ = node;
    } else {
        for (succ=node->right; succ->left!=null; succ=succ->left)
     /* void */;
    }

    child = succ->left != null ? succ->left : succ->right;
    parent = succ->parent;
    child->parent = parent;

    if (parent != null) {
  if (parent->left == succ)
      parent->left = child;
        else
     parent->right = child;
    } else
        tree->root = child;

    if (succ != node) {
        succ->parent = node->parent;
        succ->left = node->left;
        succ->right = node->right;
        succ->color = node->color;

        parent = node->parent;
        if (parent != null) {
     if (parent->left==node)
          parent->left=succ;
     else
      parent->right=succ;
        }
        if (node->left != null)
     node->left->parent = succ;;
        if (node->right != null)
      node->right->parent = succ;

        if (tree->root == node)
     tree->root = succ;
    }

    if (succ->color == PJ_RBCOLOR_BLACK) {
  if (child != null)
      delete_fixup(tree, child);
        tree->null->color = PJ_RBCOLOR_BLACK;
    }

    --tree->size;
    return node;
}


unsigned pj_rbtree_max_height( pj_rbtree *tree,
                     pj_rbtree_node *node )
{
    unsigned l, r;

    PJ_CHECK_STACK();

    if (node==NULL)
  node = tree->root;

    l = node->left != tree->null ? pj_rbtree_max_height(tree,node->left)+1 : 0;
    r = node->right != tree->null ? pj_rbtree_max_height(tree,node->right)+1 : 0;
    return l > r ? l : r;
}

unsigned pj_rbtree_min_height( pj_rbtree *tree,
                     pj_rbtree_node *node )
{
    unsigned l, r;

    PJ_CHECK_STACK();

    if (node==NULL)
  node=tree->root;

    l = (node->left != tree->null) ? pj_rbtree_max_height(tree,node->left)+1 : 0;
    r = (node->right != tree->null) ? pj_rbtree_max_height(tree,node->right)+1 : 0;
    return l > r ? r : l;
}