#include "trees.hpp"
#include <iostream>
#include <memory>
#include <string>
#include <vector>

namespace tree_implementation
{

    TreeNode::TreeNode(std::string word)
    {
        key = word;
        color = "red";
        leftChild = nullptr;
        rightChild = nullptr;
        this->parent = nullptr;
        return;
    }

    TreeList::TreeList(void)
    {
        head = nullptr;
        return;
    }

    TreeInterface::TreeInterface(void)
    {
        return;
    }

    // Inserts a node into a BST
    void TreeInterface::Insert(std::shared_ptr<TreeNode> z)
    {
        std::shared_ptr<TreeNode> y = nullptr;
        std::shared_ptr<TreeNode> x = tree.head;
        while (x)
        {
            y = x;
            if (z.get()->key < x.get()->key)
                x = x.get()->leftChild;
            else
                x = x.get()->rightChild;
        }
        z.get()->parent = y;
        if (!y)
            tree.head = z;
        else if (z.get()->key < y.get()->key)
            y->leftChild = z;
        else
            y->rightChild = z;
        return;
    }

    // Searches for the given word in a tree
    std::shared_ptr<TreeNode> TreeInterface::Search(std::string wordToFind)
    {
        return Search(tree.head, wordToFind);
    }

    bool TreeInterface::IsSearchSuccessful(std::shared_ptr<TreeNode> foundNode, std::string key)
    {
        if (foundNode)
            return true;
        std::cout << "No node found with key " << key << std::endl;
        return false;
    }

    // Prints tree while traversing it
    void TreeInterface::InOrderTreeTraversal(std::shared_ptr<TreeNode> viewedNode)
    {
        if (viewedNode)
            InOrderTreeTraversal(viewedNode->leftChild);
            std::cout << viewedNode.get()->key << '\n';
            InOrderTreeTraversal(viewedNode->rightChild);
        return;
    }

    // Prints the given word's parent's word if found
    void TreeInterface::PrintParentKey(std::string key)
    {
        std::shared_ptr<TreeNode> foundNode = Search(key);
        if (!IsSearchSuccessful(foundNode, key)) return;
        std::cout << "The parent's word is " << foundNode->parent->key << '\n';
        return;
    }

    // Prints the given word's left child's word if found
    void TreeInterface::PrintLeftChild(std::string key)
    {
        std::shared_ptr<TreeNode> foundNode = Search(key);
        if (!IsSearchSuccessful(foundNode, key)) return;
        std::cout << "The left child's word is " << foundNode->leftChild->key << '\n';
        return;
    }

    // Prints the given word's right child's word if found
    void TreeInterface::PrintRightChild(std::string key)
    {
        std::shared_ptr<TreeNode> foundNode = Search(key);
        if (!IsSearchSuccessful(foundNode, key)) return;
        std::cout << "The right child's word is " << foundNode->rightChild->key << '\n';
        return;
    }

    // Recursive search function for Search()
    std::shared_ptr<TreeNode> TreeInterface::Search(std::shared_ptr<TreeNode> viewedNode, std::string wordToFind)
    {
        if ((!viewedNode) || (wordToFind == viewedNode.get()->key))
            return viewedNode;
        if (wordToFind < viewedNode.get()->key)
            return Search(viewedNode->leftChild, wordToFind);
        return Search(viewedNode->rightChild, wordToFind);
    }

    // Insert a node into a BST
    void BinarySearchTree::Insert(std::string keyToInsert)
    {
        std::shared_ptr<TreeNode> newNode = std::make_shared<TreeNode> (keyToInsert);
        TreeInterface::Insert(newNode);
        return;
    }

    // Insert the entire vector word list into tree
    void BinarySearchTree::InsertWordList(std::vector<std::string>* newWordList)
    {
        for (int i = 0; i < newWordList->size(); i++)
        {
            Insert(newWordList->at(i));
        }
    }

    // Prints the path to root in a BST
    void BinarySearchTree::PrintPathToRoot(std::string key)
    {
        std::shared_ptr<TreeNode> selectedNode = Search(key);
        if (!IsSearchSuccessful(selectedNode, key)) return;
        int timesPrintedOnLine = 0;
        std::cout << "Path from:" << selectedNode->key << std::endl;
        while (selectedNode = selectedNode.get()->parent)
        {
            std::cout << " -> " << selectedNode.get()->key;
            if (timesPrintedOnLine < 5)
                timesPrintedOnLine++;
            else
            {
                timesPrintedOnLine = 0;
                std::cout << std::endl;
            }
        } 
        std::cout << std::endl;
        return;
    }

    // Insert a node into a RBT
    void RedBlackTree::Insert(std::string keyToInsert)
    {
        std::shared_ptr<TreeNode> newNode = std::make_shared<TreeNode> (keyToInsert);
        TreeInterface::Insert(newNode);
        InsertFixup(newNode);
        return;
    }

