#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"

#pragma kernel TemporalDenoise

#pragma multi_compile HALF_RES FULL_RES

uint PackHistoryData(float AO, float depth, float mvLen)
{
    uint packedVal = 0;
    packedVal = BitFieldInsert(0x000000ff, UnpackInt(AO, 8), packedVal);
    packedVal = BitFieldInsert(0x0000ff00, UnpackInt(mvLen, 8) << 8, packedVal);
    packedVal = BitFieldInsert(0xffff0000, UnpackInt(depth, 16) << 16, packedVal);
    return packedVal;
}

void UnpackHistoryData(uint historyData, out float AO, out float depth, out float mvLen)
{
    AO = UnpackUIntToFloat(historyData, 0, 8);
    mvLen = UnpackUIntToFloat(historyData, 8, 8);
    depth = UnpackUIntToFloat(historyData, 16, 16);
}

RW_TEXTURE2D_X(uint, _AOOutputHistory);

TEXTURE2D_X(_AOPackedBlurred);
TEXTURE2D_X_UINT(_AOPackedHistory);

#if HALF_RES
RW_TEXTURE2D_X(float, _OcclusionTexture);
#else
RW_TEXTURE2D_X(float, _OcclusionTexture);
#endif

float3 FindMinMaxAvgAO(float2 centralPos)
{
    float minAO = 2.0f;
    float maxAO = -2.0f;
    float avg = 0;
    for (int i = -1; i <= 1; ++i)
    {
        for (int j = -1; j <= 1; ++j)
        {
            float data = LOAD_TEXTURE2D_X(_AOPackedBlurred, (uint2)centralPos + uint2(i, j)).x;
            float currAO, ignored;
            UnpackData(data, currAO, ignored);
            avg += currAO;
            minAO = min(minAO, currAO);
            maxAO = max(maxAO, currAO);
        }
    }

    return float3(minAO, maxAO, avg/9);
}

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

    float currFrameData = _AOPackedBlurred[COORD_TEXTURE2D_X(dispatchThreadId.xy)].x;
    float currDepth, currAO;
    UnpackData(currFrameData, currAO, currDepth);

    PositionInputs posInputs = GetPositionInput(float2(dispatchThreadId.xy), _ScreenSize.zw, uint2(8, 8));

#if HALF_RES
    float2 closest = posInputs.positionSS * 2;
#else
    float2 closest = posInputs.positionSS;
#endif
    float2 motionVector;
    DecodeMotionVector(LOAD_TEXTURE2D_X(_CameraMotionVectorsTexture, closest), motionVector);
    float motionVecLength = length(motionVector);

    float2 uv = (dispatchThreadId.xy + 0.5) * _AOBufferSize.zw;
    float2 prevFrameNDC = uv - motionVector;

    uint prevData = asuint(_AOPackedHistory[COORD_TEXTURE2D_X((prevFrameNDC) * _AOHistorySize.xy)].x);
    float prevMotionVecLen, prevAO, prevDepth;
    UnpackHistoryData(prevData, prevAO, prevDepth, prevMotionVecLen);

    float motionVecWeighting = saturate(motionVecLength * 100.0);
    float velWeight = 1.0f - saturate((abs(prevMotionVecLen - motionVecWeighting)) * 3.0f);

    float3 minMax = FindMinMaxAvgAO(float2(dispatchThreadId.xy));
    float minAO = minMax.x;
    float maxAO = minMax.y;
    float avg = minMax.z;
    float nudge = lerp(_AOTemporalUpperNudgeLimit, _AOTemporalLowerNudgeLimit, motionVecWeighting) * abs(avg - currAO);
    minAO -= nudge;
    maxAO += nudge;

    float diff = abs(currAO - prevAO) / Max3(prevAO, currAO, 0.1f);
    float weight = 1.0 - diff;
    float feedback = lerp(0.85, 0.95, weight * weight);

    prevAO = clamp(prevAO, minAO, maxAO);

    float depth_similarity = saturate(pow(prevDepth / currDepth, 1) + 0.01);

    float newAO = (lerp(currAO, prevAO, feedback * depth_similarity * velWeight));

    _AOOutputHistory[COORD_TEXTURE2D_X(dispatchThreadId.xy)] = PackHistoryData(newAO, currDepth, motionVecWeighting);

#ifdef HALF_RES
    // The conversion and the pow happens in upsample phase.
    _OcclusionTexture[COORD_TEXTURE2D_X(dispatchThreadId.xy)] = PackAOOutput(newAO, currDepth);
#else
    _OcclusionTexture[COORD_TEXTURE2D_X(dispatchThreadId.xy)] = OutputFinalAO(newAO);
#endif

}