//  Generated on <$_SCRIBE_DATE$>
//
//  TAA.slh
//  Common component needed by TemporalAntialiasing fragment shader
//
//  Created by Sam Gateau on 8/17/2017
//  Copyright 2017 High Fidelity, Inc.
//  Copyright 2024 Overte e.V.
//
//  Distributed under the Apache License, Version 2.0.
//  See the accompanying file LICENSE or http://www.apache.org/licenses/LICENSE-2.0.html
//

<@include DeferredTransform.slh@>
<$declareDeferredFrameTransform()$>

<@include render-utils/ShaderConstants.h@>
<@include gpu/Color.slh@>

LAYOUT(binding=RENDER_UTILS_TEXTURE_TAA_HISTORY) uniform sampler2D historyMap;
LAYOUT(binding=RENDER_UTILS_TEXTURE_TAA_SOURCE) uniform sampler2D sourceMap;
LAYOUT(binding=RENDER_UTILS_TEXTURE_TAA_VELOCITY) uniform sampler2D velocityMap;
LAYOUT(binding=RENDER_UTILS_TEXTURE_TAA_DEPTH) uniform sampler2D depthMap;
LAYOUT(binding=RENDER_UTILS_TEXTURE_TAA_NEXT) uniform sampler2D nextMap;
LAYOUT(binding=RENDER_UTILS_TEXTURE_TAA_INTENSITY) uniform sampler2D intensityMap;

struct TAAParams
{
    float none;
    float blend;
    float covarianceGamma;
    float debugShowVelocityThreshold;
    ivec4 flags;
    vec4 pixelInfo_orbZoom;
    vec4 regionInfo;
};

LAYOUT_STD140(binding=RENDER_UTILS_BUFFER_TAA_PARAMS) uniform taaParamsBuffer {
    TAAParams params;
};

#define GET_BIT(bitfield, bitIndex) bool((bitfield) & (1 << (bitIndex)))

bool taa_isDebugEnabled() {
    return GET_BIT(params.flags.x, 0);
}

bool taa_showDebugCursor() {
    return GET_BIT(params.flags.x, 1);
}

bool taa_showClosestFragment() {
    return GET_BIT(params.flags.x, 3);
}

bool taa_bicubicHistoryFetch() {
    return GET_BIT(params.flags.y, 0);
}

bool taa_constrainColor() {
    return GET_BIT(params.flags.y, 1);
}

bool taa_sharpenOutput() {
    return GET_BIT(params.flags.y, 2);
}

bool taa_feedbackColor() {
    return GET_BIT(params.flags.y, 4);
}

vec2 taa_getDebugCursorTexcoord() {
    return params.pixelInfo_orbZoom.xy;
}

float taa_getDebugOrbZoom() {
    return params.pixelInfo_orbZoom.z;
}

vec2 taa_getRegionDebug() {
    return params.regionInfo.xy;
}

vec2 taa_getRegionFXAA() {
    return params.regionInfo.zw;
}
#define USE_YCOCG 1

vec2 taa_getImageSize() {
    vec2 imageSize = getWidthHeight(0);
    imageSize.x *= 1.0 + float(isStereo());
    return imageSize;
}

vec2 taa_getTexelSize() {
    vec2 texelSize = getInvWidthHeight();
    texelSize.x *= 1.0 - 0.5 * float(isStereo());
    return texelSize;
}

vec4 taa_transformColor(vec4 c) {
    // Apply rapid pseudo tonemapping as TAA is applied to a tonemapped image, using luminance as weight, as proposed in
    // https://de45xmedrsdbp.cloudfront.net/Resources/files/TemporalAA_small-59732822.pdf
    float lum = dot(vec3(0.3, 0.5, 0.2), c.rgb);
    c.rgb = c.rgb / (1.0 + lum);
#if USE_YCOCG
    return vec4(color_LinearToYCoCg(c.rgb), c.a);
#else
    return c;
#endif
}

