#[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 Params 
{
  vec2 rShift;
  vec2 gShift;
  vec2 bShift;
  vec2 padding0;

  float amount;
  float shiftAll;
  float unshiftPower;
  float unshiftOffset;

} params;


float mirrorValue( float x ) 
{
    float t = mod( x, 2.0 );
    return ( t <= 1.0 ) ? t : 2.0 - t;
}

vec2 mirrorUV( vec2 uv ) 
{    
  return vec2( mirrorValue( uv.x ), mirrorValue( uv.y ) );
}


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

    vec2 uv = ( vec2( texel_coord ) + 0.5 ) / vec2( size );

    float d = length( uv - vec2( 0.5 ) ) / 0.5;
    d = pow( max( 0.0, d - params.unshiftOffset ), params.unshiftPower );
    

    vec2 uvR = uv + params.rShift * params.shiftAll * d;
    vec2 uvG = uv + params.gShift * params.shiftAll * d;
    vec2 uvB = uv + params.bShift * params.shiftAll * d;

    uvR = mirrorUV( uvR );
    uvG = mirrorUV( uvG );
    uvB = mirrorUV( uvB ); 
    

    float r = imageLoad( inputImage, ivec2( uvR * size ) ).r;
    float g = imageLoad( inputImage, ivec2( uvG * size ) ).g;
    float b = imageLoad( inputImage, ivec2( uvB * size ) ).b;

    vec4 chromaticColor = vec4( r, g, b, 1 );

    vec4 color = imageLoad( inputImage, ivec2( uv * size ) );
    vec4 blended = mix( color, chromaticColor, params.amount );


    imageStore( outputImage, texel_coord, blended );
}