#pragma only_renderers d3d11 xboxseries ps5

// Temporal Filtering kernels
#pragma kernel BilateralFilterHSingle       BILATERAL_FILTER=BilateralFilterHSingle     SINGLE_CHANNEL
#pragma kernel BilateralFilterVSingle       BILATERAL_FILTER=BilateralFilterVSingle     FINAL_PASS SINGLE_CHANNEL

#pragma kernel BilateralFilterHColor        BILATERAL_FILTER=BilateralFilterHColor
#pragma kernel BilateralFilterVColor        BILATERAL_FILTER=BilateralFilterVColor     FINAL_PASS


#pragma only_renderers d3d11 xboxseries ps5

#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/Material/NormalBuffer.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/RaytracingSampling.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/Builtin/BuiltinData.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/PostProcessing/Shaders/TemporalAntialiasing.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/Denoising/BilateralFilter.hlsl"

// Tile size of this compute
#define SPATIO_TEMPORAL_TILE_SIZE 8

TEXTURE2D_X(_DenoiseInputTexture);
int _DenoiserFilterRadius;
#if SINGLE_CHANNEL
    RW_TEXTURE2D_X(float, _DenoiseOutputTextureRW);
#else
    RW_TEXTURE2D_X(float4, _DenoiseOutputTextureRW);
#endif

// Separated bilateral filter (two passes, each with 2*Radius taps)
[numthreads(SPATIO_TEMPORAL_TILE_SIZE, SPATIO_TEMPORAL_TILE_SIZE, 1)]
void BILATERAL_FILTER(uint3 dispatchThreadId : SV_DispatchThreadID, uint2 groupThreadId : SV_GroupThreadID, uint2 groupId : SV_GroupID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    // Fetch the current pixel coordinate
    uint2 centerCoord = groupId * SPATIO_TEMPORAL_TILE_SIZE + groupThreadId;

    #if SINGLE_CHANNEL
    float colorSum = 0.0;
    #else
    float3 colorSum = float3(0.0, 0.0, 0.0);
    #endif

    float wSum = 0.0;

    #if FINAL_PASS
    const uint2 passIncr = uint2(1, 0);
    #else
    const uint2 passIncr = uint2(0, 1);
    #endif

    const float sigma = 0.5 * _DenoiserFilterRadius;
    const int effectiveRadius = min(sigma * 2.0, _DenoiserFilterRadius);

    const BilateralData center = TapBilateralData(centerCoord);

    uint2 tapCoord = centerCoord - effectiveRadius * passIncr;
    for (int r = -effectiveRadius; r <= effectiveRadius; ++r, tapCoord += passIncr)
    {
        // We should not tap outside of the screen (given that its a unit, if we go below zero we wrap around)
        if ((int)tapCoord.x > _ScreenSize.x || (int)tapCoord.y > _ScreenSize.y)
            continue;

        // Compute the weight (skip computation for the center)
        const float w = r ? gaussian(r, sigma) * ComputeBilateralWeight(center, TapBilateralData(tapCoord)) : 1.0;

        #if SINGLE_CHANNEL
        colorSum += LOAD_TEXTURE2D_X(_DenoiseInputTexture, tapCoord).x * w;
        #else
        colorSum += LOAD_TEXTURE2D_X(_DenoiseInputTexture, tapCoord).xyz * w;
        #endif
        wSum += w;
    }

    // Store the intermediate result
    #if SINGLE_CHANNEL
    float finalColor = colorSum / wSum;
    _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(centerCoord)] = finalColor;
    #else
    float3 finalColor = colorSum / wSum;
    _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(centerCoord)] = float4(finalColor, 1.0);
    #endif
}