#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Filtering.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/PostProcessing/Shaders/UberPostFeatures.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/PostProcessing/Shaders/BloomCommon.hlsl"

#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#pragma kernel Uber

#pragma multi_compile _ CHROMATIC_ABERRATION
#pragma multi_compile _ VIGNETTE
#pragma multi_compile _ LENS_DISTORTION
#pragma multi_compile _ ENABLE_ALPHA

#include "Packages/com.unity.render-pipelines.high-definition/Runtime/PostProcessing/Shaders/PostProcessDefines.hlsl"

TEXTURE2D_X(_InputTexture);
TEXTURE2D_X(_BloomTexture);
TEXTURE2D(_BloomDirtTexture);
TEXTURE2D(_ChromaSpectralLut);
TEXTURE3D(_LogLut3D);
TEXTURE2D(_VignetteMask);

RW_TEXTURE2D_X(CTYPE, _OutputTexture);

SAMPLER(sampler_LinearClamp);
SAMPLER(sampler_ChromaSpectralLut);
SAMPLER(sampler_LogLut3D);
SAMPLER(sampler_VignetteMask);

CBUFFER_START(cb0)
    float4 _ChromaParams;
    float4 _VignetteParams1;
    float4 _VignetteParams2;
    float4 _VignetteColor;
    float4 _DistortionParams1;
    float4 _DistortionParams2;
    float4 _LogLut3D_Params;        // x: 1 / lut_size, y: lut_size - 1, z: postexposure, w: unused
    float4 _BloomParams;
    float4 _BloomThreshold;
    float4 _BloomTint;
    float4 _BloomDirtScaleOffset;
    float4 _BloomBicubicParams;
    float4 _DebugFlags;
    float4 _AlphaScaleBias;
CBUFFER_END

#define DistCenter              _DistortionParams1.xy
#define DistAxis                _DistortionParams1.zw
#define DistTheta               _DistortionParams2.x
#define DistSigma               _DistortionParams2.y
#define DistScale               _DistortionParams2.z
#define DistIntensity           _DistortionParams2.w

#define ChromaAmount            _ChromaParams.x
#define ChromaMaxSamples        int(_ChromaParams.y)

#define VignetteCenter          _VignetteParams1.xy
#define VignetteMode            uint(_VignetteParams1.z)
#define VignetteIntensity       _VignetteParams2.x
#define VignetteSmoothness      _VignetteParams2.y
#define VignetteRoundness       _VignetteParams2.z
#define VignetteRounded         _VignetteParams2.w
#define VignetteColor           _VignetteColor.xyz
#define VignetteOpacity         _VignetteColor.w

#define PostExposure            _LogLut3D_Params.z

#define BloomTint               _BloomTint.xyz
#define BloomIntensity          _BloomParams.x
#define DirtIntensity           _BloomParams.y
#define BloomEnabled            _BloomParams.z
#define DirtEnabled             _BloomParams.w
#define DirtScale               _BloomDirtScaleOffset.xy
#define DirtOffset              _BloomDirtScaleOffset.zw

#define OutputLogEnabled        _DebugFlags.x

#define AlphaScale              _AlphaScaleBias.x
#define AlphaBias               _AlphaScaleBias.y

float2 DistortUV(float2 uv)
{
    // Lens distortion
    // Note: this variant should never be set with XR
#ifdef LENS_DISTORTION
    uv = (uv - 0.5) * DistScale + 0.5;
    float2 ruv = DistAxis * (uv - 0.5 - DistCenter);
    float ru = length(float2(ruv));

    UNITY_BRANCH
    if (DistIntensity > 0.0)
    {
        float wu = ru * DistTheta;
        ru = tan(wu) * (rcp(ru * DistSigma));
        uv = uv + ruv * (ru - 1.0);
    }
    else
    {
        ru = rcp(ru) * DistTheta * atan(ru * DistSigma);
        uv = uv + ruv * (ru - 1.0);
    }
#endif

    return uv;
}

#define GROUP_SIZE 8

[numthreads(GROUP_SIZE, GROUP_SIZE, 1)]
void Uber(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);
    PositionInputs posInputs = GetPositionInput(float2(dispatchThreadId.xy), _ScreenSize.zw, uint2(GROUP_SIZE, GROUP_SIZE));
    float2 uv = posInputs.positionNDC;
    float2 uvDistorted = DistortUV(uv);
    CTYPE color = 0.0;
    CTYPE inputColor = 0.0;

    // Chromatic aberration
    // Inspired by the method described in "Rendering Inside" [Playdead 2016]
    // https://twitter.com/pixelmager/status/717019757766123520
