#version 450

/*
 * Kiyo data
 * - WORKGROUP_SIZE and NUM_IMAGES are provided by the engine
 */

layout ( local_size_x = WORKGROUP_SIZE, local_size_y = WORKGROUP_SIZE, local_size_z = 1 ) in;
layout( binding = 0, rgba8 ) uniform image2D images[NUM_IMAGES];
layout( push_constant ) uniform PushConstants
{
    float time;
    int in_image;
    int out_image;
} constants;

/*
 * User data
 */

void main()
{
    ivec2 p = ivec2( gl_GlobalInvocationID.xy );
    ivec2 screenSize = imageSize( images[constants.in_image] );
    if( p.x > screenSize.x || p.y > screenSize.y )
    {
        return;
    }

    float kernel[9];
    // Top row
    kernel[0] = 1.f;
    kernel[1] = 2.f;
    kernel[2] = 1.f;

    // Middle row
    kernel[3] = 2.f;
    kernel[4] = 4.f;
    kernel[5] = 2.f;

    // Bottom row
    kernel[6] = 1.f;
    kernel[7] = 2.f;
    kernel[8] = 1.f;

    float sum = 0.f;
    for(int i=0; i<9; ++i)
    {
        sum += kernel[i];
    }
    for(int i=0; i<9; ++i)
    {
        kernel[i] /= sum;
    }

    // Blur
    vec4 c = vec4( 0.0f );
    int range = 1;
    for( int x = -range; x <= range; x++ )
    {
        for( int y = -range; y <= range; y++ )
        {
            ivec2 offset = ivec2( x, y );
            float k = kernel[ x + range + ( y + range ) * 3 ];
            vec4 neighbor = imageLoad( images[constants.in_image], p + offset );
            c += neighbor;
        }
    }

    c /= pow( range * 2 + 1, 2 );

    imageStore( images[constants.out_image], p, c );
}