/* ══════════════════════════════════════════════════════════════════════════
 * SP17 Phase 3.1 (2026-05-08) — A_var_ema producer kernel.
 *
 * Per-branch EMA of the over-actions variance of the centered advantage
 *   A_centered[i, a, z] = A[i, a, z] − (1/n_d) Σ_a' A[i, a', z]
 * where the per-atom mean is taken over the b_d direction-of-branch
 * actions (b0_size for dir, b1_size for mag, ...).
 *
 * Writes the per-branch variance EMA to four ISV slots:
 *   ISV[A_VAR_EMA_DIR_INDEX=474]  (branch 0 — dir)
 *   ISV[A_VAR_EMA_MAG_INDEX=475]  (branch 1 — mag)
 *   ISV[A_VAR_EMA_ORD_INDEX=476]  (branch 2 — ord)
 *   ISV[A_VAR_EMA_URG_INDEX=477]  (branch 3 — urg)
 *
 * Var_d  = (1 / (B × n_d × NA)) Σ_{i, a, z} A_centered[i, a, z]^2
 *        (since E[A_centered]=0 by construction over actions, the second
 *         moment IS the variance up to f32 rounding.)
 *
 * Producer cadence: cold-path (epoch boundary) — single launch per
 * HEALTH_DIAG emit. Mirrors the v_a_means_diag_kernel.cu cadence and
 * dispatch pattern.
 *
 * ── Pearls + invariants ────────────────────────────────────────────────
 *  - `feedback_no_atomicadd.md` — block tree-reduce only; no atomicAdd.
 *  - `pearl_first_observation_bootstrap.md` — sentinel = SENTINEL_A_VAR_EMA
 *    (=0.0); first valid observation REPLACES directly. After bootstrap,
 *    blend with caller-supplied α (the host derives a Wiener-α with floor
 *    WELFORD_ALPHA_MIN_F=0.4 per the SP16 T3 fix-up; for cold-path
 *    cadence the host passes the steady-state floor directly because the
 *    variance signal is stable enough that a Welford accumulator state
 *    machine inside the kernel would add 6 ISV slots × 4 branches = 24
 *    slots for a cold-path cadence diagnostic — not worth the complexity.
 *    Instead, the floor IS the α used after Pearl-A bootstrap.).
 *  - `pearl_wiener_alpha_floor_for_nonstationary` — α is the structural
 *    floor 0.4 (NOT zero); preserves catch-up bandwidth when the policy
 *    drifts.
 *  - `pearl_no_host_branches_in_captured_graph.md` — kernel runs
 *    cold-path; no host branches inside.
 *  - `feedback_isv_for_adaptive_bounds.md` — α is the structural
 *    bandwidth floor (NOT a tuned constant — it's the
 *    `pearl_wiener_alpha_floor_for_nonstationary` MIN value), and the
 *    sentinel/bootstrap is the cold-start contract.
 *
 * Algorithm (per launch, four branches in parallel via gridDim.x = 4):
 *
 *   Per branch d (block d, threads 256):
 *     ── Pass 1: per-atom mean of A across n_d actions ──
 *       For each atom z in [0, NA):
 *         m_a[z] = (1 / (B × n_d)) Σ_{i, a} A[i, a, z]
 *
 *     ── Pass 2: variance of centered A ──
 *       Var_d = (1 / (B × n_d × NA)) Σ_{i, a, z} (A[i, a, z] - m_a[z])^2
 *
 *     ── Pass 3: Pearl-A bootstrap + α blend (thread 0 only) ──
 *       cur = isv[var_idx_d];
 *       if (fabsf(cur - SENTINEL_A_VAR_EMA) < EPS_F) {
 *         next = Var_d;          // first observation REPLACES
 *       } else {
 *         next = (1 - alpha) * cur + alpha * Var_d;
 *       }
 *       isv[var_idx_d] = next;
 *       __threadfence_system();
 *
 * Launch contract:
 *   grid_dim  = (4, 1, 1)             — one block per branch
 *   block_dim = (SP17_VA_BLOCK_SIZE, 1, 1)  — 256 threads
 *   shared mem (dynamic): (BLOCK_SIZE + NA) × sizeof(float)
 *                         = (256 + NA) × 4 bytes
 *                         = 1024 bytes for the tree-reduce tile
 *                         + NA × 4 bytes for the per-atom mean cache.
 *
 * Args:
 *   b_logits     — Online branch logits, BRANCH-MAJOR layout
 *                  [B×b0×NA | B×b1×NA | B×b2×NA | B×b3×NA]. Same buffer
 *                  `compute_expected_q` reads from. f32 device ptr.
 *   isv          — [ISV_TOTAL_DIM] f32 device ptr (mapped-pinned bus).
 *   B            — batch size.
 *   NA           — C51 atoms per action.
 *   b0_size, b1_size, b2_size, b3_size — per-branch action counts.
 *   var_idx_dir, var_idx_mag, var_idx_ord, var_idx_urg —
 *                  ISV slot indices for each branch (474..478).
 *   alpha        — EMA blend rate post-bootstrap. Caller passes
 *                  WELFORD_ALPHA_MIN (=0.4) per
 *                  `pearl_wiener_alpha_floor_for_nonstationary`.
 *   sentinel_a_var — SENTINEL_A_VAR_EMA = 0.0 (cold-start sentinel).
 * ══════════════════════════════════════════════════════════════════════════ */

