// crates/ml/src/cuda_pipeline/q_disagreement_update_kernel.cu
//
// SP14 Layer B Task B.3 (2026-05-05): Earned Gradient Flow producer kernel.
// Per-step compute the K=4↔K=2 mapped argmax mismatch between Q-head's
// 4-way direction action (DIR_SHORT/DIR_HOLD/DIR_LONG/DIR_FLAT) and the
// aux head's 2-way next-bar prediction (down/up), update fast + slow
// EMAs of the disagreement rate, and update a Welford-style variance
// EMA on the same signal in a single launch.
//
// ── Mapping (from `state_layout.cuh`) ───────────────────────────────────
//   DIR_SHORT (=0)  → aux down (=0)        contributes to numerator
//   DIR_HOLD  (=1)  → masked (no contribution; Q is keeping prior position
//                     so there is no NEW directional commitment to score
//                     against the aux head's directional prediction)
//   DIR_LONG  (=2)  → aux up   (=1)        contributes to numerator
//   DIR_FLAT  (=3)  → masked (Q is closing all positions; same rationale
//                     as Hold — no new directional commitment)
//
// Two of the four Q classes carry directional intent (Short/Long); the
// other two (Hold/Flat) are explicitly ABOUT NOT taking a new directional
// position, so it would be incorrect to penalise the aux head for not
// matching them. Masked rows do not contribute to numerator OR denominator.
//
// ── Output (ISV slots — must match `sp14_isv_slots.rs`) ─────────────────
//   ISV[Q_DISAGREEMENT_SHORT_EMA_INDEX  = 383]  fast EMA (α_short)
//   ISV[Q_DISAGREEMENT_LONG_EMA_INDEX   = 384]  slow EMA (α_long)
//   ISV[Q_DISAGREEMENT_VARIANCE_EMA_INDEX = 389] Welford variance EMA
//
// The slot indices are hardcoded below as `#define`s. They MUST match
// the constants in `crates/ml/src/cuda_pipeline/sp14_isv_slots.rs`. Any
// edit to the .rs constants requires a matching edit here. This is the
// standard pattern for SP4/SP5/SP13 ISV producer kernels — the layout
// fingerprint regression test (`fingerprint_bumped_from_pre_b1_1a` and
// the SP14-equivalent that lands when SP14 closes) catches the case
// where a slot is renamed but kernels are not rebuilt.
//
// ── Pearl-A first-observation bootstrap ────────────────────────────────
// When `prev_short == 0.5f` AND `prev_long == 0.5f` (the registry-
// prescribed sentinel from `SENTINEL_Q_DISAGREEMENT`) AND the batch has
// at least one valid (non-masked) contribution, both EMAs are replaced
// directly with the first observation per
// `pearl_first_observation_bootstrap.md`. The `0.5f` exact-match is safe
// because 0.5 is exactly representable in IEEE 754 (single-precision
// mantissa = 1.0, exponent = -1).
//
// ── No atomicAdd, no host branches ─────────────────────────────────────
// Block tree-reduce on numerator + count separately, then divide and
// update EMAs in a single thread (tid == 0). Per
// `feedback_no_atomicadd.md`. Pure GPU compute per
// `feedback_no_cpu_compute_strict.md`.
//
// ── Launch contract ────────────────────────────────────────────────────
// Single block, blockDim.x = 256, gridDim = (1, 1, 1). The kernel uses
// dynamic shared memory: `extern __shared__ float smem[]` is split into
// two `blockDim.x`-sized arrays (numerator + count). The launcher MUST
// specify `shared_mem_bytes = 2 * blockDim.x * sizeof(float)` =
// `2 * 256 * 4 = 2048` for the standard 256-thread launch. A
// `shared_mem_bytes = 0` launch will read garbage and corrupt the EMA
// — the SP14 Rust launcher (B.7+) and oracle tests both pass 2048.

#include "state_layout.cuh"  // DIR_SHORT (=0), DIR_HOLD (=1), DIR_LONG (=2), DIR_FLAT (=3)

