#pragma kernel TemporalAccumulationFullRes TEMPORAL_ACCUMULATION=TemporalAccumulationFullRes
#pragma kernel TemporalAccumulationHalfRes TEMPORAL_ACCUMULATION=TemporalAccumulationHalfRes HALF_RESOLUTION
#pragma kernel CopyHistory
#pragma kernel BilateralFilterH_FR      BILATERAL_FILTER=BilateralFilterH_FR
#pragma kernel BilateralFilterV_FR      BILATERAL_FILTER=BilateralFilterV_FR    FINAL_PASS
#pragma kernel BilateralFilterH_HR      BILATERAL_FILTER=BilateralFilterH_HR    HALF_RESOLUTION
#pragma kernel BilateralFilterV_HR      BILATERAL_FILTER=BilateralFilterV_HR    FINAL_PASS  HALF_RESOLUTION

#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"
#define BILATERAL_ROUGHNESS
#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/Denoising/DenoisingUtils.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/ShaderVariablesRaytracing.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/RayTracingCommon.hlsl"

// Tile size of this compute
#define REFLECTION_FILTER_TILE_SIZE 8

// #pragma enable_d3d11_debug_symbols

// Thereshold at which we decide to reject the reflection history
#define REFLECTION_HISTORY_REJECTION_THRESHOLD 0.75
// Threshold at which we go from accumulation to clampiong
#define ROUGHNESS_ACCUMULATION_THRESHOLD 0.5

// Input textures
TEXTURE2D_X(_DenoiseInputTexture);
TEXTURE2D_X(_HistoryBuffer);
TEXTURE2D_X(_HistoryDepthTexture);
// Value that tells us if the current history should be discarded based on scene-level data
float _HistoryValidity;
// Current inverse resolution of the history buffer
float2 _HistoryBufferSize;
// Resolution at which the effect is rendered (Half the _Screensize if half res)
float4 _CurrentEffectResolution;
float _PixelSpreadAngleTangent;
int _AffectSmoothSurfaces;
int _SingleReflectionBounce;
// Contains history buffer size in xy and the uv scale in zw
float4 _HistorySizeAndScale;

// Output texture
RW_TEXTURE2D_X(float4, _DenoiseOutputTextureRW);
RW_TEXTURE2D_X(float, _SampleCountTextureRW);

[numthreads(REFLECTION_FILTER_TILE_SIZE, REFLECTION_FILTER_TILE_SIZE, 1)]
void TEMPORAL_ACCUMULATION(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 mainCoord = groupId * REFLECTION_FILTER_TILE_SIZE + groupThreadId;

    // The source coordinate for the normal and depth depends if we are in half res or not
    #ifdef HALF_RESOLUTION
    uint2 geometryCoords = ComputeSourceCoordinates(mainCoord, _RayTracingCheckerIndex);
    #else
    uint2 geometryCoords = mainCoord;
    #endif

    // Fetch the depth
    float depth = LOAD_TEXTURE2D_X(_DepthTexture, geometryCoords).x;

    // If the history is flagged as invalid or this is a a background pixel write the noisy value and leave right away
    if (_HistoryValidity == 0.0 || depth == UNITY_RAW_FAR_CLIP_VALUE)
    {
        _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(mainCoord)] = LOAD_TEXTURE2D_X(_DenoiseInputTexture, mainCoord);
        _SampleCountTextureRW[COORD_TEXTURE2D_X(mainCoord)] = 0;
        return;
    }

    // Fetch the position of the current pixel
    PositionInputs posInputs = GetPositionInput(geometryCoords, _ScreenSize.zw, depth, UNITY_MATRIX_I_VP, GetWorldToViewMatrix());

    // Read the depth and velocity vectors
    float2 velocity;
    DecodeMotionVector(LOAD_TEXTURE2D_X(_CameraMotionVectorsTexture, geometryCoords), velocity);

    // Real the normal data for this pixel
    NormalData normalData;
    DecodeFromNormalBuffer(geometryCoords, normalData);

    // Compute the current and history UV coordinates
    float2 currentUV = (mainCoord + 0.5f) * _ScreenSize.zw;

    // Fetch the current and history values and apply the exposition to it.
    float3 currentColor = LOAD_TEXTURE2D_X(_DenoiseInputTexture, mainCoord).xyz * GetCurrentExposureMultiplier();
    if (_AffectSmoothSurfaces == 0 && _SingleReflectionBounce == 1 && PerceptualRoughnessToPerceptualSmoothness(normalData.perceptualRoughness) > 0.99)
    {
        _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(mainCoord)] = float4(currentColor * GetInverseCurrentExposureMultiplier(), LOAD_TEXTURE2D_X(_DenoiseInputTexture, mainCoord).w);
        _SampleCountTextureRW[COORD_TEXTURE2D_X(mainCoord)] = 1;
        return;
    }

    // Compute the pixel coordinate for the history tapping
    float2 historyUVUnscaled = posInputs.positionNDC - velocity;

    #ifdef HALF_RESOLUTION
    float2 historyTapCoord = (float2)((historyUVUnscaled * _HistorySizeAndScale.xy + 0.5) * 0.5);
    float2 historyUV = historyTapCoord / _HistorySizeAndScale.xy * _HistorySizeAndScale.zw;
    #else
    float2 historyTapCoord = (float2)(historyUVUnscaled * _HistorySizeAndScale.xy);
    float2 historyUV = historyUVUnscaled * _HistorySizeAndScale.zw;
    #endif

    float4 historyRaw = SAMPLE_TEXTURE2D_X_LOD(_HistoryBuffer, s_linear_clamp_sampler, historyUV, 0);
    float3 history = historyRaw.xyz * GetCurrentExposureMultiplier();

    // If the pixel was outside of the screen during the previous frame, invalidate the history
    bool canBeReprojected = true;
    if (historyTapCoord.x >= _CurrentEffectResolution.x || historyTapCoord.x < 0
        || historyTapCoord.y >= _CurrentEffectResolution.y || historyTapCoord.y < 0)
        canBeReprojected = false;

