////////////////////////////////////////// main.cpp //////////////////////////////////////////
#include <iostream>
#include <cassert>
#include <string>

#include "Tree.hpp"

using KeyType = size_t;
using ValueType = std::string;
using TreeType = Tree<KeyType, ValueType>;

int main() {
    TreeType tree;
    tree.insert(5, "");
    tree.insert(3, "");
    tree.insert(6, "");
    tree.insert(7, "");

    // Проверка удаления листа (не корня)
    tree.erase(7);
    assert(tree.size() == 3U);
    assert((tree.keys() == std::vector<KeyType>{ 3, 5, 6 }));
    assert(tree.balanceFactor() == 0);

    // Проверка удаления корня + проверка балансировки
    tree.insert(7, "");
    tree.erase(5);
    assert(tree.size() == 3U);
    assert((tree.keys() == std::vector<KeyType>{ 3, 6, 7 }));
    assert(tree.balanceFactor() == 0);

    return 0;
}

////////////////////////////////////////// Tree.hpp //////////////////////////////////////////
#pragma once
#include <vector>
#include <functional>

#include "TreeNode.hpp"
#include "TreeForwardIterator.hpp"
#include "TreeReverseIterator.hpp"


//! Класс, реализующий бинарное дерево поиска
//!
//! @tparam KeyT       Тип ключа
//! @tparam ValueT     Тип значения
template <typename KeyT,
          typename ValueT>
class Tree {
public:
    using TreeNode = TreeNode<KeyT, ValueT>;
    using ForwardIterator = TreeForwardIterator<TreeNode>;
    using ReverseIterator = TreeReverseIterator<TreeNode>;

    //! Конструктор
    Tree() = default;

    //! Конструктор копирования
    //!
    //! @param  other   Другой объект дерева
    Tree(const Tree &other) {
        using CopyFunctionT = std::function<TreeNode*(const TreeNode*)>;

        const CopyFunctionT copy = [&](const TreeNode* source) {
            TreeNode* dist = new TreeNode(source->key, source->value);
            if (source->left) {
                dist->left = copy(source->left);
            }
            if (source->right) {
                dist->right = copy(source->right);
            }
            return dist;
        };
        _root = copy(other._root);
    }

    //! Деструктор
    ~Tree() {
        clearNode(_root);
    }

    //! Размер дерева
    //!
    //! @return Количество нод в дереве
    std::size_t size() const {
        return _size;
    }

    //! Отчистка дерева
    void clear() {
        clearNode(_root);
        _root = nullptr;
    }

    //! Проверка дерева на пустоту
    bool empty() const {
        return _size == 0;
    }
    
    //! Доступ к данным по ключу.
    //!
    //! @param  key     Ключ элемента в дереве
    //! @return         Значение элемента внутри дерева
    ValueT& operator()(const KeyT& key) const {
        TreeNode* node = _find(key, _root);
        return node->value;
    }

    //! Добавление элемента в дерево
    //!
    //! @param  key       Ключ
    //! @param  value     Значение
    void insert(const KeyT& key, const ValueT& value) {
        TreeNode** indirect = &_root;  // Чтобы обобщить вставку
        std::vector<TreeNode**> path;  // 

        while (*indirect != nullptr) {
            path.push_back(indirect);

            if ((*indirect)->key > key)
                indirect = &((*indirect)->left);
            else
                indirect = &((*indirect)->right);
        }

        *indirect = new TreeNode(key, value);
        path.push_back(indirect);

        _balance(path);
        _size++;
    }

