#pragma once

#include "internal/common.hpp"
#include <algorithm>
#include <cmath>
#include <limits>
#include <memory>
#include <vector>

namespace cavc
{
template <typename Real, std::size_t NodeSize = 16>
class StaticSpatialIndex
{
public:
    StaticSpatialIndex(std::size_t numItems)
    {
        CAVC_ASSERT(numItems > 0, "number of items must be greater than 0");
        static_assert(NodeSize >= 2 && NodeSize <= 65535, "node size must be between 2 and 65535");
        // calculate the total number of nodes in the R-tree to allocate space for
        // and the index of each tree level (used in search later)
        m_numItems           = numItems;
        std::size_t n        = numItems;
        std::size_t numNodes = numItems;

        m_numLevels      = computeNumLevels(numItems);
        m_levelBounds    = std::unique_ptr<std::size_t[]>(new std::size_t[m_numLevels]);
        m_levelBounds[0] = n * 4;
        // now populate level bounds and numNodes
        std::size_t i = 1;
        do {
            n = static_cast<std::size_t>(std::ceil(static_cast<float>(n) / NodeSize));
            numNodes += n;
            m_levelBounds[i] = numNodes * 4;
            i += 1;
        } while (n != 1);

        m_numNodes = numNodes;
        m_boxes    = std::unique_ptr<Real[]>(new Real[numNodes * 4]);
        m_indices  = std::unique_ptr<std::size_t[]>(new std::size_t[numNodes]);
        m_pos      = 0;
        m_minX     = std::numeric_limits<Real>::infinity();
        m_minY     = std::numeric_limits<Real>::infinity();
        m_maxX     = -std::numeric_limits<Real>::infinity();
        m_maxY     = -std::numeric_limits<Real>::infinity();
    }

    Real minX() const { return m_minX; }
    Real minY() const { return m_minY; }
    Real maxX() const { return m_maxX; }
    Real maxY() const { return m_maxY; }

    void add(Real minX, Real minY, Real maxX, Real maxY)
    {
        std::size_t index = m_pos >> 2;
        m_indices[index]  = index;
        m_boxes[m_pos++]  = minX;
        m_boxes[m_pos++]  = minY;
        m_boxes[m_pos++]  = maxX;
        m_boxes[m_pos++]  = maxY;

        if (minX < m_minX)
            m_minX = minX;
        if (minY < m_minY)
            m_minY = minY;
        if (maxX > m_maxX)
            m_maxX = maxX;
        if (maxY > m_maxY)
            m_maxY = maxY;
    }

