stdtree.hpp

#pragma once

#include <iostream>
#include <ostream>
#include <string>

#include "treenode.hpp"


template<typename T, typename N>
class TreeView
{
    using NodeType = TreeNode<T, N>;
    using NodePtr = NodeType*;
    using NodeUniquePtr = std::unique_ptr<NodeType>;

    public:

    // Constructor to initialize TreeView with a unique pointer to the root node
    TreeView(NodeUniquePtr root) : root_(std::move(root)) {}
    
    // Method to get the root node
    NodePtr getRoot() const { return root_.get(); }

    // TreeView resolves nodes by breadth-first level and sibling index.
    NodePtr getNode(N level, N index) const
    {
        if (!root_) return nullptr;
        if (level == 0 && index == 0) return root_.get();
        
        std::vector<NodePtr> currentLevelNodes;
        currentLevelNodes.push_back(root_.get());
        
        for (N currentLevel = 0; currentLevel < level; ++currentLevel)
        {
            std::vector<NodePtr> nextLevelNodes;
            for (auto node : currentLevelNodes)
            {
                for (const auto& child : node->children)
                {
                    nextLevelNodes.push_back(child.get());
                }
            }
            currentLevelNodes = std::move(nextLevelNodes);
        }
        
        if (index < currentLevelNodes.size())
        {
            return currentLevelNodes[index];
        }
        
        return nullptr;
    }
    
    // Additional methods to get parent, child, and sibling nodes by level and index
    NodePtr getParent(N level, N index) const
    {
        NodePtr node = getNode(level, index);
        return node ? node->parent : nullptr;
    }

    // Method to get a child node by level, index, and child index
    NodePtr getChild(N level, N index, N childIndex) const
    {
        NodePtr node = getNode(level, index);
        if (node && childIndex < node->children.size())
        {
            return node->children[childIndex].get();
        }
        return nullptr;
    }

    // Method to get a sibling node by level, index, and sibling index
    NodePtr getSibling(N level, N index, N siblingIndex) const
    {
        NodePtr node = getNode(level, index);
        if (node && node->parent)
        {
            auto& siblings = node->parent->children;
            if (siblingIndex < siblings.size())
            {
                return siblings[siblingIndex].get();
            }
        }
        return nullptr;
    }

    // Method to get the next sibling node by level and index
    NodePtr getNextSibling(N level, N index) const
    {
        NodePtr node = getNode(level, index);
        if (node && node->parent)
        {
            auto& siblings = node->parent->children;
            for (N i = 0; i < siblings.size(); ++i)
            {
                if (siblings[i].get() == node && i + 1 < siblings.size())
                {
                    return siblings[i + 1].get();
                }
            }
        }
        return nullptr;
    }

    // Method to get the previous sibling node by level and index
    NodePtr getPreviousSibling(N level, N index) const
    {
        NodePtr node = getNode(level, index);
        if (node && node->parent)
        {
            auto& siblings = node->parent->children;
            for (N i = 0; i < siblings.size(); ++i)
            {
                if (siblings[i].get() == node && i > 0)
                {
                    return siblings[i - 1].get();
                }
            }
        }
        return nullptr;
    }

    // Method to get the first child node by level and index
    NodePtr getFirstChild(N level, N index) const
    {
        NodePtr node = getNode(level, index);
        if (node && !node->children.empty())
        {
            return node->children[0].get();
        }
        return nullptr;
    }

    // Method to get the last child node by level and index
    NodePtr getLastChild(N level, N index) const
    {
        NodePtr node = getNode(level, index);
        if (node && !node->children.empty())
        {
            return node->children.back().get();
        }
        return nullptr;
    }

    // Depth-first traversal over the tree for full-node inspection.
    std::vector<NodePtr> getAllNodes() const
    {
        std::vector<NodePtr> allNodes;
        if (!root_) return allNodes;

        std::vector<NodePtr> stack;
        stack.push_back(root_.get());

        while (!stack.empty())
        {
            NodePtr currentNode = stack.back();
            stack.pop_back();
            allNodes.push_back(currentNode);

            for (const auto& child : currentNode->children)
            {
                stack.push_back(child.get());
            }
        }

        return allNodes;
    }

    // printTree method for visualizing the tree structure
    void printTree(std::ostream& out = std::cout, NodePtr node = nullptr, std::string prefix = "", bool isLast = true) const
    {
        if (!node)
            node = root_.get();

        if (!node)
            return;

        out << prefix;
        out << (isLast ? "└──" : "├──");
        out << node->value << '\n';

        std::string childPrefix = prefix + (isLast ? "   " : "│  ");
        for (size_t index = 0; index < node->children.size(); ++index)
        {
            printTree(out, node->children[index].get(), childPrefix, index == node->children.size() - 1);
        }
    }


    private:
    
    NodeUniquePtr root_;
    
    
};