    //! Удаление элемента из дерева
    //!
    //! @param  key       Ключ
    void erase(const KeyT& key) {
        TreeNode** indirect = &_root;  // to generalize insertion
        std::vector<TreeNode**> path;  // to update height values

        while (*indirect != nullptr && (*indirect)->key != key) {
            path.push_back(indirect);

            if ((*indirect)->key > key)
                indirect = &((*indirect)->left);
            else
                indirect = &((*indirect)->right);
        }

        // В дереве нет ноды с таким ключем
        if (*indirect == nullptr) {
            throw std::runtime_error("Key doesn't exist");
        } else {
            path.push_back(indirect);
        }

        std::size_t index = path.size();

        // Удаляемая нода является листом - нет поддеревьев
        if ((*indirect)->left == nullptr || (*indirect)->right == nullptr) {
            delete *indirect;  // Просто удаляем ноду
            *indirect = nullptr;
            path.pop_back();
        } else if ((*indirect)->right == nullptr) {  // Существует только левое поддерево
            TreeNode *remove = *indirect;

            (*indirect) = (*indirect)->left;
            delete remove;

            path.pop_back();
        } else {  // Существует только правое поддерево
            TreeNode **successor = &((*indirect)->right);

            while ((*successor)->left != nullptr) {
                path.push_back(successor);
                successor = &((*successor)->left);
            }

            if (*successor == (*indirect)->right) {
                (*successor)->left = (*indirect)->left;

                TreeNode *toRemove = *indirect;
                *indirect = *successor;
                delete toRemove;
            } else {
                TreeNode *tmp = *path.back();
                TreeNode *suc = *successor;

                tmp->left = (*successor)->right;
                suc->left = (*indirect)->left;
                suc->right = (*indirect)->right;

                delete *indirect;
                *indirect = suc;
                path[index] = &(suc->right);
            }
        }

        _balance(path);
        _size--;
    }

    //! Список ключей в дереве
    std::vector<KeyT> keys() {
        std::vector<KeyT> result;
        for (auto it = begin(); it != end(); ++it) {
             result.push_back(it->key);
        }
        return result;
    }

    //! Доп. операция в варианте задания - определение критерия сбалансированности
    int balanceFactor() const {
        if (!_root) return 0;
        return _root->balanceFactor();
    }
    
    //! Получаем итератор на первый элемент в дереве
    //!
    //! @return Итератор, указывающий на первый элемент в дереве
    ForwardIterator begin() {
        return ForwardIterator(_root);
    }

    //! Получаем обратный итератор на последний элемент в дереве
    //!
    //! @return Обратный итератор, указывающий на последний элемент в дереве
    ReverseIterator rbegin() {
        return ReverseIterator(_root);
    }

    //! Получаем итератор, указывающий на элемент, после последнего
    //!
    //! @return Итератор, указывающий на элемент после последнего
    ForwardIterator end() {
        return ForwardIterator(nullptr);
    }

    //! Получаем обратный итератор, указывающий на элемент, перед первым
    //!
    //! @return Обратный итератор, указывающий на элемент, перед первым
    ReverseIterator rend() {
        return ReverseIterator(nullptr);
    }

private:
    //! Рекурсивный поиск ноды по ключу
    TreeNode* _find(const KeyT& key, TreeNode* node) const {
        if (!node) {
            return nullptr;
        }
        if (key == node->key) {
            return node;
        }
        if (key < node->key) {
            return _find(key, node->left);
        } else {
            return _find(key, node->right);
        }
    }

    //! Процедура балансировки дерева
    //! https://github.com/KadirEmreOto/AVL-Tree
    void _balance(std::vector<TreeNode **> path) {
        // Начинаем балансировать с корня к листу
        std::reverse(path.begin(), path.end());

        for (auto indirect : path) {
            TreeNode* node = *indirect;
            TreeNode* leftLeaf  = node->left;
            TreeNode* rightLeaf = node->right;

            node->updateValues();

            if (node->balanceFactor() >= 2 || (leftLeaf && leftLeaf->balanceFactor() > 0)) { // left - left
                *indirect = (*indirect)->rotateR();
            } else if (node->balanceFactor() >= 2) {  // left - right
                (*indirect)->left = (*indirect)->left->rotateL();
                *indirect = (*indirect)->rotateR();
            } else if (node->balanceFactor() <= -2 || (rightLeaf && rightLeaf->balanceFactor() < 0)) { // right - right
            } else if (node->balanceFactor() <= -2 || (rightLeaf && rightLeaf->balanceFactor() < 0)) { // right - right
                *indirect = (*indirect)->rotateL();
            } else if (node->balanceFactor() <= -2) {  // right - left
                (*indirect)->right = ((*indirect)->right)->rotateR();
                *indirect = (*indirect)->rotateL();
            }
        }
    }

    //! Удаляем поддерево
    void clearNode(TreeNode* node) {
        if (node->left)  clearNode(node->left);
        if (node->right) clearNode(node->right);
        delete node;
    };

private: 
    TreeNode* _root = nullptr;
    size_t    _size = 0;
};

////////////////////////////////////////// TreeForwardIterator.hpp //////////////////////////////////////////
#pragma once
#include <iterator>
#include <stack>