    void finish()
    {
        CAVC_ASSERT(m_pos >> 2 == m_numItems, "added item count should equal static size given");

        // if number of items is less than node size then skip sorting since each node of boxes must be
        // fully scanned regardless and there is only one node
        if (m_numItems <= NodeSize) {
            m_indices[m_pos >> 2] = 0;
            // fill root box with total extents
            m_boxes[m_pos++] = m_minX;
            m_boxes[m_pos++] = m_minY;
            m_boxes[m_pos++] = m_maxX;
            m_boxes[m_pos++] = m_maxY;
            return;
        }

        Real width  = m_maxX - m_minX;
        Real height = m_maxY - m_minY;
        std::unique_ptr<std::uint32_t[]> hilbertValues(new std::uint32_t[m_numItems]);

        std::size_t pos = 0;

        for (std::size_t i = 0; i < m_numItems; ++i) {
            pos       = 4 * i;
            Real minX = m_boxes[pos++];
            Real minY = m_boxes[pos++];
            Real maxX = m_boxes[pos++];
            Real maxY = m_boxes[pos++];

            // hilbert max input value for x and y
            const Real hilbertMax = static_cast<Real>((1 << 16) - 1);
            // mapping the x and y coordinates of the center of the box to values in the range
            // [0 -> n - 1] such that the min of the entire set of bounding boxes maps to 0 and the max of
            // the entire set of bounding boxes maps to n - 1 our 2d space is x: [0 -> n-1] and
            // y: [0 -> n-1], our 1d hilbert curve value space is d: [0 -> n^2 - 1]
            Real x           = std::floor(hilbertMax * ((minX + maxX) / 2 - m_minX) / width);
            std::uint32_t hx = static_cast<std::uint32_t>(x);
            Real y           = std::floor(hilbertMax * ((minY + maxY) / 2 - m_minY) / height);
            std::uint32_t hy = static_cast<std::uint32_t>(y);
            hilbertValues[i] = hilbertXYToIndex(hx, hy);
        }

        // sort items by their Hilbert value (for packing later)
        sort(&hilbertValues[0], &m_boxes[0], &m_indices[0], 0, m_numItems - 1);

        // generate nodes at each tree level, bottom-up
        pos = 0;
        for (std::size_t i = 0; i < m_numLevels - 1; i++) {
            auto end = m_levelBounds[i];

            // generate a parent node for each block of consecutive <nodeSize> nodes
            while (pos < end) {
                auto nodeMinX  = std::numeric_limits<Real>::infinity();
                auto nodeMinY  = std::numeric_limits<Real>::infinity();
                auto nodeMaxX  = -1 * std::numeric_limits<Real>::infinity();
                auto nodeMaxY  = -1 * std::numeric_limits<Real>::infinity();
                auto nodeIndex = pos;

                // calculate bbox for the new node
                for (std::size_t j = 0; j < NodeSize && pos < end; j++) {
                    auto minX = m_boxes[pos++];
                    auto minY = m_boxes[pos++];
                    auto maxX = m_boxes[pos++];
                    auto maxY = m_boxes[pos++];
                    if (minX < nodeMinX)
                        nodeMinX = minX;
                    if (minY < nodeMinY)
                        nodeMinY = minY;
                    if (maxX > nodeMaxX)
                        nodeMaxX = maxX;
                    if (maxY > nodeMaxY)
                        nodeMaxY = maxY;
                }

                // add the new node to the tree data
                m_indices[m_pos >> 2] = nodeIndex;
                m_boxes[m_pos++]      = nodeMinX;
                m_boxes[m_pos++]      = nodeMinY;
                m_boxes[m_pos++]      = nodeMaxX;
                m_boxes[m_pos++]      = nodeMaxY;
            }
        }
    }

    // Visit all the bounding boxes in the spatial index. Visitor function has the signature
    // bool(std::size_t level, Real xmin, Real ymin, Real xmax, Real ymax, std::size_t level).
    // Visiting stops early if false is returned.
    template <typename F>
    void visitBoundingBoxes(F&& visitor) const
    {
        std::size_t nodeIndex = 4 * m_numNodes - 4;
        std::size_t level     = m_numLevels - 1;

        std::vector<std::size_t> stack;
        stack.reserve(16);

        bool done = false;
        while (!done) {
            auto end = std::min(nodeIndex + NodeSize * 4, m_levelBounds[level]);
            for (std::size_t pos = nodeIndex; pos < end; pos += 4) {
                auto index = m_indices[pos >> 2];
                if (!visitor(level, m_boxes[pos], m_boxes[pos + 1], m_boxes[pos + 2], m_boxes[pos + 3])) {
                    return;
                }

                if (nodeIndex >= m_numItems * 4) {
                    stack.push_back(index);
                    stack.push_back(level - 1);
                }
            }

            if (stack.size() > 1) {
                level = stack.back();
                stack.pop_back();
                nodeIndex = stack.back();
                stack.pop_back();
            }
            else {
                done = true;
            }
        }
    }

    // Visit only the item bounding boxes in the spatial index. Visitor function has the signature
    // bool(std::size_t index, Real xmin, Real ymin, Real xmax, Real ymax). Visiting stops early if
    // false is returned.
    template <typename F>
    void visitItemBoxes(F&& visitor) const
    {
        for (std::size_t i = 0; i < m_levelBounds[0]; i += 4) {
            if (!visitor(m_indices[i >> 2], m_boxes[i], m_boxes[i + 1], m_boxes[i + 2], m_boxes[i + 3])) {
                return;
            }
        }
    }

