// Binary Search Tree implementation using zybooks algorithms
//
// Rafael Arce Nazario 2019

#define CATCH_CONFIG_MAIN  // This tells Catch to provide a main() - only do this in one cpp file
#include "catch.hpp"
#include <queue>

#include <iostream>
using namespace std;

class BTNode {
public:
    int key;
    BTNode *left, *right;
    BTNode(int k = 0, BTNode *l = NULL, BTNode *r = NULL) : key(k), left(l), right(r) {}
};

class BST {
private:
    BTNode *root;

public:
    BST(): root(NULL) {}
    void insert(int k);
    void insertRec(int k);
    void insertRec(int k, BTNode *r);

    void remove(int k);
    void remove(int k, BTNode *r);
    void recDisplay(ostream &out) const;
    void recDisplay(ostream &out, BTNode *cur, int dist) const;

    BTNode *search(int k) const;
    BTNode *searchRec(int k) const;
    BTNode *searchRec(int k, BTNode *r) const;


    int getHeight() const;
    int getHeight(BTNode *r) const;
    int sum() const;
    int sum(BTNode *r) const;
    int leafQty() const;
    int leafQty(BTNode *n) const;
    int parentsOfTwo() const;
    int parentsOfTwo(BTNode *n) const;

    string InOrder() const;
    void InOrder(BTNode *n, string &st) const;

    string BFS() const;
};



// function search:
//    Given k, a key to search, returns a pointer to the first
//    node that contains such key. If not found returns a NULL pointer.

BTNode* BST::search(int k) const {
    BTNode* cur = root;
    while (cur) {
        if (k == cur->key)     {return cur;}
            // Found
        else if (k < cur->key) { cur = cur->left; }
        else                   { cur = cur->right; }
    }
    return NULL;
}


// function insert:
//    Given k, a key to insert, inserts into the BST.
//    In this implementation if the key already exists in the tree
//    it will be inserted in the RIGHT subtree of the existing key.

void BST::insert(int k) {
    BTNode *cur;
    BTNode *n = new BTNode(k);

    if (!root) {
        // This tree is empty, the new node is the root!
        root = n;
    }
    else {
        // Lets search the tree for a place to insert....
        cur = root;
        while (cur) {
            if (k < cur->key){
                // The value to insert is less than current node
                if (cur->left == NULL) {
                    // If we have reached a node who lacks a left child
                    // then our new node will be the left child
                    cur->left = n;
                    cur = NULL;
                }
                else {
                    // keep looking for a place to insert, on the left subtree
                    cur = cur->left;
                }
            }
            else {
                // The value to insert is greater or equal than current node
                if (cur->right == NULL){
                    // If we have reached a node who lacks a right child
                    // then our new node will be the right child
                    cur->right = n;
                    cur = NULL;
                }
                else {
                    // keep looking for a place to insert, on the right subtree
                    cur = cur->right;
                }
            }
        }
    }
}

// functions recDisplay:
//    Print the tree to the standard output.

void BST::recDisplay(ostream &out) const {
    recDisplay(out, root, 0);
}

void BST::recDisplay(ostream &out, BTNode *cur, int dist) const{
    if (cur) {
        if (cur->right) recDisplay(out, cur->right, dist + 1);
        for (int i = 0; i < dist; i++) cout << "    ";
        out << cur->key << endl;
        if (cur->left) recDisplay(out, cur->left, dist + 1);
    }
}

// functions remove:
//    Given k, a key to remove, removes the first node that it finds with
//    such key.


void BST::remove(int k) {
    remove(k, root);
}