// Article: inorder depth-first traversal
// http://mike.eshva.ru/dev/obhod-binarnogo-dereva-s-pomoschyu-iteratora
template <typename TreeNode>
class TreeForwardIterator {
public:
    explicit TreeForwardIterator(TreeNode* root) : _node(root) {
        if (!_node) return;

        _stack.push(nullptr);

        while (_node->left) {
            _stack.push(_node);
            _node = _node->left;
        }
    };

    bool operator==(const TreeForwardIterator& other) const {
        // Равны, если оба итератора указывают на nullptr
        if (!_node && !other._node) return true;

        // Не равны, если один из них nullptr
        if ((!_node && other._node) || (_node && !other._node)) return false;

        // Равны, если указывают на одинаковый ключ (и значение?)
        if (_node->key == other._node->key) return true;
        return false;
    }

    bool operator!=(const TreeForwardIterator& other) const {
        return !operator==(other);
    }

    TreeNode* operator->() {
        if (!_node) {
            throw std::runtime_error("Iterator not dereferenceable");
        }
        return _node;
    }

    TreeForwardIterator& operator++() {
        if (!_node) return *this;

        if (_node->right) {
            _node = _node->right;
            while (_node->left) {
                _stack.push(_node);
                _node = _node->left;
            }
        } else {
            _node = _stack.top();
            _stack.pop();
        }

        return *this;
    }

private:
    std::stack<TreeNode*> _stack;
    TreeNode* _node;
};

////////////////////////////////////////// TreeReverseIterator.hpp //////////////////////////////////////////
#pragma once
#include <iterator>
#include <stack>

// Article: inorder depth-first traversal
// http://mike.eshva.ru/dev/obhod-binarnogo-dereva-s-pomoschyu-iteratora
template <typename TreeNode>
class TreeReverseIterator {
public:
    explicit TreeReverseIterator(TreeNode* root) : _node(root) {
        if (!_node) return;

        _stack.push(nullptr);

        while (_node->right) {
            _stack.push(_node);
            _node = _node->right;
        }
    };

    bool operator==(const TreeReverseIterator& other) const {
        // Равны, если оба итератора указывают на nullptr
        if (!_node && !other._node) return true;

        // Не равны, если один из них nullptr
        if ((!_node && other._node) || (_node && !other._node)) return false;

        // Равны, если указывают на одинаковый ключ (и значение?)
        if (_node->key == other._node->key) return true;
        return false;
    }

    bool operator!=(const TreeReverseIterator& other) const {
        return !operator==(other);
    }

    TreeNode* operator->() {
        if (!_node) {
            throw std::runtime_error("Iterator not dereferenceable");
        }
        return _node;
    }

    TreeReverseIterator& operator++() {
        if (!_node) return *this;

        if (_node->left) {
            _node = _node->left;
            while (_node->right) {
                _stack.push(_node);
                _node = _node->right;
            }
        } else {
            _node = _stack.top();
            _stack.pop();
        }

        return *this;
    }

private:
    std::stack<TreeNode*> _stack;
    TreeNode* _node;
};

////////////////////////////////////////// TreeNode.hpp //////////////////////////////////////////
#pragma once


//! Класс, реализующий узел бинарного дерева поиска
//!
//! @tparam KeyT       Тип ключа
//! @tparam ValueT     Тип значения
template <typename KeyT,
          typename ValueT>
struct TreeNode {
    // Данные, хранящиеся в элементе
    KeyT   key;
    ValueT value;

    // Высота элемента
    size_t height = 1;

    // Количество элементов в поддереве
    size_t count = 1;

    // Указатель на дочерние элементы
    TreeNode* left = nullptr;
    TreeNode* right = nullptr;

    void updateValues() {
        const auto leftCount  = (left  != nullptr) ? left->count  : 0;
        const auto rightCount = (right != nullptr) ? right->count : 0;
        count = leftCount + rightCount + 1;

        const auto leftHeight  = (left  != nullptr) ? left->height  : 0;
        const auto rightHeight = (right != nullptr) ? right->height : 0;
        height = leftHeight + rightHeight + 1;
    }
    int balanceFactor() const {
        const auto leftHeight  = (left  != nullptr) ? left->height  : 0;
        const auto rightHeight = (right != nullptr) ? right->height : 0;
        return leftHeight - rightHeight;
    }

