#pragma once

#include <Geometry/BoundingBox.h>

#include <deque>
#include <set>
#include <algorithm>

namespace Geometry
{

    class BVH
    {
    private:
        class BVHNode
        {
        private:
            static bool triangleSortCriteria(const Triangle * t1, const Triangle * t2, int axis)
            {
                return t1->center()[axis] < t2->center()[axis];
            }

        protected:
            BoundingBox * m_boundingBox;
            std::vector<Triangle *> m_triangles;
            BVHNode * m_firstSon;
            BVHNode * m_secondSon;
            bool isLeaf;

        public:

            BVHNode()
            {
                m_boundingBox = new BoundingBox();
                this->isLeaf = false;
                m_firstSon = nullptr;
                m_secondSon = nullptr;
            }

            ~BVHNode()
            {
                delete m_boundingBox;
                delete m_firstSon;
                delete m_secondSon;
            }

            // adds a deque of triangles to the node
            void addTriangles(std::vector<Triangle *> & triangles)
            {
                for (Triangle * t : triangles)
                {
                    this->addTriangle(t);
                }
            }

            // adds one triangle to the node
            void addTriangle(Triangle * t)
            {
                m_triangles.push_back(t);
                m_boundingBox->update(*t);
            }

            // adds a deque of triangle to bounding box
            void addTrianglesToBoundingBox(const std::vector<Triangle *> & triangles)
            {
                for (Triangle * t : triangles)
                {
                    addTriangleToBoundingBox(t);
                }
            }

            // adds one triangle to bounding box
            void addTriangleToBoundingBox(const Triangle * t)
            {
                if (m_boundingBox->isEmpty())
                {
                    Geometry g;
                    g.addTriangle(*t);
                    m_boundingBox->set(g);
                }
                m_boundingBox->update(*t);
            }

            const BoundingBox & getBoundingBox()
            {
                return *m_boundingBox;
            }

            // This computes the BVH before the first render of the scene
            void compute(std::vector<Triangle *> & triangles, int depth)
            {
                //std::cout << "\n\ndepth : " << depth << "\nnumber of triangles : " << triangles.size() << std::endl;
                m_firstSon = new BVHNode();
                m_secondSon = new BVHNode();

                // adds the triangles to the node
                this->addTrianglesToBoundingBox(triangles);

                if (triangles.size() < 8)
                {
                    //std::cout << "added " << triangles.size() << " triangles to leaf node !" << std::endl;
                    this->isLeaf = true;
                    this->addTriangles(triangles);
                    return;
                }

                std::vector<Triangle *> trianglesL;
                std::vector<Triangle *> trianglesR;

                // sorts the triangles
                sortTriangles(triangles, trianglesL, trianglesR);

                // computes the sons of this node
                this->m_firstSon->compute(trianglesL, depth + 1);
                this->m_secondSon->compute(trianglesR, depth + 1);
            }

            // Sorts triangles relative to an axis using SAH (the axis selected is the one on which the bounding box is the largest)
            void sortTriangles(std::vector<Triangle *> & triangles, std::vector<Triangle *> & trianglesL, std::vector<Triangle *> & trianglesR)
            {
                int axis = 0;
                for (int i = 1; i < 3; i++)
                {
                    double max = this->m_boundingBox->max()[i];
                    double min = this->m_boundingBox->min()[i];
                    if (max - min >= this->m_boundingBox->max()[axis] - this->m_boundingBox->min()[axis])
                        axis = i;
                }

                // lambda function to sort the triangles using std::sort
                auto triangleSort = [=](const Triangle * t1, const Triangle * t2) { return triangleSortCriteria(t1, t2, axis); };

                std::sort(triangles.begin(), triangles.end(), triangleSort);

                //std::cout << "sorted array size : " << triangles.size() << std::endl;
                int size = triangles.size();

                int indexSAH = getSAHIndex(triangles);

                // Dispatching triangles in sub sections (Left & Right)
                for(int i = size - 1; i >= 0; i--)
                {
                    if (i < indexSAH + 1)
                    {
                        trianglesL.push_back(triangles.back());
                        triangles.pop_back();
                    }
                    else
                    {
                        trianglesR.push_back(triangles.back());
                        triangles.pop_back();
                    }
                };

                //std::cout << "Left array size : " << trianglesL.size() << std::endl;
                //std::cout << "Right array size : " << trianglesR.size() << std::endl;
            }

