#pragma kernel RayBinning           RAY_BINNING=RayBinning
#pragma kernel RayBinningHalf       RAY_BINNING=RayBinningHalf HALF_RESOLUTION

#pragma only_renderers d3d11 xboxseries ps5

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariablesGlobal.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/NormalBuffer.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/RayTracingCommon.hlsl"

#define RAY_BINNING_TILE_SIZE 16
#define BINNING_TILE_SIZE 16

TEXTURE2D_X(_RaytracingDirectionBuffer);
RWStructuredBuffer<uint>            _RayBinResult;
RWStructuredBuffer<uint>            _RayBinSizeResult;
uint                                _RayBinTileCountX;
uint                                _RayBinViewOffset;
uint                                _RayBinTileViewOffset;

groupshared uint gs_binSize[BINNING_TILE_SIZE * BINNING_TILE_SIZE + 1];
groupshared uint gs_binOffset[BINNING_TILE_SIZE * BINNING_TILE_SIZE];

[numthreads(RAY_BINNING_TILE_SIZE, RAY_BINNING_TILE_SIZE, 1)]
void RAY_BINNING(uint2 groupThreadId : SV_GroupThreadID, uint3 dispatchThreadId : SV_DispatchThreadID, uint2 groupId : SV_GroupID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    // Compute the pixel position to process
    uint2 currentCoord = groupId * RAY_BINNING_TILE_SIZE + groupThreadId.xy;

    #if HALF_RESOLUTION
    currentCoord = ComputeSourceCoordinates(currentCoord, _RayTracingCheckerIndex);
    #endif

    // Initialize the invalid counter to 0
    if(groupThreadId.x == 0 && groupThreadId.y == 0)
    {
        gs_binSize[256] = 0;
    }

    // Every tile initializes a counter to 0
    gs_binSize[groupThreadId.y * RAY_BINNING_TILE_SIZE + groupThreadId.x] = 0;

    // Sync all threads
    GroupMemoryBarrierWithGroupSync();

    // is this sample valid
    uint binIndex = 256;

    // Is this sample valid ? otherwise its bin index is 256
    if (LOAD_TEXTURE2D_X(_RaytracingDirectionBuffer, currentCoord).w > 0.0 && (float)currentCoord.x < _ScreenSize.x && (float)currentCoord.y < _ScreenSize.y)
    {
        // Fetch the current direction
        const float3 currentDirection = LOAD_TEXTURE2D_X(_RaytracingDirectionBuffer, currentCoord).xyz;

        // Generate its bin index
        float2 octahedralDirection = PackNormalOctQuadEncode(currentDirection);
        uint2 binCoordinates = (uint2)((octahedralDirection * 0.5 + 0.5) * BINNING_TILE_SIZE);
        binIndex = binCoordinates.y * BINNING_TILE_SIZE + binCoordinates.x;
    }

    // Increment the bin size of the bin where this sample goes
    int rayBinIndex = 0;
    UNITY_BRANCH if (binIndex != 256)
    {
        InterlockedAdd(gs_binSize[binIndex], 1, rayBinIndex);
    }

    // Sync all threads
    GroupMemoryBarrierWithGroupSync();

    if(groupThreadId.x == 0 && groupThreadId.y == 0)
    {
        // Build the offset list of the bins
        gs_binOffset[0] = 0;
        for(int i = 1; i < 256; ++i)
        {
            gs_binOffset[i] = gs_binOffset[i - 1] + gs_binSize[i - 1];
        }
    }

    // Sync all threads
    GroupMemoryBarrierWithGroupSync();

    // We only want to store it if it's bin is valid
    if(binIndex < 256)
    {
        // Output the indices of the original pixels
        uint groupdIndex = groupId.y * _RayBinTileCountX + groupId.x;
        uint globalOffset = groupdIndex * RAY_BINNING_TILE_SIZE * RAY_BINNING_TILE_SIZE + gs_binOffset[binIndex] + rayBinIndex + _RayBinViewOffset * dispatchThreadId.z;
        _RayBinResult[globalOffset] = ((currentCoord.x & 0xffff) << 16) + (currentCoord.y & 0xffff);
    }

    // Then output the size of every bin
    if(groupThreadId.x == 0 && groupThreadId.y == 0)
    {
        uint groupdIndex = groupId.y * _RayBinTileCountX + groupId.x + _RayBinTileViewOffset * dispatchThreadId.z;
        _RayBinSizeResult[groupdIndex] = gs_binOffset[255] + gs_binSize[255];
    }
}