#pragma only_renderers d3d11 xboxseries ps5

#pragma kernel BilateralFilterHSingleDirectional        BILATERAL_FILTER=BilateralFilterHSingleDirectional      SINGLE_CHANNEL DIRECTIONAL_LIGHT
#pragma kernel BilateralFilterVSingleDirectional        BILATERAL_FILTER=BilateralFilterVSingleDirectional      FINAL_PASS SINGLE_CHANNEL DIRECTIONAL_LIGHT

#pragma kernel BilateralFilterHColorDirectional         BILATERAL_FILTER=BilateralFilterHColorDirectional       DIRECTIONAL_LIGHT
#pragma kernel BilateralFilterVColorDirectional         BILATERAL_FILTER=BilateralFilterVColorDirectional       FINAL_PASS DIRECTIONAL_LIGHT

#pragma kernel BilateralFilterHSinglePoint              BILATERAL_FILTER=BilateralFilterHSinglePoint            SINGLE_CHANNEL POINT_LIGHT
#pragma kernel BilateralFilterVSinglePoint              BILATERAL_FILTER=BilateralFilterVSinglePoint            FINAL_PASS SINGLE_CHANNEL POINT_LIGHT

#pragma kernel BilateralFilterHSingleSpot               BILATERAL_FILTER=BilateralFilterHSingleSpot             SINGLE_CHANNEL SPOT_LIGHT
#pragma kernel BilateralFilterVSingleSpot               BILATERAL_FILTER=BilateralFilterVSingleSpot             FINAL_PASS SINGLE_CHANNEL SPOT_LIGHT

#pragma kernel BilateralFilterHSinglePyramid            BILATERAL_FILTER=BilateralFilterHSinglePyramid          SINGLE_CHANNEL PYRAMID_LIGHT
#pragma kernel BilateralFilterVSinglePyramid            BILATERAL_FILTER=BilateralFilterVSinglePyramid          FINAL_PASS SINGLE_CHANNEL PYRAMID_LIGHT

#pragma kernel BilateralFilterHSingleBox                BILATERAL_FILTER=BilateralFilterHSingleBox              SINGLE_CHANNEL BOX_LIGHT
#pragma kernel BilateralFilterVSingleBox                BILATERAL_FILTER=BilateralFilterVSingleBox              FINAL_PASS SINGLE_CHANNEL BOX_LIGHT

#pragma multi_compile _ DISTANCE_BASED_DENOISER

#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"
#define BILATERLAL_UNLIT
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/Denoising/BilateralFilter.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/Shadows/RayTracingShadowUtilities.hlsl"

#pragma only_renderers d3d11 xboxseries ps5

// #pragma enable_d3d11_debug_symbols

// Tile size of this compute
#define SHADOW_DENOISER_TILE_SIZE 8

// Ray tracing input textures
TEXTURE2D_X(_DenoiseInputTexture);
TEXTURE2D_X(_DistanceTexture);

// Generic denoiser inputs
int _DenoiserFilterRadius;

// Camera FOV
float _CameraFOV;

// Inputs for sphere lights (point and spot)
float3  _SphereLightPosition;

// Output texture
#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(SHADOW_DENOISER_TILE_SIZE, SHADOW_DENOISER_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
    int2 centerCoord = groupId * SHADOW_DENOISER_TILE_SIZE + groupThreadId;

    // Read the center pixel0
    const BilateralData center = TapBilateralData(centerCoord);

#if SPOT_LIGHT || PYRAMID_LIGHT || BOX_LIGHT || POINT_LIGHT
    // Fetch the data of the area light
    LightData lightData = _LightDatas[_RaytracingTargetLight];
#endif

#if defined(DISTANCE_BASED_DENOISER)
#if DIRECTIONAL_LIGHT
    // For the directonal light, the solid angle can be used directly
    float lightSolidAngle = _RaytracingLightAngle;
#else
    // Compute the light distance 
#if BOX_LIGHT
    float lightPointDistance = dot(_SphereLightPosition - center.position, normalize(-lightData.forward));
#else
    float lightPointDistance = length(_SphereLightPosition - center.position);
#endif
    // Evaluate the normalized solid angle of the light
    float lightSolidAngle = atan(_RaytracingLightRadius / lightPointDistance);