    // See other overloads for details.
    void query(Real minX, Real minY, Real maxX, Real maxY, std::vector<std::size_t>& results) const
    {
        auto visitor = [&](std::size_t index) {
            results.push_back(index);
            return true;
        };

        visitQuery(minX, minY, maxX, maxY, visitor);
    }

    // Query the spatial index adding indexes to the results vector given. This overload accepts an
    // existing vector to use as a stack and takes care of clearing the stack before use.
    void query(Real minX, Real minY, Real maxX, Real maxY, std::vector<std::size_t>& results,
               std::vector<std::size_t>& stack) const
    {
        auto visitor = [&](std::size_t index) {
            results.push_back(index);
            return true;
        };

        visitQuery(minX, minY, maxX, maxY, visitor, stack);
    }

    // See other overloads for details.
    template <typename F>
    void visitQuery(Real minX, Real minY, Real maxX, Real maxY, F&& visitor) const
    {
        std::vector<std::size_t> stack;
        stack.reserve(16);
        visitQuery(minX, minY, maxX, maxY, std::forward<F>(visitor), stack);
    }

    // Query the spatial index, invoking a visitor function for each index that overlaps the bounding
    // box given. Visitor function has the signature bool(std::size_t index), if visitor returns false
    // the query stops early, otherwise the query continues. This overload accepts an existing vector
    // to use as a stack and takes care of clearing the stack before use.
    template <typename F>
    void visitQuery(Real minX, Real minY, Real maxX, Real maxY, F&& visitor,
                    std::vector<std::size_t>& stack) const
    {
        CAVC_ASSERT(m_pos == 4 * m_numNodes, "data not yet indexed - call Finish() before querying");

        auto nodeIndex = 4 * m_numNodes - 4;
        auto level     = m_numLevels - 1;

        stack.clear();

        auto done = false;

        while (!done) {
            // find the end index of the node
            auto end = std::min(nodeIndex + NodeSize * 4, m_levelBounds[level]);

            // search through child nodes
            for (std::size_t pos = nodeIndex; pos < end; pos += 4) {
                auto index = m_indices[pos >> 2];
                // check if node bbox intersects with query bbox
                if (maxX < m_boxes[pos] || maxY < m_boxes[pos + 1] || minX > m_boxes[pos + 2] ||
                    minY > m_boxes[pos + 3]) {
                    // no intersect
                    continue;
                }

                if (nodeIndex < m_numItems * 4) {
                    done = !visitor(index);
                    if (done) {
                        break;
                    }
                }
                else {
                    // push node index and level for further traversal
                    stack.push_back(index);
                    stack.push_back(level - 1);
                }
            }

            if (stack.size() > 1) {
                level = stack.back();
                stack.pop_back();
                nodeIndex = stack.back();
                stack.pop_back();
            }
            else {
                done = true;
            }
        }
    }

