#[compute]
#version 450

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

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

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

layout( set = 2, binding = 0 ) 
uniform sampler2D zSampler;

layout( set = 3, binding = 0 ) 
uniform sampler2D zTresholdSampler;

layout( set = 4, binding = 0 ) 
uniform sampler2D edgeIntensitySampler;

layout( set = 5, binding = 0 ) 
uniform sampler2D normalRoughnessSampler;

layout( set = 6, binding = 0 ) 
uniform sampler2D outlineWidthSampler;

layout( push_constant, std430 ) 
uniform Parameters 
{
  float amount;
  float edgeR;
  float edgeG;
  float edgeB;

  float zMin;
  float zMax;
  float edgeThreshold;
  float edgeIntensity;

  float adaptiveScaleAmount;
  float adaptiveScaleNormalizer;
  float zEdgeAmount;
  float normalEdgeAmount;

  float normalEdgeZMin;
  float normalEdgeZMax;
  float rimOffsetX;
  float rimOffsetY;

  float fillR;
  float fillG;
  float fillB;
  float fillA;  

  float rimContrast;
  float rimStrength;
  float outlineWidth;
  float edgeDistanceFade;


} parameters;

vec3 SRGBtoLINEAR( vec3 sRGB )
{
  return mix( pow( (sRGB + vec3( 0.055 )) * ( 1.0 / ( 1.0 + 0.055 )),vec3( 2.4 )),sRGB * ( 1.0 / 12.92 ),lessThan( sRGB,vec3( 0.04045 )) );
}

vec4 SRGBtoLINEAR( vec4 sRGB )
{
  return vec4( SRGBtoLINEAR( sRGB.rgb ), sRGB.a );
}

vec3 LINEARtoSRGB( vec3 linear )
{
  vec3 color = linear;
  const vec3 a = vec3(0.055f);
	return mix((vec3(1.0f) + a) * pow(color.rgb, vec3(1.0f / 2.4f)) - a, 12.92f * color.rgb, lessThan(color.rgb, vec3(0.0031308f)));
}

vec4 LINEARtoSRGB( vec4 linear )
{
  return vec4( LINEARtoSRGB( linear.rgb ), linear.a );
}

float randomFloat( vec2 uv ) 
{
  return fract( sin( dot( uv.xy, vec2( 12.9898,78.233 ) ) ) * 43758.5453123 );
}

vec2 random( vec2 uv ) 
{
	float x = randomFloat( uv );
  float y = randomFloat( uv + vec2( 10002, -23589 ) );
 
  return vec2( x, y ) * 2.0 - vec2( 1.0, 1.0 );
}

float computeGaussianWeight( float x, float y, float sigma )
{
  return exp( - ( x * x + y * y ) / ( 2.0 * sigma * sigma ));
}

float sampleNormalizedZ( vec2 uv, int x, int y, vec2 stepSize )
{
  return textureLod( zSampler, uv + vec2( float( x ), float( y ) ) * stepSize, 0 ).r;
}

float sampleZ( vec2 uv, int x, int y, vec2 stepSize )
{
  float z = sampleNormalizedZ( uv, x, y, stepSize );

  return z * ( parameters.zMax - parameters.zMin ) + parameters.zMin; 
}

vec3 sampleNormal( vec2 uv, int x, int y, vec2 stepSize )
{
  vec3 normalRGB = textureLod( normalRoughnessSampler, uv + vec2( float( x ), float( y ) ) * stepSize, 0 ).rgb;
  vec3 normalXYZ = normalRGB * 2.0 - vec3( 1.0 );
  return normalize( normalXYZ );
}

float clamp01( float value )
{
  return clamp( value, 0.0, 1.0 );
}

float normalizeToRange( float value, float min, float max )
{
  return ( value - min ) / ( max - min );
}

float normalizeToRange01( float value, float min, float max )
{
  return clamp01( normalizeToRange( value, min, max ) ); 
}

float map( float value, float inMin, float inMax, float outMin, float outMax )
{
  return mix( outMin, outMax, normalizeToRange( value, inMin, inMax ) );
}

