//  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()$>

const int SEQUENCE_LENGTH = 16;
struct JitterSequence {
    vec4 offsets[(SEQUENCE_LENGTH / 2) + 1];
};
layout(std140) uniform taaJitterBuffer {
    JitterSequence sequence;
};

vec2 taa_getJitterSample(int index) {
    return vec2((bool(index & 0x01) ? sequence.offsets[index >> 1].zw : sequence.offsets[index >> 1].xy));
}

int taa_getJitterSequenceLength() {
    return int(sequence.offsets[(SEQUENCE_LENGTH / 2)].z);
}

int taa_getCurrentJitterIndex() {
    return int(sequence.offsets[(SEQUENCE_LENGTH / 2)].w);
}

vec2 taa_getCurrentJitterSample() {
    return taa_getJitterSample(taa_getCurrentJitterIndex());
}

<@include gpu/Color.slh@>

uniform sampler2D depthMap;
uniform sampler2D sourceMap;
uniform sampler2D historyMap;
uniform sampler2D velocityMap;
uniform sampler2D nextMap;

struct TAAParams
{
	float debugX;
	float blend;
	float motionScale;
	float debugShowVelocityThreshold;
    ivec4 debug;
    vec4 pixelInfo_orbZoom;
};

layout(std140) uniform taaParamsBuffer {
    TAAParams params;
};

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

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

bool taa_showJitterSequence() {
    return GET_BIT(params.debug.x, 2);
}

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


bool taa_unjitter() {
    return GET_BIT(params.debug.y, 0);
}
bool taa_constrainColor() {
    return GET_BIT(params.debug.y, 1);
}

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

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

vec4 taa_fetchColor(sampler2D map, vec2 uv) {
#if USE_YCOCG
		vec4 c = texture(map, uv);
		return vec4(color_LinearToYCoCg(c.rgb), c.a);
#else
		return texture(map, uv);
#endif
}

vec3 taa_resolveColor(vec3 color) {
#if USE_YCOCG
		return color_YCoCgToLinear(color);
#else
		return color;
#endif
}

vec4 taa_fetchSourceMap(vec2 uv) {
#if USE_YCOCG
		vec4 c = texture(sourceMap, uv);
		return vec4(color_LinearToYCoCg(c.rgb), c.a);
#else
		return texture(sourceMap, uv);
#endif
}

vec4 taa_fetchHistoryMap(vec2 uv) {
#if USE_YCOCG
		vec4 c = texture(historyMap, uv);
		return vec4(color_LinearToYCoCg(c.rgb), c.a);
#else
		return texture(historyMap, uv);
#endif
}

vec4 taa_fetchNextMap(vec2 uv) {
#if USE_YCOCG
		vec4 c = texture(nextMap, uv);
		return vec4(color_LinearToYCoCg(c.rgb), c.a);
#else
		return texture(nextMap, uv);
#endif
}

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)

vec3 taa_findClosestFragment3x3(vec2 uv)
{
	vec2 dd = abs(getInvWidthHeight());
	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_fetchSourceAndHistory(vec2 fragUV, vec2 fragVelocity, out vec3 sourceColor, out vec3 historyColor) {
    sourceColor = taa_fetchSourceMap(fragUV).xyz;

    vec2 prevFragUV = fragUV - fragVelocity;
    historyColor = sourceColor;
    if (!(any(lessThan(prevFragUV, vec2(0.0))) || any(greaterThan(prevFragUV, 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;
}

vec3 taa_temporalReprojection(vec3 sourceColor, vec3 historyColor, vec2 fragUV, vec2 fragVelocity, float fragZe, vec2 fragJitterPix)
{
    vec4 texel1 = vec4(historyColor, 1.0);
    vec4 texel0 = vec4(sourceColor, 1.0);

    vec2 imageSize = getWidthHeight(0);
    vec2 texelSize = getInvWidthHeight();

  /*  if (taa_unjitter()) {
        fragUV += fragJitterPix * texelSize;
    }*/

		const float _SubpixelThreshold = 0.5;
		const float _GatherBase = 0.5;
		const float _GatherSubpixelMotion = 0.1666;
        const float _FeedbackMin = 0.1;
        const float _FeedbackMax = 0.9;

		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);
		vec4 c00 = taa_fetchSourceMap(fragUV - ss_offset11);
		vec4 c10 = taa_fetchSourceMap(fragUV - ss_offset01);
		vec4 c01 = taa_fetchSourceMap(fragUV + ss_offset01);
		vec4 c11 = taa_fetchSourceMap(fragUV + ss_offset11);

		vec4 cmin = min(c00, min(c10, min(c01, c11)));
		vec4 cmax = max(c00, max(c10, max(c01, c11)));

		#if USE_YCOCG || USE_CLIPPING
			vec4 cavg = (c00 + c10 + c01 + c11) / 4.0;
		#endif


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

		// clamp to neighbourhood of current sample
	#if USE_CLIPPING
		texel1 = clip_aabb(cmin.xyz, cmax.xyz, clamp(cavg, cmin, cmax), texel1);
	#else
		texel1 = clamp(texel1, cmin, cmax);
	#endif

		// feedback weight from unbiased luminance diff (t.lottes)
	#if USE_YCOCG
		float lum0 = texel0.r;
		float lum1 = texel1.r;
	#else
		float lum0 = Luminance(texel0.rgb);
		float lum1 = Luminance(texel1.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);

		// output
    vec3 nextColor =  mix(texel1, texel0, k_feedback).xyz;

    return taa_resolveColor(nextColor);
}

<$declareColorWheel()$>

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

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