#pragma once

#include <numeric>
#include <algorithm>
#include <vector>
#include <cstddef>
#include "Geometry.h"

struct TriangleSegmentIntersectRes
{
    bool hit;
    float t;
    TriangleSegmentIntersectRes(bool h, float tt) : hit(h), t(tt) {} // VS2008
};

TriangleSegmentIntersectRes
triangleSegmentIntersection(const LineSegment &segment, const Vec3f &a, const Vec3f &b, const Vec3f &c)
{
    Vec3f e1 = b - a;
    Vec3f e2 = c - a;
    Vec3f s = segment.start - a;
    Vec3f s1 = segment.getDir().cross(e2);
    Vec3f s2 = s.cross(e1);
    float s1_dot_e1 = s1.dot(e1);
    float b1 = s1.dot(s) / s1_dot_e1;
    float b2 = s2.dot(segment.getDir()) / s1_dot_e1;
    if (b1 >= 0. && b2 >= 0. && b1 + b2 <= 1.)
    {
        // hit
        //        return {true, s2.dot(e2) / s1_dot_e1};
        return TriangleSegmentIntersectRes(true, s2.dot(e2) / s1_dot_e1); // VS2008
    }
    //    return {false, 0.};
    return TriangleSegmentIntersectRes(false, 0.f); // VS2008
}

class BVHNode
{
public:
    size_t left;
    size_t right;
    size_t parent;
    AABB boundingBox;

    BVHNode(size_t l, size_t r, size_t p, const AABB &aabb) : left(l), right(r), parent(p), boundingBox(aabb) {}
    BVHNode() : left(0), right(0), parent(0), boundingBox() {}
};

class BVH_intersection
{
public:
    const Mesh &mesh;
    std::vector<Vec3f> faceCenters;
    std::vector<BVHNode> nodes;

    BVH_intersection(const Mesh &mesh_) : mesh(mesh_)
    {
        // ??????????????????
        faceCenters.reserve(mesh.indices.size());
        for (std::vector<Vec3u>::const_iterator it = mesh.indices.begin(); it != mesh.indices.end(); ++it)
        {
            const Vec3u &face = *it;
            Vec3f center = (mesh.vertices[face[0]] + mesh.vertices[face[1]] + mesh.vertices[face[2]]) / 3.0f;
            faceCenters.push_back(center);
        } // VS2008

        // indicesList ????0,1,2...?C++11iota??
        // std::iota(indicesList.begin(), indicesList.end(), 0);
        std::vector<size_t> indicesList(mesh.indices.size());
        for (size_t i = 0; i < indicesList.size(); ++i)
        {
            indicesList[i] = i;
        }

        // ??BVH????
        try
        {
            nodes.resize(2 * mesh.indices.size() - 1);
        }
        catch (const std::length_error &e)
        {
            std::cout << "vector.length:" << nodes.size() << " resize_pa:" << (2 * mesh.indices.size() - 1) << std::endl;
            std::cerr << "Caught a length_error: " << e.what() << std::endl;
        }


        size_t nowIdx = 0;
        dfsBuild(indicesList, computeAABB(indicesList), nowIdx);
    }
    
    // BVH???
    size_t dfsBuild(std::vector<size_t> &indicesList, AABB aabb, size_t &nowIdx)
    {
        const size_t nodeIdx = nowIdx;
        nowIdx++;
        if (indicesList.size() == 1)
        {
            // leaf
            nodes[nodeIdx] = BVHNode(0, indicesList[0], 0, aabb); // VS2008
            return nodeIdx;
        }

        // longest axis
        int longAxis = -1;
        float longAxisLen = -1;
        for (int i = 0; i < 3; ++i)
        {
            const float axisLen = aabb.max[i] - aabb.min[i];
            if (axisLen > longAxisLen)
            {
                longAxisLen = axisLen;
                longAxis = i;
            }
        }

        // split indices list
        const size_t k = indicesList.size() / 2;
        //        std::nth_element(indicesList.begin(), indicesList.begin() + k - 1, indicesList.end(),
        //                         [&](const size_t &a, const size_t &b) {
        //                             return faceCenters[a][longAxis] < faceCenters[b][longAxis];
        //                         });
        nth(indicesList, k - 1, longAxis);
        std::vector<size_t> leftIndices(indicesList.begin(), indicesList.begin() + k);
        std::vector<size_t> rightIndices(indicesList.begin() + k, indicesList.end());

        const AABB leftAABB = computeAABB(leftIndices);
        const AABB rightAABB = computeAABB(rightIndices);
        const size_t leftIdx = dfsBuild(leftIndices, leftAABB, nowIdx);
        const size_t rightIdx = dfsBuild(rightIndices, rightAABB, nowIdx);
        //		 nodes[nodeIdx] = {leftIdx, rightIdx, 0, aabb};
        nodes[nodeIdx] = BVHNode(leftIdx, rightIdx, 0, aabb); // VS2008
        nodes[leftIdx].parent = nodeIdx;
        nodes[rightIdx].parent = nodeIdx;
        return nodeIdx;
    }