float mapClamped( float value, float inMin, float inMax, float outMin, float outMax )
{
  return mix( outMin, outMax, normalizeToRange01( value, inMin, inMax ) );
}

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


  
  vec2 stepSize = vec2( 1.0 ) / vec2( size ) ;
  
  vec2 uv = ( vec2( xy ) + vec2( 0.5 ) ) * stepSize;

  float z11 = sampleNormalizedZ( uv,  0,  0, stepSize ); 
  float zCenter = z11 * ( parameters.zMax * parameters.zMin ) + parameters.zMin;
  vec2 zUV = vec2( z11, 0.0 );

  float edgeDistanceFade = mix( 1.0, 1.0 - z11, parameters.edgeDistanceFade );

  float outlineWidthCurve = pow( 10.0, textureLod( outlineWidthSampler, zUV, 0 ).r * 2.0 - 1.0 );

  float outlineWidth = outlineWidthCurve * parameters.outlineWidth;

  float z00 = sampleZ( uv, -1, -1, stepSize * outlineWidth );
  float z10 = sampleZ( uv,  0, -1, stepSize * outlineWidth );
  float z20 = sampleZ( uv,  1, -1, stepSize * outlineWidth );

  float z01 = sampleZ( uv, -1,  0, stepSize * outlineWidth );
  
  float z21 = sampleZ( uv,  1,  0, stepSize * outlineWidth );

  float z02 = sampleZ( uv, -1,  1, stepSize * outlineWidth );
  float z12 = sampleZ( uv,  0,  1, stepSize * outlineWidth );
  float z22 = sampleZ( uv,  1,  1, stepSize * outlineWidth );
    
  float gx = -z00 - 2.0 * z01 - z02 + z20 + 2.0 * z21 + z22;
  float gy = -z00 - 2.0 * z10 - z20 + z02 + 2.0 * z12 + z22;

  
  vec2 rimOffset = vec2( parameters.rimOffsetX,  parameters.rimOffsetY );
  float rimZ = sampleZ( uv + rimOffset, 0, 0, stepSize );  

  float rimStrength = min( 1.0, max( 0.0, abs( zCenter - rimZ ) ) * parameters.rimContrast / zCenter  );
  rimStrength *= parameters.rimStrength;

 
  
  
  float zThresholdCurve = textureLod( zTresholdSampler, zUV, 0 ).r;
  float edgeIntensityCurve = textureLod( edgeIntensitySampler, zUV, 0 ).r;

  float adaptiveScale = zCenter / parameters.adaptiveScaleNormalizer;
  adaptiveScale = mix( 1.0, adaptiveScale, parameters.adaptiveScaleAmount );

  float zEdge = length( vec2( gx, gy ) );
  zEdge = ( zEdge - adaptiveScale * parameters.edgeThreshold * zThresholdCurve ) * parameters.edgeIntensity *  edgeIntensityCurve;
  zEdge = clamp( zEdge, 0.0, 1.0 ); 


  vec3 n10 = sampleNormal( uv,  0, -1, stepSize * outlineWidth );
  vec3 n01 = sampleNormal( uv, -1,  0, stepSize * outlineWidth );
  vec3 n11 = sampleNormal( uv,  0,  0, stepSize * outlineWidth );
  vec3 n21 = sampleNormal( uv,  1,  0, stepSize * outlineWidth );  
  vec3 n12 = sampleNormal( uv,  0,  1, stepSize * outlineWidth );  

  vec3 centerNormal = normalize( 0.5 * n11 + 0.5 * ( n10 + n01 + n21 + n12 ) / 4.0 );

  // vec3 shadowDirection = vec3( parameters.normalX, parameters.normalY, 0 );

  // float shadowAmount = max( 0, dot( centerNormal, normalize( shadowDirection ) ) );

  float normalEdge = 0.0;

  normalEdge = max( normalEdge, 1.0 - dot( n11, n01 ) );
  normalEdge = max( normalEdge, 1.0 - dot( n11, n21 ) );
  normalEdge = max( normalEdge, 1.0 - dot( n11, n10 ) );
  normalEdge = max( normalEdge, 1.0 - dot( n11, n12 ) );

  float edgeWeight = smoothstep( parameters.normalEdgeZMin, parameters.normalEdgeZMax, normalEdge );

  float edge = max( normalEdge * parameters.normalEdgeAmount, zEdge * parameters.zEdgeAmount * edgeWeight );
  edge = clamp01( edge );

  edge *= edgeDistanceFade;

  vec4 originalColor = textureLod( inputImage, uv, 0 );
  vec4 fillColor = SRGBtoLINEAR( vec4( parameters.fillR, parameters.fillG, parameters.fillB, 1.0 ) );
  originalColor = mix( originalColor, fillColor, parameters.fillA * parameters.amount );
  originalColor = mix( originalColor, vec4( 1.0 ), rimStrength * parameters.amount );

  vec4 edgeColor = SRGBtoLINEAR( vec4( parameters.edgeR, parameters.edgeG, parameters.edgeB, 1.0 ) );
  // originalColor.rgb = mix( originalColor.rgb, vec3( 0.0, 0.0, 0.0 ), shadowAmount );
  vec4 mixedColor = mix( originalColor, edgeColor, parameters.amount * edge );
  

  imageStore( outputImage, xy, mixedColor );
}