#include "depth_normal.fsh"
#include "camera.sh"
#include "idioms.sh"

#ifdef _SHADERC_PRESET
// #define TAA_DYNAMIC_SCENE
// #define TAA_HAS_PREVIOUS_DEPTHMAP
#endif


///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Defines to tweak in/out formats
//
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
#define AA_USE_INPUT_COLOR_HDR             0
#define AA_USE_HISTORY_COLOR_HDR           0

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Defines to tweak quality options (and performance)
//
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
#define AA_YCOCG_SPACE                      0
#define AA_USE_CLOSEST_MOTION_VECTOR        1
#define AA_USE_VARIANCE_CLIPPING            1
#define AA_VARIANCE_CLIPPING_SAMPLES_9      0
#define AA_USE_FILTERED                     0
#define AA_USE_FILTERED_HISTORY             0
#define AA_ANTI_FLICKERING                  0

#define HALF_MAX_MINUS1 65472.0 // (2 - 2^-9) * 2^15
#define HALF_EPSILON    0.0000152587890625

struct FTAAResolve {
    float4x4 PixelToPrevTexelMatrix;
    float4 Resolution;
    float2 Jitter;
    float4 TemporalClipBounding; // variance gamma, x: stationary, y: motional, z: motion amplification, w: packed color sharpness
    float4 FinalBlendParameters; // x: _static, y: dynamic
};

uniform FTAAResolve CBData;

uniform sampler2D ColorTexture;
uniform sampler2D HistoryTexture;
uniform sampler2D CameraDepthTexture;
uniform sampler2D LastCameraDepthTexture;
#ifdef TAA_DYNAMIC_SCENE
    uniform sampler2D CameraMotionVectorTexture;
#endif

in float2 UV_f;
layout(location = 0) out float4 fs_color;

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
///
/// Implement
///
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

const half kFrameExposureScale = 10.f;

// Difference in pixels for velocity after which the pixel is marked as no-history
const half kFrameVelocityInPixelsDiff = 64.f;  // valid for 1920x1080

const half3x3 kRGBToYCoCgTransform = float3x3(
    1.f, 2.f, -1.f,
    2.f, 0.f, 2.f,
    1.f, -2.f, -1.f
);

const half3x3 kYCoCgToRGBTransform = float3x3(
    0.25f, 0.25f, 0.25f,
    0.25f, 0.f, -0.25f,
    -0.25f, 0.25f, -0.25f
);

const int2 kOffset3x3[8] = int2[8](
    int2(-1, -1),
    int2(0, -1),
    int2(1, -1),
    int2(-1, 0),
    int2(1, 0),
    int2(-1, 1),
    int2(0, 1),
    int2(1, 1)
);

float2 Linear01Depth(float2 depth) {
    return zParams.x / ( zParams.y - depth / zParams.w );
}

half3 RGBToYCoCg(half3 RGB) {
    return kRGBToYCoCgTransform * RGB;
}

half3 YCoCgToRGB(half3 YCoCg) {
    return kYCoCgToRGBTransform * YCoCg;
}

// Faster but less accurate luma computation. 
// Luma includes a scaling by 4.
half Luma4(half3 color) {
    return color.g * 2.0 + (color.r + color.b);
}

half HdrWeight4(half3 color, float exposure) {
    return rcp(Luma4(color) * exposure + 4.0);
}

// Optimized HDR weighting function.
half HdrWeightY(half color, half exposure) {
    return rcp(color * exposure + 4.0);
}

half3 ClipToAABB(half3 history, half3 filtered, half3 neighborMin, half3 neighborMax) {

    const half3 RayDirEps = float3(HALF_EPSILON);

    half3 rayOrigin = history;
    half3 rayDir = filtered - history;
    rayDir = select(step(RayDirEps, abs(rayDir)), rayDir, RayDirEps);
    half3 invRayDir = rcp(rayDir);
    half3 tmin = (neighborMin - rayOrigin) * invRayDir;
    half3 tmax = (neighborMax - rayOrigin) * invRayDir;
    half3 t3 = min(tmin, tmax);
    half t = max(max(t3.x, t3.y), t3.z);
    t = saturate(t);
    return rayOrigin + t * rayDir;
}