    bool intersectWithLineSegment(const LineSegment &lineSegment) const
    {
        return recursiveLineSegIntersection(lineSegment, 0);
    }

    // private:
    //     AABB computeAABB(const std::vector<size_t> &indices) {
    //         AABB aabb;
    //         for (const size_t &idx: indices) {
    //             const Vec3u &face = mesh.indices[idx];
    //             for (int i = 0; i < 3; ++i) {
    //                 const Vec3f &vertex = mesh.vertices[face[i]];
    //                 for (int j = 0; j < 3; ++j) {
    //                     aabb.min[j] = std::min(aabb.min[j], vertex[j]);
    //                    aabb.max[j] = std::max(aabb.max[j], vertex[j]);
    //                 }
    //             }
    //         }
    //         return {aabb.min, aabb.max};
    //     }
private:
    // ??????????????AABB??????????AABB
    AABB computeAABB(const std::vector<size_t> &indices)
    {
        AABB aabb;
        for (std::vector<size_t>::const_iterator it = indices.begin(); it != indices.end(); ++it)
        {
            const Vec3u &face = mesh.indices[*it];
            for (int i = 0; i < 3; ++i)
            {
                const Vec3f &vertex = mesh.vertices[face[i]];
                for (int j = 0; j < 3; ++j)
                {
                    aabb.min[j] = min(aabb.min[j], vertex[j]);
                    aabb.max[j] = max(aabb.max[j], vertex[j]);
                }
            }
        }
        return aabb;
    }

    bool recursiveLineSegIntersection(const LineSegment &lineSegment, size_t nodeIdx) const
    {
        // segment-box intersection test
        const AABB &aabb = nodes[nodeIdx].boundingBox;
        const Vec3f &dir = lineSegment.getDir();
        bool hit = false;
        for (int i = 0; !hit && i < 3; ++i)
        {
            float t_min = (aabb.min[i] - lineSegment.start[i]) / dir[i];
            float t_max = (aabb.max[i] - lineSegment.start[i]) / dir[i];
            if (t_min > t_max)
            {
                std::swap(t_min, t_max);
            }
            if (t_min > lineSegment.getLength() || t_max < 0)
                return false;
            Vec3f hitPt = lineSegment.start + dir * t_min;
            int otherPlane1 = (i + 1) % 3, otherPlane2 = (i + 2) % 3;
            if (hitPt[otherPlane1] >= aabb.min[otherPlane1] && hitPt[otherPlane1] <= aabb.max[otherPlane1] &&
                hitPt[otherPlane2] >= aabb.min[otherPlane2] && hitPt[otherPlane2] <= aabb.max[otherPlane2])
            {
                hit = true;
            }
        }
        if (hit)
        {
            if (nodes[nodeIdx].left == 0)
            {
                // leaf
                Vec3u face = mesh.indices[nodes[nodeIdx].right];
                TriangleSegmentIntersectRes res = triangleSegmentIntersection(lineSegment, mesh.vertices[face[0]], mesh.vertices[face[1]],
                                                                              mesh.vertices[face[2]]);
                if (!res.hit || res.t < 0 || res.t > lineSegment.getLength())
                    return false;
                return true;
            }
            else
            {
                // check children
                return recursiveLineSegIntersection(lineSegment, nodes[nodeIdx].left) ||
                       recursiveLineSegIntersection(lineSegment, nodes[nodeIdx].right);
            }
        }
        return false;
    }

    // implement nth, without std::nth_element
    // kth is 0-based
    void nth(std::vector<size_t> &indicesList, size_t kth, int longAxis)
    {
        recursiveChoose(indicesList, 0, indicesList.size() - 1, kth, longAxis);
    }
    void recursiveChoose(std::vector<size_t> &indicesList, size_t begin, size_t end, size_t kth, int longAxis)
    {
        if (begin >= end)
            return;
        int i = begin, j = end;
        while (i < j)
        {
            while (i < j && faceCenters[indicesList[j]][longAxis] >= faceCenters[indicesList[begin]][longAxis])
                j--;
            while (i < j && faceCenters[indicesList[i]][longAxis] <= faceCenters[indicesList[begin]][longAxis])
                i++;
            std::swap(indicesList[i], indicesList[j]);
        }
        std::swap(indicesList[begin], indicesList[i]);
        if (i == kth)
            return;
        if (i < kth)
            recursiveChoose(indicesList, i + 1, end, kth, longAxis);
        else
            recursiveChoose(indicesList, begin, i - 1, kth, longAxis);
    }
};