#[compute]
#version 450

// Invocations in the (x, y, z) dimension
layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

// Our textures.
layout(r32f, set = 0, binding = 0) uniform restrict readonly image2D current_image;
layout(r32f, set = 1, binding = 0) uniform restrict writeonly image2D output_image;

// PushConstants
layout(push_constant, std430) uniform Params {
    vec2 texture_size;
    float decay_factor;
    float dissipation_factor;
} params;

float get_dissipation_from(ivec2 uv) {
    float v = imageLoad(current_image, uv).r;
    return v * params.dissipation_factor;
}

void store_value(ivec2 uv, vec4 v) {
    imageStore(output_image, uv, v);
}

// The code we want to execute in each invocation
void main() {

    ivec2 size = ivec2(params.texture_size.x - 1, params.texture_size.y - 1);
    ivec2 uv = ivec2(gl_GlobalInvocationID.xy);

    ivec2 tl = ivec2(0, 0);

    ivec2 up_uv = clamp(uv - ivec2(0, 1), tl, size);
    ivec2 down_uv = clamp(uv + ivec2(0, 1), tl, size);
    ivec2 left_uv = clamp(uv - ivec2(1, 0), tl, size);
    ivec2 right_uv = clamp(uv + ivec2(1, 0), tl, size);

    // Just in case the texture size is not divisable by 8.
    if ((uv.x > size.x) || (uv.y > size.y)) {
        return;
    }

    float current_v = imageLoad(current_image, uv).r;

    float new_v = current_v * params.decay_factor;
    new_v += get_dissipation_from(up_uv);
    new_v += get_dissipation_from(down_uv);
    new_v += get_dissipation_from(right_uv);
    new_v += get_dissipation_from(left_uv);


    new_v = clamp(new_v, 0, 1);

    vec4 result = vec4(new_v, new_v, new_v, 1.0);
    imageStore(output_image, uv, result);


}