#pragma kernel PreBlur

#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/Material/NormalBuffer.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/Denoising/ReBlur/ReBlur_BlurUtilities.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/ReBlurDenoiser.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/RayTracingCommon.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/ShaderVariablesRaytracing.cs.hlsl"

// Bilateral filtering
#define BILATERAL_ROUGHNESS
#define BILATERLAL_UNLIT
#define BILATERLAL_SSR
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/Denoising/BilateralFilter.hlsl"

#define REBLUR_PREPASS_TILE_SIZE 8

#pragma multi_compile _ HALF_RESOLUTION

//#pragma enable_d3d11_debug_symbols

// Constants
#define PREPASS_BLUR_RADIUS 20.0

// #define WITHOUT_LDS

// Input texture
TEXTURE2D_X(_LightingInputTexture);
TEXTURE2D_X(_DistanceInputTexture);

// Output texture
RW_TEXTURE2D_X(float4, _LightingDistanceTextureRW);

#ifdef HALF_RESOLUTION
groupshared float3 gs_cacheLighting[36];
groupshared float gs_cacheDistance[36];
groupshared float gs_cacheDepth[36];
groupshared bool gs_cacheValidity[36];

void FillRegionLDS(uint groupIndex, uint2 groupOrigin)
{
    // Define which value we will be acessing with this worker thread
    int acessCoordX = groupIndex % 6;
    int acessCoordY = groupIndex / 6;

    // The initial position of the access
    int2 originXY = groupOrigin / 2 - int2(1, 1);

    // Compute the sample position
    int2 sampleCoord = int2(clamp(originXY.x + acessCoordX, 0, _ScreenSize.x / 2 - 1), clamp(originXY.y + acessCoordY, 0, _ScreenSize.y / 2 - 1));

    // The representative coordinate to use depends if we are using the checkerboard integration pattern (or not)
    int2 representativeCoord = ComputeSourceCoordinates(sampleCoord, _RayTracingCheckerIndex);

    // Read the depth value and linearize it
    float sampleDepth = LOAD_TEXTURE2D_X(_DepthTexture, representativeCoord).x;
    gs_cacheDepth[groupIndex] = Linear01Depth(sampleDepth, _ZBufferParams);

    // Evaluate if the the pixel's data can be used (if it is not unlit, has receive SSR and has the right smoothness value)
    uint stencilValue = GetStencilValue(LOAD_TEXTURE2D_X(_StencilTexture, representativeCoord));
    bool validity = (stencilValue & STENCILUSAGE_IS_UNLIT) == 0;
    validity = validity && (stencilValue & STENCILUSAGE_TRACE_REFLECTION_RAY) != 0;
    validity = validity && (sampleDepth != UNITY_RAW_FAR_CLIP_VALUE);
    gs_cacheValidity[groupIndex] = validity;

    // Read the lighting data
    float3 sampleLighting = LOAD_TEXTURE2D_X(_LightingInputTexture, sampleCoord).xyz;
    gs_cacheLighting[groupIndex] = sampleLighting;

    // Read the distance data
    float sampleDistance = LOAD_TEXTURE2D_X(_DistanceInputTexture, sampleCoord).x;
    gs_cacheDistance[groupIndex] = sampleDistance;
}

uint OffsetToLDSAdress(uint2 groupThreadId, int2 offset)
{
    // Compute the tap coordinate in the 6x6 grid
    uint2 tapAddress = (uint2)((int2)(groupThreadId / 2 + 1) + offset);
    return clamp((uint)(tapAddress.x) % 6 + tapAddress.y * 6, 0, 35);
}
#endif

[numthreads(REBLUR_PREPASS_TILE_SIZE, REBLUR_PREPASS_TILE_SIZE, 1)]
void PreBlur(uint3 dispatchThreadId : SV_DispatchThreadID,
                uint2 groupThreadId : SV_GroupThreadID,
                int groupIndex : SV_GroupIndex,
                uint2 groupId : SV_GroupID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);
    uint2 currentCoord = dispatchThreadId.xy;

#if defined(HALF_RESOLUTION) && !defined(WITHOUT_LDS)
    // Only 36 workers of the 64 do the pre-fetching
    if (groupIndex < 36)
    {
        // Load 1 value per thread
        FillRegionLDS(groupIndex, groupId * 8);
    }
    // Make sure all values are loaded in LDS by now.
    GroupMemoryBarrierWithGroupSync();
