#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/Builtin/BuiltinData.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"

#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#pragma multi_compile _ ENABLE_MAX_BLENDING

#pragma kernel KMain

TEXTURE2D_X(_InputCoCTexture);
TEXTURE2D_X(_InputHistoryCoCTexture);

RW_TEXTURE2D_X(float, _OutputCoCTexture);

#ifdef ENABLE_MAX_BLENDING
#if !defined(UNITY_CORE_SAMPLERS_INCLUDED)
    SAMPLER(sampler_PointClamp);
#endif
    #define SAMPLER_STATE sampler_PointClamp
#else
#if !defined(UNITY_CORE_SAMPLERS_INCLUDED)
    SAMPLER(sampler_LinearClamp);
#endif
    #define SAMPLER_STATE sampler_LinearClamp
#endif

CBUFFER_START(cb0)
    float4 _Params;
    float4 _Params2;
CBUFFER_END

#define MotionBlending              _Params.x
#define ScreenToTargetScaleHistory  _Params.yz

#define GROUP_SIZE 8

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

    PositionInputs posInputs = GetPositionInput(float2(dispatchThreadId.xy), _PostProcessScreenSize.zw, uint2(GROUP_SIZE, GROUP_SIZE));
    float2 uv = posInputs.positionNDC;

#if 0
    // Using Gather4 here doesn't show any performance gain (profiled on Nvidia) compared to doing
    // four loads so it's disabled for the sake of maximum compatibility
    float4 cocTL = GATHER_TEXTURE2D_X(_InputCoCTexture, sampler_LinearClamp, uv * _RTHandlePostProcessScale.xy - _PostProcessScreenSize.zw * 0.5); // Top-Left
    float4 cocBR = GATHER_TEXTURE2D_X(_InputCoCTexture, sampler_LinearClamp, uv * _RTHandlePostProcessScale.xy + _PostProcessScreenSize.zw * 0.5); // Bottom-Right
    float coc1 = cocTL.x;
    float coc2 = cocTL.z;
    float coc3 = cocBR.x;
    float coc4 = cocBR.z;
#else
    float2 coord = ClampAndScaleUVPostProcessTextureForPoint(posInputs.positionSS);
    float coc1 = LOAD_TEXTURE2D_X(_InputCoCTexture, coord - uint2(1u, 0u)).x; // Left
    float coc2 = LOAD_TEXTURE2D_X(_InputCoCTexture, coord - uint2(0u, 1u)).x; // Top
    float coc3 = LOAD_TEXTURE2D_X(_InputCoCTexture, coord + uint2(0u, 1u)).x; // Bottom
    float coc4 = LOAD_TEXTURE2D_X(_InputCoCTexture, coord + uint2(1u, 0u)).x; // Right
#endif

    // Dejittered center sample
    float coc0 = SAMPLE_TEXTURE2D_X_LOD(_InputCoCTexture, SAMPLER_STATE, ClampAndScaleUVForBilinearPostProcessTexture(uv - _TaaJitterStrength.zw) , 0.0).x;

    // CoC dilation: determine the closest point in the four neighbors
    float3 closest = float3(0.0, 0.0, coc0);
    closest = abs(coc1) < abs(closest.z) ? float3(-1.0,  0.0, coc1) : closest;
    closest = abs(coc2) < abs(closest.z) ? float3( 0.0, -1.0, coc2) : closest;
    closest = abs(coc3) < abs(closest.z) ? float3( 0.0,  1.0, coc3) : closest;
    closest = abs(coc4) < abs(closest.z) ? float3( 1.0,  0.0, coc4) : closest;

    // Sample the history buffer with the motion vector at the closest point
    float2 motionVector;
    float2 motionVecUV = FromOutputPosSSToPreupsampleUV(max(int2(posInputs.positionSS) + int2(closest.xy), 0));
    motionVecUV = ClampAndScaleUVForPoint(motionVecUV);

    float4 encodedMV = SAMPLE_TEXTURE2D_X_LOD(_CameraMotionVectorsTexture, s_point_clamp_sampler, motionVecUV, 0);
    DecodeMotionVector(encodedMV, motionVector);
    float cocHis = SAMPLE_TEXTURE2D_X_LOD(_InputHistoryCoCTexture, SAMPLER_STATE, (uv - motionVector) * ScreenToTargetScaleHistory.xy, 0.0).x;

    // Neighborhood clamping
#ifdef ENABLE_MAX_BLENDING
    float cocMin = closest.z;
#else
    float cocMin = min(closest.z, 1.0f);
#endif
    float cocMax = Max3(Max3(coc0, coc1, coc2), coc3, coc4);
    cocHis = clamp(cocHis, cocMin, cocMax);

    float outputCoC = coc0;
    if (sign(coc0) == sign(cocHis))
        outputCoC = lerp(coc0, cocHis, MotionBlending);

#if defined(SHADER_API_XBOXONE)
    // In some cases, it looks like the compiler reorganizes code so that we end up at the end with a NaN in the history (disabling compiler optimizations get rid of the NaN).
    // As a failsafe, we capture the situation here and reject history when that is the case.
    if (IsNaN(cocHis) || IsInf(cocHis))
        outputCoC = coc0;
#endif
    // Blend with the history
    _OutputCoCTexture[COORD_TEXTURE2D_X(posInputs.positionSS)] = outputCoC;
}