    // Поворот дерева вокруг узла
    TreeNode* rotateL() {
        TreeNode* root = right; // Правый лист стал корнем
        right = right->left;
        root->left = this;

        this->updateValues();  // Порядок обновления важен
        root->updateValues();

        return root;
    }
    TreeNode* rotateR() {
        TreeNode* root = left; // Левый лист стал корнем
        left = left->right;
        root->right = this;

        this->updateValues();  // Порядок обновления важен
        root->updateValues();

        return root;
    }

    TreeNode(const KeyT& key, const ValueT& value)
        : key(key), value(value) {}
};

////////////////////////////////////////// TemplateTreeTest.cpp //////////////////////////////////////////
#include <CppUnitTest.h>
#include "../TemplateTree/Tree.hpp"

#include <vector>
#include <string>
#include <numeric>

using namespace Microsoft::VisualStudio::CppUnitTestFramework;

template<> inline
std::wstring Microsoft::VisualStudio::CppUnitTestFramework::ToString<std::vector<size_t>>(const std::vector<size_t>& vector) {
    return std::accumulate(std::begin(vector), std::end(vector), std::wstring{},
                           [](std::wstring &ss, const unsigned int &s) {
                               return ss.empty() ? std::to_wstring(s) : ss + L"," + std::to_wstring(s);
                           });

}

TEST_CLASS(TemplateTreeTest) {
private:
    using KeyType   = size_t;
    using ValueType = std::string;
    using TreeType  = Tree<KeyType, ValueType>;

public:
	TEST_METHOD(SimpleTests) {
        TreeType tree;
        Assert::IsTrue(tree.size() == 0);
        Assert::IsTrue(tree.balanceFactor() == 0);
        Assert::IsTrue(tree.keys() == std::vector<KeyType>{});

        // Элемент добавился в корень дерева
        tree.insert(5, "пять");
        Assert::IsTrue(tree.size() == 1U);
        Assert::IsTrue(tree.balanceFactor() == 0);
        Assert::IsTrue(tree.keys() == std::vector<KeyType>{ 5 });

        // Элемент добавился в левый лист
        tree.insert(3, "три");
        Assert::IsTrue(tree.size() == 2U);
        Assert::IsTrue(tree.balanceFactor() == 1);
        Assert::IsTrue(tree.keys() == std::vector<KeyType>{ 3, 5 });

        // Элемент добавился в правый лист
        tree.insert(6, "шесть");
        Assert::IsTrue(tree.size() == 3U);
        Assert::IsTrue(tree.balanceFactor() == 0);
        Assert::IsTrue(tree.keys() == std::vector<KeyType>{ 3, 5, 6 });

        tree.insert(7, "семь");
        Assert::IsTrue(tree.size() == 4U);
        Assert::IsTrue(tree.balanceFactor() == -1);
        Assert::IsTrue(tree.keys() == std::vector<KeyType>{ 3, 5, 6, 7 });
	}

    TEST_METHOD(ReverseIteratorTest) {
        TreeType tree;
        tree.insert(5, "");
        tree.insert(3, "");
        tree.insert(6, "");
        tree.insert(7, "");

        // Проверка обратного итератора
        auto it = tree.rbegin();
        Assert::IsTrue(it->key == 7U);
        ++it;
        Assert::IsTrue(it->key == 6U);
        ++it;
        Assert::IsTrue(it->key == 5U);
        ++it;
        Assert::IsTrue(it->key == 3U);
        ++it;
        Assert::IsTrue(it == tree.rend());
    }

    TEST_METHOD(RemoveTest) {
        TreeType tree;
        tree.insert(5, "");
        tree.insert(3, "");
        tree.insert(6, "");
        tree.insert(7, "");

        // Проверка удаления не вершины
        tree.erase(7);
        Assert::IsTrue(tree.size() == 3U);
        Assert::IsTrue(tree.balanceFactor() == 0);
        Assert::IsTrue(tree.keys() == std::vector<KeyType>{ 3, 5, 6 });

        // Проверка удаления вершины + проверка балансировки
        tree.insert(7, "");
        tree.erase(5);
        Assert::IsTrue(tree.size() == 3U);
        Assert::IsTrue(tree.balanceFactor() == 0);
        Assert::IsTrue(tree.keys() == std::vector<KeyType>{ 3, 6, 7 });
    }
};