#[compute]
#version 450


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

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

vec2 random_v2( vec2 uv ) 
{
  uv = vec2
  ( 
    dot(uv, vec2( 127.1,311.7 ) ),
    dot(uv, vec2( 269.5,183.3 ) ) 
  );
  
  return -1.0 + 2.0 * fract( sin( uv ) * 43758.5453123 );
}


vec3 random_v3( vec3 uvw )
{
	
  uvw = vec3( dot(uvw, vec3(127.1,311.7, 513.7) ),
              dot(uvw, vec3(269.5,183.3, 396.5) ),
        dot(uvw, vec3(421.3,314.1, 119.7) ) );
    
  return -1.0 + 2.0 * fract(sin(uvw) * 43758.5453123);
}

float perlin(vec2 uv) 
{
  vec2 uv_index = floor(uv);
  vec2 uv_fract = fract(uv);

  vec2 blur = smoothstep(0.0, 1.0, uv_fract);

  return mix( mix( dot( random_v2(uv_index + vec2(0.0,0.0) ), uv_fract - vec2(0.0,0.0) ),
                    dot( random_v2(uv_index + vec2(1.0,0.0) ), uv_fract - vec2(1.0,0.0) ), blur.x),
              mix( dot( random_v2(uv_index + vec2(0.0,1.0) ), uv_fract - vec2(0.0,1.0) ),
                    dot( random_v2(uv_index + vec2(1.0,1.0) ), uv_fract - vec2(1.0,1.0) ), blur.x), blur.y) + 0.5;
}

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



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( rgba16f, set = 2, binding = 0 ) 
// uniform image2D noiseImage;

layout( push_constant, std430 ) 
uniform Parameters 
{
  float distortionAmount;
  float blendAmount;  
  float horizontalEdge;
  float verticalEdge;

  float redScale;
  float greenScale;
  float blueScale;  

  float steps;
  float smear;

  float horizontalEdgePower;
  float verticalEdgePower;
  float horizontalAmount;
  float verticalAmount;
  float p0;
  float p1;
  float p2;

} parameters;


vec4 sampleChromatic( vec2 uv, vec2 dir, ivec2 size, int steps, vec3 shifts, float smear )
{
  vec4 combinedColor = vec4( 0.0 );

  float weights = 0.0;

  for ( int i = -steps; i <= steps; i++ )
  {
    float t = float( i ) / float( steps );
    float w = 1.0 - abs( t );
    vec2 stepDir = dir * pow( smear, t );

    vec2 rUV = uv + stepDir * shifts.r;
    rUV = mirrorUV( rUV );

    vec2 gUV = uv + stepDir * shifts.g;
    gUV = mirrorUV( gUV );

    vec2 bUV = uv + stepDir * shifts.b;
    bUV = mirrorUV( bUV );

    vec4 colorR = imageLoad( inputImage, ivec2( rUV * size ) );
    vec4 colorG = imageLoad( inputImage, ivec2( gUV * size ) );
    vec4 colorB = imageLoad( inputImage, ivec2( bUV * size ) );

    combinedColor.r += colorR.r * w;
    combinedColor.g += colorG.g * w;
    combinedColor.b += colorB.b * w;
    combinedColor.a += colorG.a * w;

    weights += w;
  }

  combinedColor /= weights;

  return combinedColor;
}

float triangle( float value )
{
  float inverted = 1.0 - value;

  return min( value, inverted ) * 2.0;
}

float ringFunction( float value, float roundness  )
{
  float t = triangle( mod( value, 1.0 ) );
  float c =  0.5 - 0.5 * cos( 6.28318531 * value );

  return clamp( mix( t, c, roundness ), 0.0, 1.0 );
}

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

  

  vec2 dirX = vec2( 0.0, 0.0 );
  vec2 dirY = vec2( 0.0, 0.0 );

  float xAmount = 0.0;
  float yAmount = 0.0;

  if ( uv.x < parameters.horizontalEdge )
  {
    dirX.x = 1.0;
    xAmount = 1.0 - uv.x / parameters.horizontalEdge;
  }
  else if ( uv.x > ( 1.0 - parameters.horizontalEdge ) )
  {
    dirX.x = -1.0;
    xAmount = 1.0 - ( 1.0 - uv.x ) / parameters.horizontalEdge;
  }

  if ( uv.y < parameters.verticalEdge )
  {
    dirY.y = 1.0;
    yAmount = 1.0 - uv.y / parameters.verticalEdge;
  }
  else if ( uv.y > ( 1.0 - parameters.verticalEdge ) )
  {
    dirY.y = -1.0;
    yAmount = 1.0 - ( 1.0 - uv.y ) / parameters.verticalEdge;
  }

  xAmount = pow( xAmount, parameters.horizontalEdgePower ) * parameters.horizontalAmount;
  yAmount = pow( yAmount, parameters.verticalEdgePower ) * parameters.verticalAmount;
  
  vec2 uvOffsetX =  dirX * xAmount; 
  vec2 uvOffsetY =  dirY * yAmount; 
  vec3 distortion = vec3( parameters.redScale, parameters.greenScale, parameters.blueScale ) * parameters.distortionAmount;

  int stepsI = min( 32, int( parameters.steps ) );  

  vec4 originalColor = imageLoad( inputImage, texel_coord );
  
  vec2 uvOffset = uvOffsetX;

  if ( yAmount > 0 )
  {
    uvOffset = ( xAmount * uvOffsetX + yAmount * uvOffsetY ) / ( yAmount + xAmount );
  }

  vec4 color = sampleChromatic( uv, uvOffset, size, stepsI, distortion, parameters.smear );


  color = mix( originalColor, color, parameters.blendAmount );
  imageStore( outputImage, texel_coord, color );
}