vec4 taa_fetchColor(sampler2D map, vec2 uv) {
    vec4 c = texture(map, uv);
    return taa_transformColor(c);
}

vec3 taa_resolveColor(vec3 color) {
#if USE_YCOCG
    color = max(vec3(0), color_YCoCgToUnclampedLinear(color));
#endif
    // Apply rapid inverse tonemapping, using luminance as weight, as proposed in
    // https://de45xmedrsdbp.cloudfront.net/Resources/files/TemporalAA_small-59732822.pdf
    float lum = dot(vec3(0.3, 0.5, 0.2),color.rgb);
    color = color / (1.0 - lum);
    return color;
}

vec4 taa_fetchSourceMap(vec2 uv) {
    vec4 source = taa_fetchColor(sourceMap, uv);
    // Store AA intensity in alpha
    source.a = texture(intensityMap, uv).r;
    return source;
}

vec4 taa_fetchHistoryMap(vec2 uv) {
    vec4 result;

    if (taa_bicubicHistoryFetch()) {
        // Perform a Catmull-Rom interpolation for a (possibly) sharper result.
        // Optimized 9-tap fetch based on http://vec3.ca/bicubic-filtering-in-fewer-taps/
        vec2 samplePixelPos = uv * taa_getImageSize() - 0.5;
        vec2 samplePixelIntPos = floor(samplePixelPos);

        vec2 samplePixelFracPos = samplePixelPos - samplePixelIntPos;
        vec2 samplePixelFracPos2 = samplePixelFracPos * samplePixelFracPos;

        // Catmull-Rom interpolation weights
        vec2 w0 = samplePixelFracPos * (samplePixelFracPos * (1.0 - 0.5*samplePixelFracPos) - 0.5);
        vec2 w1 = 1.0 + samplePixelFracPos2 * (1.5*samplePixelFracPos - 2.5);
        vec2 w2 = samplePixelFracPos * (samplePixelFracPos * (2.0 - 1.5*samplePixelFracPos) + 0.5);
        vec2 w3 = samplePixelFracPos2 * (0.5 * samplePixelFracPos - 0.5);

        vec2 w12 = w1 + w2;
        vec2 offset12 = w2 / (w1 + w2);

        vec2 sampleUV0 = samplePixelIntPos - vec2(1.0);
        vec2 sampleUV3 = samplePixelIntPos + vec2(2.0);
        vec2 sampleUV12 = samplePixelIntPos + offset12;

        vec2 texelSize = taa_getTexelSize();
        sampleUV0 = (sampleUV0 + 0.5) * texelSize;
        sampleUV3 = (sampleUV3 + 0.5) * texelSize;
        sampleUV12 = (sampleUV12 + 0.5) * texelSize;

        result = texture(historyMap, vec2(sampleUV0.x,  sampleUV0.y)) * w0.x * w0.y;
        result += texture(historyMap, vec2(sampleUV12.x, sampleUV0.y)) * w12.x * w0.y;
        result += texture(historyMap, vec2(sampleUV3.x,  sampleUV0.y)) * w3.x * w0.y;

        result += texture(historyMap, vec2(sampleUV0.x,  sampleUV12.y)) * w0.x * w12.y;
        result += texture(historyMap, vec2(sampleUV12.x, sampleUV12.y)) * w12.x * w12.y;
        result += texture(historyMap, vec2(sampleUV3.x,  sampleUV12.y)) * w3.x * w12.y;

        result += texture(historyMap, vec2(sampleUV0.x,  sampleUV3.y)) * w0.x * w3.y;
        result += texture(historyMap, vec2(sampleUV12.x, sampleUV3.y)) * w12.x * w3.y;
        result += texture(historyMap, vec2(sampleUV3.x,  sampleUV3.y)) * w3.x * w3.y;

        result.a = clamp(result.a, 0.0, 1.0);
    } else {
        result = texture(historyMap, uv);
    }
    return taa_transformColor(result);
}

