#pragma once

#include <numeric>
#include "cmath"
#include "algorithm"
#include "Geometry.hpp"

struct IntersectRes {
    bool hit;
    float t;
};

IntersectRes
triangleRayIntersection(const Ray &ray, const Vec3f &a, const Vec3f &b, const Vec3f &c) {
    Vec3f e1 = b - a;
    Vec3f e2 = c - a;
    Vec3f s = ray.start - a;
    Vec3f s1 = ray.dir.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(ray.dir) / s1_dot_e1;
    if (b1 >= 0. && b2 >= 0. && b1 + b2 <= 1.) {
        // hit
        return {true, s2.dot(e2) / s1_dot_e1};
    }
    return {false, 0.};
}

bool isClose(double a, double b, double rel_tol = 1e-9, double abs_tol = 0.0) {
    return std::fabs(a - b) <= std::max(rel_tol * std::max(std::fabs(a), std::fabs(b)), abs_tol);
}

int sign(double x) {
    return (x > 0) - (x < 0);
}

float dot2(Vec3f a) {
    return a.dot(a);
}

float pointTriangleDistance(const Vec3f &p, const Vec3f &a, const Vec3f &b, const Vec3f &c) {
    Vec3f ab = b - a, ap = p - a;
    Vec3f bc = c - b, bp = p - b;
    Vec3f ca = a - c, cp = p - c;
    Vec3f nor = ab.cross(ca);

    return sqrt(
            (sign(ab.cross(nor).dot(ap)) +
             sign(bc.cross(nor).dot(bp)) +
             sign(ca.cross(nor).dot(cp)) < 2.0)
            ?
            fmin(fmin(
                         dot2(ab * std::clamp(ab.dot(ap) / dot2(ab), 0.0f, 1.0f) - ap),
                         dot2(bc * std::clamp(bc.dot(bp) / dot2(bc), 0.0f, 1.0f) - bp)),
                 dot2(ca * std::clamp(ca.dot(cp) / dot2(ca), 0.0f, 1.0f) - cp))
            :
            nor.dot(ap) * nor.dot(ap) / dot2(nor));
}

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 {
public:
    const Mesh &mesh;
    std::vector<Vec3f> faceCenters;
    std::vector<BVHNode> nodes;

    BVH(const Mesh &mesh_) : mesh(mesh_) {
        faceCenters.reserve(mesh.faces.size());
        for (const Vec3u &face: mesh.faces) {
            Vec3f center = (mesh.vertices[face[0]] + mesh.vertices[face[1]] + mesh.vertices[face[2]]) / 3.0f;
            faceCenters.push_back(center);
        }
        std::vector<size_t> indicesList(mesh.faces.size());
        std::iota(indicesList.begin(), indicesList.end(), 0);
        nodes.resize(2 * mesh.faces.size() - 1);
        size_t nowIdx = 0;
        dfsBuild(indicesList, computeAABB(indicesList), nowIdx);
    }

    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] = {0, indicesList[0], 0, aabb};
            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[leftIdx].parent = nodeIdx;
        nodes[rightIdx].parent = nodeIdx;
        return nodeIdx;
    }

    unsigned int intersectWithRay(const Ray &ray) const {
        return recursiveRayIntersection(ray, 0);
    }

    void writeBVHAsObj(const std::string &path) {
        std::ofstream file(path);
        for (const BVHNode &node: nodes) {
            Vec3f vertex;
            for (int i = 0; i < 2; ++i) {
                vertex[0] = i == 0 ? node.boundingBox.min[0] : node.boundingBox.max[0];
                for (int j = 0; j < 2; ++j) {
                    vertex[1] = j == 0 ? node.boundingBox.min[1] : node.boundingBox.max[1];
                    for (int k = 0; k < 2; ++k) {
                        vertex[2] = k == 0 ? node.boundingBox.min[2] : node.boundingBox.max[2];
                        file << "v " << vertex[0] << " " << vertex[1] << " " << vertex[2] << std::endl;
                    }
                }
            }
        }
        for (size_t i = 0; i < nodes.size(); ++i) {
            const BVHNode &node = nodes[i];

            file << "l " << i * 8 + 1 << " " << i * 8 + 2 << std::endl;
            file << "l " << i * 8 + 3 << " " << i * 8 + 4 << std::endl;
            file << "l " << i * 8 + 5 << " " << i * 8 + 6 << std::endl;
            file << "l " << i * 8 + 7 << " " << i * 8 + 8 << std::endl;
            file << "l " << i * 8 + 1 << " " << i * 8 + 3 << std::endl;
            file << "l " << i * 8 + 2 << " " << i * 8 + 4 << std::endl;
            file << "l " << i * 8 + 5 << " " << i * 8 + 7 << std::endl;
            file << "l " << i * 8 + 6 << " " << i * 8 + 8 << std::endl;
            file << "l " << i * 8 + 1 << " " << i * 8 + 5 << std::endl;
            file << "l " << i * 8 + 2 << " " << i * 8 + 6 << std::endl;
            file << "l " << i * 8 + 3 << " " << i * 8 + 7 << std::endl;
            file << "l " << i * 8 + 4 << " " << i * 8 + 8 << std::endl;

        }
        file.close();
    }

    void writeLeavesAsObj(const std::string &path) {
        std::ofstream file(path);
        for (const BVHNode &node: nodes) {
            if (node.left == 0) {
                Vec3f vertex;
                for (int i = 0; i < 2; ++i) {
                    vertex[0] = i == 0 ? node.boundingBox.min[0] : node.boundingBox.max[0];
                    for (int j = 0; j < 2; ++j) {
                        vertex[1] = j == 0 ? node.boundingBox.min[1] : node.boundingBox.max[1];
                        for (int k = 0; k < 2; ++k) {
                            vertex[2] = k == 0 ? node.boundingBox.min[2] : node.boundingBox.max[2];
                            file << "v " << vertex[0] << " " << vertex[1] << " " << vertex[2] << std::endl;
                        }
                    }
                }
            }
        }
        for (size_t i = 0; i < nodes.size(); ++i) {
            const BVHNode &node = nodes[i];
            if (node.left == 0) {
                file << "l " << i * 8 + 1 << " " << i * 8 + 2 << std::endl;
                file << "l " << i * 8 + 3 << " " << i * 8 + 4 << std::endl;
                file << "l " << i * 8 + 5 << " " << i * 8 + 6 << std::endl;
                file << "l " << i * 8 + 7 << " " << i * 8 + 8 << std::endl;
                file << "l " << i * 8 + 1 << " " << i * 8 + 3 << std::endl;
                file << "l " << i * 8 + 2 << " " << i * 8 + 4 << std::endl;
                file << "l " << i * 8 + 5 << " " << i * 8 + 7 << std::endl;
                file << "l " << i * 8 + 6 << " " << i * 8 + 8 << std::endl;
                file << "l " << i * 8 + 1 << " " << i * 8 + 5 << std::endl;
                file << "l " << i * 8 + 2 << " " << i * 8 + 6 << std::endl;
                file << "l " << i * 8 + 3 << " " << i * 8 + 7 << std::endl;
                file << "l " << i * 8 + 4 << " " << i * 8 + 8 << std::endl;
            }
        }
        file.close();
    }

