#include <psemek/gfx/effect/fxaa.hpp>
#include <psemek/gfx/program.hpp>
#include <psemek/gfx/array.hpp>
#include <psemek/gfx/error.hpp>

namespace psemek::gfx
{

	static char const fxaa_vs[] =
R"(#version 330

const vec4 vertices[6] = vec4[6](
	vec4(-1.0, -1.0, 0.0, 1.0),
	vec4( 1.0, -1.0, 0.0, 1.0),
	vec4( 1.0,  1.0, 0.0, 1.0),

	vec4(-1.0, -1.0, 0.0, 1.0),
	vec4( 1.0,  1.0, 0.0, 1.0),
	vec4(-1.0,  1.0, 0.0, 1.0)
);

out vec2 texcoord;

void main()
{
	gl_Position = vertices[gl_VertexID];

	texcoord = vertices[gl_VertexID].xy * 0.5 + vec2(0.5);
}
)";


// FXAA implementation taken from https://github.com/libretro/glsl-shaders/blob/master/anti-aliasing/shaders/fxaa.glsl
// TODO: refactor

	static char const fxaa_fs[] =
R"(#version 330

uniform sampler2D u_input;
uniform vec2 u_d;

in vec2 texcoord;

out vec4 out_color;

vec3 FxaaLerp3(vec3 a, vec3 b, float amountOfA) {
	return (vec3(-amountOfA) * b) + ((a * vec3(amountOfA)) + b);
}

vec4 FxaaTexOff(sampler2D tex, vec2 pos, ivec2 off, vec2 rcpFrame) {
	float x = pos.x + float(off.x) * rcpFrame.x;
	float y = pos.y + float(off.y) * rcpFrame.y;
	return texture(tex, vec2(x, y));
}

float FxaaLuma(vec3 c)
{
	return dot(vec3(0.2126, 0.7152, 0.0722), c);
}

#define FXAA_EDGE_THRESHOLD      (1.0/8.0)
#define FXAA_EDGE_THRESHOLD_MIN  (1.0/16.0)
#define FXAA_SEARCH_STEPS        16
#define FXAA_SEARCH_THRESHOLD    (1.0/4.0)
#define FXAA_SUBPIX_CAP          (3.0/4.0)
#define FXAA_SUBPIX_TRIM         (1.0/4.0)
#define FXAA_SUBPIX_TRIM_SCALE 1.0 / (1.0 - FXAA_SUBPIX_TRIM)