vec4 taa_fetchNextMap(vec2 uv) {
    return taa_fetchColor(nextMap, uv);
}

vec2 taa_fetchVelocityMap(vec2 uv) {
    return texture(velocityMap, uv).xy;
}

float taa_fetchDepth(vec2 uv) {
    return -texture(depthMap, vec2(uv), 0.0).x;
}


#define ZCMP_GT(a, b) float(a > b)

vec3 taa_findClosestFragment3x3(vec2 uv)
{
    vec2 dd = taa_getTexelSize();
    vec2 du = vec2(dd.x, 0.0);
    vec2 dv = vec2(0.0, dd.y);
    vec2 dm = vec2(-dd.x, dd.y);

    vec3 dtl = vec3(-1, -1, taa_fetchDepth(uv - dd));
    vec3 dtc = vec3( 0, -1, taa_fetchDepth(uv - dv));
    vec3 dtr = vec3( 1, -1, taa_fetchDepth(uv - dm));

    vec3 dml = vec3(-1, 0, taa_fetchDepth(uv - du));
    vec3 dmc = vec3( 0, 0, taa_fetchDepth(uv));
    vec3 dmr = vec3( 1, 0, taa_fetchDepth(uv + du));

    vec3 dbl = vec3(-1, 1, taa_fetchDepth(uv + dm));
    vec3 dbc = vec3( 0, 1, taa_fetchDepth(uv + dv));
    vec3 dbr = vec3( 1, 1, taa_fetchDepth(uv + dd));

    vec3 dmin = dtl;
    dmin = mix(dmin, dtc, ZCMP_GT(dmin.z, dtc.z));
    dmin = mix(dmin, dtr, ZCMP_GT(dmin.z, dtr.z));

    dmin = mix(dmin, dml, ZCMP_GT(dmin.z, dml.z));
    dmin = mix(dmin, dmc, ZCMP_GT(dmin.z, dmc.z));
    dmin = mix(dmin, dmr, ZCMP_GT(dmin.z, dmr.z));

    dmin = mix(dmin, dbl, ZCMP_GT(dmin.z, dbl.z));
    dmin = mix(dmin, dbc, ZCMP_GT(dmin.z, dbc.z));
    dmin = mix(dmin, dbr, ZCMP_GT(dmin.z, dbr.z));

    return vec3(uv + dd.xy * dmin.xy, dmin.z);
}

vec2 taa_fetchVelocityMapBest(vec2 uv) {
    vec2 dd = taa_getTexelSize();
    vec2 du = vec2(dd.x, 0.0);
    vec2 dv = vec2(0.0, dd.y);
    vec2 dm = vec2(dd.x, -dd.y);

    vec2 dtl = taa_fetchVelocityMap(uv - dd);
    vec2 dtc = taa_fetchVelocityMap(uv - dv);
    vec2 dtr = taa_fetchVelocityMap(uv - dm);

    vec2 dml = taa_fetchVelocityMap(uv - du);
    vec2 dmc = taa_fetchVelocityMap(uv);
    vec2 dmr = taa_fetchVelocityMap(uv + du);

    vec2 dbl = taa_fetchVelocityMap(uv + dm);
    vec2 dbc = taa_fetchVelocityMap(uv + dv);
    vec2 dbr = taa_fetchVelocityMap(uv + dd);

    vec3 best = vec3(dtl, dot(dtl, dtl));

    float testSpeed = dot(dtc, dtc);
    mix(best, vec3(dtc, testSpeed), float(testSpeed > best.z));
    testSpeed = dot(dtr, dtr);
    mix(best, vec3(dtr, testSpeed), float(testSpeed > best.z));

    testSpeed = dot(dml, dml);
    mix(best, vec3(dml, testSpeed), float(testSpeed > best.z));
    testSpeed = dot(dmc, dmc);
    mix(best, vec3(dmc, testSpeed), float(testSpeed > best.z));
    testSpeed = dot(dmr, dmr);
    mix(best, vec3(dmr, testSpeed), float(testSpeed > best.z));

    testSpeed = dot(dbl, dbl);
    mix(best, vec3(dbl, testSpeed), float(testSpeed > best.z));
    testSpeed = dot(dbc, dbc);
    mix(best, vec3(dbc, testSpeed), float(testSpeed > best.z));
    testSpeed = dot(dbr, dbr);
    mix(best, vec3(dbr, testSpeed), float(testSpeed > best.z));

    return best.xy;
}

