/*
 * iqn_quantile_ema_update — Plan 4 Task 3 (E.3) producer kernel; SP4 Layer
 * A Task A13.4 retrofit (2026-05-01).
 *
 * Reads the IQN online forward's per-quantile Q surface
 * (`save_q_online [TBA, B*Q]` col-major; populated by
 * `GpuIqnHead::execute_training_pipeline` after each step) and writes the
 * mean |Q(s,a;τ)| step observation at each off-median fixed-τ position
 * to `producer_step_scratch_buf[scratch_first..scratch_first+4)` for
 * Pearls A+D consumption:
 *
 *   FIXED_TAUS = [0.05, 0.25, 0.50, 0.75, 0.95]
 *
 *   blockIdx.x = 0 → τ_idx 0 (τ=0.05) → scratch[scratch_first+0] → ISV[99]
 *   blockIdx.x = 1 → τ_idx 1 (τ=0.25) → scratch[scratch_first+1] → ISV[100]
 *   blockIdx.x = 2 → τ_idx 3 (τ=0.75) → scratch[scratch_first+2] → ISV[101]
 *   blockIdx.x = 3 → τ_idx 4 (τ=0.95) → scratch[scratch_first+3] → ISV[102]
 *
 * The median (τ=0.50, idx=2) is intentionally skipped — it is already
 * surfaced via the existing greedy-Q diagnostic, no duplication needed.
 *
 * Reduction layout (one block per scratch slot):
 *   For target τ_idx t, walk every (sample b, action a) and read
 *   save_q_online[a + (b*Q + t) * TBA].  Sum |Q| across (B × TBA) entries,
 *   divide by N = B × TBA, write to scratch slot.  Single block,
 *   256-thread shmem-tree reduce — no atomicAdd
 *   (per `feedback_no_atomicadd.md`).
 *
 * **α dropped per SP4 — Pearls A+D adapt α from per-slot signal-vs-noise
 * variance.** The host launcher
 * (`GpuDqnTrainer::launch_iqn_quantile_ema`) follows up with the GPU
 * `apply_pearls_ad_kernel` on the same stream (single launch, n_slots=4)
 * mapping each scratch slot to its ISV bound slot (99/100/101/102) and
 * updating Wiener-EMA state for each. 2026-05-01 GPU-Pearls refactor —
 * graph-capture-compatible (no host sync).
 *
 * Cold-path-cadence: launched once per training step alongside
 * `h_s2_rms_ema_update`. Diagnostic only — no consumer kernel reads slots
 * [99..103). ISV is pinned device-mapped, so host HEALTH_DIAG sees the
 * Pearls-A+D-updated values without an explicit DtoH.
 */

#include <cuda_runtime.h>

extern "C" __global__ void iqn_quantile_ema_update(
    const float* __restrict__ save_q_online,  /* [TBA, B*Q] col-major */
    int          batch_size,                  /* B */
    int          num_quantiles,               /* Q (= 5 under fixed-τ) */
    int          tba,                         /* total branch actions */
    float*       __restrict__ scratch_buf,    /* producer_step_scratch_buf */
    int          scratch_first_index          /* slot 59 — first of 4 contiguous slots */
) {
    /* Map blockIdx.x → (τ_idx within FIXED_TAUS, target scratch slot).
     * Skip the median (idx=2) — that's the existing greedy-Q diagnostic.
     * Scratch slot order matches blockIdx.x: p05 → +0, p25 → +1, p75 → +2,
     * p95 → +3. */
    int tau_idx;
    int slot_offset;
    switch (blockIdx.x) {
        case 0: tau_idx = 0; slot_offset = 0; break;
        case 1: tau_idx = 1; slot_offset = 1; break;
        case 2: tau_idx = 3; slot_offset = 2; break;
        case 3: tau_idx = 4; slot_offset = 3; break;
        default: return;
    }

    /* Bounds check — caller must launch with grid.x = 4 and num_quantiles
     * matching FIXED_TAUS.len() = 5. The defensive guard skips the
     * scratch write if the layout is incompatible, surfacing the
     * misconfiguration to the host via a stale (zero or prior-step) value
     * rather than a silent bad-index read. */
    if (tau_idx >= num_quantiles) return;

    extern __shared__ float smem[];
    const int tid   = (int)threadIdx.x;
    const int block = (int)blockDim.x;
    const int N     = batch_size * tba;

    /* Sum |Q(s, a; τ_idx)| across (sample, action). Absolute value tracks
     * Q-magnitude consistently with the spec's risk-monitoring purpose
     * (sign of an off-median quantile alone is not a useful diagnostic;
     * its scale relative to the median is). */
    float local_sum = 0.0f;
    for (int i = tid; i < N; i += block) {
        int a = i % tba;
        int b = i / tba;
        /* Col-major: save_q_online[row=a, col=(b*Q + tau_idx)] */
        long col = (long)b * (long)num_quantiles + (long)tau_idx;
        long flat = (long)a + col * (long)tba;
        local_sum += fabsf(save_q_online[flat]);
    }
    smem[tid] = local_sum;
    __syncthreads();

    /* In-block tree reduction (no atomicAdd). */
    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    /* Thread 0 writes the step observation. The host caller's trainer
     * config requires batch_size > 0 and tba > 0, so N > 0 is invariant;
     * no defensive `(N > 0) ? N : 1` here — masking the divide would
     * silently emit 0 if the invariant ever broke; the NaN that
     * division-by-zero produces propagates loudly through Pearls A+D to
     * the ISV slot. */
    if (tid == 0) {
        const float mean_abs_q = smem[0] / (float)N;
        scratch_buf[scratch_first_index + slot_offset] = mean_abs_q;
        __threadfence_system();
    }
}
