#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/ScreenSpaceLighting/GTAOCommon.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/NormalBuffer.hlsl"

// TODO: This pass really could really use some quality improvement.

#pragma kernel SpatialDenoise

#pragma multi_compile _ TO_TEMPORAL

TEXTURE2D_X(_AOPackedData);

#ifdef TO_TEMPORAL
RW_TEXTURE2D_X(float, _AOPackedBlurred);
#else
RW_TEXTURE2D_X(float, _OcclusionTexture);
#endif

#define DEBUG_VISUALIZE_BILATERAL_WEIGHTS 0

#define BILATERAL_EPSILON 0.01

float BilateralWeight(float sampleDepth, float linearCentralDepth)
{
    float linearSample = LinearEyeDepth(sampleDepth, _ZBufferParams);
    float delta = abs(linearSample - linearCentralDepth);
    float w = saturate(1.0f - (_AOSpatialBilateralAggressiveness * delta + BILATERAL_EPSILON));

    return w;
}

float Blur(float2 centralPos, out float centralDepth, out float centralAO)
{
    float4 UnpackedAOs, UnpackedDepths;

    float2 UV = ClampAndScaleUVForBilinear((centralPos + float2(0.0, 0.0)) * _AOBufferSize.zw, _AOBufferSize.zw);
    GatherAOData(_AOPackedData, UV, UnpackedAOs, UnpackedDepths);

    centralDepth = UnpackedDepths.y;
    float linearCentralDepth = LinearEyeDepth(centralDepth, _ZBufferParams);
    centralAO = UnpackedAOs.y;
    float total = UnpackedAOs.y;
    float totalWeight = 1;

    // This manual unrolling is horrible looking, but I found it hard to please the PS4 compiler otherwise. TODO: Make this nicer.

    // First set of gathered data.
    float weight = BilateralWeight(UnpackedDepths.x, linearCentralDepth);
    total += weight * UnpackedAOs.x;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.z, linearCentralDepth);
    total += weight * UnpackedAOs.z;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.w, linearCentralDepth);
    total += weight * UnpackedAOs.w;
    totalWeight += weight;

    // Second set of gathered data.
    UV = ClampAndScaleUVForBilinear((centralPos + float2(2.0, 0.0)) * _AOBufferSize.zw, _AOBufferSize.zw);
    GatherAOData(_AOPackedData, UV, UnpackedAOs, UnpackedDepths);

    weight = BilateralWeight(UnpackedDepths.x, linearCentralDepth);
    total += weight * UnpackedAOs.x;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.y, linearCentralDepth);
    total += weight * UnpackedAOs.y;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.z, linearCentralDepth);
    total += weight * UnpackedAOs.z;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.w, linearCentralDepth);
    total += weight * UnpackedAOs.w;
    totalWeight += weight;


    // Third set of gathered data.
    UV = ClampAndScaleUVForBilinear((centralPos + float2(0.0, 2.0)) * _AOBufferSize.zw, _AOBufferSize.zw);
    GatherAOData(_AOPackedData, UV, UnpackedAOs, UnpackedDepths);

    weight = BilateralWeight(UnpackedDepths.x, linearCentralDepth);
    total += weight * UnpackedAOs.x;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.y, linearCentralDepth);
    total += weight * UnpackedAOs.y;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.z, linearCentralDepth);
    total += weight * UnpackedAOs.z;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.w, linearCentralDepth);
    total += weight * UnpackedAOs.w;
    totalWeight += weight;

    // Fourth set of gathered data.
    UV = ClampAndScaleUVForBilinear((centralPos + float2(2.0, 2.0)) * _AOBufferSize.zw, _AOBufferSize.zw);
    GatherAOData(_AOPackedData, UV, UnpackedAOs, UnpackedDepths);

    weight = BilateralWeight(UnpackedDepths.x, linearCentralDepth);
    total += weight * UnpackedAOs.x;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.y, linearCentralDepth);
    total += weight * UnpackedAOs.y;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.z, linearCentralDepth);
    total += weight * UnpackedAOs.z;
    totalWeight += weight;

    weight = BilateralWeight(UnpackedDepths.w, linearCentralDepth);
    total += weight * UnpackedAOs.w;
    totalWeight += weight;


    total /= totalWeight;

#if DEBUG_VISUALIZE_BILATERAL_WEIGHTS
    return totalWeight / 16.0f;
#endif
    return total;
}

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

    float centralDepth, centralAO;
    float final = Blur(dispatchThreadId.xy, centralDepth, centralAO);

#if TO_TEMPORAL
    _AOPackedBlurred[COORD_TEXTURE2D_X(dispatchThreadId.xy)] = PackAOOutput(final, centralDepth);
#else
    _OcclusionTexture[COORD_TEXTURE2D_X(dispatchThreadId.xy)] = OutputFinalAO(final);
#endif
}