void main()
{
	vec2 pos = texcoord;
	vec2 rcpFrame = u_d;

	vec3 rgbN = FxaaTexOff(u_input, pos.xy, ivec2( 0,-1), rcpFrame).xyz;
	vec3 rgbW = FxaaTexOff(u_input, pos.xy, ivec2(-1, 0), rcpFrame).xyz;
	vec3 rgbM = FxaaTexOff(u_input, pos.xy, ivec2( 0, 0), rcpFrame).xyz;
	vec3 rgbE = FxaaTexOff(u_input, pos.xy, ivec2( 1, 0), rcpFrame).xyz;
	vec3 rgbS = FxaaTexOff(u_input, pos.xy, ivec2( 0, 1), rcpFrame).xyz;

	float lumaN = FxaaLuma(rgbN);
	float lumaW = FxaaLuma(rgbW);
	float lumaM = FxaaLuma(rgbM);
	float lumaE = FxaaLuma(rgbE);
	float lumaS = FxaaLuma(rgbS);
	float rangeMin = min(lumaM, min(min(lumaN, lumaW), min(lumaS, lumaE)));
	float rangeMax = max(lumaM, max(max(lumaN, lumaW), max(lumaS, lumaE)));

	float range = rangeMax - rangeMin;
	if(range < max(FXAA_EDGE_THRESHOLD_MIN, rangeMax * FXAA_EDGE_THRESHOLD))
	{
		out_color = vec4(rgbM, 1.0);
		return;
	}

	vec3 rgbL = rgbN + rgbW + rgbM + rgbE + rgbS;

	float lumaL = (lumaN + lumaW + lumaE + lumaS) * 0.25;
	float rangeL = abs(lumaL - lumaM);
	float blendL = max(0.0, (rangeL / range) - FXAA_SUBPIX_TRIM) * FXAA_SUBPIX_TRIM_SCALE;
	blendL = min(FXAA_SUBPIX_CAP, blendL);

	vec3 rgbNW = FxaaTexOff(u_input, pos.xy, ivec2(-1,-1), rcpFrame).xyz;
	vec3 rgbNE = FxaaTexOff(u_input, pos.xy, ivec2( 1,-1), rcpFrame).xyz;
	vec3 rgbSW = FxaaTexOff(u_input, pos.xy, ivec2(-1, 1), rcpFrame).xyz;
	vec3 rgbSE = FxaaTexOff(u_input, pos.xy, ivec2( 1, 1), rcpFrame).xyz;
	rgbL += (rgbNW + rgbNE + rgbSW + rgbSE);
	rgbL *= vec3(1.0/9.0);

	float lumaNW = FxaaLuma(rgbNW);
	float lumaNE = FxaaLuma(rgbNE);
	float lumaSW = FxaaLuma(rgbSW);
	float lumaSE = FxaaLuma(rgbSE);

	float edgeVert =
		abs((0.25 * lumaNW) + (-0.5 * lumaN) + (0.25 * lumaNE)) +
		abs((0.50 * lumaW ) + (-1.0 * lumaM) + (0.50 * lumaE )) +
		abs((0.25 * lumaSW) + (-0.5 * lumaS) + (0.25 * lumaSE));
	float edgeHorz =
		abs((0.25 * lumaNW) + (-0.5 * lumaW) + (0.25 * lumaSW)) +
		abs((0.50 * lumaN ) + (-1.0 * lumaM) + (0.50 * lumaS )) +
		abs((0.25 * lumaNE) + (-0.5 * lumaE) + (0.25 * lumaSE));

	bool horzSpan = edgeHorz >= edgeVert;
	float lengthSign = horzSpan ? -rcpFrame.y : -rcpFrame.x;

	if(!horzSpan)
	{
		lumaN = lumaW;
		lumaS = lumaE;
	}

	float gradientN = abs(lumaN - lumaM);
	float gradientS = abs(lumaS - lumaM);
	lumaN = (lumaN + lumaM) * 0.5;
	lumaS = (lumaS + lumaM) * 0.5;

	if (gradientN < gradientS)
	{
		lumaN = lumaS;
		lumaN = lumaS;
		gradientN = gradientS;
		lengthSign *= -1.0;
	}

	vec2 posN;
	posN.x = pos.x + (horzSpan ? 0.0 : lengthSign * 0.5);
	posN.y = pos.y + (horzSpan ? lengthSign * 0.5 : 0.0);

	gradientN *= FXAA_SEARCH_THRESHOLD;

	vec2 posP = posN;
	vec2 offNP = horzSpan ? vec2(rcpFrame.x, 0.0) : vec2(0.0, rcpFrame.y);
	float lumaEndN = lumaN;
	float lumaEndP = lumaN;
	bool doneN = false;
	bool doneP = false;
	posN += offNP * vec2(-1.0, -1.0);
	posP += offNP * vec2( 1.0,  1.0);

	for(int i = 0; i < FXAA_SEARCH_STEPS; i++) {
		if(!doneN)
		{
			lumaEndN = FxaaLuma(texture(u_input, posN.xy).xyz);
		}
		if(!doneP)
		{
			lumaEndP = FxaaLuma(texture(u_input, posP.xy).xyz);
		}

		doneN = doneN || (abs(lumaEndN - lumaN) >= gradientN);
		doneP = doneP || (abs(lumaEndP - lumaN) >= gradientN);

		if(doneN && doneP)
		{
			break;
		}
		if(!doneN)
		{
			posN -= offNP;
		}
		if(!doneP)
		{
			posP += offNP;
		}
	}

	float dstN = horzSpan ? pos.x - posN.x : pos.y - posN.y;
	float dstP = horzSpan ? posP.x - pos.x : posP.y - pos.y;
	bool directionN = dstN < dstP;
	lumaEndN = directionN ? lumaEndN : lumaEndP;

	if(((lumaM - lumaN) < 0.0) == ((lumaEndN - lumaN) < 0.0))
	{
		lengthSign = 0.0;
	}


	float spanLength = (dstP + dstN);
	dstN = directionN ? dstN : dstP;
	float subPixelOffset = (0.5 + (dstN * (-1.0/spanLength))) * lengthSign;
	vec3 rgbF = texture(u_input, vec2(
		pos.x + (horzSpan ? 0.0 : subPixelOffset),
		pos.y + (horzSpan ? subPixelOffset : 0.0))).xyz;
	out_color = vec4(FxaaLerp3(rgbL, rgbF, blendL), 1.0);
}
)";

	struct fxaa::impl
	{
		gfx::program program{fxaa_vs, fxaa_fs};
		gfx::array array;
	};

	fxaa::fxaa()
		: pimpl_{std::make_unique<struct impl>()}
	{
		impl().program.bind();
		impl().program["u_input"] = 0;
	}

	fxaa::~fxaa() = default;

	void fxaa::invoke(texture_2d const & src, render_target const & dst)
	{
		gl::Disable(gl::BLEND);
		gl::Disable(gl::DEPTH_TEST);
		gl::Disable(gl::CULL_FACE);

		dst.bind();
		gl::ActiveTexture(gl::TEXTURE0);
		src.bind();
		impl().array.bind();
		impl().program.bind();
		impl().program["u_d"] = geom::vector{1.f / src.width(), 1.f / src.height()};
		gl::DrawArrays(gl::TRIANGLES, 0, 6);
	}

}