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

vec3 applyMatrix( vec3 v, mat4 m )
{
  return ( m * vec4( v, 1.0 ) ).xyz;
}

vec3 applyMatrixWithoutTranslation( vec3 v, mat4 m )
{
  return ( m * vec4( v, 0.0 ) ).xyz;
}

vec3 applyMatrix( vec3 v, mat3 m )
{
  return ( m * v ).xyz;
}


vec3 localToWorld( vec3 local, mat4 _MODEL_MATRIX )
{
  return applyMatrix( local, _MODEL_MATRIX );
}

vec3 localToWorldDirection( vec3 local, mat4 _MODEL_MATRIX )
{
  return applyMatrixWithoutTranslation( local, _MODEL_MATRIX );
}

vec3 localToView( vec3 local, mat4 _MODELVIEW_MATRIX )
{
  return applyMatrix( local, _MODELVIEW_MATRIX );
}

vec3 localToViewDirection( vec3 local, mat4 _MODELVIEW_MATRIX )
{
  return applyMatrixWithoutTranslation( local, _MODELVIEW_MATRIX );
}

vec3 localToViewDirection( vec3 local, mat3 _MODELVIEW_NORMAL_MATRIX )
{
  return applyMatrix( local, _MODELVIEW_NORMAL_MATRIX );
}

vec3 worldToLocal( vec3 world, mat4 _MODEL_MATRIX )
{
  mat4 inversedMatrix = inverse( _MODEL_MATRIX );
  return applyMatrix( world, inversedMatrix );
}

vec3 worldToLocalDirection( vec3 world, mat4 _MODEL_MATRIX )
{
  mat4 inversedMatrix = inverse( _MODEL_MATRIX );
  return applyMatrixWithoutTranslation( world, inversedMatrix );
}


vec3 worldToView( vec3 world, mat4 _VIEW_MATRIX )
{  
  return applyMatrix( world, _VIEW_MATRIX );
}

vec3 worldToViewDirection( vec3 worldDirection, mat4 _VIEW_MATRIX )
{
  return applyMatrixWithoutTranslation( worldDirection, _VIEW_MATRIX );
}

vec3 viewToWorld( vec3 view, mat4 _INV_VIEW_MATRIX )
{
  return applyMatrix( view, _INV_VIEW_MATRIX );
}

vec3 viewToWorldDirection( vec3 viewDirection, mat4 _INV_VIEW_MATRIX )
{
  return applyMatrixWithoutTranslation( viewDirection, _INV_VIEW_MATRIX );
}

vec3 viewToLocal( vec3 view, mat4 _INV_VIEW_MATRIX, mat4 _MODEL_MATRIX )
{
  vec3 world = viewToWorld( view, _INV_VIEW_MATRIX );
  return worldToLocal( world, _MODEL_MATRIX );
}

vec3 viewToLocal( vec3 view, mat4 _MODELVIEW_MATRIX )
{
  mat4 inversedMatrix = inverse( _MODELVIEW_MATRIX );
  return applyMatrix( view, inversedMatrix );
}

vec3 viewToLocalDirection( vec3 viewDirection, mat4 _MODELVIEW_MATRIX )
{
  mat4 inversedMatrix = inverse( _MODELVIEW_MATRIX );
  return applyMatrixWithoutTranslation( viewDirection, inversedMatrix );
}

vec3 viewToLocalDirection( vec3 view, mat4 _INV_VIEW_MATRIX, mat4 _MODEL_MATRIX )
{
  vec3 world = viewToWorldDirection( view, _INV_VIEW_MATRIX );
  return worldToLocalDirection( world, _MODEL_MATRIX );
}

vec4 viewToClip( vec3 view, mat4 _PROJECTION_MATRIX )
{
  vec4 clip = _PROJECTION_MATRIX * vec4( view, 1.0 ); 
  clip /= clip.w;
  return clip;
}

vec2 clipToScreen( vec4 clip )
{
  return ( clip.xy / clip.w ) * 0.5 + vec2( 0.5 );
}

vec4 screenToClip( vec2 screen, float z )
{
  return vec4( screen * 2.0 - 1.0, z, 1.0 );
}

vec4 clipToView( vec4 clip, mat4 _INV_PROJECTION_MATRIX )
{
  vec4 view = _INV_PROJECTION_MATRIX * clip;
  view /= view.w;
  return view;
}

vec4 screenToView( vec2 screen, float z, mat4 _INV_PROJECTION_MATRIX )
{
  vec4 clip = screenToClip( screen, z );
  vec4 view = clipToView( clip, _INV_PROJECTION_MATRIX );

  return view;
}

