#version 430 core

layout(local_size_x = 16, local_size_y = 16) in;

layout(binding = 0, rgba32f) uniform readonly image2D u_input;
layout(binding = 1, rgba32f) uniform writeonly image2D u_output;

uniform int u_radius; // 1 => 3x3, 2 => 5x5

// Gaussian weight based on distance
float gaussian_weight(float dist, float sigma) {
    return exp(-0.5 * dist * dist / (sigma * sigma));
}

// Bilateral filter: considers both spatial distance and color similarity
void main() {
    ivec2 p = ivec2(gl_GlobalInvocationID.xy);
    ivec2 size = imageSize(u_output);
    if (p.x >= size.x || p.y >= size.y) return;

    vec3 center_color = imageLoad(u_input, p).rgb;
    
    // Sigma values for bilateral filter
    float sigma_space = float(u_radius);      // spatial sigma
    float sigma_color = 0.3;                   // color sigma (adjust for more/less smoothing)

    vec3 sum = vec3(0.0);
    float weight_sum = 0.0;

    for (int dy = -u_radius; dy <= u_radius; ++dy) {
        for (int dx = -u_radius; dx <= u_radius; ++dx) {
            ivec2 q = clamp(p + ivec2(dx, dy), ivec2(0), size - ivec2(1));
            vec3 sample_color = imageLoad(u_input, q).rgb;
            
            // Spatial weight (Gaussian based on distance)
            float spatial_dist = length(vec2(float(dx), float(dy)));
            float spatial_weight = gaussian_weight(spatial_dist, sigma_space);
            
            // Color weight (Gaussian based on color difference)
            float color_dist = length(sample_color - center_color);
            float color_weight = gaussian_weight(color_dist, sigma_color);
            
            // Combined bilateral weight
            float weight = spatial_weight * color_weight;
            
            sum += sample_color * weight;
            weight_sum += weight;
        }
    }

    vec3 out_color = sum / max(weight_sum, 1e-6);
    imageStore(u_output, p, vec4(out_color, 1.0));
}