#include <cuda_runtime.h>

#define SP17_VA_BLOCK_SIZE 256
#define SP17_EPS_F         1e-6f

extern "C" __global__ void sp17_a_var_ema_update(
    const float* __restrict__ b_logits,
    float*       __restrict__ isv,
    int          B,
    int          NA,
    int          b0_size,
    int          b1_size,
    int          b2_size,
    int          b3_size,
    int          var_idx_dir,
    int          var_idx_mag,
    int          var_idx_ord,
    int          var_idx_urg,
    float        alpha,
    float        sentinel_a_var)
{
    /* One block per branch (gridDim.x == 4). */
    const int branch = (int)blockIdx.x;
    if (branch >= 4) return;

    /* Branch dispatch: select per-branch action count + ISV slot index +
     * BRANCH-MAJOR offset (in float elements) into b_logits. */
    int  n_branch;
    int  var_idx;
    long long off_branch;
    const long long Bll  = (long long)B;
    const long long NAll = (long long)NA;
    const long long n0   = (long long)b0_size;
    const long long n1   = (long long)b1_size;
    const long long n2   = (long long)b2_size;
    const long long off0 = 0;
    const long long off1 = off0 + Bll * n0 * NAll;
    const long long off2 = off1 + Bll * n1 * NAll;
    const long long off3 = off2 + Bll * (long long)b2_size * NAll;
    if (branch == 0)      { n_branch = b0_size; var_idx = var_idx_dir; off_branch = off0; }
    else if (branch == 1) { n_branch = b1_size; var_idx = var_idx_mag; off_branch = off1; }
    else if (branch == 2) { n_branch = b2_size; var_idx = var_idx_ord; off_branch = off2; }
    else                  { n_branch = b3_size; var_idx = var_idx_urg; off_branch = off3; }

    const int tid  = (int)threadIdx.x;
    const int bdim = (int)blockDim.x;

    /* Degenerate guard: zero-sized branch or empty buffers — skip the
     * EMA update entirely. Leaves the ISV slot at its current value
     * (sentinel on first launch, prior EMA on subsequent launches). */
    if (B <= 0 || NA <= 0 || n_branch <= 0) {
        return;
    }

    /* Shared memory layout (dynamic):
     *   shmem[0..bdim)              — tree-reduce tile (per pass)
     *   shmem[bdim..bdim+NA)        — per-atom mean a_mean[z] cached
     *                                  across pass 1 → pass 2.
     */
    extern __shared__ float shmem[];
    float* s_reduce = shmem;
    float* s_a_mean = shmem + bdim;

    const long long total = Bll * (long long)n_branch * NAll;

    /* ── Pass 1: per-atom mean over (B × n_branch) for each atom z ─────
     * Loop over atoms; per-atom block-wide tree reduction. Cost ≤ 51
     * reductions per branch × 4 branches; cold-path cadence makes it
     * irrelevant.
     */
    for (int z = 0; z < NA; ++z) {
        float local = 0.0f;
        const long long count_ia = Bll * (long long)n_branch;
        /* Linearise (i, a) → idx, read b_logits[off_branch + idx*NA + z]. */
        for (long long idx = (long long)tid; idx < count_ia; idx += (long long)bdim) {
            const long long off = off_branch + idx * NAll + (long long)z;
            local += b_logits[off];
        }
        s_reduce[tid] = local;
        __syncthreads();
        for (int s = bdim / 2; s > 0; s >>= 1) {
            if (tid < s) s_reduce[tid] += s_reduce[tid + s];
            __syncthreads();
        }
        if (tid == 0) {
            const float total_count = (float)count_ia;
            s_a_mean[z] = (total_count > 0.0f) ? (s_reduce[0] / total_count) : 0.0f;
        }
        __syncthreads();
    }

    /* ── Pass 2: variance over (B × n_branch × NA) of (A − a_mean[z]) ─
     * Single strided pass with per-thread sum-of-squared-deviations. */
    float local_sumsq = 0.0f;
    for (long long idx = (long long)tid; idx < total; idx += (long long)bdim) {
        const long long ia  = idx / NAll;
        const long long z   = idx - ia * NAll;
        const long long off = off_branch + ia * NAll + z;
        const float a   = b_logits[off];
        const float dev = a - s_a_mean[z];
        local_sumsq += dev * dev;
    }
    s_reduce[tid] = local_sumsq;
    __syncthreads();
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) s_reduce[tid] += s_reduce[tid + s];
        __syncthreads();
    }

    /* ── Pass 3: Pearl-A bootstrap + α blend → ISV write ──────────── */
    if (tid == 0) {
        const float total_f = (float)total;
        const float var_d = (total_f > 0.0f) ? (s_reduce[0] / total_f) : 0.0f;

        const float current = isv[var_idx];
        float next;
        if (fabsf(current - sentinel_a_var) < SP17_EPS_F) {
            /* Pearl-A first-observation bootstrap: sentinel detected →
             * REPLACE directly (no blend). */
            next = var_d;
        } else {
            /* Steady-state: blend with caller-supplied α (host passes
             * WELFORD_ALPHA_MIN=0.4 per
             * `pearl_wiener_alpha_floor_for_nonstationary`). */
            next = (1.0f - alpha) * current + alpha * var_d;
        }
        isv[var_idx] = next;
        __threadfence_system();   /* PCIe-visible write for mapped-pinned ISV */
    }
}
