#[compute]
#version 450

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

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

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


layout( push_constant, std430 )
uniform Parameters 
{
  float amount;
  float edgeThreshold;
  float edgeIntensity;
  float blurMix;

  float blurType;
  float spread;
  int pad;
  int pad2;

} parameters;

float sampleDepth( ivec2 xy, vec2 pixelSize ) 
{
  vec2 uv = ( vec2( xy ) + 0.5 ) * pixelSize;
  return texture( depthSampler, uv ).r;
}

void main( ) 
{
	ivec2 xy   = ivec2( gl_GlobalInvocationID.xy );
	ivec2 size = imageSize( outputImage );

	if ( xy.x >= size.x || xy.y >= size.y ) 
  {
		return;
	}
  
  vec2 pixelSize = vec2( 1.0 ) / vec2( size );

  float centerDepth = sampleDepth( xy, pixelSize );
  
  float topLeft     = sampleDepth( xy + ivec2( -1, -1 ), pixelSize );
  float top         = sampleDepth( xy + ivec2(  0, -1 ), pixelSize );
  float topRight    = sampleDepth( xy + ivec2(  1, -1 ), pixelSize );

  float left        = sampleDepth( xy + ivec2( -1, 0 ), pixelSize );
  float right       = sampleDepth( xy + ivec2(  1, 0 ), pixelSize );

  float bottomLeft  = sampleDepth( xy + ivec2( -1, 1 ), pixelSize );
  float bottom      = sampleDepth( xy + ivec2(  0, 1 ), pixelSize );
  float bottomRight = sampleDepth( xy + ivec2(  1, 1 ), pixelSize );

  
  float gx = -topLeft - 2.0 * left - bottomLeft + topRight   + 2.0 * right  + bottomRight;
  float gy = -topLeft - 2.0 * top  - topRight   + bottomLeft + 2.0 * bottom + bottomRight;
    

  float edge = length( vec2( gx, gy ) );
  edge = ( edge - parameters.edgeThreshold ) * parameters.edgeIntensity;
  edge = clamp( edge, 0.0, 1.0 ); 


	vec4 color = imageLoad( outputImage, xy );

  int spread = int( parameters.spread );
  vec4 color0 = imageLoad( outputImage, xy + ivec2( -1,  0 ) * spread );
  vec4 color1 = imageLoad( outputImage, xy + ivec2(  1,  0 ) * spread );
  
  ivec2 dir2 = int( parameters.blurType ) == 0 ? ivec2( 2, 0 ) : ivec2( 0, 1 );

  vec4 color2 = imageLoad( outputImage, xy - dir2 * spread );
  vec4 color3 = imageLoad( outputImage, xy + dir2 * spread );

  vec4 blurColor =  0.35 * ( color0 + color1 ) + ( color2 + color3 ) * 0.15;

  vec4 smoothedColor = mix( color, blurColor, parameters.blurMix );

  color.rgb = mix( color.rgb, smoothedColor.rgb, edge * parameters.amount );

	imageStore( outputImage, xy, color );
}