            // this returns the index of the best triangle to dispatch the triangles into two nodes according to the SAH
            int getSAHIndex(std::vector<Triangle *> & triangles)
            {
                // size of the triangles array
                int size = triangles.size();

                // vectors used to compute the area of the bounding box of the present node
                Math::Vector3f min = this->m_boundingBox->min();
                Math::Vector3f max = this->m_boundingBox->max();
                Math::Vector3f diff = max - min;

                // surface of this node's bounding box
                double surface;

                // surfaces of the sons nodes for each triangles (helps to compute the best split)
                double * surfaceL = new double[size];
                double * surfaceR = new double[size];

                // vectors used to compute the area of the sons nodes
                Math::Vector3f minL(Math::makeVector(std::numeric_limits<double>::max(), std::numeric_limits<double>::max(), std::numeric_limits<double>::max()));
                Math::Vector3f minR(Math::makeVector(std::numeric_limits<double>::max(), std::numeric_limits<double>::max(), std::numeric_limits<double>::max()));
                Math::Vector3f maxL(Math::makeVector(std::numeric_limits<double>::lowest(), std::numeric_limits<double>::lowest(), std::numeric_limits<double>::lowest()));
                Math::Vector3f maxR(Math::makeVector(std::numeric_limits<double>::lowest(), std::numeric_limits<double>::lowest(), std::numeric_limits<double>::lowest()));
                Math::Vector3f diffL, diffR;

                // computation of the area for each possible split of triangles in the sons nodes
                for (int i = 0; i < size; i++)
                {
                    // computes the min/max vectors as in a bounding box for the given triangle
                    for (int j = 0; j < 3; j++)
                    {
                        minL = minL.simdMin(triangles.at(i)->vertex(j));
                        maxL = maxL.simdMax(triangles.at(i)->vertex(j));
                        minR = minR.simdMin(triangles.at(size - i - 1)->vertex(j));
                        maxR = maxR.simdMax(triangles.at(size - i - 1)->vertex(j));
                    }

                    // used to compute the area (diff[0] = dx, diff[1] = dy, diff[2] = dz)
                    diffL = maxL - minL;
                    diffR = maxR - minR;

                    // computation of the surface with the triangles [0, i] and [size - i, size]
                    surfaceL[i] = 2 * (diffL[0] * diffL[1] + diffL[0] * diffL[2] + diffL[1] * diffL[2]);
                    surfaceR[size - i - 1] = 2 * (diffR[0] * diffR[1] + diffR[0] * diffR[2] + diffR[1] * diffR[2]);

                    //std::cout << "L : " << surfaceL[i] << "\nR : " << surfaceR[size - i - 1] << std::endl;
                }

                // computation of the present node bounding box
                surface = 2 * (diff[0] * diff[1] + diff[0] * diff[2] + diff[1] * diff[2]);

                double cost = 0;
                double bestCost;

                // definition of the traversal cost and the intersect cost
                double cost_traversal = 1;
                double cost_intersect = 1;

                int indexSAH = 0;

                // computation of the cost if the split is [0, 0] | [1, size]
                bestCost = cost_traversal + surfaceL[0] / surface * 1 * cost_intersect + surfaceR[0] / surface * (size - 1) * cost_intersect;
                //std::cout << "initial cost :" << bestCost << std::endl;
                //std::cout << "node surface : " << surface << " | surface L : " << surfaceL[0] << " | surface R : " << surfaceR[1] << " | cost : " << bestCost << std::endl;

                // computation of the best cost
                for (int i = 1; i < size - 1; i++)
                {
                    //std::cout << "node surface : " << surface << " | surface L : " << surfaceL[i] << " | surface R : "  << surfaceR[i] << " | cost : " << cost << std::endl;
                    cost = cost_traversal + surfaceL[i] / surface * (i + 1) * cost_intersect + surfaceR[i + 1] / surface * (size - i - 1) * cost_intersect;
                    if (cost <= bestCost)
                    {
                        bestCost = cost;
                        indexSAH = i;
                    }
                    if (cost > size * surface * cost_intersect)
                        break;
                }

                delete[] surfaceL;
                delete[] surfaceR;

                return indexSAH;
            }

