//  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.
//
//  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;

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_constrainColor() {
    return GET_BIT(params.flags.y, 1);
}

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

vec4 taa_fetchColor(sampler2D map, vec2 uv) {
	vec4 c = texture(map, uv);
	// 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
}

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) {
	return taa_fetchColor(sourceMap, uv);
}

vec4 taa_fetchHistoryMap(vec2 uv) {
	return taa_fetchColor(historyMap, uv);
}

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).x;
}


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

vec2 taa_getImageSize() {
    vec2 imageSize = getWidthHeight(0);
    if (isStereo()) {
        imageSize.x *= 2.0;
    }
    return imageSize;
}

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

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

	vec3 dtl = vec3(-1, -1, taa_fetchDepth(uv - dv - du));
	vec3 dtc = vec3( 0, -1, taa_fetchDepth(uv - dv));
	vec3 dtr = vec3( 1, -1, taa_fetchDepth(uv - dv + du));

	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 + dv - du));
	vec3 dbc = vec3( 0, 1, taa_fetchDepth(uv + dv));
	vec3 dbr = vec3( 1, 1, taa_fetchDepth(uv + dv + du));

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

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

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

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

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

    vec2 dtl = taa_fetchVelocityMap(uv - dv - du);
    vec2 dtc = taa_fetchVelocityMap(uv - dv);
    vec2 dtr = taa_fetchVelocityMap(uv - dv + du);

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

    vec2 dbl = taa_fetchVelocityMap(uv + dv - du);
    vec2 dbc = taa_fetchVelocityMap(uv + dv);
    vec2 dbr = taa_fetchVelocityMap(uv + dv + du);

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

    float testSpeed = dot(dtc,dtc);
    if (testSpeed > best.z) { best = vec3(dtc, testSpeed); }
    testSpeed = dot(dtr,dtr);
    if (testSpeed > best.z) { best = vec3(dtr, testSpeed); }

    testSpeed = dot(dml,dml);
    if (testSpeed > best.z) { best = vec3(dml, testSpeed); }
    testSpeed = dot(dmc,dmc);
    if (testSpeed > best.z) { best = vec3(dmc, testSpeed); }
    testSpeed = dot(dmr,dmr);
    if (testSpeed > best.z) { best = vec3(dmr, testSpeed); }

    testSpeed = dot(dbl,dbl);
    if (testSpeed > best.z) { best = vec3(dbl, testSpeed); }
    testSpeed = dot(dbc,dbc);
    if (testSpeed > best.z) { best = vec3(dbc, testSpeed); }
    testSpeed = dot(dbr,dbr);
    if (testSpeed > best.z) { best = vec3(dbr, testSpeed); }

    return best.xy;
}

vec2 taa_fromFragUVToEyeUVAndSide(vec2 fragUV, out int stereoSide) {
    vec2 eyeUV = fragUV;
    stereoSide = 0;
    if (isStereo()) {
        if (eyeUV.x > 0.5) {
            eyeUV.x -= 0.5;
            stereoSide = 1;
        }
        eyeUV.x *= 2.0;
    }
    return eyeUV;
}

vec2 taa_fromEyeUVToFragUV(vec2 eyeUV, int stereoSide) {
    vec2 fragUV = eyeUV;
    if (isStereo()) {
        fragUV.x *= 0.5;
        fragUV.x += stereoSide*0.5;
    }
    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);
}

vec2 taa_fetchSourceAndHistory(vec2 fragUV, vec2 fragVelocity, out vec3 sourceColor, out vec3 historyColor) {
    vec2 prevEyeUV;
    vec2 prevFragUV = taa_computePrevFragAndEyeUV(fragUV, fragVelocity, prevEyeUV);
    sourceColor = taa_fetchSourceMap(fragUV).xyz;

    historyColor = sourceColor;
    if (!(any(lessThan(prevEyeUV, vec2(0.0))) || any(greaterThan(prevEyeUV, vec2(1.0))))) {
        historyColor = taa_fetchHistoryMap(prevFragUV).xyz;
    }
    return prevFragUV;
}

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 du = vec2(texelSize.x, 0.0);
	vec2 dv = vec2(0.0, texelSize.y);

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

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

    sampleColor = taa_fetchSourceMap(fragUV - dv + du).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 + dv - du).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;

    sampleColor = taa_fetchSourceMap(fragUV + dv).rgb;
    sumSamples += sampleColor;
    sumSamples2 += sampleColor * sampleColor;
    
    sampleColor = taa_fetchSourceMap(fragUV + dv + du).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)));
    
       float gamma = params.covarianceGamma;
       vec3 cmin = mu - gamma * sigma;
       vec3 cmax = mu + gamma * 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);
}

vec3 taa_evalFeedbackColor(vec3 sourceColor, vec3 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);

    
    vec3 nextColor =  mix(historyColor, sourceColor, k_feedback * blendFactor).xyz;
    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
    if (lumaB < lumaMin || lumaB > lumaMax) {
        return rgbA;
    } else {
        return rgbB;
    }
}