#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;
layout(binding = 2, rgba32f) uniform readonly image2D u_history;

uniform int u_radius;               // 1 => 3x3, 2 => 5x5
uniform float u_temporal_weight;     // 0 = no temporal, 1 = full history
uniform bool u_has_history;          // Whether history texture is valid

// 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 denoised_color = sum / max(weight_sum, 1e-6);
    
    // Temporal accumulation: blend with history
    if (u_has_history && u_temporal_weight > 0.0) {
        vec3 history_color = imageLoad(u_history, p).rgb;
        
        // Simple exponential moving average
        // temporal_weight controls how much history to keep (higher = smoother but more ghosting)
        vec3 final_color = mix(denoised_color, history_color, u_temporal_weight);
        
        imageStore(u_output, p, vec4(final_color, 1.0));
    } else {
        imageStore(u_output, p, vec4(denoised_color, 1.0));
    }
}
