// Arne Andersson tree implementation
// By: Nick Cash
// Computer Science III

// For more information on this data type please visit Wikipedia

// Note: This implementation is a string to integer map

// Includes
#include <iostream>
#include "aa-tree.h"

// Constructor
cAATree::cAATree()
{
   Root = NULL;
}

// Desturctor
cAATree::~cAATree()
{
}

// Used to lookup keys, and if they don't exist, add them
int cAATree::Lookup( const std::string & key )
{
   cNode* node = new cNode;
   
   // Set default values
   node->key = key;
   node->count = 0;
   node->level = 1;
   node->Parent = node->Left = node->Right = NULL;
   
   // Start recursion
   node = Insert(Root,node);
    
   // Return number of occurances so far
   return node->count;
}

// Perform a right rotation when inserting a left red link.
void cAATree::Skew(cNode* x)
{
   cNode* y = x->Left;
   
   // Move X from its parent
   if ( x->Parent->Left == x )
     x->Parent->Left = y;
   else
     x->Parent->Right = y;
     
   // Set parents
   y->Parent = x->Parent;
   x->Parent = y;

   x->Left = y->Right;

   // X's left's parent becomes X
   if ( x->Left != NULL )
    x->Left->Parent = x;
    
   // Hang X on Y's right
   y->Right = x;
   
   // Set level
   x->level = (x->Left ? x->Left->level+1 : 1);
}

// Do a left rotation (if needed) if an insert creatures two red links right in a row
bool cAATree::Split(cNode* x)
{
   cNode* y = x->Right;
   
   // Perform left rotation if Y has a child Right and its level is == X's level
   if ( y && y->Right && (y->Right->level == x->level))
   {
     // Move X from its parent
     if ( x->Parent->Left == x )
      x->Parent->Left = y;
     else
      x->Parent->Right = y;
      
     // Set parents
     y->Parent = x->Parent;
     x->Parent = y;
     
     x->Right = y->Left;
     
     // X's right's parent becomes X
     if ( x->Right != NULL )
      x->Right->Parent = x;
      
     // Hang X on Y's left
     y->Left = x;
     
     // Set level
     y->level = x->level+1;
     
     // Let the rebalance function know we did a left rotation
     return true;
   }
   
   // Let the rebalance function know we did -NOT- do a left rotation
   return false;
}

// Attempt to rebalance the tree if need be.
// X is always a newly added leaf node
void cAATree::Rebalance(cNode* x)
{
  x->level = 1;
  x->Left = x->Right = NULL;
  
  // Climmb up the tree and balance
  for ( x = x->Parent; x != Root; x = x->Parent )
  {
     // check for right rotation
     if ( x->level != (x->Left ? x->Left->level + 1 : 1 ))
     {
        Skew(x);
        
        if ( (x->Right == NULL) || (x->level != x->Right->level))
         x = x->Parent;
     }
     
     // check for left rotation, end rebalance if no rotation was made
     if ( (x->Parent != Root) && (Split(x->Parent) == false) )
      break;
  }
}

// Recursive function to add a node to the tree, or update its word count
cAATree::cNode* cAATree::Insert(cNode* x, cNode* insert)
{
     // Check for NULL Root, our special case
     if ( Root == NULL )
     {
       insert->count = 1;
       insert->Parent = insert->Left = insert->Right = NULL;

       Root = insert;

       return Root;
     }

     // Go left?
     if ( insert->key < x->key )
     {
        if ( x->Left ) // if we have a left, call Insert again
         return Insert( x->Left, insert );

        // otherwise add the node left
        insert->count = 1;
        x->Left = insert;
        insert->Parent = x;
        Rebalance(insert);

        return insert;
     }

     // Go right?
     if ( insert->key > x->key )
     {
        if ( x->Right ) // if we have a right, call Insert again
         return Insert( x->Right, insert );

        // otherwise add the node left
        insert->count = 1;
        x->Right = insert;
        insert->Parent = x;
        Rebalance(insert);

        return insert;
     }

     // So x isn't NULL, and insert->key == x->key
     // We fonud what we were looking for, incrememnt reference
     x->count++;

     // Get rid of insert, the node we would have hung here
     delete insert;

     // Give back the node we found
     return x;
}

// Start recursive function that performs an in-order traversal of the tree
// And prints each element, along with its count
void cAATree::Print()
{
   PrintNode(Root);
}

// Recursive function that prints out keys and values.
void cAATree::PrintNode(cNode* node)
{
  // If we hit a leaf return
  if ( !node )
   return;
   
  // Print all left first
  PrintNode(node->Left);

  // Print this element next
  std::cout << node->key << ' ' << node->count << std::endl;
  
  // Print right
  PrintNode(node->Right);
}

// Start recursive function that counts the number of nodes in the tree
unsigned int cAATree::TotalNodes()
{
  return CountNode(Root);
}

// Recursive function
unsigned int cAATree::CountNode(cNode* node)
{
  // Leaves don't count
  if ( !node )
   return 0;
   
  unsigned int count = 1; // count variable, set to 1 to count this node
  
  // Count children and so on
  count += CountNode(node->Left);
  count += CountNode(node->Right);
  
  // Return our gathered number
  return count;
}