#[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 inputImage;
layout( rgba16f, set = 1, binding = 0 ) 
uniform image2D outputImage;


layout( push_constant, std430 ) 
uniform Parameters 
{
  vec4 searchColor;
  vec4 replacementColor;

  float amount;
  float colorDistanceMax;
  float colorDistanceMin;
  float colorSpace;

} 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 );
}

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

    vec4 color = imageLoad( inputImage, texel_coord ); 

    vec4 colorA = LINEARtoSRGB( color ) ;
    vec4 colorB = parameters.searchColor;

    float dist = length( colorA.rgb - colorB.rgb );
    dist = max( 0.0, dist - parameters.colorDistanceMin );
    dist = min( 1.0, dist / parameters.colorDistanceMax );

    float replaceAmount = 1.0 - dist;

    vec4 replacedColor = SRGBtoLINEAR( parameters.replacementColor );
    vec4 blended = mix( color, replacedColor, parameters.amount * replaceAmount );
    // blended = vec4( replaceAmount, replaceAmount, replaceAmount, 1.0 );

    imageStore( outputImage, texel_coord, blended );
}