#ifdef CHROMATIC_ABERRATION
    float2 coords = 2.0 * uv - 1.0;
    float2 end = uv - coords * dot(coords, coords) * ChromaAmount;
    float2 diff = end - uv;
    int samples = clamp(int(length(_ScreenSize.xy * diff / 2.0)), 3, ChromaMaxSamples);
    float2 delta = diff / samples;
    float2 pos = uv;
    float3 sum = 0.0, filterSum = 0.0;

    for (int i = 0; i < samples; i++)
    {
        float t = (i + 0.5) / samples;
        float3 s = SAMPLE_TEXTURE2D_X_LOD(_InputTexture, sampler_LinearClamp, ClampAndScaleUVForBilinear(DistortUV(pos)), 0.0).xyz;
        float3 filter = SAMPLE_TEXTURE2D_LOD(_ChromaSpectralLut, sampler_ChromaSpectralLut, float2(t, 0.0), 0).xyz;

        sum += s * filter;
        filterSum += filter;
        pos += delta;
    }

    color.xyz = sum / filterSum;
#ifdef ENABLE_ALPHA
    inputColor = SAMPLE_TEXTURE2D_X_LOD(_InputTexture, sampler_LinearClamp, ClampAndScaleUVForBilinear(uvDistorted), 0.0);
    color.w = inputColor.w;
#endif

#else
    color = SAMPLE_TEXTURE2D_X_LOD(_InputTexture, sampler_LinearClamp, ClampAndScaleUVForBilinear(uvDistorted), 0.0).CTYPE_SWIZZLE;
    inputColor = color;
#endif

    // Bloom
    UNITY_BRANCH
    if (BloomEnabled)
    {
        #if 0 // Bilinear
        float3 bloom = SAMPLE_TEXTURE2D_X_LOD(_BloomTexture, sampler_LinearClamp, ClampAndScaleUVForBilinear(uvDistorted), 0.0).xyz;
        #else
        float3 bloom = SampleTexture2DBicubic(TEXTURE2D_X_ARGS(_BloomTexture, sampler_LinearClamp), uvDistorted * _RTHandleScale.xy, _BloomBicubicParams, _RTHandleScale.xy, unity_StereoEyeIndex).xyz;
        #endif

        float3 thresholdedColor = QuadraticThreshold(color.xyz, _BloomThreshold.x, _BloomThreshold.yzw);
        color.xyz = lerp(color.xyz, (color.xyz - thresholdedColor) + (bloom * BloomTint), BloomIntensity);

        UNITY_BRANCH
        if (DirtEnabled)
        {
            // UVs for the dirt texture should be DistortUV(uv * DirtScale + DirtOffset) but
            // considering we use a cover-style scale on the dirt texture the difference isn't massive
            // so we chose to save a few ALUs here instead in case lens distortion is active
            float3 dirt = SAMPLE_TEXTURE2D_LOD(_BloomDirtTexture, sampler_LinearClamp, uvDistorted * DirtScale + DirtOffset, 0.0).xyz;
            color.xyz += bloom * dirt * DirtIntensity;
        }

#ifdef ENABLE_ALPHA
        // Bloom should also spread in areas with zero alpha, so we save the image with bloom here to do the mixing at the end of the shader
        inputColor.xyz = color.xyz;
#endif
    }

    // Vignette
#ifdef VIGNETTE

    UNITY_BRANCH
    if (VignetteMode == 0u) // Procedural
    {
        float2 d = abs(uvDistorted - VignetteCenter) * VignetteIntensity;
        d.x *= lerp(1.0, _ScreenSize.x / _ScreenSize.y, VignetteRounded);
        d = pow(saturate(d), VignetteRoundness);
        float vfactor = pow(saturate(1.0 - dot(d, d)), VignetteSmoothness);
        color.xyz *= lerp(VignetteColor, (1.0).xxx, vfactor);
    }
    else // Masked
    {
        float vfactor = SAMPLE_TEXTURE2D_LOD(_VignetteMask, sampler_VignetteMask, uvDistorted, 0).w;
        vfactor = FastSRGBToLinear(vfactor);
        float3 newColor = color.xyz * lerp(VignetteColor, (1.0).xxx, vfactor);
        color.xyz = lerp(color.xyz, newColor, VignetteOpacity);
    }

#endif

    // Grading, tonemapping
    // The branch is only used for frame settings & exr log export - else grading is always enabled
    UNITY_BRANCH
    if (OutputLogEnabled)
    {
        // Output in log space for debug & exr export (external grading)
        color.xyz = saturate(LinearToLogC(color.xyz));
    }
    else
    {
        // Artist request to fine tune exposure in post without affecting bloom, dof etc
        color.xyz *= PostExposure;

        // Move from linear to LogC
        float3 colorLutSpace = saturate(LinearToLogC(color.xyz));

        // Color lookup in the LogC lut
        color.xyz = ApplyLut3D(TEXTURE3D_ARGS(_LogLut3D, sampler_LogLut3D), colorLutSpace, _LogLut3D_Params.xy);
    }

    // Alpha mask
#ifdef ENABLE_ALPHA
        // Post processing is not applied on pixels with zero alpha
        // The alpha scale and bias control how steep is the transition between the post-processed and plain regions
        float alpha = inputColor.a * AlphaScale + AlphaBias;
        // Saturate is necessary to avoid issues when additive blending pushes the alpha over 1.
        color.xyz = lerp(inputColor.xyz, color.xyz, saturate(alpha));
#endif

    // Done
    _OutputTexture[COORD_TEXTURE2D_X(posInputs.positionSS)] = color;
}