    static std::uint32_t hilbertXYToIndex(std::uint32_t x, std::uint32_t y)
    {
        std::uint32_t a = x ^ y;
        std::uint32_t b = 0xFFFF ^ a;
        std::uint32_t c = 0xFFFF ^ (x | y);
        std::uint32_t d = x & (y ^ 0xFFFF);

        std::uint32_t A = a | (b >> 1);
        std::uint32_t B = (a >> 1) ^ a;
        std::uint32_t C = ((c >> 1) ^ (b & (d >> 1))) ^ c;
        std::uint32_t D = ((a & (c >> 1)) ^ (d >> 1)) ^ d;

        a = A;
        b = B;
        c = C;
        d = D;
        A = (a & (a >> 2)) ^ (b & (b >> 2));
        B = (a & (b >> 2)) ^ (b & ((a ^ b) >> 2));
        C ^= (a & (c >> 2)) ^ (b & (d >> 2));
        D ^= (b & (c >> 2)) ^ ((a ^ b) & (d >> 2));

        a = A;
        b = B;
        c = C;
        d = D;
        A = (a & (a >> 4)) ^ (b & (b >> 4));
        B = (a & (b >> 4)) ^ (b & ((a ^ b) >> 4));
        C ^= (a & (c >> 4)) ^ (b & (d >> 4));
        D ^= (b & (c >> 4)) ^ ((a ^ b) & (d >> 4));

        a = A;
        b = B;
        c = C;
        d = D;
        C ^= ((a & (c >> 8)) ^ (b & (d >> 8)));
        D ^= ((b & (c >> 8)) ^ ((a ^ b) & (d >> 8)));

        a = C ^ (C >> 1);
        b = D ^ (D >> 1);

        std::uint32_t i0 = x ^ y;
        std::uint32_t i1 = b | (0xFFFF ^ (i0 | a));

        i0 = (i0 | (i0 << 8)) & 0x00FF00FF;
        i0 = (i0 | (i0 << 4)) & 0x0F0F0F0F;
        i0 = (i0 | (i0 << 2)) & 0x33333333;
        i0 = (i0 | (i0 << 1)) & 0x55555555;

        i1 = (i1 | (i1 << 8)) & 0x00FF00FF;
        i1 = (i1 | (i1 << 4)) & 0x0F0F0F0F;
        i1 = (i1 | (i1 << 2)) & 0x33333333;
        i1 = (i1 | (i1 << 1)) & 0x55555555;

        return (i1 << 1) | i0;
    }

private:
    Real m_minX;
    Real m_minY;
    Real m_maxX;
    Real m_maxY;
    std::size_t m_numItems;
    std::size_t m_numLevels;
    // using std::unique_ptr arrays for uninitialized memory optimization
    std::unique_ptr<std::size_t[]> m_levelBounds;
    std::size_t m_numNodes;
    std::unique_ptr<Real[]> m_boxes;
    std::unique_ptr<std::size_t[]> m_indices;
    std::size_t m_pos;

    static std::size_t computeNumLevels(std::size_t numItems)
    {
        std::size_t n               = numItems;
        std::size_t levelBoundsSize = 1;
        do {
            n = static_cast<std::size_t>(std::ceil(static_cast<float>(n) / NodeSize));
            levelBoundsSize += 1;
        } while (n != 1);

        return levelBoundsSize;
    }

    // quicksort that partially sorts the bounding box data alongside the Hilbert values
    static void sort(std::uint32_t* values, Real* boxes, std::size_t* indices, std::size_t left,
                     std::size_t right)
    {
        CAVC_ASSERT(left <= right, "left index should never be past right index");

        // check against NodeSize (only need to sort down to NodeSize buckets)
        if (left / NodeSize >= right / NodeSize) {
            return;
        }

        auto pivot = values[(left + right) >> 1];
        auto i     = left - 1;
        auto j     = right + 1;

        while (true) {
            do
                i++;
            while (values[i] < pivot);
            do
                j--;
            while (values[j] > pivot);
            if (i >= j)
                break;
            swap(values, boxes, indices, i, j);
        }

        sort(values, boxes, indices, left, j);
        sort(values, boxes, indices, j + 1, right);
    }

    static void swap(std::uint32_t* values, Real* boxes, std::size_t* indices, std::size_t i,
                     std::size_t j)
    {
        auto temp = values[i];
        values[i] = values[j];
        values[j] = temp;

        auto k = 4 * i;
        auto m = 4 * j;

        auto a       = boxes[k];
        auto b       = boxes[k + 1];
        auto c       = boxes[k + 2];
        auto d       = boxes[k + 3];
        boxes[k]     = boxes[m];
        boxes[k + 1] = boxes[m + 1];
        boxes[k + 2] = boxes[m + 2];
        boxes[k + 3] = boxes[m + 3];
        boxes[m]     = a;
        boxes[m + 1] = b;
        boxes[m + 2] = c;
        boxes[m + 3] = d;

        auto e     = indices[i];
        indices[i] = indices[j];
        indices[j] = e;
    }
};
}  // namespace cavc
