include "pipelines/common.glsl"

compute_shader [[
	layout(std140, binding = 4) uniform Data {
		vec2 u_size;
    };

	layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
	layout (binding = 0) uniform sampler2D u_history;
	layout (binding = 1) uniform sampler2D u_depthbuf;
	layout (binding = 2) uniform sampler2D u_current;
	layout (binding = 3) uniform sampler2D u_motion_vectors;
    layout (rgba8, binding = 3) uniform image2D u_output;

	// https://gist.github.com/TheRealMJP/c83b8c0f46b63f3a88a5986f4fa982b1
	vec4 catmullRom(sampler2D tex, vec2 uv, vec2 texSize)
	{
		vec2 samplePos = uv * texSize;
		vec2 texPos1 = floor(samplePos - 0.5) + 0.5;

		vec2 f = samplePos - texPos1;

		vec2 w0 = f * (-0.5 + f * (1.0f - 0.5 * f));
		vec2 w1 = 1.0f + f * f * (-2.5 + 1.5 * f);
		vec2 w2 = f * (0.5 + f * (2.0f - 1.5 * f));
		vec2 w3 = f * f * (-0.5 + 0.5 * f);

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

		vec2 texPos0 = texPos1 - 1;
		vec2 texPos3 = texPos1 + 2;
		vec2 texPos12 = texPos1 + offset12;

		texPos0 /= texSize;
		texPos3 /= texSize;
		texPos12 /= texSize;

		vec4 result = vec4(0.0);
		result += textureLod(tex, vec2(texPos0.x, texPos0.y), 0.0f) * w0.x * w0.y;
		result += textureLod(tex, vec2(texPos12.x, texPos0.y), 0.0f) * w12.x * w0.y;
		result += textureLod(tex, vec2(texPos3.x, texPos0.y), 0.0f) * w3.x * w0.y;

		result += textureLod(tex, vec2(texPos0.x, texPos12.y), 0.0f) * w0.x * w12.y;
		result += textureLod(tex, vec2(texPos12.x, texPos12.y), 0.0f) * w12.x * w12.y;
		result += textureLod(tex, vec2(texPos3.x, texPos12.y), 0.0f) * w3.x * w12.y;

		result += textureLod(tex, vec2(texPos0.x, texPos3.y), 0.0f) * w0.x * w3.y;
		result += textureLod(tex, vec2(texPos12.x, texPos3.y), 0.0f) * w12.x * w3.y;
		result += textureLod(tex, vec2(texPos3.x, texPos3.y), 0.0f) * w3.x * w3.y;

		return result;
	}

	// https://www.iquilezles.org/www/articles/texture/texture.htm
	vec4 getTexel(sampler2D tex, vec2 p, vec2 res) {
		#if 1
			return catmullRom(tex, p, res);
		#else
			p = p * res + 0.5;

			vec2 i = floor(p);
			vec2 f = p - i;
			f = f * f * f * (f * (f * 6.0 - 15.0) + 10.0);
			p = i + f;

			p = (p - 0.5) / res;
			return textureLod(tex, p, 0);
		#endif
	}

	// https://www.elopezr.com/temporal-aa-and-the-quest-for-the-holy-trail/
	void main()
	{
		ivec2 ij = ivec2(gl_GlobalInvocationID.xy);
		if (any(greaterThanEqual(ij, ivec2(u_size.xy)))) return;

		vec2 uv = (vec2(ij) + 0.5) / u_size.xy;
		float depth = textureLod(u_depthbuf, uv, 0).x;
		vec2 motionvec = textureLod(u_motion_vectors, uv, 0).xy * 2 - 1;
		
		vec2 uv_prev = uv + motionvec;

		vec4 current = textureLod(u_current, uv /*- Global.pixel_jitter*/, 0);
		if (all(lessThan(uv_prev, vec2(1))) && all(greaterThan(uv_prev, vec2(0))) ) {
			vec4 prev = getTexel(u_history, uv_prev, u_size);

			#if 1 // color clamping
				vec4 minColor = vec4(9001.0), maxColor = vec4(-9001.0);

				#define ITER(x, y) \
				{ \
					vec4 color = textureLodOffset(u_current, uv, 0, ivec2(x, y)); \
					minColor = min(minColor, color); \
					maxColor = max(maxColor, color); \
				}				
				ITER(-1, -1)
				ITER(-1, 0)
				ITER(-1, 1)
				ITER(0, -1)
				ITER(0, 1)
				ITER(1, -1)
				ITER(1, 0)
				ITER(1, 1)
				prev = clamp(prev, minColor, maxColor);
			#endif

			float lum0 = luminance(current.rgb);
			float lum1 = luminance(prev.rgb);

			float d = 1 - abs(lum0 - lum1) / max(lum0, max(lum1, 0.1));
			float k_feedback = mix(0.9, 0.99, d * d);

			current = mix(current, prev, k_feedback);
    	}

		imageStore(u_output, ij, current);
	}
]]