/* ══════════════════════════════════════════════════════════════════════════
 * SP17 Phase 3.2 (2026-05-08) — adaptive advantage_clip_bound producer.
 *
 * Computes an adaptive upper-bound for |A_centered| by tracking
 *   advantage_clip_bound ← p99(|A_centered|) × ADVANTAGE_CLIP_SAFETY_FACTOR=1.5
 *
 * Mirrors the SP14 P0-A REWARD_POS_CAP producer pattern (also p99 × 1.5).
 * The kernel uses `sp4_histogram_p99` (the canonical block-tree-reduce
 * percentile estimator from SP4) — NO atomicAdd per
 * `pearl_fused_per_group_statistics_oracle` and `feedback_no_atomicadd`.
 *
 * The bound is observability-only in this commit (Phase 3 is the
 * diagnostic chain); the actual clipping wire-up is Phase 5 follow-up.
 *
 * Cadence: cold-path (epoch boundary). EMA α = ADVANTAGE_CLIP_EMA_ALPHA
 * = 0.01 (slow per-fold — host-passed).
 *
 * ── Pearls + invariants ────────────────────────────────────────────────
 *  - `feedback_no_atomicadd.md` — uses sp4_histogram_p99 (block
 *    tree-reduce + per-warp tile binning, no atomicAdd).
 *  - `pearl_first_observation_bootstrap.md` — sentinel
 *    SENTINEL_ADVANTAGE_CLIP_BOUND = 1.0; cold-start REPLACE on first
 *    observation.
 *  - `pearl_symmetric_clamp_audit.md` — bilateral clamp to
 *    [ADVANTAGE_CLIP_BOUND_MIN=0.1, ADVANTAGE_CLIP_BOUND_MAX=100.0]
 *    after EMA blend.
 *  - `feedback_isv_for_adaptive_bounds.md` — bound is producer-tracked,
 *    NEVER hardcoded.
 *  - `pearl_no_host_branches_in_captured_graph.md` — pure on-device.
 *
 * Algorithm (single-block, BLOCK=256):
 *
 *   For each branch d in [0, 4):
 *     Pass 1: per-atom mean of raw A over (B × n_branch) — written to
 *             `s_a_mean[NA]` (block-shared, recomputed per branch).
 *     Pass 2: write |A_raw[i, a, z] - a_mean[z]| into
 *             `abs_a_centered_scratch` at the BRANCH-MAJOR offset (so
 *             the final scratch is a flat |A_centered| over the full
 *             (B × Σ_d n_d × NA) tensor).
 *
 *   p99: call `sp4_histogram_p99<256>(scratch, total_count)` →
 *        p99(|A_centered|) over the full flat tensor.
 *
 *   Blend: bound = p99 × SAFETY × α + cur × (1-α), bilateral clamp,
 *          ISV write.
 *
 * Args:
 *   b_logits                  — BRANCH-MAJOR advantage logits.
 *   abs_a_centered_scratch    — pre-allocated scratch buffer of size
 *                                ≥ Σ_d B × b_d × NA floats. Caller-owned;
 *                                produced/consumed only inside this kernel.
 *   scratch_capacity           — sentinel (host-passed) for defensive
 *                                guard; kernel returns early if total >
 *                                capacity.
 *   isv                        — ISV bus.
 *   B, NA                      — shapes.
 *   b0..b3 sizes               — per-branch action counts.
 *   bound_idx                  — ADVANTAGE_CLIP_BOUND_INDEX = 482.
 *   alpha                      — ADVANTAGE_CLIP_EMA_ALPHA = 0.01.
 *   safety_factor              — ADVANTAGE_CLIP_SAFETY_FACTOR = 1.5.
 *   bound_min, bound_max       — bilateral clamp rails [0.1, 100.0].
 *   sentinel_bound             — SENTINEL_ADVANTAGE_CLIP_BOUND = 1.0.
 *
 * Launch: grid=(1, 1, 1), block=(256, 1, 1). Dynamic shared memory =
 * `(BLOCK_SIZE / 32) × 256 × sizeof(int) = 8192 bytes` for
 * `sp4_histogram_p99`'s per-warp tile array (extern __shared__ int
 * s_warp_tiles[]). Plus the static shared memory declared inside the
 * kernel (`s_a_mean[NA_MAX]` and `s_reduce[BLOCK_SIZE]`) — those are
 * `__shared__` static, NOT extern, so they don't compete with the
 * dynamic sp4 allocation.
 * ══════════════════════════════════════════════════════════════════════════ */

