#include "lib/shared/point.hlsl"
#include "lib/shared/quat-functions.hlsl"
#include "lib/shared/pbr.hlsl"

cbuffer Params : register(b0)
{
    float StepCount;
    float DecayW;
    float Extend;
    float SpreadColor;
    float SpreadColorShift;
}

StructuredBuffer<Point> SourcePoints : t0;
StructuredBuffer<PbrVertex> Vertices: t1;
StructuredBuffer<int3> Indices: t2;

RWStructuredBuffer<Point> ResultPoints : u0;


// Casual Moller-Trumbore GPU Ray-Triangle Intersection Routine
// bool intersectMT(
//     float3 orig, float3 dir,
//     float3 v0, float3 v1, float3 v2,
//     out float3 baryzentricUVW,
//     out float t)
// {
//     float3 e1 = v1 -  v0;
//     float3 e2 = v2 -  v0;
//     float3 normal = normalize(cross(e1, e2));
//     float b = dot(normal, dir);
//     float3 w0 = orig -  v0;
//     float a = -dot(normal, w0);
//     t = a / b;
//     float3 p = orig + t * dir;
//     float uu, vv, uv, wu, wv, inverseD;
//     //float2 baryzentricUV = 0;

//     uu = dot(e1, e1);
//     uv = dot(e1, e2);
//     vv = dot(e2, e2);
//     float3 w = p -  v0;
//     wu = dot(w, e1);
//     wv = dot(w, e2);
//     inverseD = uv * uv -  uu * vv;
//     inverseD = 1.0f / inverseD;
//     float u = (uv * wv -  vv * wu) * inverseD;

//     if (u < 0.0f || u > 1.0f) 
//         //return -1.0f;
//         return false;

//     float v = (uv * wu -  uu * wv) * inverseD;
//     if (v < 0.0f || (u + v) > 1.0f)
//         return false;
//         //return -1.0f;

//     baryzentricUVW = float3(u,v, 1-u-v).xzz;
//     return true;
// }

static const float kEpsilon = 0.0001;

// From https://graphicscodex.courses.nvidia.com/app.html?page=_rn_rayCst#section4.2
bool intersect(
    float3 orig, float3 dir,
    float3 v0, float3 v1, float3 v2,
    out float3 b,
    out float t)
{
    // Edge vectors
    float3 e_1 = v1 - v0;
    float3 e_2 = v2 - v0;

    // Face normal
    float3 n = normalize(cross(e_1, e_2));

    float3 q = cross(dir, e_2);
    float a = dot(e_1, q);

    // Backfacing / nearly parallel, or close to the limit of precision?
    if ((dot(n, dir) >= 0) || (abs(a) <= kEpsilon)) return false;

    float3 s = (orig - v0) / a;
    float3 r = cross(s, e_1);

    b[0] = dot(s, q);
    b[1] = dot(r, dir);
    b[2] = 1.0f - b[0] - b[1];

    t = dot(e_2,r);

    // Intersected inside triangle?
    return ((b[0] >= 0) && (b[1] >= 0) && (b[2] >= 0) && (t >= 0));
}

static const float NaN = sqrt(-1);

static const int RAY_THREAD_COUNT = 8;
static const int FACE_THREAD_COUNT = 512/RAY_THREAD_COUNT;

groupshared int BestHitIntDistances[RAY_THREAD_COUNT];
groupshared int BestHitIndices[RAY_THREAD_COUNT];
groupshared float3 BestHitPositions[RAY_THREAD_COUNT];
groupshared float2 BestHitBaryUV[RAY_THREAD_COUNT];