half3 ClipToAABBFast(half3 history, half3 neighborMin, half3 neighborMax) {
    half3 center = inl_half(0.5f) * (neighborMin + neighborMax);
    half3 extents = inl_half(0.5f) * (neighborMax - neighborMin);

    half3 offset = history - center;

    half3 ts = abs(extents / (offset + inl_half(0.0001)));
    half t = saturate(min(ts.x, min(ts.y, ts.z)));
    return center + offset * t;
}

// Samples a texture with Catmull-Rom filtering, using 9 texture fetches instead of 16.
//      https://therealmjp.github.io/
// Some optimizations from here:
//      http://vec3.ca/bicubic-filtering-in-fewer-taps/ for more details
// Optimized to 5 taps by removing the corner samples
// And modified for mediump support
vec4 SampleTextureCatmullRom(const sampler2D tex, const highp vec2 uv, const highp vec4 texSizeParams) {
    // We're going to sample a a 4x4 grid of texels surrounding the target UV coordinate. We'll do this by rounding
    // down the sample location to get the exact center of our "starting" texel. The starting texel will be at
    // location [1, 1] in the grid, where [0, 0] is the top left corner.

    highp vec2 samplePos = uv * texSizeParams.xy;
    highp vec2 texPos1 = floor(samplePos - 0.5) + 0.5;

    // Compute the fractional offset from our starting texel to our original sample location, which we'll
    // feed into the Catmull-Rom spline function to get our filter weights.
    highp vec2 f = samplePos - texPos1;
    highp vec2 f2 = f * f;
    highp vec2 f3 = f2 * f;

    // Compute the Catmull-Rom weights using the fractional offset that we calculated earlier.
    // These equations are pre-expanded based on our knowledge of where the texels will be located,
    // which lets us avoid having to evaluate a piece-wise function.
    vec2 w0 = f2 - 0.5 * (f3 + f);
    vec2 w1 = 1.5 * f3 - 2.5 * f2 + 1.0;
    vec2 w3 = 0.5 * (f3 - f2);
    vec2 w2 = 1.0 - w0 - w1 - w3;

    // Work out weighting factors and sampling offsets that will let us use bilinear filtering to
    // simultaneously evaluate the middle 2 samples from the 4x4 grid.
    vec2 w12 = w1 + w2;

    // Compute the final UV coordinates we'll use for sampling the texture
    highp vec2 texPos0 = texPos1 - vec2(1.0);
    highp vec2 texPos3 = texPos1 + vec2(2.0);
    highp vec2 texPos12 = texPos1 + w2 / w12;

    highp vec2 invTexSize = texSizeParams.zw;
    texPos0  *= invTexSize;
    texPos3  *= invTexSize;
    texPos12 *= invTexSize;

    float k0 = w12.x * w0.y;
    float k1 = w0.x  * w12.y;
    float k2 = w12.x * w12.y;
    float k3 = w3.x  * w12.y;
    float k4 = w12.x * w3.y;

    vec4 result =   textureLod(tex, vec2(texPos12.x, texPos0.y),  0.0) * k0
                  + textureLod(tex, vec2(texPos0.x,  texPos12.y), 0.0) * k1
                  + textureLod(tex, vec2(texPos12.x, texPos12.y), 0.0) * k2
                  + textureLod(tex, vec2(texPos3.x,  texPos12.y), 0.0) * k3
                  + textureLod(tex, vec2(texPos12.x, texPos3.y),  0.0) * k4;

    result *= 1.0 / (k0 + k1 + k2 + k3 + k4);

    return result;
}