private:
    AABB computeAABB(const std::vector<size_t> &indices) {
        AABB aabb;
        for (const size_t &idx: indices) {
            const Vec3u &face = mesh.faces[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};
    }

    unsigned int recursiveRayIntersection(const Ray &ray, size_t nodeIdx) const {
        // segment-box intersection test
        const AABB &aabb = nodes[nodeIdx].boundingBox;
        bool hit = false;
        for (int i = 0; !hit && i < 3; ++i) {
            float t_min = (aabb.min[i] - ray.start[i]) / ray.dir[i];
            float t_max = (aabb.max[i] - ray.start[i]) / ray.dir[i];
            if (t_min > t_max) {
                std::swap(t_min, t_max);
            }
            if (t_max < 0) return 0;
            Vec3f hitPt = ray.start + ray.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.faces[nodes[nodeIdx].right];
                IntersectRes res = triangleRayIntersection(ray, mesh.vertices[face[0]], mesh.vertices[face[1]],
                                                           mesh.vertices[face[2]]);
                if (!res.hit || res.t < 0) return 0;
                return 1;
            } else {
                // check children
                return recursiveRayIntersection(ray, nodes[nodeIdx].left) +
                       recursiveRayIntersection(ray, nodes[nodeIdx].right);
            }
        }
        return 0;
    }

    // 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);
    }
};