#include <cuda_runtime.h>
#include "sp4_histogram_p99.cuh"

#define SP17_VA_BLOCK_SIZE 256
#define SP17_NA_MAX        51       /* matches NUM_ATOMS_MAX in state_layout.cuh */
#define SP17_EPS_F         1e-6f

extern "C" __global__ void sp17_advantage_clip_bound_update(
    const float* __restrict__ b_logits,
    float*       __restrict__ abs_a_centered_scratch,
    int          scratch_capacity,
    float*       __restrict__ isv,
    int          B,
    int          NA,
    int          b0_size,
    int          b1_size,
    int          b2_size,
    int          b3_size,
    int          bound_idx,
    float        alpha,
    float        safety_factor,
    float        bound_min,
    float        bound_max,
    float        sentinel_bound)
{
    /* Single-block kernel — guard against accidental multi-block launch. */
    if (blockIdx.x != 0) return;

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

    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 n3   = (long long)b3_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;
    const long long total = off3 + Bll * n3 * NAll;

    if (B <= 0 || NA <= 0 || total <= 0) return;
    if ((long long)scratch_capacity < total) return;   /* defensive guard */

    /* Static shared memory: per-atom mean cache (small, fixed) +
     * per-pass tree-reduce tile. These do NOT consume the extern
     * `__shared__` dynamic allocation that sp4_histogram_p99 owns. */
    __shared__ float s_a_mean[SP17_NA_MAX];
    __shared__ float s_reduce[SP17_VA_BLOCK_SIZE];

    /* Branch passes — for each branch d, compute per-atom A mean, then
     * write |A_raw - mean[z]| into scratch at the branch's offset. After
     * all 4 branches, the scratch buffer is a flat |A_centered| tensor of
     * size `total` floats. */
    for (int branch = 0; branch < 4; ++branch) {
        int n_branch;
        long long off_branch;
        if (branch == 0)      { n_branch = b0_size; off_branch = off0; }
        else if (branch == 1) { n_branch = b1_size; off_branch = off1; }
        else if (branch == 2) { n_branch = b2_size; off_branch = off2; }
        else                  { n_branch = b3_size; off_branch = off3; }
        if (n_branch <= 0) continue;

        /* Pass 1: per-atom mean over (B × n_branch) at each atom z. */
        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 2: write |A_raw - mean[z]| into scratch at BRANCH-MAJOR
         * offset. Each thread handles a strided slice. */
        const long long branch_total = Bll * (long long)n_branch * NAll;
        for (long long idx = (long long)tid; idx < branch_total; idx += (long long)bdim) {
            const long long ia  = idx / NAll;
            const long long z   = idx - ia * NAll;
            const long long off = off_branch + idx;
            const float a       = b_logits[off];
            abs_a_centered_scratch[off] = fabsf(a - s_a_mean[z]);
        }
        __syncthreads();   /* ensure all writes visible before next branch reuses s_reduce */
    }

    /* ── p99(|A_centered|) over the full flat scratch ─────────────────
     * sp4_histogram_p99 uses its OWN extern __shared__ int per-warp
     * tiles + 256-int static + 1-float static (total 1028 bytes static
     * + 8192 bytes dynamic for BLOCK=256). Our static `s_a_mean` and
     * `s_reduce` arrays do not collide because they are statically
     * declared (different `__shared__` sub-region). The extern dynamic
     * shmem region is the caller-passed `shared_mem_bytes` and is
     * exclusively owned by sp4_histogram_p99.
     */
    const float p99 = sp4_histogram_p99<SP17_VA_BLOCK_SIZE>(
        abs_a_centered_scratch, (int)total);

    /* ── Pearl-A bootstrap + α blend + bilateral clamp → ISV write ── */
    if (tid == 0) {
        const float target = p99 * safety_factor;
        const float current = isv[bound_idx];
        float next;
        if (fabsf(current - sentinel_bound) < SP17_EPS_F) {
            /* Pearl-A: cold-start sentinel → REPLACE directly. */
            next = target;
        } else {
            next = (1.0f - alpha) * current + alpha * target;
        }
        /* `pearl_symmetric_clamp_audit` — bilateral clamp. */
        next = fmaxf(bound_min, fminf(next, bound_max));
        isv[bound_idx] = next;
        __threadfence_system();
    }
}
