#ifndef TREESET_PRIVATE_HPP_INCLUDED
#define TREESET_PRIVATE_HPP_INCLUDED 1

#include <iostream>
#include <cassert>

using namespace std;

/*
 * Public member functions.  Most of these call recursive private helper
 * functions on the root of the TreeSet.
 */

template <typename T>
TreeSet<T>::TreeSet()
    : size_(0), rotations_(0), root_(NULL)
{
    // Nothing else to do!
}

template <typename T>
TreeSet<T>::~TreeSet()
{
    destroySubtree(root_);
}

template <typename T>
size_t TreeSet<T>::size() const
{
    return size_;
}

template <typename T>
void TreeSet<T>::insert(const T& item)
{
    insert(item, root_);
    ++size_;
}

template <typename T>
bool TreeSet<T>::exists(const T& item) const
{
    return exists(item, root_);
}

template <typename T>
void TreeSet<T>::printStatistics()
{
    cerr << "height " << height() << ", " <<  rotations_ << " rotations, "
         << "size " << size() << endl;
}

template <typename T>
ostream& TreeSet<T>::print(ostream& out) const
{
    return print(out, root_);
}

template <typename T>
size_t TreeSet<T>::height()
{
    return height(root_);
}

/*
 * Recursive private helper functions.
 */

template <typename T>
void TreeSet<T>::destroySubtree(Node* here)
{
    if (here != NULL) {
        destroySubtree(here->left_);
        destroySubtree(here->right_);
    }
    
    delete here;
}

template <typename T>
void TreeSet<T>::insert(const T& item, Node*& here)
{
    if (here == NULL)
        here = new Node(item);
    else {
        here->maybeRecolor();
        
        // Don't make root_ red!
        if (here == root_)
            here->isRed_ = false;
        
        if (item < here->value_)
            insert(item, here->left_);
        else
            insert(item, here->right_);
        
        rotations_ += Node::maybeRotate(here);
    }
}

template <typename T>
bool TreeSet<T>::exists(const T& item, const Node* const& here) const
{
    if (here == NULL)
        return false;
    
    if (item == here->value_)
        return true;
    
    if (item < here-> value_)
        return exists(item, here->left_);
    else
        return exists(item, here->right_);
}

template <typename T>
ostream& TreeSet<T>::print(ostream& out, const Node* const& here) const
{
    if (here == NULL)
        out << "-";
    else {
        out << "(";
        print(out, here->left_);
        out << ", " << here->value_ << ", ";
        print(out, here->right_);
        out << ")";
    }
    
    return out;
}

template <typename T>
int TreeSet<T>::height(Node*& here)
{
    if (here == NULL)
        return -1;
    else
        return 1 + max(height(here->left_), height(here->right_));
}

/*
 * Public node member functions
 */

template <typename T>
TreeSet<T>::Node::Node(const T& item)
    : isRed_(true), value_(item), left_(NULL), right_(NULL)
{
    // Nothing else to do!
}


// A node needs rotating if one of its children and one of that child's children
// are both red. It should be rotated in the opposite direction of the red child.
// If the grandchild is an inner child (a.k.a. right's left or
// left's right, as opposed to right's right or left's left), then the child
// also needs rotating, in the opposite direction of the red grandchild.
// In addition, the original node should become red and the new root of the
// subtree should become black.

template <typename T>
size_t TreeSet<T>::Node::maybeRotate(Node*& rotateThis)
{
    size_t rotations = 0;
   
    if (isRed(rotateThis->left_) && hasRedKids(rotateThis->left_)) {
        rotateThis->isRed_ = true;
        
        if (isRed(rotateThis->left_->right_)) {
            rotateLeft(rotateThis->left_);
            ++rotations;
        }
        
        rotateRight(rotateThis);
        ++rotations;
        
        // Set the new subtree root to be black
        rotateThis->isRed_ = false;
    }
    
    if (isRed(rotateThis->right_) && hasRedKids(rotateThis->right_)) {
        rotateThis->isRed_ = true;
        
        if (isRed(rotateThis->right_->left_)) {
            rotateRight(rotateThis->right_);
            ++rotations;
        }
        
        rotateLeft(rotateThis);
        ++rotations;
        
        // Set the new subtree root to be black
        rotateThis->isRed_ = false;
    }
    // If we've rotated more than twice, the tree was unbalanced.
    assert(rotations <= 2);
        
    return rotations;
}

template <typename T>
void TreeSet<T>::Node::maybeRecolor()
{
    if (isRed(left_) && isRed(right_) && !isRed_) {
        left_->isRed_ = false;
        right_->isRed_ = false;
        isRed_ = true;
    }
}

/* 
 * Private Node member functions
 */
 
template <typename T>
void TreeSet<T>::Node::rotateLeft(Node*& rotateThis)
{
    
    assert(rotateThis->right_ != NULL);
    
    // Move pointers around
    Node* newTop       = rotateThis->right_;
    rotateThis->right_ = newTop->left_;
    newTop->left_      = rotateThis;
    
    // The top of the subtree is what used to be rotateThis->right_
    rotateThis = newTop;
}

template <typename T>
void TreeSet<T>::Node::rotateRight(Node*& rotateThis)
{
    
    assert(rotateThis->left_ != NULL);
    
    // Move pointers around
    Node* newTop       = rotateThis->left_;
    rotateThis->left_  = newTop->right_;
    newTop->right_     = rotateThis;
    
    // The top of the subtree is what used to be rotateThis->left_
    rotateThis = newTop;
}

template <typename T>
bool TreeSet<T>::Node::isRed(Node*& here)
{
    if (here == NULL)
        return false;
    
    return here->isRed_;
}

template <typename T>
bool TreeSet<T>::Node::hasRedKids(Node*& here)
{
    if (here == NULL)
        return false;
    
    return (isRed(here->left_) || isRed(here->right_));
}

#endif