#endif
    // Compute the distances we need for our filtering
    const float distanceCameraToPlane = length(GetWorldSpaceViewDir(center.position));
    const float distancePlaneToObject = LOAD_TEXTURE2D_X(_DistanceTexture, centerCoord).x;

    // Compute the cone footprint on the image reflection plane for this configuration
    const float brdfConeRadius = tan(lightSolidAngle * 0.5) * distancePlaneToObject * 2.0f;

    // We need to compute the view cone radius
    const float viewConeRadius = brdfConeRadius * distanceCameraToPlane / (distancePlaneToObject + distanceCameraToPlane);

    // Compute the view cone's half angle. This matches the FOV angle to see exactly the half of the cone (The tangent could be precomputed in the table)
    const float viewConeHalfAngle = FastATanPos(viewConeRadius / distanceCameraToPlane);
    // Given the camera's fov and pixel resolution convert the viewConeHalfAngle to a number of pixels
    const float pixelDistance = viewConeHalfAngle / _CameraFOV * _ScreenSize.x;

    // Evaluate the radius that should be used for the filter
    #if FINAL_PASS
        const float radius = clamp(pixelDistance, 1, _DenoiserFilterRadius);
    #else
        const float radius = clamp(pixelDistance, 1, _DenoiserFilterRadius);
    #endif
    // Compute the sigma value for our filter
    const float sigma = 0.5 * radius;
#else
    const float sigma = 0.5 * _DenoiserFilterRadius;
    const int radius = min(sigma * 2.0, _DenoiserFilterRadius);
#endif

    // Initialize the accumulation variables
    #if SINGLE_CHANNEL
        float colorSum = 0.0;
    #else
        float3 colorSum = float3(0.0, 0.0, 0.0);
    #endif
    float wSum = 0.0;

    // Based on if we are horizontal of vertical, define our increment direction
    #if FINAL_PASS
        const int2 passIncr = int2(1, 0);
    #else
        const int2 passIncr = int2(0, 1);
    #endif


    // Evaluate the filter
    int2 tapCoord = centerCoord - radius * passIncr;
    for (int r = -radius; r <= radius; ++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 (tapCoord.x >= _ScreenSize.x || tapCoord.y >= _ScreenSize.y || tapCoord.x < 0 || tapCoord.y < 0)
            continue;

        // Grab the data of the tap pixel
        const BilateralData tapData = TapBilateralData(tapCoord);

        // Compute the weight (skip computation for the center)
        float w = r ? gaussian(r, sigma) * ComputeBilateralWeight(center, tapData) : 1.0;
        w = tapData.isUnlit ? 0.0 : w;

    #if POINT_LIGHT
        float dist2 = DistSqrToLight(lightData, tapData.position);
        w = PositionInPointRange(lightData, dist2) ? w : 0.0;
    #endif

    #if SPOT_LIGHT
        float dist2 = DistSqrToLight(lightData, tapData.position);
        w = PositionInSpotRange(lightData, _RaytracingLightAngle, tapData.position, dist2) ? w : 0.0;
    #endif

    #if PYRAMID_LIGHT
        float dist2 = DistSqrToLight(lightData, tapData.position);
        w = PositionInPyramidRange(lightData, _RaytracingLightSizeX, _RaytracingLightSizeY, tapData.position, dist2) ? w : 0.0;
    #endif

    #if BOX_LIGHT
        float dist = dot(lightData.positionRWS - tapData.position, normalize(-lightData.forward));
        w = PositionInBoxRange(lightData, _RaytracingLightSizeX, _RaytracingLightSizeY, tapData.position, dist * dist) ? w : 0.0;
    #endif

        // Accumulate the color
        #if SINGLE_CHANNEL
            colorSum += LOAD_TEXTURE2D_X(_DenoiseInputTexture, tapCoord).x * w;
        #else
            colorSum += LOAD_TEXTURE2D_X(_DenoiseInputTexture, tapCoord).xyz * w;
        #endif

        // Accumulate its weight
        wSum += w;
    }

    // Normalize and store the value
    #if SINGLE_CHANNEL
        float finalColor = wSum != 0.0 ? colorSum / wSum : 0.0;
        _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(centerCoord)] = finalColor;
    #else
        float3 finalColor = wSum != 0.0 ? colorSum / wSum : 0.0;
        _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(centerCoord)] = float4(finalColor, 1.0);
    #endif
}