#[compute]
#version 450

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

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

layout( set = 1, binding = 0 ) 
uniform sampler2D depth_sampler;


layout( push_constant, std430 )
uniform Params 
{
	vec2 rasterSize;
	vec2 amounts;
  float effectStrength;
  float edgeThreshold;
  float edgeIntensity;
  float contrastThreshold;
  float contrastIntensity;
  float debugView;

} params;

float sampleDepth( ivec2 coord ) 
{
    coord = clamp( coord, ivec2( 0 ), ivec2( params.rasterSize ) - ivec2( 1 ) );

    vec2 uv = ( vec2( coord ) + 0.5 ) / params.rasterSize;
    return texture( depth_sampler, uv ).r;
}

void main( ) 
{
	ivec2 uv   = ivec2( gl_GlobalInvocationID.xy );
	ivec2 size = ivec2( params.rasterSize );

	if ( uv.x >= size.x || uv.y >= size.y ) 
  {
		return;
	}

  float centerDepth = sampleDepth( uv );
  
  float top_left = sampleDepth( uv + ivec2( -1, -1 ) );
  float top = sampleDepth( uv + ivec2( 0, -1 ) );
  float top_right = sampleDepth( uv + ivec2( 1, -1 ) );
  float left = sampleDepth( uv + ivec2( -1, 0 ) );
  float right = sampleDepth( uv + ivec2( 1, 0 ) );
  float bottom_left = sampleDepth( uv + ivec2( -1, 1 ) );
  float bottom = sampleDepth( uv + ivec2( 0, 1 ) );
  float bottom_right = sampleDepth( uv + ivec2( 1, 1 ) );

  
  float gx = -top_left - 2.0 * left - bottom_left + top_right   + 2.0 * right  + bottom_right;
  float gy = -top_left - 2.0 * top  - top_right   + bottom_left + 2.0 * bottom + bottom_right;
    
  float edgeStrength = length( vec2( gx, gy ) );

  // float edgeStrength = abs( left - right );

  edgeStrength = clamp(  ( edgeStrength - params.edgeThreshold ) * params.edgeIntensity, 0.0, 1.0 );
  float edge = edgeStrength;


	vec4 color = imageLoad( color_image, uv );

  vec4 color_left          = imageLoad( color_image, uv + ivec2( -1, 0 ) );
  vec4 color_left2         = imageLoad( color_image, uv + ivec2( -2, 0 ) );
  vec4 color_right         = imageLoad( color_image, uv + ivec2( 1, 0 ) );
  vec4 color_right2        = imageLoad( color_image, uv + ivec2( 2, 0 ) );

  float c = ( length( color_left - color_right ) - params.contrastThreshold ) * params.contrastIntensity;

  float d = min( 1.0, c );

  edge = max( d, edge );

  vec4 smoothedColor = color_left + color_left2 +
                       2.0 * color +
                       color_right + color_right2;

  smoothedColor /= 6.0;

	float gray = color.r * 0.2125 + color.g * 0.7154 + color.b * 0.0721;
  vec3 debugViewColor = vec3( params.amounts.x * gray, params.amounts.y * centerDepth, edge );

  color.rgb = mix( color.rgb, smoothedColor.rgb, edge * params.effectStrength );
	color.rgb = mix( color.rgb, debugViewColor, params.debugView );

	imageStore( color_image, uv, color );
}