#include "KdTree.h"
#include "GLViewer.h"
#include "Mesh.h"
#include "Scene.h"

#undef near
#undef far


KdTreeNode::KdTreeNode(std::vector<unsigned int> indices, const Mesh& mesh, unsigned int depth, const BoundingBox& box) :
    left(0),
    right(0),
    axis(depth % 3),
    split(0.0f),
    bbox(box),
    mesh(mesh)
{
    // End recursion case: not enought triangles to continue splitting, or max depth encountered
    if(indices.size() <= kdTreeMinSize || depth >= kdTreeMaxDepth)
    {
        tris = indices;
        return;
    }

    // Optimal split
    std::vector<unsigned int> atLeft, atRight;
    findOptimalSplit(indices, atLeft, atRight);
    BoundingBox cLeft, cRight;
    box.split(split, axis, cLeft, cRight);

    // Build children tree
    if(atLeft.size())
        left = new KdTreeNode(atLeft, mesh, depth + 1, cLeft);
    if(atRight.size())
        right = new KdTreeNode(atRight, mesh, depth + 1, cRight);
}

KdTreeNode::~KdTreeNode()
{
    if(left) delete left;
    if(right) delete right;
}

void KdTreeNode::splitTriangles(float splitPos, const std::vector<unsigned int>& indices, std::vector<unsigned int>& tLeft, std::vector<unsigned int>& tRight) const
{
    // Add triangles in a part they intersect. A triangle can be in both right and left
    for(unsigned int t = 0; t < indices.size(); t++)
    {
        Triangle tr = mesh.getTriangles()[indices[t]];

        // Number of vertices of the triangle in the left part
        unsigned int nbInLeft = 0;
        for(unsigned int j = 0; j < 3; j++)
            if(mesh.getVertices()[tr.getVertex(j)].getPos()[axis] < splitPos)
                nbInLeft++;

        // Add to the parts the triangle intersects
        if(nbInLeft > 0)
            tLeft.push_back(indices[t]);
        if(nbInLeft < 3)
            tRight.push_back(indices[t]);
    }
}

void KdTreeNode::findOptimalSplit(const std::vector<unsigned int>& indices, std::vector<unsigned int>& tLeft, std::vector<unsigned int>& tRight)
{
    // Surface Area Heuristic : do not use, only 27% intersect speed improvement, but takes waaaaay longer to build (100x slower)
    /*
    float mincost = -1, splitPos = 0.0f;
    BoundingBox bLeft, bRight;
    std::vector<unsigned int> atLeft, atRight;

    for(unsigned int t = 0; t < indices.size(); t++)
    {
        float tmin = triangles[indices[t]].minPos(axis, vertices);
        float tmax = triangles[indices[t]].maxPos(axis, vertices);

        // Check first extreme
        if(box.split(tmin, axis, bLeft, bRight))
        {
            // Split triangles
            splitTriangles(tmin, indices, atLeft, atRight);

            // Compute cost
            float cost = bLeft.getArea() * atLeft.size() + bRight.getArea() * atRight.size();
            if(mincost < 0.0f || cost < mincost)
            {
                tLeft = atLeft;
                tRight = atRight;
                mincost = cost;
                splitPos = tmin;
            }

            // Clear split vectors
            atLeft.clear();
            atRight.clear();
        }

        // Check second extreme
        if(box.split(tmax, axis, bLeft, bRight))
        {
            // Split triangles
            splitTriangles(tmax, indices, atLeft, atRight);

            // Compute cost
            float cost = bLeft.getArea() * atLeft.size() + bRight.getArea() * atRight.size();
            if(mincost < 0.0f || cost < mincost)
            {
                tLeft = atLeft;
                tRight = atRight;
                mincost = cost;
                splitPos = tmax;
            }

            // Clear split vectors
            atLeft.clear();
            atRight.clear();
        }
    }

    // If no split position is found
    if(mincost < 0.0f)
    {
        split = (box.getMax()[axis] - box.getMin()[axis]) / 2;
        splitTriangles(split, indices, tLeft, tRight);
    }
    else // Triangles already split
        split = splitPos;

    /*/
    float min = 1000000.0f, max = -100000000.0f;
    for(unsigned int t = 0; t < indices.size(); t++)
        for(unsigned int j = 0; j < 3; j++)
        {
            float v = mesh.getVertices()[mesh.getTriangles()[indices[t]].getVertex(j)].getPos()[axis];
            if(v > max) max = v;
            if(v < min) min = v;
        }

    split = (min + max) / 2;
    splitTriangles(split, indices, tLeft, tRight);
    //*/
}

