// crates/ml/src/cuda_pipeline/sp4_histogram_p99.cuh
//
// SP4 dynamic-range linear-histogram p99 estimator.
//
// Single-block, BLOCK_SIZE-thread, three-pass:
//   Pass 1: block-wide max-reduce of |buf| → step_max
//   Pass 2: linear-spaced [0, step_max] binning into 256 bins via per-warp
//           tiles in shared memory (no atomicAdd, per
//           feedback_no_atomicadd.md)
//   Pass 3: cumulative-from-top → p99 = bin upper-edge
//
// All passes use __syncthreads()/__syncwarp(); no atomicAdd anywhere.
// Returns 0.0f if step_max == 0.0f (degenerate: caller skips ISV update).
//
// Linear spacing is chosen over log because SP4 producers care about
// resolution at the *top* of the |buf| distribution (the 99th percentile).
// Top bin width = step_max / 256 ≈ 0.39% of step_max — well within the
// 1% quantile precision budget.
//
// Shared-memory contract for the caller:
//   - __shared__ static: 256 × sizeof(int) for s_bins (and float reduce
//     storage in pass 1 via __float_as_int reuse) + 1 × sizeof(float)
//     for s_step_max → 1028 bytes static.
//   - __shared__ dynamic (caller passes via launch config):
//     ≥ (BLOCK_SIZE / 32) × 256 × sizeof(int) for the per-warp tiles.
//     With BLOCK_SIZE=256, that is 8 warps × 256 ints × 4 bytes = 8192
//     bytes.
//
// Caller MUST launch with grid_dim = (1, 1, 1), block_dim = (BLOCK_SIZE,
// 1, 1). BLOCK_SIZE must be a positive multiple of warpSize (32).

#pragma once

#define SP4_HIST_BINS 256

template <int BLOCK_SIZE>
__device__ float sp4_histogram_p99(const float* __restrict__ buf, int count) {
    static_assert(BLOCK_SIZE >= 32 && (BLOCK_SIZE % 32) == 0,
        "BLOCK_SIZE must be a positive multiple of 32 (warp size)");

    __shared__ int   s_bins[SP4_HIST_BINS];
    __shared__ float s_step_max;

    const int tid = threadIdx.x;

    // ── Pass 1: block-wide max-reduce of |buf| ──────────────────────────
    // Each thread folds its strided slice into local_max, then a shared-
    // memory tree reduces across the block. We reuse s_bins[0..BLOCK_SIZE)
    // as a `float`-via-`int`-reinterpret store (BLOCK_SIZE ≤ 256 = bin
    // count, so it fits).
    float local_max = 0.0f;
    for (int i = tid; i < count; i += BLOCK_SIZE) {
        local_max = fmaxf(local_max, fabsf(buf[i]));
    }
    s_bins[tid] = __float_as_int(local_max);
    __syncthreads();
    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            float a = __int_as_float(s_bins[tid]);
            float b = __int_as_float(s_bins[tid + s]);
            s_bins[tid] = __float_as_int(fmaxf(a, b));
        }
        __syncthreads();
    }
    if (tid == 0) s_step_max = __int_as_float(s_bins[0]);
    __syncthreads();

    if (s_step_max == 0.0f) return 0.0f;   // degenerate: no signal magnitude

    const float step_max  = s_step_max;
    const float bin_width = step_max / (float)SP4_HIST_BINS;

    // ── Pass 2: linear-bin |buf| via per-warp tiles (no atomicAdd) ─────
    // Each warp owns its own 256-int tile in dynamic shared memory; lanes
    // within the warp may collide, but the rate is bounded by 1/(256×32)
    // per binning event (≈0.012% expected count loss for typical signals),
    // well within the 1% quantile precision budget. Cross-warp collisions
    // are eliminated by the per-warp tiling.
    extern __shared__ int s_warp_tiles[];
    const int warp_id = tid >> 5;
    const int lane    = tid & 31;
    const int warps   = BLOCK_SIZE >> 5;

    // Zero this warp's tile (lanes cooperate: 32 lanes × 8 ints = 256).
    for (int b = lane; b < SP4_HIST_BINS; b += 32) {
        s_warp_tiles[warp_id * SP4_HIST_BINS + b] = 0;
    }
    __syncwarp();

    // Each thread bins its strided slice into its warp's tile.
    for (int i = tid; i < count; i += BLOCK_SIZE) {
        int bin_idx = (int)floorf(fabsf(buf[i]) / bin_width);
        if (bin_idx >= SP4_HIST_BINS) bin_idx = SP4_HIST_BINS - 1;
        s_warp_tiles[warp_id * SP4_HIST_BINS + bin_idx] += 1;
    }
    __syncthreads();

    // Tree-reduce warp tiles into s_bins[].
    for (int b = tid; b < SP4_HIST_BINS; b += BLOCK_SIZE) {
        int sum = 0;
        for (int w = 0; w < warps; ++w) sum += s_warp_tiles[w * SP4_HIST_BINS + b];
        s_bins[b] = sum;
    }
    __syncthreads();

    // ── Pass 3: cumulative-from-top → p99 (bin upper-edge) ─────────────
    if (tid == 0) {
        const int target = (count + 99) / 100;   // ceil(count / 100); top 1%
        int cumul   = 0;
        int p99_bin = SP4_HIST_BINS - 1;
        for (int b = SP4_HIST_BINS - 1; b >= 0; --b) {
            cumul += s_bins[b];
            if (cumul >= target) { p99_bin = b; break; }
        }
        // p99 = upper edge of p99_bin = (p99_bin + 1) × bin_width.
        s_step_max = (float)(p99_bin + 1) * bin_width;
    }
    __syncthreads();
    return s_step_max;
}