vec2 taa_fromFragUVToEyeUVAndSide(vec2 fragUV, out int stereoSide) {
    vec2 eyeUV = fragUV;

    float check = float(isStereo());
    float check2 = float(eyeUV.x > 0.5);
    eyeUV.x -= check * check2 * 0.5;
    stereoSide = int(check * check2);
    eyeUV.x *= 1.0 + check;
    return eyeUV;
}

vec2 taa_fromEyeUVToFragUV(vec2 eyeUV, int stereoSide) {
    vec2 fragUV = eyeUV;
    float check = float(isStereo());
    fragUV.x += check * float(stereoSide);
    fragUV.x *= 1.0 - 0.5 * check;
    return fragUV;
}

vec2 taa_computePrevFragAndEyeUV(vec2 fragUV, vec2 fragVelocity, out vec2 prevEyeUV) {
    int stereoSide = 0;
    vec2 eyeUV = taa_fromFragUVToEyeUVAndSide(fragUV, stereoSide);
    prevEyeUV = eyeUV - fragVelocity;
    return taa_fromEyeUVToFragUV(prevEyeUV, stereoSide);
}

bool taa_fetchSourceAndHistory(vec2 fragUV, vec2 fragVelocity, out vec4 sourceColor, out vec4 historyColor) {
    sourceColor.rgb = texture(sourceMap, fragUV).rgb;
    // Store AA intensity in alpha
    sourceColor.a = texture(intensityMap, fragUV).r;

    // If velocity is 0 then don't fetch history, just return the source. This means there is no jitter on this pixel
    if (any(notEqual(fragVelocity, vec2(0.0)))) {
        vec2 prevEyeUV;
        vec2 prevFragUV = taa_computePrevFragAndEyeUV(fragUV, fragVelocity, prevEyeUV);
        sourceColor = taa_transformColor(sourceColor);
        historyColor = mix(sourceColor, taa_fetchHistoryMap(prevFragUV), float(!(any(lessThan(prevEyeUV, vec2(0.0))) || any(greaterThan(prevEyeUV, vec2(1.0))))));
        return true;
    } else {
        historyColor = sourceColor;
        return false;
    }
}

float Luminance(vec3 rgb) {
    return rgb.x/4.0 + rgb.y/2.0 + rgb.z/4.0;
}

#define MINMAX_3X3_ROUNDED 1

mat3 taa_evalNeighbourColorVariance(vec3 sourceColor, vec2 fragUV, vec2 fragVelocity) {
    vec2 texelSize = taa_getTexelSize();

    vec2 dd = texelSize;
    vec2 du = vec2(dd.x, 0.0);
    vec2 dv = vec2(0.0, dd.y);
    vec2 dm = vec2(dd.x, -dd.y);

    vec3 sampleColor = taa_fetchSourceMap(fragUV - dd).rgb;
    vec3 sumSamples = sampleColor;
    vec3 sumSamples2 = sampleColor * sampleColor;

    sampleColor = taa_fetchSourceMap(fragUV - dv).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;

    sampleColor = taa_fetchSourceMap(fragUV - dm).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;

    sampleColor = taa_fetchSourceMap(fragUV - du).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;

    sampleColor = sourceColor; //taa_fetchSourceMap(fragUV).rgb; // could resuse the same osurce sampleColor isn't it ?
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;

    sampleColor = taa_fetchSourceMap(fragUV + du).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;

    sampleColor = taa_fetchSourceMap(fragUV + dm).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;

    sampleColor = taa_fetchSourceMap(fragUV + dv).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;

    sampleColor = taa_fetchSourceMap(fragUV + dd).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;


    vec3 mu = sumSamples / vec3(9.0);
    vec3 sigma = sqrt(max(sumSamples2 / vec3(9.0) - mu * mu, vec3(0.0)));

    sigma *= params.covarianceGamma;
    vec3 cmin = mu - sigma;
    vec3 cmax = mu + sigma;

    return mat3(cmin, cmax, mu);
}

