#ifndef RBTREE_HPP
#define RBTREE_HPP

#include "Exception/ValueExistsException.hpp"

#include <boost/optional.hpp>

#define CONSISTENCY_CHECKING 1

#if CONSISTENCY_CHECKING
#define RBTREE_ASSERT(X) assert(X)
#else
#define RBTREE_ASSERT(X)
#endif

#define CONSISTENCY_CHECK(n) RBTREE_ASSERT(!n || (n)->consistent())


template<typename Data>
class RBTree {
private:
  class RBTreeNode {
    Data data;
    RBTreeNode* left;
    RBTreeNode *right;
    RBTreeNode *parent;
    bool red;

    enum Color { RED, BLACK };

  public:
    RBTreeNode(const Data& data) : left(0), right(0), parent(0), red(true)
    {
      this->data = data;
    }

    ~RBTreeNode()
    {
      delete left;
      delete right;
    }

    void setColor(Color color)
    {
      red = (color == RED);
    }

    Color color() const
    {
      return red ? RED : BLACK;
    }

    RBTreeNode* grandparent()
    {
      if ((this != NULL) && (this->parent != NULL))
	return this->parent->parent;
      else
	return NULL;
    }

    RBTreeNode* uncle()
    {
      RBTreeNode* g = grandparent();
      if (g == NULL)
	return NULL; // No grandparent means no uncle
      if (this->parent == g->left)
	return g->right;
      else
	return g->left;
    }

    RBTreeNode* getRootNode()
    {
      RBTreeNode* p = this;
      while(p->parent)
	p = p->parent;
      return p;
    }

    void rotate_left()
    {
      RBTreeNode* Q = this->right;
      RBTREE_ASSERT(Q);
      //      Set Q to be the new root.
      if (this->parent) {
	if (this->parent->left == this) {
	  this->parent->left = Q;
	}
	else {
	  RBTREE_ASSERT(this->parent->right == this);
	  this->parent->right = Q;
	}
      }
      this->right = Q->left;
      if (this->right)
	this->right->parent = this;
      Q->left = this;
      Q->parent = this->parent;
      this->parent = Q;
    }

    void rotate_right()
    {
      RBTreeNode* P = this->left;
      RBTREE_ASSERT(P);
      //Set P to be the new root.
      if (this->parent) {
	if (this->parent->left == this) {
	  this->parent->left = P;
	}
	else {
	  RBTREE_ASSERT(this->parent->right == this);
	  this->parent->right = P;
	}
      }
      this->left = P->right;
      if (this->left)
	this->left->parent = this;
      P->right = this;
      P->parent = this->parent;
      this->parent = P;
    }

    void insert_case5()
    {
      RBTreeNode* g = grandparent();
      CONSISTENCY_CHECK(this);
      CONSISTENCY_CHECK(g);

      this->parent->setColor(BLACK);
      g->setColor(RED);
      if (this == this->parent->left)
	g->rotate_right();
      else {
	RBTREE_ASSERT (this == this->parent->right);
	g->rotate_left();
      }
    }

    void insert_case4()
    {
      RBTreeNode* g = grandparent();

      CONSISTENCY_CHECK(this);

      if (!g)
	return;
      CONSISTENCY_CHECK(g);

      RBTreeNode* n = this;

      if ((this == n->parent->right) && (n->parent == g->left)) {
	n->parent->rotate_left();
	n = n->left;
      } else if ((n == n->parent->left) && (n->parent == g->right)) {
	n->parent->rotate_right();
	n = n->right;
      }
      n->insert_case5();
    }


    void insert_case3()
    {
      RBTreeNode *u = uncle();
      CONSISTENCY_CHECK(this);
      CONSISTENCY_CHECK(u);

      if ((u != NULL) && (u->color() == RED)) {
	this->parent->setColor(BLACK);
	u->setColor(BLACK);
	RBTreeNode* g;
	g = grandparent();
	g->setColor(RED);
	g->insert_case1();
      } else {
	insert_case4();
      }
    }

    void insert_case2()
    {
      CONSISTENCY_CHECK(this);
      if (this->parent->color() == BLACK)
	return; /* Tree is still valid */
      else
	insert_case3();
    }

    void insert_case1()
    {
      CONSISTENCY_CHECK(this);

      if (this->parent == NULL)
	this->setColor(BLACK);
      else
	insert_case2();
    }

    void insert(const Data& data)
    {
      if (data < this->data) {
	if (!left) {
	  left = new RBTreeNode(data);
	  left->parent = this;
	  left->insert_case1();
	}
	else
	  left->insert(data);
      }
      else if (data > this->data) {
	if (!right) {
	  right = new RBTreeNode(data);
	  right->parent = this;
	  right->insert_case1();
	}
	else
	  right->insert(data);
      }
      else {
	throw ValueExistsException();
      }
      CONSISTENCY_CHECK(this);
    }

    boost::optional<Data> find(const Data& data)
    {
      if (this->data == data) {
	return this->data;
      }
      else if (data < this->data) {
	if (left)
	  return left->find(data);
	else
	  return boost::optional<Data>();
      }
      else {
	if (right)
	  return right->find(data);
	else
	  return boost::optional<Data>();
      }
      return boost::optional<Data>();
    }

    uint size() const
    {
      return 1 + (left ? left->size() : 0) + (right ? right->size() : 0);
    }

    uint depth() const
    {
      return 1 + std::max((left ? left->depth() : 0), (right ? right->depth() : 0));
    }

    uint total_depth(uint n) const
    {
      return n + (left ? left->total_depth(n + 1) : 0) + (right ? right->total_depth(n + 1) : 0);
    }

    bool consistent() const
    {
      bool retVal = true;
      if (left)
	retVal = retVal && (left->parent == this) && left->consistent();
      if (right)
	retVal = retVal && (right->parent == this) && right->consistent();
      return retVal;
    }

    friend class RBTree;
  };

  RBTreeNode* node;
public:
  RBTree() : node(0)
  {
  }

  ~RBTree()
  {
    delete node;
  }

  bool consistent() const
  {
    if (node == 0)
      return true;
    if (node->parent == 0 && node->consistent())
      return true;
    return false;
  }

  void insert(const Data& data)
  {
    if (node == 0) {
      node = new RBTreeNode(data);
    }
    else {
      node->insert(data);
    }
    node = node->getRootNode();
    CONSISTENCY_CHECK(this);
  }

  boost::optional<Data> find(const Data& data)
  {
    if (node) {
      return node->find(data);
    }
    else
      return boost::optional<Data>();
  }

  uint size() const
  {
    if (!node)
      return 0;
    else
      return node->size();
  }

  uint depth() const
  {
    if (!node)
      return 0;
    return node->depth();
  }

  float avg_depth() const
  {
    if (!node)
      return 0;
    return double(total_depth())/size();
  }

  uint total_depth() const
  {
    if (!node)
      return 0;
    return node->total_depth(1);
  }

  uint optimal_depth() const
  {
    int s = size();
    if (s == 0)
      return 0;
    int d = s;
    int targetlevel = 1;
    while (d >>= 1)
      ++targetlevel;
    return targetlevel;
  }
};

#endif //RBTREE_HPP