/* ══════════════════════════════════════════════════════════════════════════
 * SP18 v2 weight-drift diagnostic + ISV-adaptive shrink-and-perturb producer
 * (2026-05-09).
 *
 * Original cadence (commit `aab13a83f`): cold-path producer of the per-epoch
 * weight-drift diagnostic. Reads `params_flat [n]` and `best_params_snapshot
 * [n]`; emits two floats into a `MappedF32Buffer<4>`:
 *
 *   out[0] = ||params_flat - best_params||₂                    (L2 distance)
 *   out[1] = out[0] / max(||best_params||₂, EPS)               (relative)
 *
 * Surfaces HD2→HD5 weight drift across epochs so the post-deploy reviewer
 * can correlate Sharpe-peak decay with actual parameter trajectory in
 * `(W_best, W_curr)` space.
 *
 * SP18 v2 ISV-adaptive shrink-and-perturb close-out (this commit, 2026-05-09):
 * extended to ALSO write the adaptive α and σ for the next fold's
 * `shrink_and_perturb` call into ISV[505] / ISV[506] using the values
 * already computed by the two block-tree-reductions. No new kernel — this
 * is the single producer for both diagnostics and the controller (fewer
 * kernels = same atomicity per the spec's "Watch for: Drift diagnostic
 * timing — have the drift kernel ALSO emit the adaptive α/σ as a 4-element
 * output — fewer kernels, same atomicity"). Per `feedback_isv_for_adaptive_
 * bounds` and `feedback_adaptive_not_tuned`. Outputs:
 *
 *   out[2] = α_adaptive  (clamped to [SHRINK_ALPHA_MIN, SHRINK_ALPHA_MAX])
 *   out[3] = σ_adaptive  (clamped to [SHRINK_SIGMA_MIN, SHRINK_SIGMA_MAX])
 *
 * AND simultaneously:
 *
 *   isv[SHRINK_ALPHA_ADAPTIVE_INDEX=505] = α_adaptive
 *   isv[SHRINK_SIGMA_ADAPTIVE_INDEX=506] = σ_adaptive
 *
 * Formulas (single-thread-0 cold-path arithmetic post-reduction):
 *
 *   relative      = ||p − b||₂ / max(||b||₂, EPS)        — drift signal
 *   best_rms      = ||b||₂ / sqrt(n)                     — scale signal
 *
 *   α_target      = 1.0 − clamp(relative,
 *                                1.0 − SHRINK_ALPHA_MAX,
 *                                1.0 − SHRINK_ALPHA_MIN)
 *                 ≡ small drift → α near MAX (preserve the converged
 *                   best-Sharpe weights hard); large drift → α near MIN
 *                   (shrink hard to escape the divergent region).
 *
 *   σ_target      = clamp(SENTINEL_SHRINK_SIGMA × (best_rms /
 *                                                  SHRINK_SIGMA_RMS_REFERENCE),
 *                         SHRINK_SIGMA_MIN, SHRINK_SIGMA_MAX)
 *                 ≡ scale-aware noise — σ tracks `||best||_RMS` so noise
 *                   stays scale-relative as weight magnitudes evolve
 *                   across folds (LayerNorm γ drift, weight-decay
 *                   relaxation). At RMS = REFERENCE = 0.5, σ_target =
 *                   SENTINEL = 0.01 → bit-identical to prior hardcoded
 *                   value for typical healthy weights.
 *
 * Pearl-A first-observation bootstrap is enforced by the host launcher:
 * when `best_params_snapshot` is `None` (first fold pre-best-Sharpe, or
 * any fold with no improvement), the launcher writes
 * [0.0, 0.0, SENTINEL_SHRINK_ALPHA, SENTINEL_SHRINK_SIGMA] directly into
 * the mapped-pinned output buffer and skips the kernel — bit-identical
 * cold-start to prior hardcoded `α=0.8 / σ=0.01`. The ISV slots are NOT
 * touched on that path; consumers fall back to host-side sentinels read
 * via `read_isv_signal_at` (which returns the FoldReset-written sentinel
 * before the first valid kernel launch lands a real signal).
 *
 * ── Pearls + invariants ────────────────────────────────────────────────
 *  - `feedback_no_atomicadd.md` — block tree-reduce only; no atomicAdd.
 *    Two reductions (sum-of-squares of diff, sum-of-squares of best)
 *    use the same shmem tile sequentially.
 *  - `feedback_no_htod_htoh_only_mapped_pinned.md` — output is a
 *    `MappedF32Buffer<4>`; kernel writes via `__threadfence_system()`
 *    for PCIe-visible coherence; host reads after stream sync. ISV bus
 *    write is also mapped-pinned (host-readable directly).
 *  - `pearl_no_host_branches_in_captured_graph.md` — kernel runs
 *    cold-path (epoch boundary, between training-step graph
 *    invocations); no host branches inside.
 *  - `pearl_first_observation_bootstrap.md` — first valid observation
 *    REPLACES sentinel directly. The producer ALWAYS writes a clamped
 *    target value (no EMA blend), so on the first valid launch, ISV[505]
 *    and ISV[506] go from sentinel directly to the computed target.
 *    Subsequent launches REPLACE again — there is no Welford state for
 *    this controller; the underlying signal (relative drift, best_rms)
 *    is itself computed afresh each epoch from current/best weights.
 *  - `pearl_symmetric_clamp_audit.md` — bilateral
 *    `fmaxf(MIN, fminf(x, MAX))` for both α and σ; never one-sided.
 *  - `feedback_isv_for_adaptive_bounds.md` — replaces the hardcoded
 *    `α=0.8 / σ=0.01` constants formerly at
 *    `fused_training.rs::reset_for_fold` lines 1014-1015. The bounds
 *    [MIN, MAX] are themselves Category-1 dimensional safety floors —
 *    structural (never adaptive themselves; their values are the floor
 *    and ceiling).
 *  - `feedback_adaptive_not_tuned.md` — α/σ derive from the EMA-tracked
 *    drift signal + scale-aware best_rms ratio, not tuned constants.
 *
 * Algorithm (single block, 256 threads):
 *
 *   Pass 1 (sum-of-squared diffs):
 *     For each i in [0, n) striding by blockDim.x:
 *       d = params[i] - best[i]
 *       local_sumsq += d*d
 *     Block tree-reduce into s_reduce[0]; thread 0 stores to local s0.
 *
 *   Pass 2 (sum-of-squared best):
 *     For each i in [0, n) striding by blockDim.x:
 *       local_sumsq += best[i] * best[i]
 *     Block tree-reduce into s_reduce[0]; thread 0 stores to local s1.
 *
 *   Pass 3 (write outputs):
 *     thread 0 only:
 *       l2        = sqrtf(s0)
 *       norm_best = sqrtf(s1)
 *       rel       = l2 / max(norm_best, SP18_DRIFT_EPS_F)
 *       best_rms  = norm_best / sqrtf(n)                    [ NEW ]
 *       α_target  = 1.0 − clamp(rel, 1.0 − ALPHA_MAX, 1.0 − ALPHA_MIN)  [ NEW ]
 *       σ_target  = clamp(SIGMA_SENTINEL × best_rms / RMS_REF,
 *                         SIGMA_MIN, SIGMA_MAX)             [ NEW ]
 *       out[0..4] = [l2, rel, α_target, σ_target]
 *       isv[505]  = α_target                                [ NEW ]
 *       isv[506]  = σ_target                                [ NEW ]
 *       __threadfence_system()
 *
 * Launch contract:
 *   grid_dim  = (1, 1, 1)
 *   block_dim = (SP18_DRIFT_BLOCK, 1, 1)  — 256 threads
 *   shared mem (dynamic): SP18_DRIFT_BLOCK × sizeof(float)  = 1024 bytes
 *
 * Args:
 *   params       — `params_flat [n]` f32 device pointer (online weights).
 *   best         — `best_params_snapshot [n]` f32 device pointer.
 *   out          — `MappedF32Buffer<4>` device pointer for
 *                  [l2, rel, α_target, σ_target].
 *   isv          — ISV bus device pointer (mapped-pinned). Writes ISV[505]
 *                  and ISV[506].
 *   n            — element count of `params` and `best` (must match).
 * ══════════════════════════════════════════════════════════════════════════ */

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

