#[compute]
#version 450

#include "res://addons/rokojori_action_library/Runtime/Shading/Library/Colors.gdshaderinc"

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( set = 0, binding = 0 ) 
uniform sampler2D inputImage;

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



layout( push_constant, std430 ) 
uniform Parameters 
{
  float amount;
  float redScale;
  float greenScale;
  float blueScale;

  float steps;
  float smear;
  float distortionAmount;
  float distortionAngle;

  float uvOffsetX; 
  float uvOffsetY;
  float blendMode;
  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 = textureLod( inputImage, rUV, 0 );
    vec4 colorG = textureLod( inputImage, gUV, 0 );
    vec4 colorB = textureLod( inputImage, bUV, 0 );

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



vec2 rotate( vec2 uv, float angle )
{
  float s = sin( angle );
  float c = cos( angle );

    
  float x = uv.x;
  float y = uv.y; 
 
  uv.x = c * x - s * y;
  uv.y = s * x + c * y;

  return uv;
}

float rgb2luma( vec3 rgb ) 
{
	return sqrt( dot( rgb, vec3( 0.299, 0.587, 0.114 ) ) );
}

vec4 compositeHDR_glow( vec4 originalColor, vec4 light, float blendAmount )
{
  vec3 additiveColor   = compositeHDR_add( originalColor.rgb, light.rgb );
  vec3 colorBlendColor = compositeHDR_color( originalColor.rgb, light.rgb, 1.0, 5.0 );

  vec4 resultColor = vec4( mix( additiveColor, colorBlendColor, blendAmount ), originalColor.a );

  return resultColor;
}

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

  if ( any( greaterThanEqual( xy, size ) )  ) 
  {
    return;
  }

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

  vec3 distortion = vec3( parameters.redScale, parameters.greenScale, parameters.blueScale ) * parameters.distortionAmount;

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

  vec4 originalColor = imageLoad( outputImage, xy );

  float angle = parameters.distortionAngle;
  vec2 ratio = vec2( 1.0, float( size.y ) / float( size.x ) );

  vec2 uvDir = vec2( cos( angle ), sin( angle ) );



  vec2 uvOffset = vec2( parameters.uvOffsetX, parameters.uvOffsetY );

  uvOffset += -uvDir * parameters.distortionAmount * parameters.smear;

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

  // float luma = rgb2luma( color.rgb );
  // float lumaScale = max( 0.0, luma - parameters.lumaTreshold ) * parameters.lumaScale;
  // lumaScale = min( lumaScale, parameters.lumaClamp ) * parameters.lumaStrength;

  // vec3 lumaTint = vec3( parameters.lumaTintR, parameters.lumaTintG, parameters.lumaTintB );
  
  // vec4 blendedColor = originalColor;
  // blendedColor.rgb += color.rgb * lumaTint * lumaScale * parameters.amount; 

  // vec4 blendedColor = originalColor;
  // blendedColor.rgb += color.rgb * parameters.amount; 

  vec4 resultColor = compositeHDR_glow( originalColor, color * parameters.amount, parameters.blendMode );

  imageStore( outputImage, xy, resultColor );
}