vec3 screenToWorld( vec2 screen, float z, mat4 _INV_PROJECTION_MATRIX, mat4 _INV_VIEW_MATRIX  )
{
  vec4 view = screenToView( screen, z, _INV_PROJECTION_MATRIX );

  return viewToWorld( view.xyz, _INV_VIEW_MATRIX );
}

vec3 screenToLocal( vec2 screen_point, float z, mat4 _PROJECTION_MATRIX, mat4 _VIEW_MATRIX, mat4 _MODEL_MATRIX )
{
  vec4 clip = vec4( screen_point * 2.0 - 1.0, z, 1.0 );
  vec4 view4 = inverse( _PROJECTION_MATRIX ) * clip;
  view4 /= view4.w;
  vec3 view = view4.xyz;
  mat4 m = inverse( _VIEW_MATRIX * _MODEL_MATRIX );
  return ( m * vec4( view, 1.0 ) ).xyz;
}

vec2 viewToScreen( vec3 view, mat4 _PROJECTION_MATRIX )
{
  vec4 clip = viewToClip( view, _PROJECTION_MATRIX );

  return clipToScreen( clip );
}

vec2 localToScreen( vec3 local, mat4 _MODELVIEW_MATRIX, mat4 _PROJECTION_MATRIX )
{
  vec3 view = localToView( local, _MODELVIEW_MATRIX );
  return viewToScreen( view, _PROJECTION_MATRIX );
}

vec2 worldToScreen( vec3 world, mat4 _VIEW_MATRIX, mat4 _PROJECTION_MATRIX )
{
  vec3 view   = worldToView( world, _VIEW_MATRIX );
  vec4 clip   = viewToClip( view, _PROJECTION_MATRIX );
  vec2 screen = clipToScreen( clip );

  return screen;
}


vec3 extractTranslation( mat4 matrix )
{
  return vec3( matrix[ 3 ][ 0 ], matrix[ 3 ][ 1 ], matrix[ 3 ][ 2 ] );
}

vec3 extractScale( mat3 matrix )
{
  mat3 m = matrix;
  
  float x = length( vec3( m[ 0 ][ 0 ], m[ 1 ][ 0 ], m[ 2 ][ 0 ] ) ); 
  float y = length( vec3( m[ 0 ][ 1 ], m[ 1 ][ 1 ], m[ 2 ][ 1 ] ) ); 
  float z = length( vec3( m[ 0 ][ 2 ], m[ 1 ][ 2 ], m[ 2 ][ 2 ] ) ); 

  return vec3( x, y, z );
}

vec3 extractScale( mat4 matrix )
{
  mat4 m = matrix;
  
  float x = length( vec3( m[ 0 ][ 0 ], m[ 1 ][ 0 ], m[ 2 ][ 0 ] ) ); 
  float y = length( vec3( m[ 0 ][ 1 ], m[ 1 ][ 1 ], m[ 2 ][ 1 ] ) ); 
  float z = length( vec3( m[ 0 ][ 2 ], m[ 1 ][ 2 ], m[ 2 ][ 2 ] ) ); 

  return vec3( x, y, z );
}

mat3 extractRotationMatrix( mat4 m )
{
  vec3 x = normalize( m[ 0 ].xyz ); 
  vec3 y = normalize( m[ 1 ].xyz ); 
  vec3 z = normalize( m[ 2 ].xyz ); 

  return mat3( x, y, z );
}

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


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 smearSteps;
  float smear;
  float ringSteps;
  float ringSizeX;

  float ringSizeY;
  float distortionAmount;
  float blendMode;
  float p3;
  

} parameters;


vec4 sampleChromatic( vec2 uv, vec2 dir, 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 ) ) );
}



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.smearSteps ) );  
  int ringSteps = min( 32, int( parameters.ringSteps ) );  

  vec4 color = vec4( 0.0 );

  float angleStep = 6.283185307 / float( ringSteps ); 


  float combinedZ = 0;

  float ratio = float( size.y ) / float( size.x );

  for ( int i = 0; i < ringSteps; i++ )
  {
    float angle = float( i ) * angleStep;
    vec2 uvDir = vec2( cos( angle ) * ratio * parameters.ringSizeX, sin( angle )  * parameters.ringSizeY  );
    vec2 uvOffset = uvDir * parameters.distortionAmount;
 
    vec4 chromatic = sampleChromatic( uv + uvOffset, uvDir, stepsI, distortion, parameters.smear );
    color.rgb += chromatic.rgb;    
  } 	
 
  color.rgb /= float( ringSteps );  

  // vec4 blendedColor = imageLoad( outputImage, xy );
  // blendedColor.rgb += color.rgb * parameters.amount;

  vec4 originalColor = imageLoad( outputImage, xy );
  vec4 blendedColor = compositeHDR_glow( originalColor, color * parameters.amount, parameters.blendMode );

  imageStore( outputImage, xy, blendedColor );
}