#endif

    // Read the central position data
    const BilateralData center = TapBilateralData(currentCoord);

    // If this is a background pixel or an unlit one, we are done
    if (center.z01 == 1.0 || center.isUnlit || !center.hasSSR)
    {
        _LightingDistanceTextureRW[COORD_TEXTURE2D_X(currentCoord)] = float4(0.0, 0.0, 0.0, -1);
        return;
    }

    // Center distance
    #if defined(HALF_RESOLUTION)
        #if !defined(WITHOUT_LDS)
            int halfResLDSIndex = OffsetToLDSAdress(groupThreadId, int2(0, 0));
            float centerDistance = gs_cacheDistance[halfResLDSIndex];
        #else
            float centerDistance = LOAD_TEXTURE2D_X(_DistanceInputTexture, currentCoord / 2).x;
        #endif
    #else
    float centerDistance = LOAD_TEXTURE2D_X(_DistanceInputTexture, currentCoord).x;
    #endif

    // Evaluate the blur radius
    float blurRadius = ComputeBlurRadius(center.roughness, PREPASS_BLUR_RADIUS) * _ReBlurDenoiserRadius;
    blurRadius *= HitDistanceAttenuation(center.roughness, length(center.position), centerDistance);

    // Evaluate the UV coordinates of the pixel
    float2 currentUV = (currentCoord + 0.5) * _ScreenSize.zw;

    // Loop through the samples
    float4 signalSum = 0.0;
    float sumWeight = 0.0;
    for (uint sampleIndex = 0; sampleIndex < POISSON_SAMPLE_COUNT; ++sampleIndex)
    {
        // Pick the next sample value
        float3 offset = k_PoissonDiskSamples[sampleIndex];

        // Evaluate the tap uv
        float2 uv = currentUV + RotateVector(_ReBlurPreBlurRotator, offset.xy) * _ScreenSize.zw * blurRadius;

        // Evaluate the tap coordinates
        float2 tapCoord = uv * _ScreenSize.xy;

    // If we are working in half resolution, we need to pick which up-resolution pixel was used to generate the current frame data.
    #if defined(HALF_RESOLUTION)
        tapCoord = ComputeSourceCoordinates(tapCoord / 2, _RayTracingCheckerIndex);
    #endif

        // Fetch the corresponding data
        const BilateralData tapData = TapBilateralData(tapCoord);

        // Sample weights
        float w = GetGaussianWeight(offset.z);
        w *= ComputeBilateralWeight(center, tapData);
        w *= tapData.z01 != 1.0 ? 1.0 : 0.0;
        w = UVInScreen(uv) ? w : 0.0;
        w = tapData.hasSSR ? w : 0.0;

    #if defined(HALF_RESOLUTION)
        // Fetch the full resolution depth
        float3 lighting = LOAD_TEXTURE2D_X(_LightingInputTexture, tapCoord / 2).xyz;
        float distance = LOAD_TEXTURE2D_X(_DistanceInputTexture, tapCoord / 2).x;
    #else
        // Fetch the full resolution depth
        float3 lighting = LOAD_TEXTURE2D_X(_LightingInputTexture, tapCoord).xyz;
        float distance = LOAD_TEXTURE2D_X(_DistanceInputTexture, tapCoord).x;
    #endif

        // Accumulate
        signalSum += float4(lighting, distance) * w;
        sumWeight += w;
    }

    // If none of the samples picked are usable, use the central pixel
    float4 finalSignal;
    if (sumWeight != 0.0)
    {
        finalSignal = signalSum / sumWeight;
    }
    else
    {
        // If we are 
        #if defined(HALF_RESOLUTION)
        float4 fallbackLighting = 0.0;
        float weightSum = 0.0;
        for (int y = -1; y <= 1; ++y)
        {
            for (int x = -1; x <= 1; ++x)
            {
                // Evaluate the initial weight
                float radius = sqrt(x * x + y * y);
                float weight = gaussian(radius, 0.9);

                #if !defined(WITHOUT_LDS)
                    // Compute the LDS index of this half res tap
                    int tapHalfResLDSIndex = OffsetToLDSAdress(groupThreadId, int2(x, y));
                    float tapDepth = gs_cacheDepth[tapHalfResLDSIndex];
                #else
                    int2 halfResCoord = currentCoord / 2 + int2(x,y);
                    int2 representativeCoord = ComputeSourceCoordinates(halfResCoord, _RayTracingCheckerIndex);
                    float tapDepth = Linear01Depth(LOAD_TEXTURE2D_X(_DepthTexture, representativeCoord).x, _ZBufferParams);
                #endif

                // Make sure the depth is not a background depth and is similar to the final one
                weight *= ((tapDepth != 1.0) ? 1.0 : 0.0);
                weight *= (((abs(tapDepth - center.z01) < center.z01 * 0.1)) ? 1.0 : 0.0);

                // Make sure the pixel has a valid signal
                #if !defined(WITHOUT_LDS)
                    weight *= gs_cacheValidity[tapHalfResLDSIndex] ? 1.0 : 0.0;
                #else
                    uint stencilValue = GetStencilValue(LOAD_TEXTURE2D_X(_StencilTexture, representativeCoord));
                    bool validity = (stencilValue & STENCILUSAGE_IS_UNLIT) == 0;
                    validity = validity && ((stencilValue & STENCILUSAGE_TRACE_REFLECTION_RAY) != 0);
                    weight *= (validity ? 1.0 : 0.0);
                #endif

                // Add the contribution of the sample
                #if !defined(WITHOUT_LDS)
                    fallbackLighting += float4(gs_cacheLighting[tapHalfResLDSIndex], gs_cacheDistance[tapHalfResLDSIndex]) * weight;
                #else
                    fallbackLighting += float4(LOAD_TEXTURE2D_X(_LightingInputTexture, halfResCoord).xyz, LOAD_TEXTURE2D_X(_DistanceInputTexture, halfResCoord).x) * weight;
                #endif
                weightSum += weight; 
            }
        }

        // After all of this, it is still possible that we're not able to find a candidate to upscale, in that case, there is nothing we can do, so the signal is going to be black.
        finalSignal = weightSum == 0.0 ? float4(0, 0, 0, 0) : fallbackLighting / weightSum;
        #else
        finalSignal = float4(LOAD_TEXTURE2D_X(_LightingInputTexture, currentCoord).xyz, centerDistance);
        #endif
    }

    // Normalize the result
    _LightingDistanceTextureRW[COORD_TEXTURE2D_X(currentCoord)] = finalSignal;
}