/* ══════════════════════════════════════════════════════════════════════════
 * SP17 Phase 3.2 (2026-05-08) — V_share producer kernel.
 *
 * Per-branch magnitude share
 *   V_share[d] = |E[V]| / (|E[V]| + |E[A_centered, max-Q action]|)
 *
 * The picked-action defaults to argmax E[Q] across the b_d actions of the
 * branch (max-Q semantic — tractable per-batch without depending on the
 * Thompson-sampled action stored in actions_history_buf, which is
 * collector-time state and is stale relative to the cuBLAS forward
 * driving on_v_logits_buf / on_b_logits_buf at HEALTH_DIAG cadence).
 *
 * Argmax simplification: argmax_a E[Q[i, a]] = argmax_a (1/NA) Σ_z
 * (V[i, z] + A_centered[i, a, z]). Since V[i, z] is the same across
 * actions and the per-atom mean of A is also action-invariant, the
 * ranking depends only on Σ_z A[i, a, z]. So:
 *   a*(i) = argmax_a Σ_z A_raw[i, a, z]
 *
 * For each sample i in branch d:
 *   v_i  = (1/NA) Σ_z V[i, z]                            (mean V over atoms)
 *   a_i  = (1/NA) Σ_z A_centered[i, a*(i), z]            (mean A_centered over atoms)
 *
 * Aggregated across the batch (one block per branch):
 *   |E[V]|_batch    = | (1/B) Σ_i v_i |
 *   |E[A_c]|_batch  = | (1/B) Σ_i a_i |
 *
 * V_share[d] = |E[V]|_batch / (|E[V]|_batch + |E[A_c]|_batch + EPS)
 *
 * Then EMA-blend into ISV[V_SHARE_*_INDEX] with the same Pearl-A bootstrap +
 * Wiener-α-floor pattern as A_var_ema.
 *
 * ── Pearls + invariants ────────────────────────────────────────────────
 *  - `feedback_no_atomicadd.md` — block tree-reduce only.
 *  - `pearl_first_observation_bootstrap.md` — sentinel SENTINEL_V_SHARE
 *    = 0.5 ("healthy 50/50 cold-start"); first observation REPLACES.
 *  - `pearl_wiener_alpha_floor_for_nonstationary` — α floored at 0.4
 *    (caller passes WELFORD_ALPHA_MIN).
 *  - `pearl_no_host_branches_in_captured_graph.md` — pure on-device
 *    branching only (per-branch dispatch via blockIdx.x).
 *  - `pearl_symmetric_clamp_audit.md` — V_share is structurally bounded
 *    in [0, 1]; defense-in-depth bilateral clamp before ISV write.
 *
 * 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 + 1) × sizeof(float)
 *                         = tree-reduce tile + per-atom mean cache
 *                         + 1 float for batch-aggregate E[V] broadcast.
 *
 * Args:
 *   v_logits         — [B, NA] online value logits, sample-major.
 *   b_logits         — BRANCH-MAJOR online advantage logits.
 *   isv              — ISV bus.
 *   B, NA            — batch size, atoms.
 *   b0..b3 sizes     — per-branch action counts.
 *   share_idx_*      — V_SHARE_*_INDEX per branch (478..482).
 *   alpha            — Wiener-α blend rate (caller passes WELFORD_ALPHA_MIN).
 *   sentinel_v_share — SENTINEL_V_SHARE = 0.5.
 * ══════════════════════════════════════════════════════════════════════════ */

#include <cuda_runtime.h>

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