void BST::remove(int k, BTNode *r) {
    BTNode *parent = NULL;
    BTNode *cur = r;
    while (cur) { // Search for node
        if (cur->key == k) { // Node found
            if (cur->left == NULL && cur->right == NULL) {        // Remove leaf
                cout << "cur is a leaf: " << cur->key << endl;
                if (!parent) { // Node is root
                    root = NULL;
                }
                else if (parent->left == cur)
                    parent->left = NULL;
                else
                    parent->right = NULL;
                delete cur;
            }
            else if (cur->left && cur->right == NULL) {    // Remove node with only left child
                if (!parent) // Node is root
                    root = cur->left;
                else if (parent->left == cur)
                    parent->left = cur->left;
                else
                    parent->right = cur->left;
                delete cur;
            }
            else if (cur->left == NULL && cur->right) {    // Remove node with only right child
                if (!parent) // Node is root
                    root = cur->right;
                else if (parent->left == cur)
                    parent->left = cur->right;
                else
                    parent->right = cur->right;
                delete cur;
            }
            else {                                  // Remove node with two children
                // Find successor (leftmost child of right subtree)
                BTNode *suc = cur->right;

                while (suc->left )
                    suc = suc->left;
                int successorData = suc->key;
                remove(suc->key, cur);     // Remove successor
                cur->key = successorData;
            }
            return; // Node found and removed
        }
        else if (cur->key < k) { // Search right
            parent = cur;
            cur = cur->right;
        }
        else {                     // Search left
            parent = cur;
            cur = cur->left;
        }
    }
    return; // Node not found
}


// function getHeight:
//    Returns the height of the tree.

int  BST::getHeight() const {
    return getHeight(root);
}

int  BST::getHeight(BTNode *r) const {
    if (!r) return -1;

    int leftHeight = getHeight(r->left);
    int rightHeight = getHeight(r->right);
    return 1 + max(leftHeight, rightHeight);
}

// function sum:
//    Returns the sum of all keys in the tree.

int BST::sum() const { return sum(root);}

int BST::sum(BTNode *n) const {
    if (!n) return 0;
    else return n->key + sum(n->left) + sum(n->right);
}


// function leafQty:
//    Returns the number of leaves in the tree.

int BST::leafQty() const { return leafQty(root);}

int BST::leafQty(BTNode *n) const {
    if (!n) return 0;
    if (n->left == NULL && n->right == NULL) return 1;
    else return leafQty(n->left) + leafQty(n->right);
}

// function parentsOfTwo:
//    Returns the number of nodes that are parents of two children.

int BST::parentsOfTwo() const {
    return parentsOfTwo(root);
}
int BST::parentsOfTwo(BTNode *n) const {
    if (!n || (!n->left && !n->right)) return 0;
    int resultLeft = parentsOfTwo(n->left);
    int resultRight = parentsOfTwo(n->right);
    if (n->left && n->right) return 1 + resultLeft + resultRight;
    else return resultLeft + resultRight;
}

// function InOrder:
//    Will return a string containing the sequence of visited keys
//    during an in-order traversal (BFS) of the tree.

string BST::InOrder() const {
    string st;
    InOrder(root, st);
    return st;
}
void BST::InOrder(BTNode *n, string &st) const {
    if (n) {
        InOrder(n->left, st);
        st = st + to_string(n->key) + " ";
        InOrder(n->right, st);
    }
}

// function BFS:
//    Will return a string containing the sequence of visited keys
//    during a breadth-first traversal of the tree.

string BST::BFS() const {
    string st;
    queue<BTNode*> q;

    if(root!=NULL) q.push(root);

    while(!q.empty()) {
        BTNode *f = q.front();
        st.append(to_string(f->key));
        st.append(" ");
        if (f->left  != NULL) q.push(f->left);
        if (f->right != NULL) q.push(f->right);
        q.pop();
    }

    return st;
}


// function searchRec:
//    A recursive version of the search algorithm.

BTNode* BST::searchRec(int k) const { return searchRec(k,root); }

BTNode* BST::searchRec(int k, BTNode *r) const {
    if (r) {
        if (k == r->key) return r;
        else if (k < r->key)
            return searchRec(k, r->left);
        else
            return searchRec(k, r->right);
    }
    return NULL;
}

// function insertRec:
//    A recursive version of the insert algorithm.

void BST::insertRec(int k) {
    if (!root) root = new BTNode(k);
    else insertRec(k, root);
}
void BST::insertRec(int k, BTNode *r) {
    if (k < r->key) {
        if (r->left == NULL) r->left = new BTNode(k);
        else insertRec(k, r->left);
    }
    else { // k is greater than r->key
        if (r->right == NULL) r->right = new BTNode(k);
        else insertRec(k, r->right);
    }
}


TEST_CASE( "BST is tested", "[BST]" ) {
    vector<int> v = {8, 9, 5, 2, 4, 10, 1};
    BST B;
    for (auto e: v) B.insert(e);
    REQUIRE (B.InOrder() == "1 2 4 5 8 9 10 ");
    REQUIRE (B.BFS() == "8 5 9 2 10 1 4 ");
}