// ── Multi-sub-buffer variant (SP4 Task A7 fix-up #2) ─────────────────────
//
// `sp4_histogram_p99_multi` treats the union of `n_sub` sub-buffers as one
// logical distribution. Sub-buffer pointers + counts are read from device
// arrays (typically mapped-pinned tables populated host-side) so the
// caller can dispatch with `n_sub=1` (the contiguous-buffer common case;
// SP4 groups 0-6) or `n_sub=4` (Curiosity's `[w1, b1, w2, b2]` layout).
//
// Pass structure mirrors the single-buffer template:
//   - Pass 1: max-reduce across all `n_sub` sub-buffers.
//   - Pass 2: linear-bin all `n_sub` sub-buffers into the same per-warp
//             tile space; binned counts reflect the union distribution.
//   - Pass 3: cumulative-from-top sees `total_count` (sum of sub_counts);
//             yields the same p99 semantics as the single-buffer template
//             when n_sub == 1.
//
// `sub_buf_ptrs`/`sub_counts` must be readable from the device (mapped-
// pinned device-pointer is the canonical case). `n_sub` is a small
// integer (≤ 8 in practice) so the per-pass `for (s)` loops add no
// material overhead vs the inner per-element loops.
template <int BLOCK_SIZE>
__device__ float sp4_histogram_p99_multi(
    const unsigned long long* __restrict__ sub_buf_ptrs,
    const int*                __restrict__ sub_counts,
    int                                    n_sub,
    int                                    total_count
) {
    static_assert(BLOCK_SIZE >= 32 && (BLOCK_SIZE % 32) == 0,
        "BLOCK_SIZE must be a positive multiple of 32 (warp size)");

    __shared__ int   s_bins[SP4_HIST_BINS];
    __shared__ float s_step_max;

    const int tid = threadIdx.x;

    // ── Pass 1: block-wide max-reduce of |buf| across all sub-buffers ──
    float local_max = 0.0f;
    for (int s = 0; s < n_sub; ++s) {
        const float* buf_s = reinterpret_cast<const float*>(sub_buf_ptrs[s]);
        const int    cnt_s = sub_counts[s];
        for (int i = tid; i < cnt_s; i += BLOCK_SIZE) {
            local_max = fmaxf(local_max, fabsf(buf_s[i]));
        }
    }
    s_bins[tid] = __float_as_int(local_max);
    __syncthreads();
    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            float a = __int_as_float(s_bins[tid]);
            float b = __int_as_float(s_bins[tid + s]);
            s_bins[tid] = __float_as_int(fmaxf(a, b));
        }
        __syncthreads();
    }
    if (tid == 0) s_step_max = __int_as_float(s_bins[0]);
    __syncthreads();

    if (s_step_max == 0.0f) return 0.0f;

    const float step_max  = s_step_max;
    const float bin_width = step_max / (float)SP4_HIST_BINS;

    // ── Pass 2: linear-bin all sub-buffers into per-warp tiles ─────────
    extern __shared__ int s_warp_tiles[];
    const int warp_id = tid >> 5;
    const int lane    = tid & 31;
    const int warps   = BLOCK_SIZE >> 5;

    for (int b = lane; b < SP4_HIST_BINS; b += 32) {
        s_warp_tiles[warp_id * SP4_HIST_BINS + b] = 0;
    }
    __syncwarp();

    for (int s = 0; s < n_sub; ++s) {
        const float* buf_s = reinterpret_cast<const float*>(sub_buf_ptrs[s]);
        const int    cnt_s = sub_counts[s];
        for (int i = tid; i < cnt_s; i += BLOCK_SIZE) {
            int bin_idx = (int)floorf(fabsf(buf_s[i]) / bin_width);
            if (bin_idx >= SP4_HIST_BINS) bin_idx = SP4_HIST_BINS - 1;
            s_warp_tiles[warp_id * SP4_HIST_BINS + bin_idx] += 1;
        }
    }
    __syncthreads();

    // Tree-reduce warp tiles into s_bins[].
    for (int b = tid; b < SP4_HIST_BINS; b += BLOCK_SIZE) {
        int sum = 0;
        for (int w = 0; w < warps; ++w) sum += s_warp_tiles[w * SP4_HIST_BINS + b];
        s_bins[b] = sum;
    }
    __syncthreads();

    // ── Pass 3: cumulative-from-top → p99 over the union distribution ──
    if (tid == 0) {
        const int target = (total_count + 99) / 100; // ceil(total_count / 100)
        int cumul   = 0;
        int p99_bin = SP4_HIST_BINS - 1;
        for (int b = SP4_HIST_BINS - 1; b >= 0; --b) {
            cumul += s_bins[b];
            if (cumul >= target) { p99_bin = b; break; }
        }
        s_step_max = (float)(p99_bin + 1) * bin_width;
    }
    __syncthreads();
    return s_step_max;
}
