#include "KDTree.h"
#include "KDAxis.h"
#include "KDNode.h"
#include <set>

//#define RAYTRY_KDTREE_DEBUG

namespace raytry::KD {

    void Tree::insertSorted(triangleValueVector &sortedTriangles, float value, const Triangle &triangle) {
        bool emplaced = false;
        for (auto mpos = sortedTriangles.begin(); !emplaced && mpos != sortedTriangles.end(); ++mpos) {
            if ((*mpos).first > value) {
                sortedTriangles.emplace(mpos, value, &triangle);
                emplaced = true;
            }
        }
        if (!emplaced) {
            sortedTriangles.emplace_back(value, &triangle);
        }
    }

    void Tree::insertSortedReversed(triangleValueVector &sortedTriangles, float value, const Triangle &triangle) {
        bool emplaced = false;
        for (auto mpos = sortedTriangles.begin(); !emplaced && mpos != sortedTriangles.end(); ++mpos) {
            if ((*mpos).first < value) {
                sortedTriangles.emplace(mpos, value, &triangle);
                emplaced = true;
            }
        }
        if (!emplaced) {
            sortedTriangles.emplace_back(value, &triangle);
        }
    }

    float Tree::calculateCost(float costConstant, float bndMin, float bndMax, unsigned int count) {
        return costConstant + ((bndMax - bndMin) * static_cast<float>(count));
    }

    Tree::Tree(const Tree &other)
        : nodes{other.nodes}, leafs{other.leafs}, boundaries{other.boundaries},
          randomDevice{}, randomGenerator{randomDevice()}, randomProbes{other.randomProbes} {}

    Tree &Tree::operator=(Tree other) {
        std::swap(nodes, other.nodes);
        std::swap(leafs, other.leafs);
        std::swap(boundaries, other.boundaries);
        std::swap(randomProbes, other.randomProbes);
        return *this;
    }

    void Tree::separateAxisData(AxisData &data, std::array<triangleValueVector, Z + 1> &separatingData, std::array<triangleValueVector, Z + 1> &otherData, Axis axis, float tVal, bool separatorIsEnd) {
        auto separator = separatingData[axis].begin();
        for (; separator != separatingData[axis].end(); ++separator) {
            if (separatorIsEnd == separator->first > tVal) {
                break;
            }
        }
        std::set<const Triangle *> trianglesToRemove{};
        for (auto removeIter = separator; removeIter != separatingData[axis].end(); ++removeIter) {
            trianglesToRemove.insert(removeIter->second);
        }
        separatingData[axis].erase(separator, separatingData[axis].end());
        separatingData[axis].shrink_to_fit();
        for (unsigned int a = X; a <= Z; ++a) {
            if (a != axis) {
                separatingData[a].erase(
                        std::remove_if(separatingData[a].begin(), separatingData[a].end(), [&](const std::pair<float, const Triangle *> &pair) {
                            return trianglesToRemove.contains(pair.second);
                        }),
                        separatingData[a].end());
                separatingData[a].shrink_to_fit();
            }
            otherData[a].erase(
                    std::remove_if(otherData[a].begin(), otherData[a].end(), [&](const std::pair<float, const Triangle *> &pair) {
                        return trianglesToRemove.contains(pair.second);
                    }),
                    otherData[a].end());
            otherData[a].shrink_to_fit();
        }
    }

    float Tree::probeBestT(const Axis &axis, const AxisData &data) {
        float bestT = NAN;
        if (data.boundaries[axis].first <= data.boundaries[axis].second) {
            float bestTCost = calculateCost(0, data.boundaries[axis].first, data.boundaries[axis].second, data.mins[axis].size());
            size_t availableSamples = data.mins[axis].size() + data.maxs[axis].size();
            assert(availableSamples > 0u);
            std::uniform_int_distribution<size_t> dist{1u, availableSamples};
            const auto& samples = std::min(availableSamples, randomSamples);
            for (size_t i = 1; i <= samples; ++i) {
                size_t idx;
                if (availableSamples > randomSamples) {
                    idx = dist(randomGenerator);
                    ++randomProbes;
                } else {
                    idx = i;
                }
                unsigned int tVertexCountFirst = 0;
                unsigned int tVertexCountSecond = 0;
                float tVal;
                if (idx > data.mins[axis].size()) {
                    idx -= data.mins[axis].size();
                    tVal = data.maxs[axis][idx - 1u].first;
                    tVertexCountSecond = idx;
                    for (const auto &item: data.mins[axis]) {
                        if (item.first > tVal) {
                            break;
                        }
                        ++tVertexCountFirst;
                    }
                } else {
                    assert(idx > 0u);
                    tVal = data.mins[axis][idx - 1u].first;
                    tVertexCountFirst = idx;
                    for (const auto &item: data.maxs[axis]) {
                        if (item.first < tVal) {
                            break;
                        }
                        ++tVertexCountSecond;
                    }
                }
#ifdef RAYTRY_KDTREE_DEBUG
                qDebug() << "T:" << tVal << ", VC1:" << tVertexCountFirst << ", VC2:" << tVertexCountSecond;
#endif
                float tCost = calculateCost(traversalCost, data.boundaries[axis].first, tVal, tVertexCountFirst) +
                              calculateCost(traversalCost, tVal, data.boundaries[axis].second, tVertexCountSecond);
                if (tCost < bestTCost) {
#ifdef RAYTRY_KDTREE_DEBUG
                    qDebug() << "Found better T:" << tVal << "with cost" << tCost << "Better than T:" << bestT << "with cost" << bestTCost;
#endif
                    bestT = tVal;
                    bestTCost = tCost;
                }
#ifdef RAYTRY_KDTREE_DEBUG
                else {
                    qDebug() << "Found worse T:" << tVal << "with cost" << tCost << "Worse than T:" << bestT << "with cost" << bestTCost;
                }
#endif
            }
        }
        return bestT;
    }

