#include <metal_stdlib>
using namespace metal;

/// Bilateral filter with shared memory tiling.
/// Spatial Gaussian * range Gaussian weighting for edge-preserving smoothing.

struct BilateralParams {
    uint   width;
    uint   height;
    float  sigma_space;
    float  sigma_range;
};

// Tile size for shared memory caching
constant uint TILE_SIZE = 16;
// Maximum filter radius (determined by sigma_space, capped for shared memory)
constant uint MAX_RADIUS = 8;
// Shared memory tile dimension includes apron
constant uint SHARED_DIM = TILE_SIZE + 2 * MAX_RADIUS;

kernel void bilateral_filter(
    device const float*        input    [[buffer(0)]],
    device float*              output   [[buffer(1)]],
    constant BilateralParams&  params   [[buffer(2)]],
    threadgroup float*         tile     [[threadgroup(0)]],
    uint2 gid      [[thread_position_in_grid]],
    uint2 tid      [[thread_position_in_threadgroup]],
    uint2 tg_pos   [[threadgroup_position_in_grid]],
    uint2 tg_size  [[threads_per_threadgroup]]
) {
    uint w = params.width;
    uint h = params.height;
    int radius = (int)clamp((int)ceil(params.sigma_space * 3.0f), 1, (int)MAX_RADIUS);
    float inv_2sigma_s2 = -0.5f / (params.sigma_space * params.sigma_space);
    float inv_2sigma_r2 = -0.5f / (params.sigma_range * params.sigma_range);

    // Load tile with apron into shared memory
    int tile_origin_x = (int)(tg_pos.x * TILE_SIZE) - radius;
    int tile_origin_y = (int)(tg_pos.y * TILE_SIZE) - radius;

    uint shared_dim = TILE_SIZE + 2 * (uint)radius;
    uint total_shared = shared_dim * shared_dim;
    uint linear_tid = tid.y * tg_size.x + tid.x;
    uint threads_in_group = tg_size.x * tg_size.y;

    for (uint i = linear_tid; i < total_shared; i += threads_in_group) {
        int sy = tile_origin_y + (int)(i / shared_dim);
        int sx = tile_origin_x + (int)(i % shared_dim);
        sx = clamp(sx, 0, (int)w - 1);
        sy = clamp(sy, 0, (int)h - 1);
        tile[i] = input[(uint)sy * w + (uint)sx];
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    if (gid.x >= w || gid.y >= h) return;

    float center = tile[(tid.y + (uint)radius) * shared_dim + (tid.x + (uint)radius)];
    float weight_sum = 0.0f;
    float value_sum  = 0.0f;

    for (int dy = -radius; dy <= radius; dy++) {
        for (int dx = -radius; dx <= radius; dx++) {
            uint sy = (uint)((int)tid.y + radius + dy);
            uint sx = (uint)((int)tid.x + radius + dx);
            float neighbor = tile[sy * shared_dim + sx];

            float spatial_dist = (float)(dx * dx + dy * dy);
            float range_dist  = (neighbor - center) * (neighbor - center);

            float weight = exp(spatial_dist * inv_2sigma_s2 + range_dist * inv_2sigma_r2);
            weight_sum += weight;
            value_sum  += weight * neighbor;
        }
    }

    output[gid.y * w + gid.x] = value_sum / weight_sum;
}