#ifdef HALF_RESOLUTION
    float2 historyDepthUV = ((float2)(mainCoord * 2.0 + 0.5f) - velocity * _CurrentEffectResolution.xy * 2.0f) * _HistoryBufferSize.xy;
#else
    float2 historyDepthUV = historyUV;
#endif

    // Fetch the depth of the history pixel. If the history position was a background point, invalidate the history
    float historyDepth = SAMPLE_TEXTURE2D_X_LOD(_HistoryDepthTexture, s_linear_clamp_sampler, historyDepthUV, 0).r;
    if (historyDepth == UNITY_RAW_FAR_CLIP_VALUE)
        canBeReprojected = false;

    // Compute the world space position
    PositionInputs posInput = GetPositionInput(geometryCoords, _ScreenSize.zw, depth, UNITY_MATRIX_I_VP, UNITY_MATRIX_V);

    // Compute the world space position (from previous frame)
    float3 historyPositionWS = ComputeWorldSpacePosition(posInput.positionNDC - velocity, historyDepth, UNITY_MATRIX_PREV_I_VP);

    // Compute the max world radius that we consider acceptable for history reprojection
    float maxRadius = ComputeMaxReprojectionWorldRadius(posInput.positionWS, normalData.normalWS, _PixelSpreadAngleTangent);

    // Is it too far from the current position?
    if (length(historyPositionWS - posInput.positionWS) > maxRadius)
        canBeReprojected = false;

    // Depending on the roughness of the surface run one or the other temporal reprojection
    float3 result;
    if (normalData.perceptualRoughness > ROUGHNESS_ACCUMULATION_THRESHOLD)
    {


        float sampleCount = historyRaw.w;
        if (canBeReprojected && sampleCount != 0.0)
        {
            float accumulationFactor = sampleCount >= 8.0 ? 0.93 : (sampleCount / (sampleCount + 1.0));
            result = (currentColor * (1.0 - accumulationFactor) + history * accumulationFactor);
            sampleCount = max(sampleCount + 1.0, 8.0);
        }
        else
        {
            result = currentColor;
            sampleCount = 1.0;
        }

        // Update the sample count
        historyRaw.w = sampleCount;
    }
    else
    {
        float exposureMultiplier = GetCurrentExposureMultiplier();
        float3 topLeft = LOAD_TEXTURE2D_X(_DenoiseInputTexture, clamp(mainCoord - int2(1, 1), int2(0, 0), _CurrentEffectResolution.xy)).xyz;
        float3 bottomRight = LOAD_TEXTURE2D_X(_DenoiseInputTexture, clamp(mainCoord + int2(1, 1), int2(0, 0), _CurrentEffectResolution.xy)).xyz;
        topLeft *= exposureMultiplier;
        bottomRight *= exposureMultiplier;

        float3 corners = 4.0 * (topLeft + bottomRight) - 2.0 * currentColor;

        float3 color = clamp(currentColor, 0.0, CLAMP_MAX);

        float3 average = FastTonemap((corners + color) / 7.0);

        topLeft = FastTonemap(topLeft);
        bottomRight = FastTonemap(bottomRight);
        color = FastTonemap(color);

        float colorLuma = Luminance(color);
        float averageLuma = Luminance(average);
        float velocityLength = length(velocity);
        float nudge = lerp(4.0, 0.25, saturate(velocityLength * 100.0)) * abs(averageLuma - colorLuma);

        float3 minimum = min(bottomRight, topLeft) - nudge;
        float3 maximum = max(topLeft, bottomRight) + nudge;

        history = FastTonemap(history);

        // Clip history samples
        history = DirectClipToAABB(history, minimum, maximum);

        // Blend color & history
        // Feedback weight from unbiased luminance diff (Timothy Lottes)
        float historyLuma = Luminance(history);
        float diff = abs(colorLuma - historyLuma) / Max3(colorLuma, historyLuma, 0.2);
        float weight = 1.0 - diff;
        const float feedbackMin = 0.96;
        const float feedbackMax = 0.91;
        float feedback = lerp(feedbackMin, feedbackMax, weight * weight);

        color = FastTonemapInvert(lerp(color, history, feedback));
        result = canBeReprojected ? clamp(color, 0.0, CLAMP_MAX) : currentColor;
        historyRaw.w = 0.0;
    }

    _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(mainCoord)] = float4(result * GetInverseCurrentExposureMultiplier(), LOAD_TEXTURE2D_X(_DenoiseInputTexture, mainCoord).w);
    _SampleCountTextureRW[COORD_TEXTURE2D_X(mainCoord)] = historyRaw.w;
}

