// crates/ml/src/cuda_pipeline/plasticity_injection_kernel.cu
//
// SP15 Phase 3.5.4 — plasticity injection trigger + warm-up tracker.
// SP15 Wave 4.2 / Phase 3.5.4.b — cuRAND Kaiming-He weight reset (now wired).
//
// TWO-STEP recovery per spec §9.2 (3.5.4) post-amendment-2 fix:
//   1. On fire: set PLASTICITY_FIRED_THIS_FOLD = 1.0, set
//      PLASTICITY_WARM_BARS_REMAINING = M_warm, AND reset the last 10%
//      of advantage-head weights to Kaiming-He init (gain=√2 for
//      ReLU/GLU activations) via cuRAND `curand_normal()`. The
//      action-selection layer (consumer-wiring follow-up) reads the
//      PLASTICITY_FIRED_THIS_FOLD == 1 flag transition AND emits Flat
//      that bar.
//   2. Per-bar decrement of PLASTICITY_WARM_BARS_REMAINING; the
//      consumer-wiring follow-up uses
//      `effective_cooldown = max(COOLDOWN_BARS_REMAINING,
//      PLASTICITY_WARM_BARS_REMAINING)` to force Hold while either
//      counter > 0.
//
// Wave 4.2 lands the previously-deferred Kaiming-He weight reset:
//
//   - Reset region: the last 10% of `advantage_head_weights`
//     (`reset_count = n_weights / 10`, addressed via offset
//     `advantage_head_weights + n_weights - reset_count`). Resetting
//     the trailing slice keeps the upstream feature-extraction trunk
//     intact while re-initialising the head's last layer where the
//     drawdown-correlated gradients have presumably saturated. This
//     matches the spec §9.2 (3.5.4) "last 10%" prescription.
//   - Init formula: Kaiming-He for ReLU/GLU activations,
//     `weights[i] ~ Normal(0, sqrt(2 / fan_in))` (gain = sqrt(2)).
//     `fan_in` is supplied by the launcher (host computes it from the
//     trainer's network config; for oracle tests the test sets it to
//     a representative value).
//   - Per-thread cuRAND state via `curand_init(seed, tid, 0, &state)`
//     + `curand_normal(&state)` × stddev. `seed` is a host-passed
//     `unsigned long long`; the launcher derives it deterministically
//     from a trainer RNG state (or `SystemTime` nanos in the oracle
//     tests). The same (seed, tid) pair always produces the same
//     sample, satisfying determinism for the oracle test's mean/std
//     assertions.
//   - Grid: multi-block, `[((reset_count + 255) / 256), 1, 1]` ×
//     `[256, 1, 1]`. Block 0 / thread 0 still runs the trigger +
//     warm-bars decrement (the only ISV-write code path); all other
//     threads independently re-evaluate the same fire condition from
//     the same ISV source-of-truth (ISV state is read at kernel-launch
//     time, the trigger condition is a pure function of those reads,
//     so every thread computes the same `fire_now` without any
//     cross-block synchronisation).
//
// SP15 Phase 1.3.b-followup contract change (2026-05-07): the
// persistence read migrated from ISV[404] (mean-aggregate of per-env
// DD_PERSISTENCE) to ISV[443] (DD_PERSISTENCE_MAX — max-aggregate
// across envs). The plasticity gate fires when ANY env's persistence
// exceeds the threshold, NOT when the across-env mean exceeds it: one
// set of advantage weights → global firing semantics → max-aggregate
// is the only correct rule. Mean-aggregate would silently gate the
// trigger when only a fraction of envs are in long DD (e.g. 1 env at
// persistence=200 + 7 envs at persistence=0 ⇒ mean=25 ⇒ gate stays
// closed even though one env is deep in the regime that motivates
// plasticity injection).
//
// Slot reads/writes (literal indices to keep this kernel self-contained;
// the canonical names live in `crates/ml/src/cuda_pipeline/sp15_isv_slots.rs`
// and the constructor + state_reset_registry pin those names to these
// indices via compile-time `assert_eq!` regression checks):
//   ISV[443] DD_PERSISTENCE_MAX               (read — Phase 1.3.b-followup)
//   ISV[437] PLASTICITY_PERSISTENCE_THRESHOLD (read)
//   ISV[436] PLASTICITY_FIRED_THIS_FOLD       (read-modify-write — block 0 thread 0 only)
//   ISV[438] PLASTICITY_WARM_BARS_REMAINING   (read-modify-write — block 0 thread 0 only)
//
// Hard rules:
//
//   * `feedback_no_atomicadd` — weight-reset is per-thread independent
//     (each thread writes a unique element of `advantage_head_weights`);
//     no reductions, no `atomicAdd`. The trigger logic remains
//     single-thread (block 0 / thread 0) so its ISV writes are
//     race-free without atomics.
//   * `feedback_no_partial_refactor` — kernel + launcher signature
//     change + 4 oracle tests (3 existing migrated + 1 new) land
//     atomically in this commit (Wave 4.2); consumer wiring (action
//     selection forces Hold via `effective_cooldown =
//     max(COOLDOWN_BARS_REMAINING, PLASTICITY_WARM_BARS_REMAINING)`
//     remains the separate follow-up tracked under Phase 3.5.4.c.
//   * `feedback_no_stubs` — kernel performs the actual Kaiming-He
//     reset; no `(void)` casts remain. Wave 4.2 closes the Phase
//     3.5.4-deferred weight-reset NO-OP per
//     `feedback_wire_everything_up.md`.
//   * `feedback_isv_for_adaptive_bounds` — PLASTICITY_PERSISTENCE_
//     THRESHOLD is ISV-tracked (slot 437); the producer driving it
//     from the running mean of dd_persistence_bars is deferred but
//     the kernel reads it rather than a hardcoded literal so the
//     follow-up swap is a no-op at the consumer.
//   * `feedback_no_cpu_compute_strict` — cuRAND init + sampling are
//     GPU-resident (`curand_init` / `curand_normal` are device-only
//     functions inlined by nvcc into the cubin from
//     `<curand_kernel.h>`).
//   * `pearl_no_host_branches_in_captured_graph` — `seed` is a kernel
//     argument (a `u64` value, not a host-side branch); the launcher
//     itself can be invoked from inside a graph capture region
//     because the kernel-launch shape (grid/block) is fixed at the
//     call (a function of `n_weights / 10` which is also a launch
//     argument, not a host-conditional).