            // This intersects the ray with the BVH. Not optimized yet.
            bool intersection(CastedRay & r, double entryT, double exitT, int depth)
            {
                if (this->isLeaf)
                {
                    for (Triangle * t : m_triangles)
                    {
                        r.intersect(t);
                    }
                    return r.validIntersectionFound();
                }
                else
                {
                    double entryTR, exitTR, entryTL, exitTL;

                    // epsilon to account errors in floating point operations
                    double epsilon = std::numeric_limits<double>::epsilon() * exitT;

                    // tests intersection with the sons node's bounding boxes
                    boolean left = this->m_firstSon->m_boundingBox->intersect(r, entryT-epsilon, exitT+epsilon, entryTL, exitTL);
                    boolean right = this->m_secondSon->m_boundingBox->intersect(r, entryT-epsilon, exitT+epsilon, entryTR, exitTR);

                    // checks if the two intervals are empty
                    if ((!left) && (!right))
                    {
                        return false;
                    }
                    else if (!left || !right) // checks if one of the intervals is empty
                    {
                        if (!right)
                        {
                            return this->m_firstSon->intersection(r, entryTL, exitTL, depth + 1);
                        }
                        else
                        {
                            return this->m_secondSon->intersection(r, entryTR, exitTR, depth + 1);
                        }
                    }
                    else // if none is empty
                    {
                        if (entryTL < entryTR)
                        {
                            bool lSon = this->m_firstSon->intersection(r, entryTL, exitTL, depth + 1);
                            if(!lSon)
                                return this->m_secondSon->intersection(r, entryTR, exitTR, depth + 1);
                            else if (r.intersectionFound().tRayValue() > entryTR)
                                return (this->m_secondSon->intersection(r, entryTR, r.intersectionFound().tRayValue(), depth + 1) || lSon);
                            return lSon;
                        }
                        else
                        {
                            bool rSon = this->m_secondSon->intersection(r, entryTR, exitTR, depth + 1);
                            if (!rSon)
                                return this->m_firstSon->intersection(r, entryTL, exitTL, depth + 1);
                            else if (r.intersectionFound().tRayValue() > entryTL)
                                return (this->m_firstSon->intersection(r, entryTL, r.intersectionFound().tRayValue(), depth + 1) || rSon);
                            return rSon;
                        }
                    }
                }
            }
        };

    protected:
        BVHNode * m_root;
        std::vector<Triangle *> m_triangles;

    public:

        BVH()
        {
            m_root = new BVHNode();
        }

        ~BVH()
        {
            delete m_root;
            for (Triangle * t : m_triangles)
            {
                delete t;
            }
        }

        // adds every triangle in a geometry to the BVH
        void addGeometry(Geometry & g)
        {
            for (auto it = g.getTriangles().begin(); it!= g.getTriangles().end(); it++)
            {
                Triangle * t = new Triangle(*it);
                m_triangles.push_back(t);
            }
        }

        // adds a triangle to the bVH
        void addTriangle(Triangle * t)
        {
            this->m_triangles.push_back(t);
        }

        // computes BVH tree
        void computeTree()
        {
            std::cout << "nbr de triangles total : " << m_triangles.size() << std::endl;
            this->m_root->compute(this->m_triangles, 0);
        }

        // intersects the ray with the scene
        void intersection(CastedRay & r)
        {
            double test1, test2;
            if (!this->m_root->getBoundingBox().intersect(r, 0.0, std::numeric_limits<double>::max(), test1, test2))
                return;
            this->m_root->intersection(r, 0.0, std::numeric_limits<double>::max(), 0);
        }
    };
}