// Each #kernel tells which function to compile; you can have many kernels
#pragma kernel Trace

struct Triangle
{
    float3 a;
    float3 b;
    float3 c;
};

struct BVHNode
{
    float4 l_xy;
    float4 r_xy;
    float4 lr_z;
    uint offset;
    uint primitiveCount;
    uint axis;
    uint pad;
};

RWStructuredBuffer<Triangle> _Triangles;
RWStructuredBuffer<BVHNode> _BVHNodes;
RWStructuredBuffer<uint> _BVHIndices;
RWTexture2D<float4> _LightMap;
Texture2D<float4> _RayOrigin;
Texture2D<float4> _RayDirection;
uint _Resolution;
float _Time;

#define BVH_STACK_SIZE 64

float4 RayTriangleTest(float4 o, float4 d, float3 a, float3 b, float3 c)
{
    float3 e1 = b - a;
    float3 e2 = c - a;

    float3 pvec = cross(d.xyz, e2);
    float det = dot(e1, pvec);

    if (abs(det) < 0.0001f)
    {
        return float4(0.0f, 0.0f, 0.0f, -1.0f);
    }

    float inv_det = 1.0f / det;

    float3 tvec = o.xyz - a;
    float bU = dot(tvec, pvec) * inv_det;
    if (bU < 0.0f || bU > 1.0f)
    {
        return float4(0.0f, 0.0f, 0.0f, -1.0f);
    }

    float3 qvec = cross(tvec, e1);
    float bV = dot(d.xyz, qvec) * inv_det;
    if (bV < 0.0f || bU + bV > 1.0f)
    {
        return float4(0.0f, 0.0f, 0.0f, -1.0f);
    }

    float dist = dot(e2, qvec) * inv_det;

    return float4(dist, bU, bV, 1.0f);
}

float4 RayAABBTest(float4 o, float4 d, float4 i, float3 bmin, float3 bmax)
{
    float3 v1 = (bmin - o.xyz) * i.xyz;
    float3 v2 = (bmax - o.xyz) * i.xyz;

    float3 tnear = min(v1, v2);
    float3 tfar = max(v1, v2);

    float tenter = max(tnear.x, max(tnear.y, tnear.z));
    float texit = min(tfar.x, min(tfar.y, tfar.z));

    if (tenter < texit && texit > 0.0f)
    {
        return float4(max(tenter, 0.0f), 0.0f, 0.0f, 1.0f);
    }

    return float4(0.0f, 0.0f, 0.0f, -1.0f);
}