#include <curand_kernel.h>

extern "C" __global__ void plasticity_injection_kernel(
    float* __restrict__ isv,
    float* __restrict__ advantage_head_weights,
    int                 n_weights,
    float               m_warm,                  /* default 200 bars */
    int                 fan_in,                  /* Kaiming gain denominator */
    unsigned long long  seed                     /* cuRAND seed */
) {
    // Every thread independently reads the ISV state to determine
    // `fire_now`. The trigger condition is a pure function of those
    // reads, so all threads converge on the same answer without any
    // cross-block synchronisation. The same reads also drive block 0
    // thread 0's ISV-write path below.
    //
    // Phase 1.3.b-followup (2026-05-07): the persistence read uses
    // ISV[DD_PERSISTENCE_MAX_INDEX=443] (max across envs, written by
    // `dd_state_reduce_kernel`) so the gate fires when ANY env's
    // persistence exceeds the threshold — one shared advantage-head
    // means a global firing condition.
    const float persistence = isv[/*DD_PERSISTENCE_MAX_INDEX*/ 443];
    const float threshold   = isv[/*PLASTICITY_PERSISTENCE_THRESHOLD_INDEX*/ 437];
    const float fired       = isv[/*PLASTICITY_FIRED_THIS_FOLD_INDEX*/ 436];

    // Trigger: persistence exceeds threshold AND not yet fired this
    // fold. PLASTICITY_FIRED_THIS_FOLD is reset to 0.0 at fold boundary
    // by the `sp15_plasticity_fired_this_fold` registry-arm dispatch,
    // re-arming the gate for the next fold. The `< 0.5f` half-bit
    // tolerance check matches the existing `cooldown_remaining <= 0.0f`
    // float-equality-with-tolerance pattern in `cooldown_kernel`.
    const bool fire_now = (persistence > threshold) && (fired < 0.5f);

    // Block 0 / thread 0 owns the ISV writes (fired flag + warm-bars
    // counter). Single-thread to mirror the cooldown_kernel /
    // dd_asymmetric_reward_kernel / hold_floor_kernel / regret_signal
    // / dd_penalty_kernel pattern — no reduction, no atomic.
    if (threadIdx.x == 0 && blockIdx.x == 0) {
        float warm = isv[/*PLASTICITY_WARM_BARS_REMAINING_INDEX*/ 438];

        if (fire_now) {
            isv[/*PLASTICITY_FIRED_THIS_FOLD_INDEX*/ 436] = 1.0f;
            warm = m_warm;
        }

        // Per-bar decrement of warm-bars (every call, regardless of
        // fire). The trigger above runs first so the same call that
        // fires the gate also immediately consumes one bar of warm-up
        // — bias downward by 1 is acceptable per the test's [198, 200]
        // tolerance and avoids the bookkeeping branch needed to skip
        // the decrement on the trip step (mirrors `cooldown_kernel`'s
        // identical trigger-then-decrement sequence).
        if (warm > 0.0f) {
            warm -= 1.0f;
        }
        isv[/*PLASTICITY_WARM_BARS_REMAINING_INDEX*/ 438] = warm;

        __threadfence_system();
    }

    // Weight reset: last 10% of advantage-head weights via Kaiming-He
    // sampling. Each thread writes a unique element so no atomic /
    // reduction is required. Threads beyond `reset_count` early-out;
    // when `fire_now == false` every thread early-outs and the kernel
    // is a pure trigger/decrement (the previous Phase 3.5.4 behaviour).
    if (!fire_now) return;

    const int reset_count = n_weights / 10;
    if (reset_count <= 0) return;

    const int reset_start = n_weights - reset_count;
    const int tid         = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= reset_count) return;

    // Per-thread cuRAND state: (seed, sequence=tid, offset=0). Same
    // (seed, tid) pair always produces the same sample, satisfying
    // determinism for oracle-test mean/std assertions. `fan_in` is the
    // input dimension of the advantage head's last Linear layer; for
    // ReLU/GLU activations the Kaiming-He gain is sqrt(2).
    curandState state;
    curand_init(seed, (unsigned long long)tid, 0ULL, &state);

    const float gain   = sqrtf(2.0f);
    const float stddev = gain / sqrtf((float)fan_in);
    advantage_head_weights[reset_start + tid] = curand_normal(&state) * stddev;
}
