#include <algorithm>
#include <numeric>

#include "blobtree.hpp"

EXTERN_C_BEGIN

size_t baked_blobtree_t::get_node_count() const noexcept { return nodes.size(); }

size_t baked_blobtree_t::get_primitive_count() const noexcept { return primitives.size(); }

internal::node_t& baked_blobtree_t::get_node(uint32_t index) noexcept
{
    assert(index < nodes.size());
    return nodes[index];
}

pointer_wrapper<primitive> baked_blobtree_t::get_primitive(uint32_t index) noexcept
{
    assert(index < primitives.size());
    auto& primitive_node = nodes[index];
    auto  primitive_ptr  = primitives[primitive_node.primitive_index];

    return primitive_ptr.object_ptr;
}

void recursive_bake_blobtree(baked_blobtree_t&                               baked_tree,
                             const pointer_wrapper<internal::runtime_node_t> root_node,
                             size_t&                                         max_index)
{
    auto& node       = baked_tree.nodes[max_index];
    auto  root_index = max_index;
    max_index--;
    if (root_node->is_primitive_node()) {
        auto primitive_ptr = root_node->node_data.get_ptr();

        node.type            = static_cast<uint32_t>(primitive_ptr->get_type());
        node.primitive_index = static_cast<uint32_t>(baked_tree.primitives.size());

        baked_tree.leaf_indices.emplace_back(root_index);
        auto& primitive      = baked_tree.primitives.emplace_back();
        primitive.object_ptr = make_pointer_wrapper(primitive_ptr);
        auto subfaces        = primitive_ptr->get_subfaces();
        primitive.index_mapping.reserve(subfaces.size());
        for (size_t i = 0; i < subfaces.size(); ++i) {
            primitive.index_mapping.emplace_back(static_cast<uint32_t>(baked_tree.subfaces.size()));
            auto& subface         = baked_tree.subfaces.emplace_back();
            subface.object_ptr    = make_pointer_wrapper(subfaces[i].get_ptr());
            subface.index_mapping = {node.primitive_index};
        }
    } else {
        node.operation = static_cast<uint32_t>(root_node->node_data.get_mark());

        if (!root_node->is_right_child_null()) {
            node.right_child_index   = max_index;
            auto& right_child        = baked_tree.nodes[max_index];
            right_child.parent_index = root_index;
            recursive_bake_blobtree(baked_tree, root_node->right_child, max_index);
        }
        if (!root_node->is_left_child_null()) {
            node.left_child_index   = max_index;
            auto& left_child        = baked_tree.nodes[max_index];
            left_child.parent_index = root_index;
            recursive_bake_blobtree(baked_tree, root_node->left_child, max_index);
        }
    }
}

// build post-order traversal of the blobtree
baked_blobtree_t::baked_blobtree_t(const blobtree_t& tree) noexcept
{
    if (tree.nodes.empty()) return;

    // step 1: find root node
    auto root = make_pointer_wrapper(const_cast<internal::runtime_node_t*>(&*tree.nodes.begin()));
    while (!root->is_parent_null()) root = root->parent;

    this->nodes.resize(tree.nodes.size());
    // assume full binary tree, then we should reserve half of the nodes for leaf nodes
    this->leaf_indices.reserve(tree.nodes.size() / 2 + 1);
    this->primitives.reserve(tree.nodes.size() / 2 + 1);
    this->subfaces.reserve(tree.nodes.size() / 2 + 1);
    // step 2: build the blobtree
    size_t root_index = this->nodes.size() - 1;
    recursive_bake_blobtree(*this, root, root_index);
    this->leaf_indices.shrink_to_fit();
    this->primitives.shrink_to_fit();
    std::sort(this->leaf_indices.begin(), this->leaf_indices.end());
    assert(this->leaf_indices.front() == 0);

    // step 3: sort & remove duplicates of subfaces
    {
        stl_vector_mp<uint32_t> old_to_new_mapping(this->subfaces.size());
        {
            stl_vector_mp<uint32_t> subface_indices(this->subfaces.size());
            std::iota(subface_indices.begin(), subface_indices.end(), 0);
            std::sort(subface_indices.begin(), subface_indices.end(), [this](uint32_t lhs, uint32_t rhs) {
                return this->subfaces[lhs].object_ptr < this->subfaces[rhs].object_ptr;
            });
            for (uint32_t new_index = 0; new_index < subface_indices.size(); ++new_index) {
                const auto old_index          = subface_indices[new_index];
                old_to_new_mapping[old_index] = new_index;
            }
        }
        stl_vector_mp<object_with_index_mapping<subface>> temp_subfaces(this->subfaces.size());
        for (size_t i = 0; i < this->subfaces.size(); ++i) {
            auto& subface     = this->subfaces[i];
            auto& new_subface = temp_subfaces[old_to_new_mapping[i]];
            new_subface       = std::move(subface);
        }
        std::swap(this->subfaces, temp_subfaces);
        for (auto& primitive : this->primitives) {
            auto& primitive_mapping = primitive.index_mapping;
            for (auto& index : primitive_mapping) { index = old_to_new_mapping[index]; }
        }
    }
    // merge subfaces: basically a specialization of std::unique
    {
        auto current_iter = this->subfaces.begin();
        auto end_iter     = this->subfaces.end();
        auto result_iter  = current_iter;
        while (++current_iter != end_iter) {
            if (current_iter->object_ptr != result_iter->object_ptr) {
                // unique, try to move to result
                ++result_iter;
                // iff the same iter, do not have to change elements
                if (result_iter != current_iter) *result_iter = std::move(*current_iter);
            } else {
                // duplicate, merge indices
                result_iter->index_mapping.emplace_back(current_iter->index_mapping.front());
                auto& primitive_mapping = this->primitives[current_iter->index_mapping.front()].index_mapping;
                for (auto& index : primitive_mapping) {
                    if (index == std::distance(this->subfaces.begin(), current_iter))
                        index = std::distance(this->subfaces.begin(), result_iter);
                }
            }
        }

        this->subfaces.erase(++result_iter, this->subfaces.end());
    }
    this->subfaces.shrink_to_fit();
}

EXTERN_C_END