float4 TraceRay(float4 o, float4 d)
{
    /*float tmax = 10000.0f;
    float bU = 0.0f;
    float bV = 0.0f;

    int id = -1;

    for (uint i = 0; i < 1748; i++)
    {
        float4 tri_res = RayTriangleTest(o, d, _Triangles[i].a, _Triangles[i].b, _Triangles[i].c);

        if (tri_res.w > 0.0f)
        {
            if (tri_res.x < tmax && tri_res.x > 0.0f)
            {
                tmax = tri_res.x;
                bU = tri_res.y;
                bV = tri_res.z;
                id = i;
            }
        }
    }

    return float4(bU, bV, tmax, asfloat(id));*/

    float4 inv = float4(1.0f / d.x, 1.0f / d.y, 1.0f / d.z, 1.0f);
    float4 oinv = o * inv;

    int stack[BVH_STACK_SIZE];
    int stack_ptr = 0;
    stack[stack_ptr] = 0xffffffff;
    uint node_id = 0;

    float tmin = 0.0f;
    float tmax = 10000.0f;
    float bU = 0.0f;
    float bV = 0.0f;
    float dist = tmax;

    float leaf = 0.0f;
    int id = -1;

    float step = 0.0f;

    [loop] for (int i = 0; i < 1000; i++)
    {
        step += 1.0f;

        // Interior node
        if (_BVHNodes[node_id].primitiveCount == 0)
        {
            // Fetch children bounding boxes
            float4 n0xy = _BVHNodes[node_id].l_xy;
            float4 n1xy = _BVHNodes[node_id].r_xy;
            float4 nz = _BVHNodes[node_id].lr_z;

            // Test against child AABBs
            /*float c0lox = n0xy.x * inv.x - oinv.x;
            float c0hix = n0xy.y * inv.x - oinv.x;
            float c0loy = n0xy.z * inv.y - oinv.y;
            float c0hiy = n0xy.w * inv.y - oinv.y;
            float c0loz = nz.x * inv.z - oinv.z;
            float c0hiz = nz.y * inv.z - oinv.z;
            float c1loz = nz.z * inv.z - oinv.z;
            float c1hiz = nz.w * inv.z - oinv.z;
            float c0min = max(max(min(c0lox, c0hix), min(c0loy, c0hiy)), max(min(c0loz, c0hiz), tmin));
            float c0max = min(min(max(c0lox, c0hix), max(c0loy, c0hiy)), min(max(c0loz, c0hiz), tmax));
            float c1lox = n1xy.x * inv.x - oinv.x;
            float c1hix = n1xy.y * inv.x - oinv.x;
            float c1loy = n1xy.z * inv.y - oinv.y;
            float c1hiy = n1xy.w * inv.y - oinv.y;
            float c1min = max(max(min(c1lox, c1hix), min(c1loy, c1hiy)), max(min(c1loz, c1hiz), tmin));
            float c1max = min(min(max(c1lox, c1hix), max(c1loy, c1hiy)), min(max(c1loz, c1hiz), tmax));

            bool traverseChild0 = (c0max >= c0min);
            bool traverseChild1 = (c1max >= c1min);*/

            float4 c0 = RayAABBTest(o, d, inv, float3(n0xy.x, n0xy.z, nz.x), float3(n0xy.y, n0xy.w, nz.y));
            float4 c1 = RayAABBTest(o, d, inv, float3(n1xy.x, n1xy.z, nz.z), float3(n1xy.y, n1xy.w, nz.w));

            bool traverseChild0 = (c0.w >= 0.0f);
            bool traverseChild1 = (c1.w >= 0.0f);

            float c0min = c0.x;
            float c1min = c1.x;

            // If no children was hit, get node from stack
            if (!traverseChild0 && !traverseChild1)
            {
                node_id = stack[stack_ptr];
                stack_ptr--;
            }
            else
            {
                unsigned int first_child = node_id + 1;
                unsigned int second_child = _BVHNodes[node_id].offset;
                node_id = (traverseChild0) ? first_child : second_child;

                if (traverseChild0 && traverseChild1)
                {
                    if (c1min < c0min)
                    {
                        node_id = second_child;
                        stack_ptr++;
                        stack[stack_ptr] = first_child;
                    }
                    else
                    {
                        stack_ptr++;
                        stack[stack_ptr] = second_child;
                    }
                }
            }
        }
        else
        {
            int prim_offset = _BVHNodes[node_id].offset;
            for (uint i = 0; i < _BVHNodes[node_id].primitiveCount; i++)
            {
                int tri_idx = _BVHIndices[prim_offset + i];

                float4 tri_res = RayTriangleTest(o, d, _Triangles[tri_idx].a, _Triangles[tri_idx].b, _Triangles[tri_idx].c);

                if (tri_res.w > 0.0f)
                {
                    if (tri_res.x < tmax && tri_res.x > 0.0f)
                    {
                        tmax = tri_res.x;
                        bU = tri_res.y;
                        bV = tri_res.z;
                        id = i;
                    }
                }
            }

            node_id = stack[stack_ptr];
            stack_ptr--;
        }

        if (node_id == 0xffffffff)
        {
            break;
        }
    }
    //while (node_id != 0xffffffff);

    return float4(bU, bV, tmax, asfloat(id));
}

[numthreads(8,8,1)]
void Trace(uint3 id : SV_DispatchThreadID)
{
    // Early out on pixels that are outside of texture boundaries
    if (id.x >= _Resolution || id.y >= _Resolution)
    {
        return;
    }

    float4 direction = normalize(_RayDirection.Load(id) * 2.0f - 1.0f);
    float4 origin = _RayOrigin.Load(id) + direction * 0.001f;

    if (dot(direction, direction) < 0.0002f)
    {
        _LightMap[id.xy] = float4(0.0f, 0.0f, 0.0f, 0.0f);
        return;
    }

    float4 res = TraceRay(origin, normalize(float4(-sin(_Time * 6.28f) * 0.05f, 1.0f, cos(_Time * 6.28f) * 0.05f, 1.0f)));

    float l = max(dot(direction, normalize(float4(-sin(_Time * 6.28f) * 0.05f, 1.0f, cos(_Time * 6.28f) * 0.05f, 1.0f))), 0.0f);

    if (asuint(res.w) != (uint) (-1))
    {
        l *= 0.0f;
    }

    //_LightMap[id.xy] = float4(id.x & id.y, (id.x & 15) / 15.0, (id.y & 15) / 15.0, 0.0);
    _LightMap[id.xy] = 0.99f * _LightMap[id.xy] + 0.01f * float4(l, l, l, 1.0f);

}