mat3 taa_evalNeighbourColorRegion(vec3 sourceColor, vec2 fragUV, vec2 fragVelocity, float fragZe) {
    vec2 imageSize = taa_getImageSize();
    vec2 texelSize = taa_getTexelSize();
    vec3 cmin, cmax, cavg;

#if MINMAX_3X3_ROUNDED
    vec2 du = vec2(texelSize.x, 0.0);
    vec2 dv = vec2(0.0, texelSize.y);

    vec3 ctl = taa_fetchSourceMap(fragUV - dv - du).rgb;
    vec3 ctc = taa_fetchSourceMap(fragUV - dv).rgb;
    vec3 ctr = taa_fetchSourceMap(fragUV - dv + du).rgb;
    vec3 cml = taa_fetchSourceMap(fragUV - du).rgb;
    vec3 cmc = sourceColor; //taa_fetchSourceMap(fragUV).rgb; // could resuse the same osurce sample isn't it ?
    vec3 cmr = taa_fetchSourceMap(fragUV + du).rgb;
    vec3 cbl = taa_fetchSourceMap(fragUV + dv - du).rgb;
    vec3 cbc = taa_fetchSourceMap(fragUV + dv).rgb;
    vec3 cbr = taa_fetchSourceMap(fragUV + dv + du).rgb;

    cmin = min(ctl, min(ctc, min(ctr, min(cml, min(cmc, min(cmr, min(cbl, min(cbc, cbr))))))));
    cmax = max(ctl, max(ctc, max(ctr, max(cml, max(cmc, max(cmr, max(cbl, max(cbc, cbr))))))));

    #if MINMAX_3X3_ROUNDED || USE_YCOCG || USE_CLIPPING
        cavg = (ctl + ctc + ctr + cml + cmc + cmr + cbl + cbc + cbr) / 9.0;
    #elif
        cavg = (cmin + cmax ) * 0.5;
    #endif

    #if MINMAX_3X3_ROUNDED
        vec3 cmin5 = min(ctc, min(cml, min(cmc, min(cmr, cbc))));
        vec3 cmax5 = max(ctc, max(cml, max(cmc, max(cmr, cbc))));
        vec3 cavg5 = (ctc + cml + cmc + cmr + cbc) / 5.0;
        cmin = 0.5 * (cmin + cmin5);
        cmax = 0.5 * (cmax + cmax5);
        cavg = 0.5 * (cavg + cavg5);
    #endif
#else
    const float _SubpixelThreshold = 0.5;
    const float _GatherBase = 0.5;
    const float _GatherSubpixelMotion = 0.1666;

    vec2 texel_vel = fragVelocity * imageSize;
    float texel_vel_mag = length(texel_vel) * -fragZe;
    float k_subpixel_motion = clamp(_SubpixelThreshold / (0.0001 + texel_vel_mag), 0.0, 1.0);
    float k_min_max_support = _GatherBase + _GatherSubpixelMotion * k_subpixel_motion;

    vec2 ss_offset01 = k_min_max_support * vec2(-texelSize.x, texelSize.y);
    vec2 ss_offset11 = k_min_max_support * vec2(texelSize.x, texelSize.y);
    vec3 c00 = taa_fetchSourceMap(fragUV - ss_offset11).rgb;
    vec3 c10 = taa_fetchSourceMap(fragUV - ss_offset01).rgb;
    vec3 c01 = taa_fetchSourceMap(fragUV + ss_offset01).rgb;
    vec3 c11 = taa_fetchSourceMap(fragUV + ss_offset11).rgb;

    cmin = min(c00, min(c10, min(c01, c11)));
    cmax = max(c00, max(c10, max(c01, c11)));
    cavg = (cmin + cmax ) * 0.5;

    #if USE_YCOCG || USE_CLIPPING
        cavg = (c00 + c10 + c01 + c11) / 4.0;
    #elif
        cavg = (cmin + cmax ) * 0.5;
    #endif
#endif

    // shrink chroma min-max
#if USE_YCOCG
    vec2 chroma_extent = vec2(0.25 * 0.5 * (cmax.r - cmin.r));
    vec2 chroma_center = sourceColor.gb;
    cmin.yz = chroma_center - chroma_extent;
    cmax.yz = chroma_center + chroma_extent;
    cavg.yz = chroma_center;
#endif

   return mat3(cmin, cmax, cavg);
}