extern "C" __global__ void sp17_v_share_update(
    const float* __restrict__ v_logits,
    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          share_idx_dir,
    int          share_idx_mag,
    int          share_idx_ord,
    int          share_idx_urg,
    float        alpha,
    float        sentinel_v_share)
{
    /* One block per branch (gridDim.x == 4). */
    const int branch = (int)blockIdx.x;
    if (branch >= 4) return;

    int n_branch, share_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 * n2 * NAll;
    if (branch == 0)      { n_branch = b0_size; share_idx = share_idx_dir; off_branch = off0; }
    else if (branch == 1) { n_branch = b1_size; share_idx = share_idx_mag; off_branch = off1; }
    else if (branch == 2) { n_branch = b2_size; share_idx = share_idx_ord; off_branch = off2; }
    else                  { n_branch = b3_size; share_idx = share_idx_urg; off_branch = off3; }

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

    if (B <= 0 || NA <= 0 || n_branch <= 0) return;

    /* Shared layout (dynamic):
     *   shmem[0..bdim)       — block tree-reduce tile
     *   shmem[bdim..bdim+NA) — per-atom mean of raw A across actions
     *   shmem[bdim+NA]       — batch-aggregate E[V] broadcast scalar
     */
    extern __shared__ float shmem[];
    float* s_reduce = shmem;
    float* s_a_mean = shmem + bdim;
    float* s_e_v    = shmem + bdim + NA;

    /* ── Pass A: per-atom mean of raw A (same as A_var_ema's pass 1) ── */
    for (int z = 0; z < NA; ++z) {
        float local = 0.0f;
        const long long count_ia = Bll * (long long)n_branch;
        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 B: per-sample v_i = (1/NA) Σ_z V[i, z] aggregated to E[V] ──
     * v_logits is [B, NA] sample-major (same buffer used by all branches —
     * the value head is shared across branches). */
    float local_v_sum = 0.0f;
    for (long long i = (long long)tid; i < Bll; i += (long long)bdim) {
        float v_i = 0.0f;
        const long long base = i * NAll;
        for (long long z = 0; z < NAll; ++z) {
            v_i += v_logits[base + z];
        }
        v_i /= (float)NA;
        local_v_sum += v_i;
    }
    s_reduce[tid] = local_v_sum;
    __syncthreads();
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) s_reduce[tid] += s_reduce[tid + s];
        __syncthreads();
    }
    if (tid == 0) {
        s_e_v[0] = s_reduce[0] / (float)B;
    }
    __syncthreads();
    const float e_v = s_e_v[0];

    /* ── Pass C: per-sample picked-action = argmax_a Σ_z A_raw[i, a, z] ──
     * Build a_i = (1/NA) Σ_z (A_raw[i, a*, z] − a_mean[z]) and aggregate
     * to E[A_centered, picked]. */
    float local_a_sum = 0.0f;
    for (long long i = (long long)tid; i < Bll; i += (long long)bdim) {
        float best_score   = -INFINITY;
        long long best_a   = 0;
        for (long long a = 0; a < (long long)n_branch; ++a) {
            float s_score = 0.0f;
            const long long base = off_branch + (i * (long long)n_branch + a) * NAll;
            for (long long z = 0; z < NAll; ++z) {
                s_score += b_logits[base + z];
            }
            if (s_score > best_score) {
                best_score = s_score;
                best_a     = a;
            }
        }
        float a_i_centered = 0.0f;
        const long long base_pick = off_branch + (i * (long long)n_branch + best_a) * NAll;
        for (long long z = 0; z < NAll; ++z) {
            a_i_centered += b_logits[base_pick + z] - s_a_mean[z];
        }
        a_i_centered /= (float)NA;
        local_a_sum += a_i_centered;
    }
    s_reduce[tid] = local_a_sum;
    __syncthreads();
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) s_reduce[tid] += s_reduce[tid + s];
        __syncthreads();
    }

    /* ── Pass D: V_share + Pearl-A bootstrap + α blend → ISV write ── */
    if (tid == 0) {
        const float e_a_centered = s_reduce[0] / (float)B;
        const float abs_v = fabsf(e_v);
        const float abs_a = fabsf(e_a_centered);
        const float denom = abs_v + abs_a + SP17_EPS_F;
        const float v_share = abs_v / denom;

        const float current = isv[share_idx];
        float next;
        if (fabsf(current - sentinel_v_share) < SP17_EPS_F) {
            /* Pearl-A: cold-start sentinel → REPLACE directly. */
            next = v_share;
        } else {
            next = (1.0f - alpha) * current + alpha * v_share;
        }
        /* `pearl_symmetric_clamp_audit` — V_share is structurally
         * bounded in [0, 1] by the formula, but apply bilateral
         * defense-in-depth clamp before ISV write. */
        next = fmaxf(0.0f, fminf(next, 1.0f));
        isv[share_idx] = next;
        __threadfence_system();   /* PCIe-visible write for mapped-pinned ISV */
    }
}
