#[compute]
#version 450

layout( local_size_x = 8, local_size_y = 8, local_size_z = 1 ) in;


layout( set = 0, binding = 0 ) 
uniform sampler2D sourceColor;

layout( rgba16f, set = 1, binding = 0 ) 
uniform image2D outputImage;

layout( push_constant, std430 ) 
uniform Parameters 
{
  float luminance_multiplier;
  float exposure;
  float amount;
  float pad1;

} parameters;


float QUALITY( float q ) 
{
	return ( q < 5 ? 1.0 : ( q > 5 ? ( q < 10 ? 2.0 : ( q < 11 ? 4.0 : 8.0 ) ) : 1.5 ) );
}

float rgb2luma( vec3 rgb ) 
{
	return sqrt( dot( rgb, vec3( 0.299, 0.587, 0.114 ) ) );
}

vec3 applyFXAA( vec3 color, float exposure, vec2 uv_interp, vec2 pixel_size ) 
{
	const float EDGE_THRESHOLD_MIN = 0.0312;
	const float EDGE_THRESHOLD_MAX = 0.125;
	const int ITERATIONS = 12;
	const float SUBPIXEL_QUALITY = 0.75;
	
	float lumaUp = rgb2luma( textureLodOffset( sourceColor, uv_interp, 0.0, ivec2( 0, 1 ) ).xyz * exposure * parameters.luminance_multiplier );
	float lumaDown = rgb2luma( textureLodOffset( sourceColor, uv_interp, 0.0, ivec2( 0, -1 ) ).xyz * exposure * parameters.luminance_multiplier );
	float lumaLeft = rgb2luma( textureLodOffset( sourceColor, uv_interp, 0.0, ivec2( -1, 0 ) ).xyz * exposure * parameters.luminance_multiplier );
	float lumaRight = rgb2luma( textureLodOffset( sourceColor, uv_interp, 0.0, ivec2( 1, 0 ) ).xyz * exposure * parameters.luminance_multiplier );

	float lumaCenter = rgb2luma( color );

	float lumaMin = min( lumaCenter, min( min( lumaUp, lumaDown ), min( lumaLeft, lumaRight ) ) );
	float lumaMax = max( lumaCenter, max( max( lumaUp, lumaDown ), max( lumaLeft, lumaRight ) ) );

	float lumaRange = lumaMax - lumaMin;

	if ( lumaRange < max( EDGE_THRESHOLD_MIN, lumaMax * EDGE_THRESHOLD_MAX ) ) 
  {
		return color;
	}

	float lumaDownLeft = rgb2luma( textureLodOffset( sourceColor, uv_interp, 0.0, ivec2( -1, -1 ) ).xyz * exposure * parameters.luminance_multiplier );
	float lumaUpRight = rgb2luma( textureLodOffset( sourceColor, uv_interp, 0.0, ivec2( 1, 1 ) ).xyz * exposure * parameters.luminance_multiplier );
	float lumaUpLeft = rgb2luma( textureLodOffset( sourceColor, uv_interp, 0.0, ivec2( -1, 1 ) ).xyz * exposure * parameters.luminance_multiplier );
	float lumaDownRight = rgb2luma( textureLodOffset( sourceColor, uv_interp, 0.0, ivec2( 1, -1 ) ).xyz * exposure * parameters.luminance_multiplier );

	float lumaDownUp = lumaDown + lumaUp;
	float lumaLeftRight = lumaLeft + lumaRight;

	float lumaLeftCorners = lumaDownLeft + lumaUpLeft;
	float lumaDownCorners = lumaDownLeft + lumaDownRight;
	float lumaRightCorners = lumaDownRight + lumaUpRight;
	float lumaUpCorners = lumaUpRight + lumaUpLeft;

	float edgeHorizontal = abs( -2.0 * lumaLeft + lumaLeftCorners ) + abs( -2.0 * lumaCenter + lumaDownUp ) * 2.0 + abs( -2.0 * lumaRight + lumaRightCorners );
	float edgeVertical = abs( -2.0 * lumaUp + lumaUpCorners ) + abs( -2.0 * lumaCenter + lumaLeftRight ) * 2.0 + abs( -2.0 * lumaDown + lumaDownCorners );

	bool isHorizontal = ( edgeHorizontal >= edgeVertical );

	float stepLength = isHorizontal ? pixel_size.y : pixel_size.x;

	float luma1 = isHorizontal ? lumaDown : lumaLeft;
	float luma2 = isHorizontal ? lumaUp : lumaRight;
	float gradient1 = luma1 - lumaCenter;
	float gradient2 = luma2 - lumaCenter;

	bool is1Steepest = abs( gradient1 ) >= abs( gradient2 );

	float gradientScaled = 0.25 * max( abs( gradient1 ), abs( gradient2 ) );

	float lumaLocalAverage = 0.0;
	if ( is1Steepest ) {
		stepLength = -stepLength;
		lumaLocalAverage = 0.5 * ( luma1 + lumaCenter );
	} else {
		lumaLocalAverage = 0.5 * ( luma2 + lumaCenter );
	}

	vec2 currentUv = uv_interp;
	if ( isHorizontal ) {
		currentUv.y += stepLength * 0.5;
	} else {
		currentUv.x += stepLength * 0.5;
	}

	vec2 offset = isHorizontal ? vec2( pixel_size.x, 0.0 ) : vec2( 0.0, pixel_size.y );
	vec2 uv1 = currentUv - offset * QUALITY( 0 );
	vec2 uv2 = currentUv + offset * QUALITY( 0 );

	float lumaEnd1 = rgb2luma( textureLod( sourceColor, uv1, 0.0 ).xyz * exposure * parameters.luminance_multiplier );
	float lumaEnd2 = rgb2luma( textureLod( sourceColor, uv2, 0.0 ).xyz * exposure * parameters.luminance_multiplier );
	lumaEnd1 -= lumaLocalAverage;
	lumaEnd2 -= lumaLocalAverage;

	bool reached1 = abs( lumaEnd1 ) >= gradientScaled;
	bool reached2 = abs( lumaEnd2 ) >= gradientScaled;
	bool reachedBoth = reached1 && reached2;

	if ( !reached1 ) {
		uv1 -= offset * QUALITY( 1 );
	}
	if ( !reached2 ) {
		uv2 += offset * QUALITY( 1 );
	}

	if ( !reachedBoth ) {
		for ( int i = 2; i < ITERATIONS; i++ ) {
			if ( !reached1 ) {
				lumaEnd1 = rgb2luma( textureLod( sourceColor, uv1, 0.0 ).xyz * exposure * parameters.luminance_multiplier );
				lumaEnd1 = lumaEnd1 - lumaLocalAverage;
			}
			if ( !reached2 ) {
				lumaEnd2 = rgb2luma( textureLod( sourceColor, uv2, 0.0 ).xyz * exposure * parameters.luminance_multiplier );
				lumaEnd2 = lumaEnd2 - lumaLocalAverage;
			}
			reached1 = abs( lumaEnd1 ) >= gradientScaled;
			reached2 = abs( lumaEnd2 ) >= gradientScaled;
			reachedBoth = reached1 && reached2;
			if ( !reached1 ) {
				uv1 -= offset * QUALITY( i );
			}
			if ( !reached2 ) {
				uv2 += offset * QUALITY( i );
			}
			if ( reachedBoth ) {
				break;
			}
		}
	}

	float distance1 = isHorizontal ? ( uv_interp.x - uv1.x ) : ( uv_interp.y - uv1.y );
	float distance2 = isHorizontal ? ( uv2.x - uv_interp.x ) : ( uv2.y - uv_interp.y );

	bool isDirection1 = distance1 < distance2;
	float distanceFinal = min( distance1, distance2 );

	float edgeThickness = ( distance1 + distance2 );

	bool isLumaCenterSmaller = lumaCenter < lumaLocalAverage;

	bool correctVariation1 = ( lumaEnd1 < 0.0 ) != isLumaCenterSmaller;
	bool correctVariation2 = ( lumaEnd2 < 0.0 ) != isLumaCenterSmaller;

	bool correctVariation = isDirection1 ? correctVariation1 : correctVariation2;

	float pixelOffset = -distanceFinal / edgeThickness + 0.5;

	float finalOffset = correctVariation ? pixelOffset : 0.0;

	float lumaAverage = ( 1.0 / 12.0 ) * ( 2.0 * ( lumaDownUp + lumaLeftRight ) + lumaLeftCorners + lumaRightCorners );

	float subPixelOffset1 = clamp( abs( lumaAverage - lumaCenter ) / lumaRange, 0.0, 1.0 );
	float subPixelOffset2 = ( -2.0 * subPixelOffset1 + 3.0 ) * subPixelOffset1 * subPixelOffset1;

	float subPixelOffsetFinal = subPixelOffset2 * subPixelOffset2 * SUBPIXEL_QUALITY;

	finalOffset = max( finalOffset, subPixelOffsetFinal );

	vec2 finalUv = uv_interp;
	if ( isHorizontal ) 
  {
		finalUv.y += finalOffset * stepLength;
	} 
  else 
  {
		finalUv.x += finalOffset * stepLength;
	}

	vec3 finalColor = textureLod( sourceColor, finalUv, 0.0 ).xyz * exposure * parameters.luminance_multiplier;
	return finalColor;

}


void main() 
{
  ivec2 size = imageSize( outputImage );
  ivec2 xy = ivec2( gl_GlobalInvocationID.xy );
  
  if ( any( greaterThanEqual( xy, size ) )  ) 
  {
    return;
  }

  vec2 uv = ( vec2( xy ) + vec2( 0.5 ) ) / vec2( size );
  
  float exposure = parameters.exposure;

  vec4 color = textureLod( sourceColor, uv, 0 );
  vec4 original = color;

  color.rgb = applyFXAA( color.rgb, exposure, uv, 1.0 / size );

  color = mix( original, color, parameters.amount );
  
  imageStore( outputImage, xy, color );
}