#pragma kernel RaytracingLightCluster
#pragma kernel RaytracingLightCull

#pragma only_renderers d3d11

// SRP & HDRP includes
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/HDRaytracingLightCluster.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/ShaderVariablesRaytracingLightLoop.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/GeometricTools.hlsl"

#define CLUSTER_GROUP_SIZE 8

// Light Data
StructuredBuffer<LightVolume>           _LightVolumes;
uint                                    _LightVolumeCount;

// The target data that this computer shader must fill
RWStructuredBuffer<uint>                _RaytracingLightClusterRW;
float3                                  _ClusterCellSize;

static const float3 CornerSubDirections[8] = {
    float3(-0.5f, -0.5f, -0.5f),
    float3(-0.5f, -0.5f, 0.5f),
    float3(0.5f, -0.5f, -0.5f),
    float3(0.5f, -0.5f, 0.5f),

    float3(-0.5f, 0.5f, -0.5f),
    float3(-0.5f, 0.5f, 0.5f),
    float3(0.5f, 0.5f, -0.5f),
    float3(0.5f, 0.5f, 0.5f)
};

#define CULL_GROUP_SIZE 16

// The target data that this computer shader must fill
RWStructuredBuffer<uint>                _RaytracingLightCullResult;
float3                                  _ClusterCenterPosition;
float3                                  _ClusterDimension;

[numthreads(CLUSTER_GROUP_SIZE, CLUSTER_GROUP_SIZE, CLUSTER_GROUP_SIZE)]
void RaytracingLightCluster(uint3 threadID : SV_GroupThreadID, uint3 groupId : SV_GroupID)
{
    // Fetch the coordinates of the current cell
    uint3 targetCellIndex = groupId * CLUSTER_GROUP_SIZE + threadID;

    // Get its global cell index
    uint cellIndex = targetCellIndex.z + targetCellIndex.y * 32 + targetCellIndex.x * 2048;

    // The size of a cell data-wise
    uint cellDataSize = _LightPerCellCount + 4;

    // Diagonal direction of the cluster
    float3 bottomCorner = _MinClusterPos + float3(0.5f, 0.5f, 0.5f) * _ClusterCellSize;
    // Let's compute the position of this cell
    float3 cellCenterPosition = bottomCorner + (_MaxClusterPos - _MinClusterPos) * targetCellIndex / float3(64.0, 64.0, 32.0);

    // The actual light count that intersects with this cell
    uint currentLightCount = 0;
    uint punctualLightCount = 0;
    uint areaLightCount = 0;
    uint envLightCount = 0;

    // Now let's loop through the lights and fill the cell's information
    for(uint lightIdx = 0; lightIdx < _LightVolumeCount; ++lightIdx)
    {
        // If no more lights can fit, just skip
        if(currentLightCount >= _LightPerCellCount) break;

        // Fetch the target light data
        LightVolume currentLight = _LightVolumes[lightIdx];

        // If this light should be skipped, skip it
        // if(_RaytracingLightCullResult[lightIdx] == 1) continue;
        // if(_LightVolumes[lightIdx].active == 0) continue;

        bool intersects = false;

        if (currentLight.shape == 0)
        {
            // Do a box sphere intersection
            intersects = IntersectSphereAABB(currentLight.position, currentLight.range.x, cellCenterPosition + CornerSubDirections[0] * _ClusterCellSize, cellCenterPosition + CornerSubDirections[7] * _ClusterCellSize);
        }
        else
        {
            intersects =  true;
            for(uint cIdx = 0; cIdx < 3; ++cIdx)
            {
                // Check if this corner is inside the shphere
                float minLightPos = currentLight.position[cIdx] - currentLight.range[cIdx];
                float maxLightPos = currentLight.position[cIdx] + currentLight.range[cIdx];
                float clusterMinPos = (cellCenterPosition[cIdx] - _ClusterCellSize[cIdx]);
                float clusterMaxPos = (cellCenterPosition[cIdx] + _ClusterCellSize[cIdx]);
                if (minLightPos > clusterMaxPos || maxLightPos < clusterMinPos)
                {
                    intersects = false;
                    break;
                }
            }
        }

        if(intersects)
        {
            // Flag this light in this cell and increase the light count
            _RaytracingLightClusterRW[cellIndex * cellDataSize + 4 + currentLightCount] = currentLight.lightIndex;
            currentLightCount++;

            // Also increase the matching light count
            if (currentLight.lightType == 0)
            {
                punctualLightCount++;
            }
            else if (currentLight.lightType == 1)
            {
                areaLightCount++;
            }
            else
            {
                envLightCount++;
            }
        }
    }

    // Set the light count for the cell
    _RaytracingLightClusterRW[cellIndex * cellDataSize] = currentLightCount;
    _RaytracingLightClusterRW[cellIndex * cellDataSize + 1] = punctualLightCount;
    _RaytracingLightClusterRW[cellIndex * cellDataSize + 2] = punctualLightCount + areaLightCount;
    _RaytracingLightClusterRW[cellIndex * cellDataSize + 3] = areaLightCount + punctualLightCount + envLightCount;
}


[numthreads(CULL_GROUP_SIZE, 1, 1)]
void RaytracingLightCull(uint2 threadID : SV_GroupThreadID, uint2 groupId : SV_GroupID)
{
    // Fetch the coordinates of the current cell
    uint targetLightIndex = groupId.x * CULL_GROUP_SIZE + threadID.x;

    // Reset the culling information of this light
    _RaytracingLightCullResult[targetLightIndex] = 0;

    // if this index is beyond the target index, it is done
    if(_LightVolumeCount <= targetLightIndex) return;

    // Fetch the target light data
    LightVolume currentLight = _LightVolumes[targetLightIndex];

    bool intersects = false;
    /*
    for(uint cIdx = 0; cIdx < 3; ++cIdx)
    {
        // Check if this corner is inside the shphere
        float3 minLightPos = currentLight.position - currentLight.shape[cIdx];
        float3 maxLightPos = currentLight.position + currentLight.shape[cIdx];
        float3 clusterMinPos = (_ClusterCenterPosition - CornerSubDirections[cIdx] * _ClusterDimension);
        float3 clusterMinPos = (_ClusterCenterPosition + CornerSubDirections[cIdx] * _ClusterDimension);
        if(dot(dir, dir) <= squareRange)
        {
            intersects = true;
            break;
        }
    }
    */
    // Flag this light as culled or visible
    _RaytracingLightCullResult[targetLightIndex] = intersects ? 0 : 1;
}