#ifndef BINARYTREE_H_
#define BINARYTREE_H_

#include <iostream>
#include "Collection.h"
using std::cout;
using std::endl;

/**
 * A class for Binary Tree Nodes
 */
class BinaryTreeNode {
  public:
    NodeData *data;
    BinaryTreeNode *left;
    BinaryTreeNode *right;
  
    BinaryTreeNode(NodeData *pData,
      BinaryTreeNode *pLeft = NULL, BinaryTreeNode *pRight = NULL) 
      : data(pData), left(pLeft), right(pRight) {};
    
    ~BinaryTreeNode() {
      if (left != NULL) {
        delete left;
      }
      if (right != NULL) {
        delete right;
      }
    }

    
    /**
     * Returns true if the given integer key is a member of the Binary Tree
     * rooted at this Node.
     */
    bool member(int key);

    /**
     * Removes the node that matches the given integer key from the Binary Tree
     * rooted at this Node.
     * 
     * Returns the removed BinaryTreeNode, with its left child set to the new
     * root of this subtree.
     */
    BinaryTreeNode* removeKey(int key);
    
    /**
     * Inserts the given NodeData into the subtree rooted at this BinaryTreeNode.
     */
    void insert(NodeData *newData);
    
    /**
     * Returns the first (smallest value key) NodeData of the subtree rooted at this
     * BinaryTreeNode.
     */
    NodeData* first();
    
    /**
     * Returns the size of the subtree rooted at this node.
     */
    int size();
    
    /**
     * Returns true if this Node is a leaf node.
     */
    bool isLeaf() {
      return left == NULL && right == NULL;
    }
    
    /**
     * Returns the count of immediate children at this Node.
     */
    int getChildCount() {
      if (left == NULL) {
        return right == NULL ? 0 : 1;
      }
      else {
        return right == NULL ? 1 : 2;
      }
    }
    
    /**
     * Prints the subtree rooted at this node, with the given spacing.
     */
    void print(int spacing) {
      for (int i = 0; i < spacing; i++) {
        cout << ' ';
      }
      data->print();
      cout << endl;
      if (left != NULL) {
        left->print(spacing + 1);
      }
      if (right != NULL) {
        right->print(spacing + 1);
      }
    }
  
};

/**
 * A basic implementation of our Collection ADT
 * using an ordered Binary Tree.
 */
class BinaryTree : public Collection {
  BinaryTreeNode *root;
  
  public:
    BinaryTree() : root(NULL) {};
    ~BinaryTree() {
      if (root != NULL) {
        delete root;
      }
    }
    
    /**
     * Returns the count of the number of elements in the Collection
     */
    int size() const {
      if (root == NULL) {
        return 0;
      }
      else {
        return root->size();
      }
    }
    
    /**
     * Adds the given NodeData to the Collection. Because this is
     * an ordered BinaryTree, add(NodeData) places the new node in
     * the correct position in the tree to maintain ordering over
     * keys.
     */
    void add(NodeData *data) {
      if (root == NULL) {
        root = new BinaryTreeNode(data);
      }
      else {
        root->insert(data);
      }
    }
    
    /**
     * Returns the NodeData of the first element of the Collection.
     * For the ordered BinaryTree this returns the smallest key
     * in the Collection.
     */
    NodeData* first() const {
      if (root == NULL) {
        return NULL;
      }
      else {
        return root->first();
      }
    }
    
    /**
     * Returns true if the given integer key is a member of the Collection.
     */
    bool member(int key) const {
      if (root == NULL) {
        return false;
      }
      else {
        return root->member(key);
      }
    }
    
    /**
     * Not implemented yet, do not implement for HW3
     */
    NodeData* remove(int) {
      return NULL;
    }
    
    /**
     * Removes the node corresponding to the given key.
     */
    NodeData* removeKey(int key) {
      if (root == NULL) {
        return NULL;
      }
      else {
        BinaryTreeNode *removed = root->removeKey(key);
        if (removed == NULL) {
          return NULL;
        }
        else {
          if (removed == root) {
            root = removed->left;
          }
          
          removed->left = NULL;
          removed->right = NULL;
          NodeData *data = removed->data;
          
          delete removed;
          return data;
        }
      }
    }
    
    /**
     * Prints the BinaryTree
     */
    void print() {
      if (root != NULL) {
        root->print(0);
      }
    }
};

#endif /*BINARYTREE_H_*/