    // Insert the entire vector word list into tree
    void RedBlackTree::InsertWordList(std::vector<std::string>* newWordList)
    {
        for (int i = 0; i < newWordList->size(); i++)
        {
            Insert(newWordList->at(i));
        }
    }

    // Print path to root in RBT
    // TODO: Implement printing path to root for RBT
    void RedBlackTree::PrintPathToRoot(std::string key)
    {
        std::shared_ptr<TreeNode> selectedNode = Search(key);
        if (!IsSearchSuccessful(selectedNode, key)) return;
        int timesPrintedOnLine = 0;
        std::cout << "Path from:" << selectedNode->key << std::endl;
        while (selectedNode = selectedNode.get()->parent)
        {
            std::cout << " -> " << selectedNode.get()->key;
            std::cout << '(' << selectedNode.get()->color << ')';
            if (timesPrintedOnLine < 5)
                timesPrintedOnLine++;
            else
            {
                timesPrintedOnLine = 0;
                std::cout << std::endl;
            }
        } 
        std::cout << std::endl;
        return;
    }

    // Print color of word if found in RBT
    void RedBlackTree::PrintColor(std::string key)
    {
        std::shared_ptr<TreeNode> foundNode = Search(key);
        if (!IsSearchSuccessful(foundNode, key)) return;
        std::cout << "The color is " << foundNode->color << '\n';
        return;
    }

    // Print color of word's parent if found in RBT
    void RedBlackTree::PrintParentColor(std::string key)
    {
        std::shared_ptr<TreeNode> foundNode = Search(key);
        if (!IsSearchSuccessful(foundNode, key)) return;
        std::cout << "The color is " << foundNode->parent->color << '\n';
        return;
    }

    // Print color of word's uncle if found in RBT
    void RedBlackTree::PrintUncleColor(std::string key)
    {
        std::shared_ptr<TreeNode> foundNode = Search(key);
        if (!IsSearchSuccessful(foundNode, key)) return;
        std::cout << "The color is " << GetUncleNode(foundNode).get()->color << '\n';
        return;
    }

    void RedBlackTree::InsertFixup(std::shared_ptr<TreeNode> z)
    {
        std::shared_ptr<TreeNode> y;
        while (z->parent && z->parent->color == "red")
        {
            if (z->parent == z->parent->parent->leftChild)
            {
                y = z->parent->parent->rightChild;
                if (y && y->color == "red")
                {
                    z->parent->color = "black";
                    y->color = "black";
                    z->parent->parent->color = "red";
                    z = z->parent->parent;
                } else {
                    if (z == z->parent->rightChild)
                    {
                        z = z->parent;
                        LeftRotate(z);
                    }
                    z->parent->color = "black";
                    z->parent->parent->color = "red";
                    RightRotate(z->parent->parent);
                }
            } else {
                y = z->parent->parent->leftChild;
                if (y && y->color == "red")
                {
                    z->parent->color = "black";
                    y->color = "black";
                    z->parent->parent->color = "red";
                    z = z->parent->parent;
                } else {
                    if (z == z->parent->leftChild)
                    {
                        z = z->parent;
                        RightRotate(z);
                    }
                    z->parent->color = "black";
                    z->parent->parent->color = "red";
                    LeftRotate(z->parent->parent);
                }
            }
        }
        tree.head->color = "black";
        return;
    }

    // Returns the uncle node in RBT
    std::shared_ptr<TreeNode> RedBlackTree::GetUncleNode(std::shared_ptr<TreeNode> startNode)
    {
        if (startNode->parent == startNode->parent->parent->leftChild)
            return startNode->parent->parent->rightChild;
        return startNode->parent->parent->leftChild;
    }

    // Performs left rotate on a given node
    void RedBlackTree::LeftRotate(std::shared_ptr<TreeNode> x)
    {
        std::shared_ptr<TreeNode> y = x->rightChild;
        x->rightChild = y->leftChild;
        if (y->leftChild)
            y->leftChild->parent = x;
        y->parent = x->parent;
        if (!x->parent)
            tree.head = y;
        else if (x == x->parent->leftChild)
            x->parent->leftChild = y;
        else
            x->parent->rightChild = y;
        y->leftChild = x;
        x->parent = y;
        return;
    }

    // Performs right rotate on a given node
    void RedBlackTree::RightRotate(std::shared_ptr<TreeNode> x)
    {
        std::shared_ptr<TreeNode> y = x->leftChild;
        x->leftChild = y->rightChild;
        if (y->rightChild)
            y->rightChild->parent = x;
        y->parent = x->parent;
        if (!x->parent)
            tree.head = y;
        else if (x == x->parent->leftChild)
            x->parent->leftChild = y;
        else
            x->parent->rightChild = y;
        y->rightChild = x;
        x->parent = y;
        return;
    }
}