    unsigned int Tree::recursiveNodeBuild(const unsigned int depthLeft, const Axis axis, const AxisData &data) {
        if (depthLeft <= 0 || data.mins[X].empty()) {
            unsigned int nearIdx = nodes.size();
            Leaf leaf{};
            assert(data.mins[X].size() == data.maxs[X].size());
            assert(data.mins[X].size() == data.mins[Y].size());
            assert(data.mins[Y].size() == data.maxs[Y].size());
            assert(data.mins[Y].size() == data.mins[Z].size());
            assert(data.mins[Z].size() == data.maxs[Z].size());
            for (const auto &item: data.mins[X]) {
                leaf.triangles.push_back(item.second);
            }
            leaf.boundaries = data.boundaries;
            unsigned long leafIdx = leafs.size();
            leafs.push_back(leaf);
#ifdef RAYTRY_KDTREE_DEBUG
            qDebug() << "Adding leaf node" << leafIdx << "with" << leaf.triangles.size() << "triangles";
#endif
            nodes.emplace_back(None, leafIdx, NAN);
            return nearIdx;
        }

        float bestT = probeBestT(axis, data);
        unsigned int newAxis = axis + 1u;
        unsigned int newDepth = depthLeft;
        if (newAxis > Z) {
            newAxis = X;
            --newDepth;
        }
        if (qIsNaN(bestT)) {
#ifdef RAYTRY_KDTREE_DEBUG
            qDebug() << "No partitioning with better cost on axis" << axis << "in depth" << depthLeft;
#endif
            return recursiveNodeBuild(newDepth, static_cast<Axis>(newAxis), data);
        } else {
            unsigned int nearIdx = nodes.size();
#ifdef RAYTRY_KDTREE_DEBUG
            qDebug() << "Adding node on axis:" << axis << "in depth" << depthLeft << ", node idx:" << nearIdx;
#endif
            nodes.emplace_back(axis, 0, bestT);

            AxisData nearAxisData = AxisData(data);
            nearAxisData.boundaries[axis] = std::pair<float, float>{nearAxisData.boundaries[axis].first, bestT};
            separateAxisData(nearAxisData, nearAxisData.mins, nearAxisData.maxs, axis, bestT, true);
            recursiveNodeBuild(newDepth, static_cast<Axis>(newAxis), nearAxisData);

            AxisData farAxisData = AxisData(data);
            farAxisData.boundaries[axis] = std::pair<float, float>{bestT, farAxisData.boundaries[axis].second};
            separateAxisData(farAxisData, farAxisData.maxs, farAxisData.mins, axis, bestT, false);
            unsigned int absoluteFarIdx = recursiveNodeBuild(newDepth, static_cast<Axis>(newAxis), farAxisData);
            nodes[nearIdx].setFarIdx(absoluteFarIdx - nearIdx);

            return nearIdx;
        }
    }

    Tree Tree::build(const std::vector<Triangle> &triangles) {
        if (triangles.empty()) {
            return {};
        }

        AxisData data = AxisData();
        assert(data.mins.size() > Z);
        assert(data.maxs.size() > Z);
        for (unsigned int a = X; a <= Z; ++a) {
            data.mins[a].reserve(triangles.size());
            data.maxs[a].reserve(triangles.size());
        }

        for (const auto &item: triangles) {
            for (unsigned int a = X; a <= Z; ++a) {
                Axis axis = static_cast<Axis>(a);
                const auto itemMin = std::min(item.aVec()[axis], std::min(item.bVec()[axis], item.cVec()[axis]));
                insertSorted(data.mins[a], itemMin, item);
                const auto itemMax = std::max(item.aVec()[axis], std::max(item.bVec()[axis], item.cVec()[axis]));
                insertSortedReversed(data.maxs[a], itemMax, item);
            }
        }
        for (unsigned int a = X; a <= Z; ++a) {
            data.boundaries[a].first = data.mins[a].front().first;
            data.boundaries[a].second = data.maxs[a].front().first;
        }

        Tree ret{};
        ret.boundaries = data.boundaries; // tree boundary = root boundary
        qDebug() << "Tree boundaries:" << ret.boundaries;
        ret.recursiveNodeBuild(maxDepth, X, data);
#ifdef RAYTRY_KDTREE_DEBUG
        qDebug() << nodes;
        qDebug() << leafs;
#endif
        qInfo() << "Tree built. Probed random source" << ret.randomProbes << "times";
        return ret;
    }

    const std::vector<KD::Node> &Tree::getNodes() {
        return nodes;
    }

    const std::vector<Leaf> &Tree::getLeafs() {
        return leafs;
    }

    const AABB &Tree::getBounds() {
        return boundaries;
    }

}// namespace raytry::KD