float2 ReprojectMotionVectorUV(float2 uv, out float depth) {

#if AA_USE_CLOSEST_MOTION_VECTOR
    float2 k = CBData.Resolution.zw;

    float4 neighborhood = float4(
        texture(CameraDepthTexture, uv - k).r,
        texture(CameraDepthTexture, uv + float2(k.x, -k.y)).r,
        texture(CameraDepthTexture, uv + float2(-k.x, k.y)).r,
        texture(CameraDepthTexture, uv + k).r
    );

#define COMPARE_DEPTH(a, b) step(a, b)

    float3 result = float3(0.0, 0.0, texture(CameraDepthTexture, uv).r);
    result = lerp(result, float3(-1.0, -1.0, neighborhood.x), COMPARE_DEPTH(neighborhood.x, result.z));
    result = lerp(result, float3( 1.0, -1.0, neighborhood.y), COMPARE_DEPTH(neighborhood.y, result.z));
    result = lerp(result, float3(-1.0,  1.0, neighborhood.z), COMPARE_DEPTH(neighborhood.z, result.z));
    result = lerp(result, float3( 1.0,  1.0, neighborhood.w), COMPARE_DEPTH(neighborhood.w, result.z));

#undef COMPARE_DEPTH

    depth = result.z;

    return uv + result.xy * k;
#else
    depth = texture(CameraDepthTexture, uv).r;
    return uv;
#endif /* AA_USE_CLOSEST_MOTION_VECTOR */
}

float2 GetVelocityUV(float3 texcoord) {
    float2 mv;
#ifdef TAA_DYNAMIC_SCENE
    mv = texture(CameraMotionVectorTexture, texcoord.xy).xy;
#else
    float4 prevClip = CBData.PixelToPrevTexelMatrix * float4(texcoord, 1.0);
    prevClip.xyz /= prevClip.w;
    mv = texcoord.xy - prevClip.xy;
#endif
    return mv;
}

float2 GetTexcoord() {
    return UV_f;
}

half4 GetHistory(float2 uv) {
#if AA_USE_FILTERED_HISTORY
    half4 history = SampleTextureCatmullRom( HistoryTexture, uv, CBData.Resolution );
#else
    half4 history = textureLod( HistoryTexture, uv, 0.0 );
#endif
    return history;
}

float GetSceneColorSharpness() {
    return CBData.TemporalClipBounding.w;
}

half4 TransformSceneColor(half4 rawLinearSceneColorRGBA) {
#if AA_YCOCG_SPACE
    return float4( RGBToYCoCg(rawLinearSceneColorRGBA.rgb), rawLinearSceneColorRGBA.a );
#else
    return rawLinearSceneColorRGBA;
#endif
}

half3 TransformSceneColor(half3 rawLinearSceneColorRGB) {
#if AA_YCOCG_SPACE
    return RGBToYCoCg(rawLinearSceneColorRGB);
#else
    return rawLinearSceneColorRGB;
#endif
}

half4 TransformBackToRawLinearSceneColor(half4 sceneColor) {
#if AA_YCOCG_SPACE
    return float4(YCoCgToRGB(sceneColor.rgb), sceneColor.a);
#else
    return  sceneColor;
#endif
}

half3 TransformBackToRawLinearSceneColor(half3 sceneColor) {
#if AA_YCOCG_SPACE
    return YCoCgToRGB(sceneColor);
#else
    return sceneColor;
#endif
}

half GetSceneColorLuma4(half4 sceneColor) {
#if AA_YCOCG_SPACE
    return sceneColor.x;
#else
    return Luma4(sceneColor.rgb);
#endif
}

half GetSceneColorHdrWeight(half4 sceneColor) {
#if AA_YCOCG_SPACE
    return HdrWeightY(sceneColor.x, kFrameExposureScale);
#else
    return HdrWeight4(sceneColor.rgb, kFrameExposureScale);
#endif
}

half3 ClampHistory(half3 history, half3 filtered, half3 neighborMin, half3 neighborMax) {

#if AA_USE_VARIANCE_CLIPPING
    // return ClipToAABB(history, filtered, neighborMin, neighborMax);
    return ClipToAABBFast(history, neighborMin, neighborMax);
#else
    return clamp(history, neighborMin, neighborMax);
#endif
}