[numthreads(REFLECTION_FILTER_TILE_SIZE, REFLECTION_FILTER_TILE_SIZE, 1)]
void CopyHistory(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    if (any(dispatchThreadId.xy > uint2(_ScreenSize.xy)))
        return;  // Out of bounds, discard

    float4 currentColor = LOAD_TEXTURE2D_X(_DenoiseInputTexture, dispatchThreadId.xy);
    // We need to apply a step function on the blend factor to evaluate the validity of the history (if it is stricly higher than 0.0 then its valid)
    _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(dispatchThreadId.xy)] = float4(currentColor.xyz, _SampleCountTextureRW[COORD_TEXTURE2D_X(dispatchThreadId.xy)].x);
}

int _DenoiserFilterRadius;
float _RoughnessBasedDenoising;
TEXTURE2D(_ReflectionFilterMapping);

#define BILATERAL_FILTER_SIGMA 0.9

// Separated bilateral filter (two passes, each with 2*Radius taps)
[numthreads(REFLECTION_FILTER_TILE_SIZE, REFLECTION_FILTER_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 * REFLECTION_FILTER_TILE_SIZE + groupThreadId;

    // Based on which pass of the filter we are doing, adjust the sampling direction
    #if FINAL_PASS
    const uint2 passIncr = uint2(1, 0);
    #else
    const uint2 passIncr = uint2(0, 1);
    #endif

    #if HALF_RESOLUTION
    uint2 geometrCoords = ComputeSourceCoordinates(centerCoord, _RayTracingCheckerIndex);
    #else
    uint2 geometrCoords = centerCoord;
    #endif
    // Tap the central pixel coordinates
    const BilateralData center = TapBilateralData(geometrCoords);

    // Compute the effective radius we should be using for the filtering
    const float3 viewWS = GetWorldSpaceNormalizeViewDir(center.position);
    float2 mappingUV = float2(dot(viewWS, center.normal), center.roughness);
    float2 radiusScale = lerp(1.0f, SAMPLE_TEXTURE2D_LOD(_ReflectionFilterMapping, s_linear_clamp_sampler, mappingUV, 0.0f).xy, _RoughnessBasedDenoising);
    #if FINAL_PASS
    const float radius = _DenoiserFilterRadius * radiusScale.x;
    #else
    const float radius = _DenoiserFilterRadius * radiusScale.y;
    #endif
    const float sigma =  radius * BILATERAL_FILTER_SIGMA;
    const int effectiveRadius = min(sigma * 2.0, radius);

    // Store the intermediate result
    float3 finalColor = LOAD_TEXTURE2D_X(_DenoiseInputTexture, centerCoord).xyz;

    // If this pixels does not have ray traced reflections anyway, just skip it.
    if (_RaytracingReflectionMinSmoothness <= PerceptualRoughnessToPerceptualSmoothness(center.roughness))
    {
        // Initialize variables for accumulation
        float3 colorSum = float3(0.0, 0.0, 0.0);
        float wSum = 0.0;

        int2 tapCoord = centerCoord - effectiveRadius * passIncr;
        for (int r = -effectiveRadius; r <= effectiveRadius; ++r, tapCoord += passIncr)
        {
            // Make sure the pixel coord we are trying to use is in the screen (not out of bounds)
            if (tapCoord.x >= _CurrentEffectResolution.x || tapCoord.x < 0 || tapCoord.y >= _CurrentEffectResolution.y || tapCoord.y < 0)
                continue;

            #if HALF_RESOLUTION
            uint2 tapGeometryCoords = ComputeSourceCoordinates(tapCoord, _RayTracingCheckerIndex);
            #else
            uint2 tapGeometryCoords = tapCoord;
            #endif
            // Compute the weight (skip computation for the center)
            const BilateralData tapData = TapBilateralData(tapGeometryCoords);
            float w = r ? gaussian(r, sigma) * ComputeBilateralWeight(center, tapData) : 1.0;
            w = _RaytracingReflectionMinSmoothness > PerceptualRoughnessToPerceptualSmoothness(tapData.roughness) ? 0.0 : w;
            colorSum += LOAD_TEXTURE2D_X(_DenoiseInputTexture, tapCoord).xyz * w;
            wSum += w;
        }
        // Normalize the result
        finalColor = colorSum / wSum;
    }

    _DenoiseOutputTextureRW[COORD_TEXTURE2D_X(centerCoord)] = float4(finalColor, LOAD_TEXTURE2D_X(_DenoiseInputTexture, centerCoord).w);
}