#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.core/ShaderLibrary/PhysicalCamera.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"


#pragma kernel KHistogramGen
#define GROUP_SIZE_X 16
#define GROUP_SIZE_Y 16

#define HISTOGRAM_BINS 256

// RGB: Histogram for RGB separately A: Histogram for luminance
RWStructuredBuffer<uint4> _HistogramBuffer;
TEXTURE2D_X(_SourceTexture);


uint4 GetBinsLoc(float3 rgbValue)
{
    float3 srgbVal = LinearToSRGB(rgbValue);

    return uint4(saturate(float4(srgbVal, Luminance(srgbVal))) * (HISTOGRAM_BINS - 1));
}

groupshared uint4 gs_localHistogram[HISTOGRAM_BINS];

[numthreads(GROUP_SIZE_X, GROUP_SIZE_Y, 1)]
void KHistogramGen(uint groupIndex : SV_GroupIndex,
    uint3 dispatchThreadId : SV_DispatchThreadID)
{
    if (groupIndex < HISTOGRAM_BINS)
    {
        gs_localHistogram[groupIndex] = uint4(0u, 0u, 0u, 0u);
    }

    GroupMemoryBarrierWithGroupSync();

    uint2 fullResCoords = dispatchThreadId.xy << 1u;

    if (all(fullResCoords < uint2(_ScreenSize.xy)))
    {
        float2 uv = ClampAndScaleUVForBilinear((fullResCoords + 0.5) * _ScreenSize.zw);
        float3 rgbVal = SAMPLE_TEXTURE2D_X_LOD(_SourceTexture, s_linear_clamp_sampler, uv, 0.0).xyz;

        uint4 bins = GetBinsLoc(rgbVal);

        InterlockedAdd(gs_localHistogram[bins.x].x, 1u);
        InterlockedAdd(gs_localHistogram[bins.y].y, 1u);
        InterlockedAdd(gs_localHistogram[bins.z].z, 1u);
        InterlockedAdd(gs_localHistogram[bins.w].w, 1u);
    }

    GroupMemoryBarrierWithGroupSync();

    if (groupIndex < HISTOGRAM_BINS)
    {
        InterlockedAdd(_HistogramBuffer[groupIndex].x, gs_localHistogram[groupIndex].x);
        InterlockedAdd(_HistogramBuffer[groupIndex].y, gs_localHistogram[groupIndex].y);
        InterlockedAdd(_HistogramBuffer[groupIndex].z, gs_localHistogram[groupIndex].z);
        InterlockedAdd(_HistogramBuffer[groupIndex].w, gs_localHistogram[groupIndex].w);
    }
}