bool KdTreeNode::intersectTriangles(const Ray& r, IntersectedTriangle& tri, float tmin, float tmax) const
{
    static const float EPSILON = 0.00001f;
    IntersectedTriangle tmp;
    float dist = -1.0f;

    // Find intersecting triangle closest from the origin
    for(unsigned int t = 0; t < tris.size(); t++)
    {
        if(r.intersect(mesh.getVertices()[mesh.getTriangles()[tris[t]].getVertex(0)].getPos(),
                       mesh.getVertices()[mesh.getTriangles()[tris[t]].getVertex(1)].getPos(),
                       mesh.getVertices()[mesh.getTriangles()[tris[t]].getVertex(2)].getPos(),
                       tmp.point, tmp.u, tmp.v, tmp.t)
           && tmp.t >= tmin - EPSILON && tmp.t <= tmax + EPSILON)
        {
            float d = (tmp.point - r.getOrigin()).getSquaredLength();
            if(dist < 0.0f || d < dist)
            {
                dist = d;
                tri = tmp;
                tri.tri = mesh.getTriangles()[tris[t]];
            }
        }
    }

    return dist >= 0.0f;
}

bool KdTreeNode::intersect(const Ray& r, float tmin, float tmax, IntersectedTriangle& tri) const
{
    if(isLeaf())
        return intersectTriangles(r, tri, tmin, tmax);
    else
    {
        // Find near and far subtrees
        float nearDet = axisVec().dot(r.getDirection());
        KdTreeNode* near = (nearDet > 0 ? left : right);
        KdTreeNode* far = (nearDet > 0 ? right : left);

        // Get intersection with split plane
        float t = (split - r.getOrigin()[axis]) / r.getDirection()[axis];

        // Intersect with children
        if(tmax < t)
        {
            // Only near will intersect
            if(near) return near->intersect(r, tmin, tmax, tri);
            else return false;
        }
        else if(t < tmin)
        {
            // Only tfar will intersect
            if(far) return far->intersect(r, tmin, tmax, tri);
            else return false;
        }
        else
        {
            // Both will intersect
            bool nearSuccess = near ? near->intersect(r, tmin, t, tri) : false;
            if(!nearSuccess)
                return far && far->intersect(r, t, tmax, tri);
            else
                return true;
        }
    }
}

void KdTreeNode::draw(const BoundingBox& bbox, const GLViewer* viewer, const Vec3Df& trans) const
{
    static const Color axisColors[3] = {
        Color((unsigned int)255,0,0),
        Color((unsigned int)0,255,0),
        Color((unsigned int)0,0,255)
    };
    viewer->drawBox(trans + bbox.getMin(), trans + bbox.getMax(), axisColors[axis]);

    BoundingBox cLeft, cRight;
    bbox.split(split, axis, cLeft, cRight);

    if(left) left->draw(cLeft, viewer, trans);
    if(right) right->draw(cRight, viewer, trans);
}

Vec3Df KdTreeNode::axisVec() const
{
    switch(axis)
    {
    case 0: return Vec3Df(1,0,0);
    case 1: return Vec3Df(0,1,0);
    case 2: return Vec3Df(0,0,1);
    default: return Vec3Df(0,0,0);
    }
}


KdTree::KdTree(const Mesh& m) :
    root(0),
    mesh(m)
{
    rebuild();
}

KdTree::~KdTree()
{
    clear();
}

void KdTree::clear()
{
    if(root) delete root;
    root = 0;
}

void KdTree::rebuild()
{
    const std::vector<Triangle>& triangles = mesh.getTriangles();
    clear();
    if(triangles.size() > 0)
    {
        std::vector<unsigned int> indices(triangles.size(), 0);
        for(unsigned int t = 0; t < indices.size(); t++)
            indices[t] = t;
        root = new KdTreeNode(indices, mesh, 0, mesh.getBoundingBox());
    }
}

bool KdTree::intersect(Ray r, IntersectedTriangle &tri) const
{
    float tmin=0, tmax=0;
    if(!root || !r.getRange(mesh.getBoundingBox(), tmin, tmax))
        return false;

    return root->intersect(r, tmin, tmax, tri);
}

void KdTree::draw(const GLViewer* viewer, const Vec3Df& trans) const
{
    if(root)
        root->draw(mesh.getBoundingBox(), viewer, trans);
}