//#define USE_OPTIMIZATIONS 0

vec3 taa_clampColor(vec3 colorMin, vec3 colorMax, vec3 colorSource, vec3 color) {
    const float eps = 0.00001;
    vec3 p = colorSource;
    vec3 q = color;
    // note: only clips towards aabb center (but fast!)
    vec3 p_clip = 0.5 * (colorMax + colorMin);
    vec3 e_clip = 0.5 * (colorMax - colorMin) + vec3(eps);

    vec3 v_clip = q - p_clip;
    vec3 v_unit = v_clip.xyz / e_clip;
    vec3 a_unit = abs(v_unit);
    float ma_unit = max(a_unit.x, max(a_unit.y, a_unit.z));

    if (ma_unit > 1.0) {
        return p_clip + v_clip / ma_unit;
    } else {
        return q;// point inside aabb
    }
}

vec3 taa_evalConstrainColor(vec3 sourceColor, vec2 sourceUV, vec2 sourceVel, vec3 candidateColor) {
    mat3 colorMinMaxAvg;

    colorMinMaxAvg = taa_evalNeighbourColorVariance(sourceColor, sourceUV, sourceVel);
     
    // clamp history to neighbourhood of current sample
    return taa_clampColor(colorMinMaxAvg[0], colorMinMaxAvg[1], sourceColor, candidateColor);
}

vec4 taa_evalFeedbackColor(vec4 sourceColor, vec4 historyColor, float blendFactor) {
    const float _FeedbackMin = 0.1;
    const float _FeedbackMax = 0.9;
    // feedback weight from unbiased luminance diff (t.lottes)
#if USE_YCOCG
    float lum0 = sourceColor.r;
    float lum1 = historyColor.r;
#else
    float lum0 = Luminance(sourceColor.rgb);
    float lum1 = Luminance(historyColor.rgb);
#endif
    float unbiased_diff = abs(lum0 - lum1) / max(lum0, max(lum1, 0.2));
    float unbiased_weight = 1.0 - unbiased_diff;
    float unbiased_weight_sqr = unbiased_weight * unbiased_weight;
    float k_feedback = mix(_FeedbackMin, _FeedbackMax, unbiased_weight_sqr);

    vec4 nextColor =  mix(historyColor, sourceColor, k_feedback * blendFactor);
    return nextColor;
}


<$declareColorWheel()$>

vec3 taa_getVelocityColorRelative(float velocityPixLength) {
    return colorRamp(velocityPixLength/params.debugShowVelocityThreshold);
}

vec3 taa_getVelocityColorAboveThreshold(float velocityPixLength) {
    return colorRamp((velocityPixLength - params.debugShowVelocityThreshold)/params.debugShowVelocityThreshold);
}