[numthreads(RAY_THREAD_COUNT,FACE_THREAD_COUNT,1)]
void main(uint3 i : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID)
{
    uint rayCount, stride;
    SourcePoints.GetDimensions(rayCount, stride);

    uint faceCount;
    Indices.GetDimensions(faceCount, stride);

    uint rayId = i.x;
    uint rayThreadId = GTid.x;
    uint faceThreadId = i.y;

    uint stepCount = (uint)StepCount; // including separator
    uint rayGroupStartIndex= i.x * stepCount;

    Point p = SourcePoints[i.x];

    // Write ray start and seperator
    ResultPoints[rayGroupStartIndex + 0] = p;
    ResultPoints[rayGroupStartIndex + stepCount -1].W = NaN;

    float3 rayOrigin = p.Position;
    float3 rayDirection = qRotateVec3( float3(0,0,1), p.Rotation);
    float w = p.W;

    int _bestHitIndex = -1;
    float3 _bestHitPosition = rayOrigin + rayDirection * Extend;
    float2 _bestHitBaryUv = 0;
    for(uint stepIndex=1; stepIndex < (stepCount - 1); stepIndex++ )
    {

        if(faceThreadId == 0) 
        {
            if(rayId < rayCount) 
            {
                BestHitIntDistances[rayThreadId] = 99999999;        
                BestHitIndices[rayThreadId] = -1;
                BestHitPositions[rayThreadId] = rayOrigin + rayDirection * Extend;

                _bestHitIndex = -1;
                _bestHitPosition = rayOrigin + rayDirection * Extend;
                _bestHitBaryUv = 0;
            }
        }
        GroupMemoryBarrierWithGroupSync();

        int faceGroupCount = faceCount/FACE_THREAD_COUNT;
        for(uint faceGroupStartIndex = 0 ; faceGroupStartIndex < faceCount ; faceGroupStartIndex += FACE_THREAD_COUNT ) 
        {
            uint faceId = faceThreadId + faceGroupStartIndex;
            if(faceId < faceCount) 
            {
                int3 f = Indices[faceId];
                float3 bary;
                float t;

                if(intersect(
                    rayOrigin,
                    rayDirection,
                    Vertices[f[0]].Position,
                    Vertices[f[1]].Position,
                    Vertices[f[2]].Position,
                    bary,
                    t
                )) {
                    float org;
                    int intt = t * 1000;
                    InterlockedMin(BestHitIntDistances[rayThreadId], intt, org);
                    if(org > intt) {
                        BestHitIndices[rayThreadId] = faceId;
                        BestHitBaryUV[rayThreadId] = bary.zx;
                        BestHitPositions[rayThreadId] = rayOrigin + rayDirection *t;

                        _bestHitIndex = faceId;
                        _bestHitBaryUv = bary.zx;
                        _bestHitPosition = rayOrigin + rayDirection *t;
                    }
                }
            }
            //GroupMemoryBarrierWithGroupSync();
        }
        GroupMemoryBarrierWithGroupSync();


        _bestHitIndex = BestHitIndices[rayThreadId]; 


        if(_bestHitIndex < 0)
        {
            rayOrigin += rayDirection * Extend;
            ResultPoints[rayGroupStartIndex + stepIndex].Rotation = p.Rotation;
            ResultPoints[rayGroupStartIndex + stepIndex].Position = rayOrigin;
            ResultPoints[rayGroupStartIndex + stepIndex].W = w;
        }
        else {            
            _bestHitPosition = BestHitPositions[rayThreadId];  
            rayOrigin = _bestHitPosition;
            ResultPoints[rayGroupStartIndex + stepIndex].Position = rayOrigin;
            ResultPoints[rayGroupStartIndex + stepIndex].W = w;

            float3 n0 = normalize(Vertices[Indices[_bestHitIndex][0]].Normal);
            float3 n1 = normalize(Vertices[Indices[_bestHitIndex][1]].Normal);
            float3 n2 = normalize(Vertices[Indices[_bestHitIndex][2]].Normal);

            _bestHitBaryUv = BestHitBaryUV[rayThreadId]; // <----
            float u = _bestHitBaryUv.x;
            float v = _bestHitBaryUv.y;

            float3 n = normalize(u*n0 + v*n1 + (1 - u - v)*n2);
            w *= DecayW;

            //rayDirection= reflect( rayDirection, n * 1);
            float3 I = rayDirection;
            //float3 R = I - 2* ( dot(n,I)* n);

            float phi = acos(dot( n,I ));
            phi += (w- SpreadColorShift) * SpreadColor;
            float3 R= I - 2*cos(phi) * n;
            rayDirection = R;
            //rayDirection= reflect( rayDirection, n);
        }

        GroupMemoryBarrierWithGroupSync();
    }
}
