#include <metal_stdlib>
using namespace metal;

/// Per-threadgroup histogram using shared memory atomics, then reduced to global.
/// Input:  buffer of f32 pixel values (expected range [0, 1]).
/// Output: buffer of 256 uint32 bins.

struct HistogramParams {
    uint pixel_count;
    uint _pad0;
    uint _pad1;
    uint _pad2;
};

kernel void histogram_compute(
    device const float*   input       [[buffer(0)]],
    device atomic_uint*   output      [[buffer(1)]],
    constant HistogramParams& params  [[buffer(2)]],
    threadgroup atomic_uint* local_hist [[threadgroup(0)]],
    uint tid       [[thread_index_in_threadgroup]],
    uint tg_size   [[threads_per_threadgroup]],
    uint gid       [[thread_position_in_grid]]
) {
    // Initialize threadgroup histogram bins to zero
    for (uint i = tid; i < 256; i += tg_size) {
        atomic_store_explicit(&local_hist[i], 0, memory_order_relaxed);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Each thread accumulates its pixel into the local histogram
    if (gid < params.pixel_count) {
        float val = input[gid];
        uint bin = (uint)clamp(val * 255.0f, 0.0f, 255.0f);
        atomic_fetch_add_explicit(&local_hist[bin], 1, memory_order_relaxed);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // First threads in group flush local histogram to global
    for (uint i = tid; i < 256; i += tg_size) {
        uint count = atomic_load_explicit(&local_hist[i], memory_order_relaxed);
        if (count > 0) {
            atomic_fetch_add_explicit(&output[i], count, memory_order_relaxed);
        }
    }
}