vec3 taa_evalFXAA(vec2 fragUV) {

    // vec2 texelSize = getInvWidthHeight();
    vec2 texelSize = taa_getTexelSize();

    // filter width limit for dependent "two-tap" texture samples
    float FXAA_SPAN_MAX = 8.0;

    // local contrast multiplier for performing AA
    // higher = sharper, but setting this value too high will cause near-vertical and near-horizontal edges to fail
    // see "fxaaQualityEdgeThreshold"
    float FXAA_REDUCE_MUL = 1.0 / 8.0;

    // luminance threshold for processing dark colors
    // see "fxaaQualityEdgeThresholdMin"
    float FXAA_REDUCE_MIN = 1.0 / 128.0;

    // fetch raw RGB values for nearby locations
    // sampling pattern is "five on a die" (each diagonal direction and the center)
    // computing the coordinates for these texture reads could be moved to the vertex shader for speed if needed
    vec3 rgbNW = texture(sourceMap, fragUV + (vec2(-1.0, -1.0) * texelSize)).xyz;
    vec3 rgbNE = texture(sourceMap, fragUV + (vec2(+1.0, -1.0) * texelSize)).xyz;
    vec3 rgbSW = texture(sourceMap, fragUV + (vec2(-1.0, +1.0) * texelSize)).xyz;
    vec3 rgbSE = texture(sourceMap, fragUV + (vec2(+1.0, +1.0) * texelSize)).xyz;
    vec3 rgbM  = texture(sourceMap, fragUV).xyz;
    
    // convert RGB values to luminance
    vec3 luma = vec3(0.299, 0.587, 0.114);
    float lumaNW = dot(rgbNW, luma);
    float lumaNE = dot(rgbNE, luma);
    float lumaSW = dot(rgbSW, luma);
    float lumaSE = dot(rgbSE, luma);
    float lumaM  = dot( rgbM, luma);
    
    // luma range of local neighborhood
    float lumaMin = min(lumaM, min(min(lumaNW, lumaNE), min(lumaSW, lumaSE)));
    float lumaMax = max(lumaM, max(max(lumaNW, lumaNE), max(lumaSW, lumaSE)));
    
    // direction perpendicular to local luma gradient
    vec2 dir;
    dir.x = -((lumaNW + lumaNE) - (lumaSW + lumaSE));
    dir.y =  ((lumaNW + lumaSW) - (lumaNE + lumaSE));

    // compute clamped direction offset for additional "two-tap" samples
    // longer vector = blurry, shorter vector = sharp
    float dirReduce = max((lumaNW + lumaNE + lumaSW + lumaSE) * (0.25 * FXAA_REDUCE_MUL), FXAA_REDUCE_MIN);
    float rcpDirMin = 1.0 / (min(abs(dir.x), abs(dir.y)) + dirReduce);
    dir = min(vec2(FXAA_SPAN_MAX,  FXAA_SPAN_MAX), 
          max(vec2(-FXAA_SPAN_MAX, -FXAA_SPAN_MAX), dir * rcpDirMin)) * texelSize;
        
    // perform additional texture sampling perpendicular to gradient
    vec3 rgbA = (1.0 / 2.0) * (
                texture(sourceMap, fragUV + dir * (1.0 / 3.0 - 0.5)).xyz +
                texture(sourceMap, fragUV + dir * (2.0 / 3.0 - 0.5)).xyz);
    vec3 rgbB = rgbA * (1.0 / 2.0) + (1.0 / 4.0) * (
                texture(sourceMap, fragUV + dir * (0.0 / 3.0 - 0.5)).xyz +
                texture(sourceMap, fragUV + dir * (3.0 / 3.0 - 0.5)).xyz);
    float lumaB = dot(rgbB, luma);

    // compare luma of new samples to the luma range of the original neighborhood
    // if the new samples exceed this range, just use the first two samples instead of all four
    return mix(rgbB, rgbA, float(lumaB < lumaMin || lumaB > lumaMax));
}