// ── ISV slot indices (must match crates/ml/src/cuda_pipeline/sp14_isv_slots.rs) ──
//
// Q_DISAGREEMENT_SHORT_EMA_INDEX        = 383
// Q_DISAGREEMENT_LONG_EMA_INDEX         = 384
// Q_DISAGREEMENT_VARIANCE_EMA_INDEX     = 389
//
// These differ from the original SP14 plan (which documented 381/382/387)
// because SP13 closeout shifted the SP14 base by +2 (HOLD_RATE_TARGET=381
// and HOLD_RATE_OBSERVED_EMA=382 landed after the plan was written). See
// `sp14_isv_slots.rs` lines 17-20 for the full rationale.
#define Q_DISAGREEMENT_SHORT_EMA_IDX     383
#define Q_DISAGREEMENT_LONG_EMA_IDX      384
#define Q_DISAGREEMENT_VARIANCE_EMA_IDX  389

extern "C" __global__
void q_disagreement_update_kernel(
    /* Per-bar K=2 softmax tile [B, 2] from `aux_next_bar_forward`
     * (SP13 B1.1a). Row-major; the kernel reads `softmax[b * 2 + 0]`
     * (down) and `softmax[b * 2 + 1]` (up). Same-stream barrier with
     * the producer; safe to read once the launcher has ordered the
     * ops on a single stream.
     */
    const float* __restrict__ aux_softmax,
    /* Per-bar Q-head values. Row-major with stride `q_stride`; the
     * kernel reads the FIRST 4 columns per row (direction Q-values)
     * via `q_logits[b * q_stride + k]` for k in 0..4 and argmaxes
     * over them. Stride accommodates either a tight [B, K_DIR=4]
     * direction-only buffer (q_stride = 4) or the full
     * [B, total_actions=13] q_out_buf (q_stride = 13). The kernel does
     * NOT need calibrated probabilities — only the argmax matters for
     * the disagreement rate.
     *
     * SP14 B.11 (2026-05-05): added `q_stride` parameter so the
     * trainer can pass q_out_buf [B, 13] directly without a stride-
     * stripping copy. K_DIR (the loop bound) stays 4 — the wider
     * stride just skips the magnitude/order/urgency tail per row.
     */
    const float* __restrict__ q_logits,
    /* Global ISV bus. Slots 383, 384, 389 read+written; all other
     * slots untouched.
     */
    float* __restrict__ isv,
    /* Number of bars in the batch. The strided loop covers any
     * `B > 0`; empty batch falls through to the variance update only
     * (count == 0 → batch_mean = 0/1 = 0; sentinel-bootstrap guard
     * via `total_cnt > 0` keeps the EMA from being replaced with 0
     * on cold start with empty batch).
     */
    int B,
    /* Stride between rows in `q_logits` (>= K_DIR=4). Direction-only
     * tight buffer uses 4; full q_out_buf uses 13. SP14 B.11.
     */
    int q_stride,
    /* Fast-EMA blend coefficient (typical 0.3 — ~3-step half-life). */
    float alpha_short,
    /* Slow-EMA blend coefficient (typical 0.05 — ~14-step half-life;
     * stagnation detector counterpart to slot 383).
     */
    float alpha_long,
    /* Welford-EMA blend coefficient for variance (typical 0.05). */
    float alpha_var)
{
    extern __shared__ float smem[];
    float* num_smem = smem;                  // [blockDim.x] numerator (sum of disagreements)
    float* cnt_smem = smem + blockDim.x;     // [blockDim.x] contribution count

    const int tid = threadIdx.x;
    const int K_AUX = 2;
    const int K_DIR = 4;

    // Per-thread strided accumulation across rows.
    float local_num = 0.0f;
    float local_cnt = 0.0f;

    for (int b = tid; b < B; b += blockDim.x) {
        // argmax(aux_softmax[b]) — K=2 so this is a single compare.
        const float as0 = aux_softmax[b * K_AUX + 0];
        const float as1 = aux_softmax[b * K_AUX + 1];
        const int aux_class = (as1 > as0) ? 1 : 0;

        // argmax(q_logits[b, 0..K_DIR]) over the first 4 columns of
        // each row (stride = q_stride). Strict `>` tie-break (matches
        // the C-comparator semantics used elsewhere in the codebase
        // — first-class-wins on ties, no special handling needed
        // because the production logits are ~never bit-equal across
        // classes after a forward pass).
        int q_class = 0;
        float q_max = q_logits[b * q_stride + 0];
        #pragma unroll
        for (int k = 1; k < K_DIR; ++k) {
            const float v = q_logits[b * q_stride + k];
            if (v > q_max) { q_max = v; q_class = k; }
        }

        // Mask Hold and Flat: no new directional commitment to score.
        if (q_class == DIR_HOLD || q_class == DIR_FLAT) {
            continue;
        }

        // Map Q-class to aux space: Short(=0) → down(=0); Long(=2) → up(=1).
        const int q_mapped = (q_class == DIR_SHORT) ? 0 : 1;
        local_num += (q_mapped != aux_class) ? 1.0f : 0.0f;
        local_cnt += 1.0f;
    }

    num_smem[tid] = local_num;
    cnt_smem[tid] = local_cnt;
    __syncthreads();

    // Block tree-reduce on numerator + count separately. blockDim.x is
    // expected to be a power of two (typically 256); the standard
    // pairwise-halving reduction below requires that. Per
    // `feedback_no_atomicadd.md`.
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            num_smem[tid] += num_smem[tid + s];
            cnt_smem[tid] += cnt_smem[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float total_num = num_smem[0];
        const float total_cnt = cnt_smem[0];

        // ── Empty-batch gate (post-train-6fcml fix, 2026-05-07) ────────
        // When the batch has zero non-masked contributions (every row
        // Q-picked Hold or Flat — common for the training replay buffer
        // because Hold/Flat dominate naturally), preserve all three EMAs
        // at their previous values. Without this gate, the kernel would
        // blend `batch_mean = 0/1 = 0` into the EMAs every step, which
        // exponentially decayed the post-rollout signal to zero across
        // the per-step training launches (commit 5608b866b cadence
        // migration). Per `pearl_first_observation_bootstrap.md`:
        // "no observation" preserves prior; only "first observation"
        // replaces sentinel. Decay-on-empty was inconsistent with both
        // rules. Verified via train-6fcml HEALTH_DIAG: post-rollout
        // q_dis_s=0.0595 → post-training q_dis_s=0.0000 in one epoch.
        if (total_cnt > 0.0f) {
            const float batch_mean = total_num / total_cnt;

            const float prev_short = isv[Q_DISAGREEMENT_SHORT_EMA_IDX];
            const float prev_long  = isv[Q_DISAGREEMENT_LONG_EMA_IDX];
            const float prev_var   = isv[Q_DISAGREEMENT_VARIANCE_EMA_IDX];

            // Pearl-A first-observation: sentinel = 0.5 exact (per
            // `SENTINEL_Q_DISAGREEMENT`). 0.5 is exactly representable
            // in IEEE 754 (mantissa 1.0, exponent -1) so the `==`
            // compare is safe. The first observation REPLACES the
            // sentinel directly — see
            // `pearl_first_observation_bootstrap.md`. The previous
            // `&& (total_cnt > 0.0f)` guard on `is_first` is now
            // redundant — the outer empty-batch gate above guarantees
            // total_cnt > 0 here.
            const bool is_first =
                (prev_short == 0.5f) && (prev_long == 0.5f);

            const float new_short = is_first
                ? batch_mean
                : (alpha_short * batch_mean + (1.0f - alpha_short) * prev_short);
            const float new_long = is_first
                ? batch_mean
                : (alpha_long * batch_mean + (1.0f - alpha_long) * prev_long);

            // Welford-style variance EMA: var_new = α_var × (x - mean)²
            // + (1 - α_var) × var_old. Uses `new_short` as the running
            // mean (the fast EMA tracks the mean's current best
            // estimate); the alternative would be a separate
            // mean-tracking slot, but the adaptive k_q controller in
            // B.4 only consumes the variance ratio against
            // `VARIANCE_REF_Q`, so the choice of "mean" here is mostly
            // a numerical-stability detail (small bias toward the
            // fast-EMA estimate, which is what B.4 wants anyway).
            const float diff = batch_mean - new_short;
            const float new_var = alpha_var * (diff * diff) + (1.0f - alpha_var) * prev_var;

            isv[Q_DISAGREEMENT_SHORT_EMA_IDX]    = new_short;
            isv[Q_DISAGREEMENT_LONG_EMA_IDX]     = new_long;
            isv[Q_DISAGREEMENT_VARIANCE_EMA_IDX] = new_var;
        }
        // else: total_cnt == 0 → preserve EMA state untouched. Kernel
        // is a no-op for this step; stream-ordered launches still run.
    }
}