half2 WeightedLerpFactors(half weightA, half weightB, half blend) {
    half blendA = (inl_half(1.f) - blend) * weightA;
    half blendB = blend * weightB;
    half rcpBlend = rcp(blendA + blendB);
    blendA *= rcpBlend;
    blendB *= rcpBlend;
    return float2(blendA, blendB);
}

void main() {

    float2 texcoord = GetTexcoord();
    float2 uv = texcoord - CBData.Jitter;

    float depth;
    float2 closest = ReprojectMotionVectorUV(texcoord, depth);

    float2 velocityUV = GetVelocityUV(float3(closest, depth));

    float2 prevCoord = texcoord - velocityUV;
    float4 currColor = texture(ColorTexture, uv);
    if(any(greaterThan(float4(prevCoord, 0.0, 0.0), float4(1.0, 1.0, prevCoord)))) {
        fs_color = currColor;
        return;
    }

    half4 neighborhood[8];
    neighborhood[0] = TransformSceneColor( textureOffset( ColorTexture, uv, kOffset3x3[0] ) );
    neighborhood[1] = TransformSceneColor( textureOffset( ColorTexture, uv, kOffset3x3[1] ) );
    neighborhood[2] = TransformSceneColor( textureOffset( ColorTexture, uv, kOffset3x3[2] ) );
    neighborhood[3] = TransformSceneColor( textureOffset( ColorTexture, uv, kOffset3x3[3] ) );
    neighborhood[4] = TransformSceneColor( textureOffset( ColorTexture, uv, kOffset3x3[4] ) );
    neighborhood[5] = TransformSceneColor( textureOffset( ColorTexture, uv, kOffset3x3[5] ) );
    neighborhood[6] = TransformSceneColor( textureOffset( ColorTexture, uv, kOffset3x3[6] ) );
    neighborhood[7] = TransformSceneColor( textureOffset( ColorTexture, uv, kOffset3x3[7] ) );

    half4 middleCenter = TransformSceneColor( currColor );

    half4 filtered;
#if !AA_USE_FILTERED
    filtered = middleCenter;
#endif

    float velocityLength = length(velocityUV.xy);
    float velocityWeight = saturate( velocityLength * CBData.TemporalClipBounding.z );

    half4 minColor, maxColor;
    // Variance clipping
    {
        half AABBScale = lerp( CBData.TemporalClipBounding.x, CBData.TemporalClipBounding.y, velocityWeight );

#if AA_USE_FILTERED
        half sampleWeight = GetSceneColorHdrWeight(middleCenter);
        half totalWeight = sampleWeight;

        filtered = middleCenter * sampleWeight;
#endif

        float4 m1 = middleCenter;
        float4 m2 = middleCenter * middleCenter;

#if AA_VARIANCE_CLIPPING_SAMPLES_9
        #pragma unroll
        for(int i = 0; i < 8; ++i) {
            half4 color = neighborhood[i];
            m1 += color;
            m2 += color * color;

#if AA_USE_FILTERED
            sampleWeight = GetSceneColorHdrWeight(color);
            filtered += color * sampleWeight;
            totalWeight += sampleWeight;
#endif
        }
        const float rcpDivider = 0.11111111;
#else
        const int filterIndices[4] = int[4](1, 3, 4, 6);
        #pragma unroll
        for(int i = 0; i < 4; ++i) {
            int index = filterIndices[i];
            half4 color = neighborhood[index];
            m1 += color;
            m2 += color * color;

#if AA_USE_FILTERED
            sampleWeight = GetSceneColorHdrWeight(color);
            filtered += color * sampleWeight;
            totalWeight += sampleWeight;
#endif
        }
        const float rcpDivider = 0.2;
#endif
        float4 mean = m1 * rcpDivider;
        float4 variance = sqrt(m2 * rcpDivider - mean * mean) * AABBScale;

#if AA_USE_FILTERED
        filtered /= totalWeight;
#endif
        minColor = mean - variance;
        maxColor = mean + variance;
        minColor = min(minColor, filtered);
        maxColor = max(maxColor, filtered);
    }

#if AA_USE_FILTERED

    // Sample history
    half4 history = TransformSceneColor( GetHistory(prevCoord) );

    // Clamp history
    history.xyz = ClampHistory( history.xyz, filtered.xyz, minColor.xyz, maxColor.xyz );

    // Add back in aliasing to sharpen
    {
        float2 velocityST = velocityUV * CBData.Resolution.xy;

        const float HistoryBlurAmp = 2.0;
        float historyBlur = saturate( dot(abs(velocityST), float2(HistoryBlurAmp) ) );

        float addAliasing = historyBlur * 0.5;
        #if AA_YCOCG_SPACE
        const float LumaContrastFactor = 32.0 * 4.0;
        #else
        const float LumaContrastFactor = 32.0;
        #endif

        float LumaMax = GetSceneColorLuma4(maxColor);
        float LumaMin = GetSceneColorLuma4(minColor);
        float LumaContrast = LumaMax - LumaMin;
        addAliasing = saturate( addAliasing + rcp(1.0 + LumaContrast * LumaContrastFactor) );
        filtered = lerp(filtered, middleCenter, addAliasing);
    }

    // Compute blend amount
    float blendFinal;
    {
        float LumaFiltered = GetSceneColorLuma4(filtered);
        float LumaHistory = GetSceneColorLuma4(history);

        blendFinal = 1.0 - lerp(CBData.FinalBlendParameters.x, CBData.FinalBlendParameters.y, velocityWeight);

		#if AA_ANTI_FLICKERING
		{
			// Anti-flicker
			float DistToClamp = 2 * abs(min(LumaHistory - LumaMin, LumaMax - LumaHistory) / (LumaMax - LumaMin));
			//BlendFinal *= lerp( 0, 1, saturate(4 * DistToClamp) );
			blendFinal += 0.8 * saturate(0.02 * LumaHistory / abs(LumaFiltered - LumaHistory));
			blendFinal *= (LumaMin * InExposureScale + 0.5) / (LumaMax * InExposureScale + 0.5);
		}
		#endif

        // Make sure to have at least some small contribution
        blendFinal = max(blendFinal, saturate(0.01 * LumaHistory / abs(LumaFiltered - LumaHistory)) );
    }

    // Do final blend between history and filtered color
    {
        float filterWeight = GetSceneColorHdrWeight(filtered);
        float historyWeight = GetSceneColorHdrWeight(history);

        half2 weights = WeightedLerpFactors(historyWeight, filterWeight, blendFinal);

        half4 temporalColor = history * weights.x + filtered * weights.y;

        fs_color = TransformBackToRawLinearSceneColor( temporalColor );
    }
#else
    // Sharpen output
    half4 corners = ((neighborhood[0] + neighborhood[2] + neighborhood[5] + neighborhood[7]) - filtered) * 2.f;
    filtered += (filtered - (corners * 0.166667)) * 2.718282 * GetSceneColorSharpness();
    filtered = clamp(filtered, 0.f, HALF_MAX_MINUS1);

    half4 history = TransformSceneColor( GetHistory(prevCoord) );

    history.xyz = ClampHistory( history.xyz, middleCenter.xyz, minColor.xyz, maxColor.xyz );

    float blendWeight = lerp(CBData.FinalBlendParameters.x, CBData.FinalBlendParameters.y, velocityWeight);

#if AA_ANTI_FLICKERING
    {
        // [Lottes] prevents flickering by modulating the blend weight by the difference in luma
        float LumaFiltered = GetSceneColorLuma4(filtered);
        float LumaHistory = GetSceneColorLuma4(history);
        float Diff = abs(LumaFiltered - LumaHistory) / ( 0.001 + max(LumaFiltered, LumaHistory) );

        blendWeight *= (Diff * Diff);
    }
#endif

    half4 temporalColor = lerp( filtered, history, blendWeight );

    fs_color = TransformBackToRawLinearSceneColor( temporalColor );
#endif
}