#define SP18_DRIFT_BLOCK   256
#define SP18_DRIFT_EPS_F   1e-12f

extern "C" __global__ void weight_drift_diag_update(
    const float* __restrict__ params,
    const float* __restrict__ best,
    float*       __restrict__ out,
    float*       __restrict__ isv,
    int          n)
{
    /* Single block (gridDim.x == 1). */
    if (blockIdx.x != 0) return;

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

    /* Degenerate guard: n <= 0 — write zeros + sentinels and return. */
    if (n <= 0) {
        if (tid == 0) {
            out[0] = 0.0f;
            out[1] = 0.0f;
            out[2] = SENTINEL_SHRINK_ALPHA;
            out[3] = SENTINEL_SHRINK_SIGMA;
            isv[SHRINK_ALPHA_ADAPTIVE_INDEX] = SENTINEL_SHRINK_ALPHA;
            isv[SHRINK_SIGMA_ADAPTIVE_INDEX] = SENTINEL_SHRINK_SIGMA;
            __threadfence_system();
        }
        return;
    }

    extern __shared__ float s_reduce[];

    /* ── Pass 1: ||params - best||₂² ─────────────────────────────────── */
    float local = 0.0f;
    const long long N = (long long)n;
    for (long long i = (long long)tid; i < N; i += (long long)bdim) {
        const float d = params[i] - best[i];
        local += d * d;
    }
    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();
    }
    /* Capture diff sum-of-squares before reusing s_reduce for pass 2. */
    float sumsq_diff = 0.0f;
    if (tid == 0) sumsq_diff = s_reduce[0];

    /* ── Pass 2: ||best||₂² ──────────────────────────────────────────── */
    local = 0.0f;
    for (long long i = (long long)tid; i < N; i += (long long)bdim) {
        const float b = best[i];
        local += b * b;
    }
    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();
    }

    /* ── Pass 3: write outputs (thread 0 only) ───────────────────────── */
    if (tid == 0) {
        const float sumsq_best = s_reduce[0];
        const float l2_diff    = sqrtf(sumsq_diff);
        const float norm_best  = sqrtf(sumsq_best);
        const float denom      = (norm_best > SP18_DRIFT_EPS_F)
                                 ? norm_best
                                 : SP18_DRIFT_EPS_F;
        const float rel        = l2_diff / denom;

        /* SP18 v2 ISV-adaptive shrink-and-perturb (2026-05-09):
         * compute α_target and σ_target from the same reductions. */
        const float best_rms = norm_best / sqrtf((float)n);

        /* α_target: drift-driven preservation factor. Bilateral clamp
         * on relative drift, then 1 − clamp gives the preserve fraction.
         * `pearl_symmetric_clamp_audit` — `fmaxf(LO, fminf(rel, HI))`. */
        const float rel_lo   = 1.0f - SHRINK_ALPHA_MAX;  /* 0.05 */
        const float rel_hi   = 1.0f - SHRINK_ALPHA_MIN;  /* 0.50 */
        const float rel_clamped = fmaxf(rel_lo, fminf(rel, rel_hi));
        const float alpha_target = 1.0f - rel_clamped;

        /* σ_target: scale-aware noise. ratio = best_rms / RMS_REFERENCE
         * — when weights have typical RMS, ratio = 1.0 → σ = SENTINEL.
         * As RMS grows, σ scales up proportionally, capped at MAX. */
        const float ratio        = best_rms / SHRINK_SIGMA_RMS_REFERENCE;
        const float sigma_raw    = SENTINEL_SHRINK_SIGMA * ratio;
        const float sigma_target = fmaxf(SHRINK_SIGMA_MIN,
                                          fminf(sigma_raw, SHRINK_SIGMA_MAX));

        out[0] = l2_diff;
        out[1] = rel;
        out[2] = alpha_target;
        out[3] = sigma_target;

        /* ISV bus write (mapped-pinned). Consumer (`reset_for_fold`)
         * reads via `read_isv_signal_at` after stream sync. */
        isv[SHRINK_ALPHA_ADAPTIVE_INDEX] = alpha_target;
        isv[SHRINK_SIGMA_ADAPTIVE_INDEX] = sigma_target;

        __threadfence_system();   /* PCIe-visible write for mapped-pinned outputs */
    }
}
