//! GPU-oracle invariants for [`MultiHeadPolicy`] (Phase 2A-A + 2A-B).
//!
//! Phase 2A-A — forward-pass invariants (5 tests):
//!
//!   1. `gate_probs_sum_to_one` — softmax invariant on the gating head.
//!   2. `pi_probs_mixture_sums_to_one` — exp(pi_logits[b]).sum() ≈ 1 for
//!      every batch row (the mixture log-prob output is a valid distribution).
//!   3. `k1_reduces_to_single_head` — at K=1 the mixture math collapses to
//!      the single-head softmax. With the per-head bias zeroed for K=1
//!      (`b_heads_init[0 * N_ACTIONS + 0] += 0.5` is the ONLY entry at K=1),
//!      we verify the output mixture matches the reference softmax over the
//!      bias-only logits — the single head with a +0.5 bias on action 0.
//!   4. `gate_responds_to_regime_change` — with hand-crafted gating weights
//!      (W_gate[0] indicates dim 0 strongly), changing `regime_h[0]` from
//!      large-positive to large-negative shifts the gate distribution by
//!      more than a chosen threshold (TVD > 0.1).
//!   5. `forward_is_deterministic_across_contexts` — two separate
//!      `MultiHeadPolicy::new` instances with the same seed produce
//!      bit-equal outputs on the same input. Verifies no determinism
//!      regression vs `pearl_determinism_achieved`.
//!
//! Phase 2A-B — backward-pass invariants (4 tests):
//!
//!   6. `backward_gradcheck_w_heads` — numerical gradient check on
//!      `grad_W_heads` via 5-point finite-difference at random entries.
//!      β=0 (aux KL disabled) for a clean Q-distill-chain-only test.
//!   7. `backward_gradcheck_w_gate` — same pattern for `grad_W_gate`.
//!   8. `aux_prior_grad_sums_to_zero_per_head` — the softmax-Jacobian
//!      identity Σ_a (π · (log_ratio − KL))[a] = 0 must hold per (b, k).
//!   9. `backward_is_deterministic_across_contexts` — bit-equal `grad_W_heads`
//!      + `grad_W_gate` across two fresh same-seed instances on identical
//!      inputs. Determinism regression guard.
//!
//! Per `feedback_no_cpu_test_fallbacks`: oracles are analytical (softmax
//! sum, per-head reduction, deterministic seed, finite-difference identity).
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: tests use
//! `write_slice_f32_d_pub` / `read_slice_d_pub` for all transfers; the
//! ISV bus uses `MappedF32Buffer`.
//!
//! Run with:
//!   `cargo test -p ml-alpha --test multi_head_policy_invariants --release -- --ignored --nocapture`

use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use ml_alpha::cfc::snap_features::REGIME_DIM;
use ml_alpha::heads::HIDDEN_DIM;
use ml_alpha::pinned_mem::MappedF32Buffer;
use ml_alpha::rl::common::N_ACTIONS;
use ml_alpha::rl::isv_slots::{
    RL_POLICY_AUX_PRIOR_BETA_INDEX,
    RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX,
    RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX,
    RL_POLICY_GATE_ENTROPY_EMA_INDEX,
    RL_POLICY_GATE_ENTROPY_MEAN_INDEX,
    RL_POLICY_GATE_LR_MULTIPLIER_INDEX,
    RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX,
    RL_POLICY_NUM_HEADS_INDEX,
    RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX,
    RL_SLOTS_END,
};
use ml_alpha::rl::multi_head_policy::{MultiHeadPolicy, MultiHeadPolicyConfig, MAX_K_HEADS};
use ml_alpha::trainer::integrated::{read_slice_d_pub, write_slice_f32_d_pub};
use ml_core::device::MlDevice;
use std::sync::Arc;

/// Phase 2A-D fix B1.3 — `rl_gate_lr_multiplier_controller.cu` cubin,
/// loaded directly so the four controller invariants exercise the kernel
/// in isolation (no trainer dependency).
const RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_gate_lr_multiplier_controller.cubin"
));

/// Relative tolerance for sum-to-1 invariants. fp32 single-block reductions
/// accumulate ~1 ULP per term; 1e-5 is loose enough to absorb that plus the
/// `+1e-12` log-of-zero guard in the forward kernel.
const TOL: f32 = 1e-5;

fn build_device() -> Option<MlDevice> {
    match MlDevice::cuda(0) {
        Ok(d) => Some(d),
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            None
        }
    }
}

fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    let mut d = stream.alloc_zeros::<f32>(host.len())?;
    write_slice_f32_d_pub(stream, host, &mut d)?;
    Ok(d)
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn gate_probs_sum_to_one() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 4usize;
    let k = 3usize;
    let mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0xA1A1,
        },
    )?;

    // Random regime input (range [-1, 1] — covers realistic feature mag).
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0xCAFE);
    let regime: Vec<f32> = (0..b_size * REGIME_DIM)
        .map(|_| r.gen_range(-1.0..1.0))
        .collect();
    let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
        .map(|_| r.gen_range(-0.5..0.5))
        .collect();

    let regime_d = upload_f32(&stream, &regime)?;
    let h_t_d = upload_f32(&stream, &h_t)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;

    let gate_probs = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?;
    for b in 0..b_size {
        let row = &gate_probs[b * k..(b + 1) * k];
        let sum: f32 = row.iter().sum();
        assert!(
            (sum - 1.0).abs() < TOL,
            "gate_probs[b={b}] sum = {sum}, expected ≈1.0 (row = {row:?})"
        );
        // Also assert every entry is in [0, 1] (softmax range invariant).
        for (k_idx, p) in row.iter().enumerate() {
            assert!(
                *p >= 0.0 && *p <= 1.0,
                "gate_probs[b={b}, k={k_idx}] = {p} out of [0, 1]"
            );
        }
    }
    eprintln!("PASS — gate_probs sum to 1.0 (±{TOL}) for all {b_size} batch rows");
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn pi_probs_mixture_sums_to_one() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 4usize;
    let k = 3usize;
    let mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0xB2B2,
        },
    )?;

    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0xBEEF);
    let regime: Vec<f32> = (0..b_size * REGIME_DIM)
        .map(|_| r.gen_range(-1.0..1.0))
        .collect();
    let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
        .map(|_| r.gen_range(-0.5..0.5))
        .collect();

    let regime_d = upload_f32(&stream, &regime)?;
    let h_t_d = upload_f32(&stream, &h_t)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;

    let pi_logits = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
    // pi_logits = log(pi_probs); exp + sum should equal 1 per batch row.
    for b in 0..b_size {
        let row = &pi_logits[b * N_ACTIONS..(b + 1) * N_ACTIONS];
        // log-sum-exp stable accumulation (matches the kernel's
        // numerical guard — `pi_logits = log(pi_probs + 1e-12)`).
        let mut sum = 0.0_f32;
        for &lp in row {
            sum += lp.exp();
        }
        assert!(
            (sum - 1.0).abs() < TOL,
            "exp(pi_logits[b={b}]).sum() = {sum}, expected ≈1.0 (row = {row:?})"
        );
    }
    eprintln!(
        "PASS — exp(pi_logits) sums to 1.0 (±{TOL}) for all {b_size} batch rows (mixture is a valid distribution)"
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn k1_reduces_to_single_head() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 4usize;
    let k = 1usize;
    let mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0xC3C3,
        },
    )?;

    // h_t = 0 so the policy-head matmul produces zero logits;
    // post-init, the bias is the only source of non-zero pi_logits_k.
    //
    // Post Phase 3D-A: at K=1 the head's bias is now
    //   [+0.5, 0, +log(4), 0, 0, 0, 0, 0, 0, 0, 0]
    // (action 0 = ShortLarge → +0.5 per Head 0 specialization;
    //  action 2 = Hold      → +log(4) per Phase 3D-A all-heads lift).
    let h_t = vec![0.0_f32; b_size * HIDDEN_DIM];
    let regime: Vec<f32> = (0..b_size * REGIME_DIM).map(|i| (i as f32) * 0.01).collect();
    let h_t_d = upload_f32(&stream, &h_t)?;
    let regime_d = upload_f32(&stream, &regime)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;

    let pi_logits = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;

    // At K=1: mixture = single head's softmax. With h_t=0 the per-head
    // pre-softmax logits are exactly b_heads above. Let s = exp(0.5),
    // h = exp(log(4)) = 4, and (N_ACTIONS - 2) = 9 unbiased actions.
    //   denom = s + h + 9
    //   p[0]    = s / denom    (ShortLarge bias)
    //   p[2]    = h / denom    (Hold bias)
    //   p[else] = 1 / denom
    let hold_logit = (4.0_f32).ln();
    let denom = 0.5_f32.exp() + hold_logit.exp() + (N_ACTIONS as f32 - 2.0);
    let p0_expected = 0.5_f32.exp() / denom;
    let phold_expected = hold_logit.exp() / denom;
    let pother_expected = 1.0_f32 / denom;
    // The kernel writes pi_logits = log(pi_probs + 1e-12); convert back.
    for b in 0..b_size {
        let row = &pi_logits[b * N_ACTIONS..(b + 1) * N_ACTIONS];
        let p0 = row[0].exp();
        let phold = row[2].exp();
        assert!(
            (p0 - p0_expected).abs() < TOL,
            "k1 single-head: b={b} p[0]={p0} expected {p0_expected}"
        );
        assert!(
            (phold - phold_expected).abs() < TOL,
            "k1 single-head (Phase 3D-A Hold lift): b={b} p[2]={phold} expected {phold_expected}"
        );
        for a in 0..N_ACTIONS {
            if a == 0 || a == 2 { continue; }
            let pa = row[a].exp();
            assert!(
                (pa - pother_expected).abs() < TOL,
                "k1 single-head: b={b} p[{a}]={pa} expected {pother_expected}"
            );
        }
    }
    // Phase 3D-A invariant: P(Hold) > P(any specific open action) at init.
    // 4 open actions are 0,1,5,6. With +log(4) on Hold and +0.5 only on
    // action 0, Hold dominates: phold = 4/denom, p_open_max = max(exp(0.5),1)/denom.
    let p_open_max = 0.5_f32.exp().max(1.0) / denom;
    assert!(
        phold_expected > p_open_max,
        "Phase 3D-A bias goal not met: P(hold)={phold_expected} should exceed \
         P(any open)={p_open_max}"
    );
    eprintln!(
        "PASS — K=1 Phase 3D-A bias: p[0]={p0_expected:.6}, p[hold]={phold_expected:.6}, \
         p[other]={pother_expected:.6} (hold > any open ✓)"
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn phase_3d_a_hold_bias_propagates_all_heads() -> Result<()> {
    // Phase 3D-A invariant: with h_t = 0, P(Hold | head k) should be the
    // dominant action (or tied with the specialization bias) for EVERY
    // head k ∈ [0, K). This guards against accidental removal of the
    // all-heads `+log(4)` lift in `MultiHeadPolicy::new` and ensures
    // every gate-routed head agrees on the structural Hold prior.
    let Some(dev) = build_device() else { return Ok(()); };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 2usize;
    let k = 3usize;
    let mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig { k, b_size, seed: 0xC3D4 },
    )?;

    // h_t = 0 → per-head pre-softmax logits = b_heads only.
    // pi_probs_k_d holds [B × K × N_ACTIONS] per-head softmax probs.
    let h_t = vec![0.0_f32; b_size * HIDDEN_DIM];
    let regime = vec![0.0_f32; b_size * REGIME_DIM];
    let h_t_d = upload_f32(&stream, &h_t)?;
    let regime_d = upload_f32(&stream, &regime)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;

    let pi_probs_k = read_slice_d_pub(&stream, &mhp.pi_probs_k_d, b_size * k * N_ACTIONS)?;

    // Per-head expectations at h_t=0:
    //   Head 0: b = [+0.5, 0, +log(4), 0..0]  → P(hold) > P(short_large) since 4 > exp(0.5)≈1.65
    //   Head 1: b = [0, 0, +log(4), 0, 0, +0.5, +0.5, 0, 0, 0, 0]
    //                                          → P(hold) > P(long_small)+P(long_large)? No, just P(hold) > each open individually
    //   Head 2: b = [0, 0, +0.5+log(4), 0..0] → P(hold) very dominant
    for b in 0..b_size {
        for head in 0..k {
            let row_start = (b * k + head) * N_ACTIONS;
            let row = &pi_probs_k[row_start..row_start + N_ACTIONS];
            let p_hold = row[2];
            // Hold should be the single largest action probability per head.
            for a in 0..N_ACTIONS {
                if a == 2 { continue; }
                assert!(
                    p_hold > row[a],
                    "Phase 3D-A bias violated: b={b} head={head} P(hold)={p_hold} \
                     should exceed P(a={a})={}",
                    row[a],
                );
            }
        }
    }
    eprintln!(
        "PASS — Phase 3D-A Hold-bias propagates to all K={k} heads (h_t=0 → P(Hold) dominant in every head)"
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn gate_responds_to_regime_change() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 2usize;
    let k = 3usize;
    let mut mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0xD4D4,
        },
    )?;

    // Overwrite the gating weights so head 0 strongly tracks regime[0]:
    //   W_gate[0, :] = [+5.0, 0, 0, 0, 0, 0]
    //   W_gate[1, :] = [0, 0, 0, 0, 0, 0]
    //   W_gate[2, :] = [-5.0, 0, 0, 0, 0, 0]
    // This forces a strong gate response to regime[0]: large positive
    // regime[0] routes mass to head 0, large negative routes to head 2.
    let mut w_gate_host = vec![0.0_f32; k * REGIME_DIM];
    w_gate_host[0 * REGIME_DIM + 0] = 5.0;
    w_gate_host[2 * REGIME_DIM + 0] = -5.0;
    write_slice_f32_d_pub(&stream, &w_gate_host, &mut mhp.w_gate_d)?;

    // h_t doesn't influence the gate (gate reads only regime_h).
    let h_t = vec![0.0_f32; b_size * HIDDEN_DIM];
    let h_t_d = upload_f32(&stream, &h_t)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    // Run 1: regime[0] = +1 (one batch) / +1 (other batch) → mass routes to head 0.
    let regime_pos = vec![1.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0];
    let regime_pos_d = upload_f32(&stream, &regime_pos)?;
    mhp.forward(&h_t_d, &regime_pos_d, &mut pi_logits_out)?;
    let gp_pos = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?;

    // Run 2: regime[0] = -1 / -1 → mass routes to head 2.
    let regime_neg = vec![-1.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0];
    let regime_neg_d = upload_f32(&stream, &regime_neg)?;
    mhp.forward(&h_t_d, &regime_neg_d, &mut pi_logits_out)?;
    let gp_neg = read_slice_d_pub(&stream, &mhp.gate_probs_d, b_size * k)?;

    // Total Variation Distance between the two gate distributions per batch.
    // TVD = 0.5 × Σ_k |p1[k] − p2[k]|. With ±5σ logit swing on head 0 vs
    // head 2 (gap of 10 in pre-softmax), TVD should easily exceed 0.5.
    for b in 0..b_size {
        let row_pos = &gp_pos[b * k..(b + 1) * k];
        let row_neg = &gp_neg[b * k..(b + 1) * k];
        let tvd: f32 = row_pos
            .iter()
            .zip(row_neg.iter())
            .map(|(p, q)| (p - q).abs())
            .sum::<f32>()
            * 0.5;
        assert!(
            tvd > 0.5,
            "gate not responding to regime[0]: b={b} TVD={tvd}; pos={row_pos:?} neg={row_neg:?}"
        );
        // With W_gate[0,0]=+5 and W_gate[2,0]=-5 and r0=+1: gate logits =
        // [+5, 0, -5] → after softmax p[0] dominates. Sanity-check that
        // p[0] is the modal head under positive regime.
        let modal_pos =
            row_pos.iter().enumerate().fold(
                0usize,
                |acc, (i, p)| if *p > row_pos[acc] { i } else { acc },
            );
        assert_eq!(
            modal_pos, 0,
            "under regime[0]=+1, expected head 0 to be modal; got head {modal_pos} (row = {row_pos:?})"
        );
        let modal_neg =
            row_neg.iter().enumerate().fold(
                0usize,
                |acc, (i, p)| if *p > row_neg[acc] { i } else { acc },
            );
        assert_eq!(
            modal_neg, 2,
            "under regime[0]=-1, expected head 2 to be modal; got head {modal_neg} (row = {row_neg:?})"
        );
    }
    eprintln!(
        "PASS — gate routes mass between heads 0/2 in response to regime[0] sign (TVD > 0.5 per batch)"
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn forward_is_deterministic_across_contexts() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 4usize;
    let k = 3usize;
    let seed = 0xE5E5;

    // Identical inputs for both runs.
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0xF00D);
    let regime: Vec<f32> = (0..b_size * REGIME_DIM)
        .map(|_| r.gen_range(-1.0..1.0))
        .collect();
    let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
        .map(|_| r.gen_range(-0.5..0.5))
        .collect();
    let regime_d = upload_f32(&stream, &regime)?;
    let h_t_d = upload_f32(&stream, &h_t)?;

    // Run A
    let mhp_a = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed,
        },
    )?;
    let mut out_a = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
    mhp_a.forward(&h_t_d, &regime_d, &mut out_a)?;
    let pi_a = read_slice_d_pub(&stream, &out_a, b_size * N_ACTIONS)?;

    // Run B — fresh instance, identical seed/config
    let mhp_b = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed,
        },
    )?;
    let mut out_b = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
    mhp_b.forward(&h_t_d, &regime_d, &mut out_b)?;
    let pi_b = read_slice_d_pub(&stream, &out_b, b_size * N_ACTIONS)?;

    // Bit-equal — same-seed → identical weights → identical kernel inputs
    // → identical outputs. Use `to_bits()` to catch any sub-ULP drift that
    // would invalidate the determinism contract.
    for (i, (a, b)) in pi_a.iter().zip(pi_b.iter()).enumerate() {
        assert_eq!(
            a.to_bits(),
            b.to_bits(),
            "non-deterministic output at i={i}: a={a} b={b} (bits {} vs {})",
            a.to_bits(),
            b.to_bits()
        );
    }
    eprintln!(
        "PASS — forward bit-equal across two fresh contexts with seed=0x{seed:X} ({} floats compared)",
        pi_a.len()
    );
    Ok(())
}

// ─────────────────────────────────────────────────────────────────────
// Phase 2A-B — backward-pass tests
// ─────────────────────────────────────────────────────────────────────

/// Allocate a mapped-pinned ISV bus, set `β = beta`. Returns the buffer
/// (caller keeps it alive for the lifetime of the kernel launches).
fn make_isv_with_beta(beta: f32) -> Result<MappedF32Buffer> {
    let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) }
        .map_err(|e| anyhow::anyhow!("isv MappedF32Buffer alloc: {e}"))?;
    // Zero the whole bus deterministically.
    for slot in 0..RL_SLOTS_END {
        isv.write_record(slot, 0.0);
    }
    isv.write_record(RL_POLICY_AUX_PRIOR_BETA_INDEX, beta);
    Ok(isv)
}

/// Tolerance for the finite-difference gradcheck.
///
/// ── Why ε=1e-2 ────────────────────────────────────────────────────────
/// In fp32 the gradcheck error has TWO components:
///   * truncation:    ~ε²·f'''/6   (decreases with smaller ε)
///   * fp32 noise:    ~ulp(L)/(2ε·|grad|) ≈ 2^-23·|L|/(2ε·|grad|)
///                    (INCREASES with smaller ε — catastrophic cancellation)
///
/// The forward chain `pi_logits = log(softmax_K(gate_logits) · softmax_A(W_h·h))`
/// gives |L| ~ |log(1/N_ACTIONS)| × |grad_pi_logits| ≈ 22 · 0.3 = 6.6, with
/// `|grad| ~ 1e-2`. At ε=1e-3 noise dominates: ulp(L)/(2ε·|grad|) ≈ 4%,
/// matching the observed 1.7-1.9% relative error (averaged over probes).
///
/// At ε=1e-2 the noise term drops by 10× (to ~0.4%) and truncation rises
/// by 100× — but for this chain truncation at ε=1e-2 is still well under
/// noise, so total error MINIMIZES around ε=1e-2 at ~0.5% relative.
/// Verified by `backward_gradcheck_w_heads_eps_sweep` /
/// `backward_gradcheck_w_gate_eps_sweep`:
///
///   | ε     | max_rel_err W_heads | max_rel_err W_gate |
///   |-------|---------------------|--------------------|
///   | 1e-2  | 1.6e-3              | 5.0e-3             |  ← sweet spot
///   | 5e-3  | 1.3e-3              | 1.8e-2             |
///   | 1e-3  | 1.9e-2              | 3.3e-2             |  ← old setting
///   | 5e-4  | 1.0e-2              | 4.3e-2             |
///   | 1e-4  | 1.0e-1              | 4.7e-1             |  ← cancellation
///
/// Since the error GROWS as ε shrinks, it is NOT truncation (which would
/// shrink ∝ ε²). It is fp32 cancellation in the per-output `pi_logits`
/// values. The analytical gradient is therefore correct to ≤1% — anything
/// larger would not vanish at the cancellation-minimum ε.
///
/// Tolerance 1e-2 matches `tests/backward_finite_diff.rs::max_relative`.
const GRADCHECK_REL_TOL: f32 = 1e-2;
const GRADCHECK_ABS_TOL: f32 = 5e-4;
const GRADCHECK_EPS: f32 = 1e-2;

/// Build a small reproducible `(h_t, regime_h, grad_pi_logits)` triple. The
/// `grad_pi_logits` is the "outer loss" gradient — any nonzero vector
/// will do for gradcheck as long as it's deterministic.
fn build_gradcheck_inputs(
    seed: u64,
    b_size: usize,
) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(seed);
    let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
        .map(|_| r.gen_range(-0.3..0.3))
        .collect();
    let regime: Vec<f32> = (0..b_size * REGIME_DIM)
        .map(|_| r.gen_range(-0.5..0.5))
        .collect();
    let grad_pi_logits: Vec<f32> = (0..b_size * N_ACTIONS)
        .map(|_| r.gen_range(-0.5..0.5))
        .collect();
    (h_t, regime, grad_pi_logits)
}

/// Compute the "loss" implied by an `(outputs, grad_outputs)` pair:
///   L(outputs) = Σ_i outputs[i] · grad_outputs[i]
/// This is the canonical pairing such that dL/d(outputs) = grad_outputs.
/// Finite-differencing this L through a forward pass validates the
/// analytical backward of `grad_outputs → grad_inputs`.
fn linear_loss(outputs: &[f32], grad_outputs: &[f32]) -> f64 {
    debug_assert_eq!(outputs.len(), grad_outputs.len());
    let mut s = 0.0_f64;
    for (o, g) in outputs.iter().zip(grad_outputs.iter()) {
        s += (*o as f64) * (*g as f64);
    }
    s
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn backward_gradcheck_w_heads() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 2usize;
    let k = 2usize;

    // β=0 disables the aux KL contribution so this test isolates the
    // softmax-mixture chain rule on its own. Aux is checked separately
    // by `aux_prior_grad_sums_to_zero_per_head`.
    let isv = make_isv_with_beta(0.0)?;
    let isv_dev_ptr = isv.dev_ptr;

    let mut mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0x6E_AD_AC_AC,
        },
    )?;

    let (h_t_host, regime_host, grad_pi_logits_host) =
        build_gradcheck_inputs(0xF00D_BABE, b_size);
    let h_t_d = upload_f32(&stream, &h_t_host)?;
    let regime_d = upload_f32(&stream, &regime_host)?;
    let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    // 1. Forward + backward at the baseline W_heads to get analytical grad.
    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
    mhp.backward(&grad_pi_logits_d, &h_t_d, &regime_d, isv_dev_ptr)?;

    let analytical = read_slice_d_pub(
        &stream,
        &mhp.grad_w_heads_d,
        k * N_ACTIONS * HIDDEN_DIM,
    )?;
    let w_heads_baseline = read_slice_d_pub(
        &stream,
        &mhp.w_heads_d,
        k * N_ACTIONS * HIDDEN_DIM,
    )?;

    // 2. Choose 5 deterministic perturbation indices.
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0x1234_5678);
    let n_params = k * N_ACTIONS * HIDDEN_DIM;
    let probe_idxs: Vec<usize> = (0..5).map(|_| r.gen_range(0..n_params)).collect();

    let mut max_abs_err = 0.0_f32;
    let mut max_rel_err = 0.0_f32;

    for &idx in &probe_idxs {
        // L(+ε)
        let mut w_plus = w_heads_baseline.clone();
        w_plus[idx] += GRADCHECK_EPS;
        write_slice_f32_d_pub(&stream, &w_plus, &mut mhp.w_heads_d)?;
        mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
        let pl_plus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
        let l_plus = linear_loss(&pl_plus, &grad_pi_logits_host);

        // L(-ε)
        let mut w_minus = w_heads_baseline.clone();
        w_minus[idx] -= GRADCHECK_EPS;
        write_slice_f32_d_pub(&stream, &w_minus, &mut mhp.w_heads_d)?;
        mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
        let pl_minus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
        let l_minus = linear_loss(&pl_minus, &grad_pi_logits_host);

        let num_grad = ((l_plus - l_minus) / (2.0 * GRADCHECK_EPS as f64)) as f32;
        let ana_grad = analytical[idx];
        let abs_err = (num_grad - ana_grad).abs();
        let denom = num_grad.abs().max(ana_grad.abs()).max(1e-8);
        let rel_err = abs_err / denom;

        if abs_err > max_abs_err {
            max_abs_err = abs_err;
        }
        if rel_err > max_rel_err {
            max_rel_err = rel_err;
        }

        // Accept if EITHER relative OR absolute tolerance is met.
        let pass = rel_err < GRADCHECK_REL_TOL || abs_err < GRADCHECK_ABS_TOL;
        assert!(
            pass,
            "grad_W_heads[{idx}]: analytical={ana_grad}, numerical={num_grad}, \
             abs_err={abs_err}, rel_err={rel_err} (tol rel={GRADCHECK_REL_TOL} abs={GRADCHECK_ABS_TOL})"
        );

        // Restore W_heads to baseline.
        write_slice_f32_d_pub(&stream, &w_heads_baseline, &mut mhp.w_heads_d)?;
    }

    eprintln!(
        "PASS — grad_W_heads gradcheck (5 probes, ε={GRADCHECK_EPS}): max_abs_err={max_abs_err:.3e}, max_rel_err={max_rel_err:.3e}"
    );
    Ok(())
}

/// Diagnostic: sweep ε across {1e-2, 5e-3, 1e-3, 5e-4, 1e-4} and check
/// 10 probe indices to characterize whether the gradcheck error is
/// truncation (∝ ε²) or a real bug (~constant across ε).
///
/// Also reports per-(k,a,h) localization of the maximum error to identify
/// whether errors cluster around saturated softmax / init-bias positions.
#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn backward_gradcheck_w_heads_eps_sweep() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 2usize;
    let k = 2usize;

    let isv = make_isv_with_beta(0.0)?;
    let isv_dev_ptr = isv.dev_ptr;

    let mut mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0x6E_AD_AC_AC,
        },
    )?;

    let (h_t_host, regime_host, grad_pi_logits_host) =
        build_gradcheck_inputs(0xF00D_BABE, b_size);
    let h_t_d = upload_f32(&stream, &h_t_host)?;
    let regime_d = upload_f32(&stream, &regime_host)?;
    let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    // Analytical grad (computed once).
    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
    mhp.backward(&grad_pi_logits_d, &h_t_d, &regime_d, isv_dev_ptr)?;
    let analytical = read_slice_d_pub(
        &stream,
        &mhp.grad_w_heads_d,
        k * N_ACTIONS * HIDDEN_DIM,
    )?;
    let w_heads_baseline = read_slice_d_pub(
        &stream,
        &mhp.w_heads_d,
        k * N_ACTIONS * HIDDEN_DIM,
    )?;

    // 10 deterministic probe indices — same seed as production gradcheck
    // so we land on a comparable set.
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0x1234_5678);
    let n_params = k * N_ACTIONS * HIDDEN_DIM;
    let probe_idxs: Vec<usize> = (0..10).map(|_| r.gen_range(0..n_params)).collect();

    let epsilons: [f32; 5] = [1e-2, 5e-3, 1e-3, 5e-4, 1e-4];
    eprintln!("\n── ε-sweep gradcheck on grad_W_heads ─────────────────");
    eprintln!("(k = head, a = action, h = hidden-dim)");
    eprintln!("|     ε     | max_abs_err | max_rel_err | argmax (k,a,h) | ana@argmax | num@argmax |");

    for &eps in epsilons.iter() {
        let mut max_abs_err = 0.0_f32;
        let mut max_rel_err = 0.0_f32;
        let mut argmax_idx = 0usize;
        let mut ana_at = 0.0_f32;
        let mut num_at = 0.0_f32;

        for &idx in &probe_idxs {
            let mut w_plus = w_heads_baseline.clone();
            w_plus[idx] += eps;
            write_slice_f32_d_pub(&stream, &w_plus, &mut mhp.w_heads_d)?;
            mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
            let pl_plus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
            let l_plus = linear_loss(&pl_plus, &grad_pi_logits_host);

            let mut w_minus = w_heads_baseline.clone();
            w_minus[idx] -= eps;
            write_slice_f32_d_pub(&stream, &w_minus, &mut mhp.w_heads_d)?;
            mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
            let pl_minus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
            let l_minus = linear_loss(&pl_minus, &grad_pi_logits_host);

            let num_grad = ((l_plus - l_minus) / (2.0 * eps as f64)) as f32;
            let ana_grad = analytical[idx];
            let abs_err = (num_grad - ana_grad).abs();
            let denom = num_grad.abs().max(ana_grad.abs()).max(1e-8);
            let rel_err = abs_err / denom;

            if abs_err > max_abs_err {
                max_abs_err = abs_err;
            }
            if rel_err > max_rel_err {
                max_rel_err = rel_err;
                argmax_idx = idx;
                ana_at = ana_grad;
                num_at = num_grad;
            }

            write_slice_f32_d_pub(&stream, &w_heads_baseline, &mut mhp.w_heads_d)?;
        }

        // Decode (k, a, h) from flat index. Layout: [K × N_ACTIONS × HIDDEN_DIM].
        let h_idx = argmax_idx % HIDDEN_DIM;
        let a_idx = (argmax_idx / HIDDEN_DIM) % N_ACTIONS;
        let k_idx = argmax_idx / (HIDDEN_DIM * N_ACTIONS);
        eprintln!(
            "| {eps:.0e} | {max_abs_err:.3e}    | {max_rel_err:.3e}   | ({k_idx},{a_idx:>2},{h_idx:>3})  | {ana_at:+.4e} | {num_at:+.4e} |"
        );
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn backward_gradcheck_w_gate_eps_sweep() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 2usize;
    let k = 2usize;

    let isv = make_isv_with_beta(0.0)?;
    let isv_dev_ptr = isv.dev_ptr;

    let mut mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0xAA_BB_CC_DD,
        },
    )?;

    let w_gate_init: Vec<f32> = (0..k * REGIME_DIM)
        .map(|i| ((i as f32) * 0.37).sin() * 0.5)
        .collect();
    write_slice_f32_d_pub(&stream, &w_gate_init, &mut mhp.w_gate_d)?;

    let (h_t_host, regime_host, grad_pi_logits_host) =
        build_gradcheck_inputs(0xC0DE_BEEF, b_size);
    let h_t_d = upload_f32(&stream, &h_t_host)?;
    let regime_d = upload_f32(&stream, &regime_host)?;
    let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
    mhp.backward(&grad_pi_logits_d, &h_t_d, &regime_d, isv_dev_ptr)?;
    let analytical = read_slice_d_pub(&stream, &mhp.grad_w_gate_d, k * REGIME_DIM)?;
    let w_gate_baseline = read_slice_d_pub(&stream, &mhp.w_gate_d, k * REGIME_DIM)?;

    // For W_gate (K * REGIME_DIM = 2*6 = 12 params), probe ALL entries.
    let n_params = k * REGIME_DIM;
    let probe_idxs: Vec<usize> = (0..n_params).collect();

    let epsilons: [f32; 5] = [1e-2, 5e-3, 1e-3, 5e-4, 1e-4];
    eprintln!("\n── ε-sweep gradcheck on grad_W_gate ──────────────────");
    eprintln!("(k = head, r = regime-dim)");
    eprintln!("|     ε     | max_abs_err | max_rel_err | argmax (k,r) | ana@argmax | num@argmax |");

    for &eps in epsilons.iter() {
        let mut max_abs_err = 0.0_f32;
        let mut max_rel_err = 0.0_f32;
        let mut argmax_idx = 0usize;
        let mut ana_at = 0.0_f32;
        let mut num_at = 0.0_f32;

        for &idx in &probe_idxs {
            let mut w_plus = w_gate_baseline.clone();
            w_plus[idx] += eps;
            write_slice_f32_d_pub(&stream, &w_plus, &mut mhp.w_gate_d)?;
            mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
            let pl_plus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
            let l_plus = linear_loss(&pl_plus, &grad_pi_logits_host);

            let mut w_minus = w_gate_baseline.clone();
            w_minus[idx] -= eps;
            write_slice_f32_d_pub(&stream, &w_minus, &mut mhp.w_gate_d)?;
            mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
            let pl_minus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
            let l_minus = linear_loss(&pl_minus, &grad_pi_logits_host);

            let num_grad = ((l_plus - l_minus) / (2.0 * eps as f64)) as f32;
            let ana_grad = analytical[idx];
            let abs_err = (num_grad - ana_grad).abs();
            let denom = num_grad.abs().max(ana_grad.abs()).max(1e-8);
            let rel_err = abs_err / denom;

            if abs_err > max_abs_err {
                max_abs_err = abs_err;
            }
            if rel_err > max_rel_err {
                max_rel_err = rel_err;
                argmax_idx = idx;
                ana_at = ana_grad;
                num_at = num_grad;
            }

            write_slice_f32_d_pub(&stream, &w_gate_baseline, &mut mhp.w_gate_d)?;
        }

        let r_idx = argmax_idx % REGIME_DIM;
        let k_idx = argmax_idx / REGIME_DIM;
        eprintln!(
            "| {eps:.0e} | {max_abs_err:.3e}    | {max_rel_err:.3e}   | ({k_idx},{r_idx})       | {ana_at:+.4e} | {num_at:+.4e} |"
        );
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn backward_gradcheck_w_gate() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 2usize;
    let k = 2usize;

    let isv = make_isv_with_beta(0.0)?;
    let isv_dev_ptr = isv.dev_ptr;

    let mut mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0xAA_BB_CC_DD,
        },
    )?;

    // Make W_gate moderately nonzero so the gate is genuinely sensitive to
    // perturbations (otherwise the gating gradient flows through ~uniform
    // softmax and the FD signal is too small to verify).
    let w_gate_init: Vec<f32> = (0..k * REGIME_DIM)
        .map(|i| ((i as f32) * 0.37).sin() * 0.5)
        .collect();
    write_slice_f32_d_pub(&stream, &w_gate_init, &mut mhp.w_gate_d)?;

    let (h_t_host, regime_host, grad_pi_logits_host) =
        build_gradcheck_inputs(0xC0DE_BEEF, b_size);
    let h_t_d = upload_f32(&stream, &h_t_host)?;
    let regime_d = upload_f32(&stream, &regime_host)?;
    let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
    mhp.backward(&grad_pi_logits_d, &h_t_d, &regime_d, isv_dev_ptr)?;

    let analytical = read_slice_d_pub(&stream, &mhp.grad_w_gate_d, k * REGIME_DIM)?;
    let w_gate_baseline = read_slice_d_pub(&stream, &mhp.w_gate_d, k * REGIME_DIM)?;

    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0x9999_0000);
    let n_params = k * REGIME_DIM;
    let probe_idxs: Vec<usize> = (0..5).map(|_| r.gen_range(0..n_params)).collect();

    let mut max_abs_err = 0.0_f32;
    let mut max_rel_err = 0.0_f32;

    for &idx in &probe_idxs {
        let mut w_plus = w_gate_baseline.clone();
        w_plus[idx] += GRADCHECK_EPS;
        write_slice_f32_d_pub(&stream, &w_plus, &mut mhp.w_gate_d)?;
        mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
        let pl_plus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
        let l_plus = linear_loss(&pl_plus, &grad_pi_logits_host);

        let mut w_minus = w_gate_baseline.clone();
        w_minus[idx] -= GRADCHECK_EPS;
        write_slice_f32_d_pub(&stream, &w_minus, &mut mhp.w_gate_d)?;
        mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
        let pl_minus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?;
        let l_minus = linear_loss(&pl_minus, &grad_pi_logits_host);

        let num_grad = ((l_plus - l_minus) / (2.0 * GRADCHECK_EPS as f64)) as f32;
        let ana_grad = analytical[idx];
        let abs_err = (num_grad - ana_grad).abs();
        let denom = num_grad.abs().max(ana_grad.abs()).max(1e-8);
        let rel_err = abs_err / denom;

        if abs_err > max_abs_err {
            max_abs_err = abs_err;
        }
        if rel_err > max_rel_err {
            max_rel_err = rel_err;
        }

        let pass = rel_err < GRADCHECK_REL_TOL || abs_err < GRADCHECK_ABS_TOL;
        assert!(
            pass,
            "grad_W_gate[{idx}]: analytical={ana_grad}, numerical={num_grad}, \
             abs_err={abs_err}, rel_err={rel_err} (tol rel={GRADCHECK_REL_TOL} abs={GRADCHECK_ABS_TOL})"
        );

        write_slice_f32_d_pub(&stream, &w_gate_baseline, &mut mhp.w_gate_d)?;
    }

    eprintln!(
        "PASS — grad_W_gate gradcheck (5 probes, ε={GRADCHECK_EPS}): max_abs_err={max_abs_err:.3e}, max_rel_err={max_rel_err:.3e}"
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn aux_prior_grad_sums_to_zero_per_head() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 3usize;
    let k = 3usize;

    // Strong aux β to make the contribution dominant + measurable.
    let isv = make_isv_with_beta(1.0)?;
    let isv_dev_ptr = isv.dev_ptr;

    let mut mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0x4711_4711,
        },
    )?;

    // The default `W_heads` is zero, which makes per-head softmax outputs
    // EXACTLY match the priors (by construction — priors mirror init biases),
    // so the aux KL contribution would be identically zero and the test
    // wouldn't exercise the kernel. Override `W_heads` with deterministic
    // non-zero values so π diverges from the prior at this `h_t`.
    use rand::{Rng, SeedableRng};
    use rand_chacha::ChaCha8Rng;
    let mut r = ChaCha8Rng::seed_from_u64(0xDEAD_BEEF);
    let w_heads_init: Vec<f32> = (0..k * N_ACTIONS * HIDDEN_DIM)
        .map(|_| r.gen_range(-0.1..0.1))
        .collect();
    write_slice_f32_d_pub(&stream, &w_heads_init, &mut mhp.w_heads_d)?;

    let h_t_host: Vec<f32> = (0..b_size * HIDDEN_DIM)
        .map(|_| r.gen_range(-0.5..0.5))
        .collect();
    let regime_host: Vec<f32> = (0..b_size * REGIME_DIM)
        .map(|_| r.gen_range(-0.5..0.5))
        .collect();
    // ZERO grad_pi_logits — so the only source of `grad_pi_logits_k` is the
    // aux KL kernel; the softmax-Jacobian sum-to-zero identity then must
    // hold per (b, k).
    let grad_pi_logits_host = vec![0.0_f32; b_size * N_ACTIONS];

    let h_t_d = upload_f32(&stream, &h_t_host)?;
    let regime_d = upload_f32(&stream, &regime_host)?;
    let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?;
    let mut pi_logits_out = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;

    mhp.forward(&h_t_d, &regime_d, &mut pi_logits_out)?;
    mhp.backward(&grad_pi_logits_d, &h_t_d, &regime_d, isv_dev_ptr)?;

    let grad_pi_logits_k = read_slice_d_pub(
        &stream,
        &mhp.grad_pi_logits_k_d,
        b_size * k * N_ACTIONS,
    )?;

    // Mathematical identity: with β>0 and a softmax-parameterized distribution,
    //   Σ_a (π_a · (log(π_a/prior_a) − KL_k))
    //     = Σ_a π_a · log(π_a/prior_a) − KL_k · Σ_a π_a
    //     = KL_k − KL_k = 0
    // The tolerance reflects fp32 accumulation noise across N_ACTIONS=11 terms.
    let tol = 1e-5_f32;
    let mut max_sum_abs = 0.0_f32;
    for b in 0..b_size {
        for head in 0..k {
            let base = (b * k + head) * N_ACTIONS;
            let row = &grad_pi_logits_k[base..base + N_ACTIONS];
            let s: f32 = row.iter().sum();
            if s.abs() > max_sum_abs {
                max_sum_abs = s.abs();
            }
            assert!(
                s.abs() < tol,
                "Σ_a aux_grad[b={b}, k={head}, a]={s} expected 0 (row={row:?})"
            );
        }
    }

    // Sanity: at least one entry must be NONzero — otherwise the aux
    // kernel silently no-op'd (which would mask a real bug).
    let any_nonzero = grad_pi_logits_k.iter().any(|x| x.abs() > 1e-7);
    assert!(
        any_nonzero,
        "aux prior produced all-zero grad with β=1.0 — kernel did not fire"
    );

    eprintln!(
        "PASS — aux prior softmax-Jacobian sum-to-zero identity holds per (b, k) (max |Σ|={max_sum_abs:.3e})"
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn backward_is_deterministic_across_contexts() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 4usize;
    let k = 3usize;
    let seed = 0xCAFE_F00D;

    // Both runs use the SAME ISV (same β) for parity.
    let isv = make_isv_with_beta(0.05)?;
    let isv_dev_ptr = isv.dev_ptr;

    let (h_t_host, regime_host, grad_pi_logits_host) =
        build_gradcheck_inputs(0x1357_2468, b_size);
    let h_t_d = upload_f32(&stream, &h_t_host)?;
    let regime_d = upload_f32(&stream, &regime_host)?;
    let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?;

    // Run A
    let mut mhp_a = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed,
        },
    )?;
    let mut pi_logits_out_a = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
    mhp_a.forward(&h_t_d, &regime_d, &mut pi_logits_out_a)?;
    mhp_a.backward(&grad_pi_logits_d, &h_t_d, &regime_d, isv_dev_ptr)?;
    let gw_heads_a = read_slice_d_pub(&stream, &mhp_a.grad_w_heads_d, k * N_ACTIONS * HIDDEN_DIM)?;
    let gw_gate_a = read_slice_d_pub(&stream, &mhp_a.grad_w_gate_d, k * REGIME_DIM)?;
    let ght_a = read_slice_d_pub(&stream, &mhp_a.grad_h_t_scratch_d, b_size * HIDDEN_DIM)?;

    // Run B — fresh instance, identical seed + ISV
    let mut mhp_b = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed,
        },
    )?;
    let mut pi_logits_out_b = stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
    mhp_b.forward(&h_t_d, &regime_d, &mut pi_logits_out_b)?;
    mhp_b.backward(&grad_pi_logits_d, &h_t_d, &regime_d, isv_dev_ptr)?;
    let gw_heads_b = read_slice_d_pub(&stream, &mhp_b.grad_w_heads_d, k * N_ACTIONS * HIDDEN_DIM)?;
    let gw_gate_b = read_slice_d_pub(&stream, &mhp_b.grad_w_gate_d, k * REGIME_DIM)?;
    let ght_b = read_slice_d_pub(&stream, &mhp_b.grad_h_t_scratch_d, b_size * HIDDEN_DIM)?;

    // Bit-equal across all three backward outputs.
    for (i, (a, b)) in gw_heads_a.iter().zip(gw_heads_b.iter()).enumerate() {
        assert_eq!(
            a.to_bits(),
            b.to_bits(),
            "grad_W_heads[{i}] non-deterministic: a={a} b={b}"
        );
    }
    for (i, (a, b)) in gw_gate_a.iter().zip(gw_gate_b.iter()).enumerate() {
        assert_eq!(
            a.to_bits(),
            b.to_bits(),
            "grad_W_gate[{i}] non-deterministic: a={a} b={b}"
        );
    }
    for (i, (a, b)) in ght_a.iter().zip(ght_b.iter()).enumerate() {
        assert_eq!(
            a.to_bits(),
            b.to_bits(),
            "grad_h_t_scratch[{i}] non-deterministic: a={a} b={b}"
        );
    }
    eprintln!(
        "PASS — backward bit-equal across two fresh contexts with seed=0x{seed:X} \
         (grad_W_heads={}, grad_W_gate={}, grad_h_t_scratch={})",
        gw_heads_a.len(),
        gw_gate_a.len(),
        ght_a.len()
    );
    Ok(())
}

// ─────────────────────────────────────────────────────────────────────
// Phase 2A-D — device-aggregated gate-and-head diagnostics tests
// ─────────────────────────────────────────────────────────────────────
//
// These two tests exercise the `emit_diag_stats` kernel by synthesizing
// known `gate_probs` / `pi_probs_k` distributions directly into the
// MultiHeadPolicy's persistent forward-output buffers (bypassing the
// real forward kernel) and checking that the aggregator-written ISV
// slots match the analytical expectation.
//
// The synthetic write path is the minimal way to assert correctness of
// the reduction kernel itself, independent of the forward-pass softmax
// (already covered by `gate_probs_sum_to_one` /
// `pi_probs_mixture_sums_to_one`).

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn aggregate_diag_gate_probs_mean_correct() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 8usize;
    let k = 3usize;

    let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) }
        .map_err(|e| anyhow::anyhow!("isv MappedF32Buffer alloc: {e}"))?;
    for slot in 0..RL_SLOTS_END {
        isv.write_record(slot, 0.0);
    }
    let isv_dev_ptr = isv.dev_ptr;

    let mut mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0x2A2D_0001,
        },
    )?;

    // ── Case 1: uniform gate (1/K per head) ──────────────────────────
    // gate_probs_mean[k] = 1/K for every k ∈ [0, K), 0 for tail.
    let uniform = 1.0_f32 / (k as f32);
    let mut gate_probs_host = vec![uniform; b_size * k];
    // pi_probs_k: every (b, k) row a uniform 1/N_ACTIONS distribution
    // (so per_head_entropy_mean[k] = log(N_ACTIONS) ≈ 2.398).
    let uniform_a = 1.0_f32 / (N_ACTIONS as f32);
    let pi_probs_k_host = vec![uniform_a; b_size * k * N_ACTIONS];
    write_slice_f32_d_pub(&stream, &gate_probs_host, &mut mhp.gate_probs_d)?;
    write_slice_f32_d_pub(&stream, &pi_probs_k_host, &mut mhp.pi_probs_k_d)?;

    mhp.emit_diag_stats(isv_dev_ptr)?;
    stream.synchronize()?;

    let gpm_base = RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX;
    for k_idx in 0..k {
        let got = isv.read_record(gpm_base + k_idx);
        assert!(
            (got - uniform).abs() < 1e-5,
            "uniform: gate_probs_mean[{k_idx}] = {got}, expected {uniform}"
        );
    }
    // Tail (k >= runtime K) must be zero.
    for k_idx in k..MAX_K_HEADS {
        let got = isv.read_record(gpm_base + k_idx);
        assert_eq!(
            got, 0.0,
            "tail gate_probs_mean[{k_idx}] = {got}, expected exact 0.0"
        );
    }
    // ── Case 2: asymmetric per-batch (still deterministic mean) ─────
    //
    // gate_probs[b, k] = (b + k + 1) / Σ_k'(b + k' + 1) — row-normalized
    // ramps. The analytical mean over B = mean_b [(b+k+1) / S_b] where
    // S_b = Σ_k' (b + k' + 1) = K·(b+1) + K(K-1)/2.
    gate_probs_host.fill(0.0);
    let mut expected_gpm = vec![0.0_f64; k];
    for b in 0..b_size {
        let s_b: f64 = (0..k).map(|k2| (b + k2 + 1) as f64).sum();
        for k2 in 0..k {
            let p = ((b + k2 + 1) as f64) / s_b;
            gate_probs_host[b * k + k2] = p as f32;
            expected_gpm[k2] += p / (b_size as f64);
        }
    }
    write_slice_f32_d_pub(&stream, &gate_probs_host, &mut mhp.gate_probs_d)?;
    // Re-zero ISV slots so we read fresh writes.
    for slot in gpm_base..gpm_base + MAX_K_HEADS {
        isv.write_record(slot, 0.0);
    }

    mhp.emit_diag_stats(isv_dev_ptr)?;
    stream.synchronize()?;

    let mut max_err = 0.0_f64;
    for k_idx in 0..k {
        let got = isv.read_record(gpm_base + k_idx) as f64;
        let err = (got - expected_gpm[k_idx]).abs();
        max_err = max_err.max(err);
        assert!(
            err < 1e-5,
            "asymmetric: gate_probs_mean[{k_idx}] = {got}, expected {} (err {:.2e})",
            expected_gpm[k_idx],
            err
        );
    }
    eprintln!(
        "PASS — gate_probs_mean correct under uniform + asymmetric inputs (max_abs_err = {max_err:.2e})"
    );
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn aggregate_diag_entropy_pins_top_and_bottom() -> Result<()> {
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let b_size = 16usize;
    let k = 3usize;

    let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) }
        .map_err(|e| anyhow::anyhow!("isv MappedF32Buffer alloc: {e}"))?;
    for slot in 0..RL_SLOTS_END {
        isv.write_record(slot, 0.0);
    }
    let isv_dev_ptr = isv.dev_ptr;

    let mut mhp = MultiHeadPolicy::new(
        &dev,
        MultiHeadPolicyConfig {
            k,
            b_size,
            seed: 0x2A2D_0002,
        },
    )?;

    // ── Top: uniform gate (max entropy = log(K)) ─────────────────────
    let uniform = 1.0_f32 / (k as f32);
    let gate_uniform = vec![uniform; b_size * k];
    // Uniform per-head action probs (max per_head_entropy = log(N_ACTIONS)).
    let uniform_a = 1.0_f32 / (N_ACTIONS as f32);
    let pi_uniform = vec![uniform_a; b_size * k * N_ACTIONS];
    write_slice_f32_d_pub(&stream, &gate_uniform, &mut mhp.gate_probs_d)?;
    write_slice_f32_d_pub(&stream, &pi_uniform, &mut mhp.pi_probs_k_d)?;

    mhp.emit_diag_stats(isv_dev_ptr)?;
    stream.synchronize()?;

    let log_k = (k as f32).ln();
    let log_a = (N_ACTIONS as f32).ln();
    let ge_top = isv.read_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX);
    assert!(
        (ge_top - log_k).abs() < 1e-5,
        "top: gate_entropy_mean = {ge_top}, expected log(K) = {log_k}"
    );
    let phe_base = RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX;
    for k_idx in 0..k {
        let got = isv.read_record(phe_base + k_idx);
        assert!(
            (got - log_a).abs() < 1e-5,
            "top: per_head_entropy_mean[{k_idx}] = {got}, expected log(N_ACTIONS) = {log_a}"
        );
    }
    // Tail must be zero.
    for k_idx in k..MAX_K_HEADS {
        let got = isv.read_record(phe_base + k_idx);
        assert_eq!(
            got, 0.0,
            "top tail: per_head_entropy_mean[{k_idx}] = {got}, expected 0.0"
        );
    }
    // gate_argmax_mass with strict-tie-broken-low: all rows uniform, so
    // argmax is index 0 for every batch → mass[0]=1, mass[k>=1]=0.
    let gam_base = RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX;
    let mass0 = isv.read_record(gam_base);
    assert!(
        (mass0 - 1.0).abs() < 1e-6,
        "top: gate_argmax_mass[0] = {mass0}, expected 1.0 (ties broken by lowest index)"
    );
    for k_idx in 1..MAX_K_HEADS {
        let got = isv.read_record(gam_base + k_idx);
        assert_eq!(
            got, 0.0,
            "top: gate_argmax_mass[{k_idx}] = {got}, expected 0.0"
        );
    }

    // ── Bottom: one-hot gate (min entropy = 0) ────────────────────────
    // Spread argmax mass across the K heads so we exercise the per-head
    // counter scatter.
    let mut gate_onehot = vec![0.0_f32; b_size * k];
    for b in 0..b_size {
        let chosen = b % k;
        gate_onehot[b * k + chosen] = 1.0;
    }
    // One-hot per-head action probs (zero entropy per head).
    let mut pi_onehot = vec![0.0_f32; b_size * k * N_ACTIONS];
    for b in 0..b_size {
        for k_idx in 0..k {
            pi_onehot[(b * k + k_idx) * N_ACTIONS + (k_idx % N_ACTIONS)] = 1.0;
        }
    }
    write_slice_f32_d_pub(&stream, &gate_onehot, &mut mhp.gate_probs_d)?;
    write_slice_f32_d_pub(&stream, &pi_onehot, &mut mhp.pi_probs_k_d)?;

    // Reset target slots so we read fresh writes.
    for slot in gam_base..gam_base + MAX_K_HEADS {
        isv.write_record(slot, 0.0);
    }
    for slot in phe_base..phe_base + MAX_K_HEADS {
        isv.write_record(slot, 0.0);
    }
    isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 0.0);

    mhp.emit_diag_stats(isv_dev_ptr)?;
    stream.synchronize()?;

    let ge_bot = isv.read_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX);
    assert!(
        ge_bot.abs() < 1e-6,
        "bottom: gate_entropy_mean = {ge_bot}, expected ≈0"
    );
    for k_idx in 0..k {
        let got = isv.read_record(phe_base + k_idx);
        assert!(
            got.abs() < 1e-6,
            "bottom: per_head_entropy_mean[{k_idx}] = {got}, expected ≈0"
        );
    }
    // argmax_mass should sum to 1 across the K heads (every batch gets
    // one argmax). With chosen = b % k and b_size = 16, k = 3:
    //   head 0: b ∈ {0,3,6,9,12,15} → 6/16
    //   head 1: b ∈ {1,4,7,10,13}   → 5/16
    //   head 2: b ∈ {2,5,8,11,14}   → 5/16
    let mut expected_mass = vec![0.0_f32; k];
    for b in 0..b_size {
        expected_mass[b % k] += 1.0 / (b_size as f32);
    }
    for k_idx in 0..k {
        let got = isv.read_record(gam_base + k_idx);
        assert!(
            (got - expected_mass[k_idx]).abs() < 1e-6,
            "bottom: gate_argmax_mass[{k_idx}] = {got}, expected {}",
            expected_mass[k_idx]
        );
    }
    for k_idx in k..MAX_K_HEADS {
        let got = isv.read_record(gam_base + k_idx);
        assert_eq!(
            got, 0.0,
            "bottom tail: gate_argmax_mass[{k_idx}] = {got}, expected 0.0"
        );
    }
    eprintln!(
        "PASS — entropy pins: top gate_entropy_mean = log(K) = {log_k:.4}, \
         bottom = 0; per_head_entropy top = log(N_ACTIONS) = {log_a:.4}, bottom = 0; \
         argmax_mass sums to 1.0"
    );
    Ok(())
}

// ─────────────────────────────────────────────────────────────────────
// Phase 2A-D fix B1.3 — gate-LR multiplier controller invariants.
//
// The controller reads `gate_entropy_mean` (slot 781) via an EMA →
// escalates / decays / leaves alone `RL_POLICY_GATE_LR_MULTIPLIER_INDEX`
// (slot 790) based on whether the EMA sits above 0.85·log(K),
// below 0.20·log(K), or inside the healthy band.
//
// These four tests exercise the kernel in isolation by allocating a
// fresh `MappedF32Buffer`, seeding the relevant slots, and launching
// the standalone cubin. No `MultiHeadPolicy` instance is required.
// ─────────────────────────────────────────────────────────────────────

/// Allocate a zero-initialised ISV buffer covering all RL slots.
fn isv_zero() -> Result<MappedF32Buffer> {
    let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) }
        .map_err(|e| anyhow::anyhow!("isv MappedF32Buffer alloc: {e}"))?;
    for slot in 0..RL_SLOTS_END {
        isv.write_record(slot, 0.0_f32);
    }
    Ok(isv)
}

/// Launch the standalone `rl_gate_lr_multiplier_controller` once with a
/// pre-populated ISV. Single-thread (1,1,1)/(1,1,1) — mirrors the trainer's
/// `launch_rl_gate_lr_multiplier_controller`. Synchronises before returning
/// so subsequent `isv.read_record` calls see kernel writes through the
/// mapped pinned host pointer.
fn run_gate_lr_controller_once(
    stream: &Arc<CudaStream>,
    func: &cudarc::driver::CudaFunction,
    isv: &MappedF32Buffer,
) -> Result<()> {
    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (1, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(func);
    launch.arg(&isv.dev_ptr);
    unsafe {
        launch.launch(cfg).context("gate-LR controller launch")?;
    }
    stream.synchronize().context("gate-LR controller sync")?;
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn gate_lr_controller_bootstrap_initializes_ema_to_first_observation() -> Result<()> {
    // Phase 2A-D fix B1.3: on the first launch with `bootstrap_done = 0`,
    // the EMA must be set to the current observation (NOT blended with the
    // sentinel-zero starting value) and the bootstrap-done flag must
    // latch to 1.0. This is the canonical first-observation pattern from
    // `pearl_welford_trade_count_is_step_not_trade`.
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let ctx = dev.cuda_context().context("controller ctx")?;
    let module = ctx
        .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec())
        .context("load rl_gate_lr_multiplier_controller cubin")?;
    let func = module
        .load_function("rl_gate_lr_multiplier_controller")
        .context("load rl_gate_lr_multiplier_controller fn")?;

    let isv = isv_zero()?;
    // Seed canonical bootstrap state matching trainer construction:
    //   - K = 3
    //   - multiplier bootstrap = 5.0
    //   - entropy_now = 0.5 (inside the healthy band for K=3:
    //     [0.20·log(3) ≈ 0.220, 0.85·log(3) ≈ 0.934], so the multiplier
    //     stays put — isolates the EMA-bootstrap effect)
    //   - entropy_ema = 0 (will be overwritten by first-observation)
    //   - bootstrap_done = 0 (un-bootstrapped)
    isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0);
    isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0);
    isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 0.5);
    isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 0.0);
    isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 0.0);

    run_gate_lr_controller_once(&stream, &func, &isv)?;

    let ema_after = isv.read_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX);
    let boot_after = isv.read_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX);
    let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX);
    assert!(
        (ema_after - 0.5).abs() < 1e-5,
        "first-observation EMA = {ema_after}, expected 0.5 (must replace, not blend with 0.0)"
    );
    assert!(
        (boot_after - 1.0).abs() < 1e-6,
        "bootstrap_done after first call = {boot_after}, expected 1.0 (monotonic latch)"
    );
    // Healthy-band leaves multiplier alone (5.0 unchanged within 1 ULP).
    assert!(
        (mult_after - 5.0).abs() < 1e-5,
        "healthy-band multiplier = {mult_after}, expected 5.0 (left alone)"
    );

    // Second call should EMA-blend (α = 0.02): with entropy_now still 0.5 and
    // ema = 0.5 the result is still 0.5 (idempotent fixed point) — verifies
    // we're now in the post-bootstrap branch, not re-bootstrapping.
    run_gate_lr_controller_once(&stream, &func, &isv)?;
    let ema_after_2 = isv.read_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX);
    let boot_after_2 = isv.read_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX);
    assert!(
        (ema_after_2 - 0.5).abs() < 1e-5,
        "second-call EMA = {ema_after_2}, expected 0.5 (fixed point — confirms post-bootstrap branch)"
    );
    assert!(
        (boot_after_2 - 1.0).abs() < 1e-6,
        "bootstrap_done after second call = {boot_after_2}, expected 1.0 (latch persists)"
    );
    eprintln!("PASS — first-observation bootstrap: EMA snaps to obs, flag latches to 1.0");
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn gate_lr_controller_escalates_when_entropy_above_threshold() -> Result<()> {
    // Phase 2A-D fix B1.3: when `entropy_ema > 0.85·log(K)` the controller
    // must compound the multiplier upward at +0.3 %/step (≈+50 % over 137
    // steps for K=3). Drive entropy_now = 1.05 (above 0.85·log(3) ≈ 0.934),
    // already past the bootstrap branch, and run 100 steps. Multiplier
    // must strictly increase.
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let ctx = dev.cuda_context().context("controller ctx")?;
    let module = ctx
        .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec())
        .context("load rl_gate_lr_multiplier_controller cubin")?;
    let func = module
        .load_function("rl_gate_lr_multiplier_controller")
        .context("load rl_gate_lr_multiplier_controller fn")?;

    let isv = isv_zero()?;
    isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0);
    // 1.05 > 0.85·log(3) ≈ 0.934 → escalate branch
    isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 1.05);
    // Pre-bootstrap the EMA so we skip the first-observation pin (1.05 is
    // already above threshold so the EMA stays above threshold while the
    // controller compounds).
    isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 1.05);
    isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0);
    isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0);

    let mult_before = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX);
    let n_steps = 100;
    for _ in 0..n_steps {
        run_gate_lr_controller_once(&stream, &func, &isv)?;
    }
    let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX);

    // Analytical expectation: 5.0 × 1.003^100 ≈ 5.0 × 1.3493 ≈ 6.747.
    // Allow generous slack to absorb fp32 compounding noise.
    let expected = 5.0_f32 * 1.003_f32.powi(n_steps as i32);
    assert!(
        mult_after > 6.0,
        "after {n_steps} escalate steps: mult {mult_before} → {mult_after}, expected > 6.0 (analytical {expected:.4})"
    );
    assert!(
        (mult_after - expected).abs() < 0.1,
        "after {n_steps} escalate steps: mult = {mult_after}, expected ≈ {expected:.4} (within ±0.1)"
    );
    // EMA must still be above the over-threshold (otherwise the test
    // doesn't actually exercise the escalate branch end-to-end).
    let ema_after = isv.read_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX);
    let over_thr = 0.85_f32 * 3.0_f32.ln();
    assert!(
        ema_after > over_thr,
        "EMA after {n_steps} steps = {ema_after}, expected > over_threshold {over_thr} (test ran in escalate branch)"
    );
    eprintln!("PASS — escalate branch: {mult_before} → {mult_after} (expected ≈ {expected:.4}) over {n_steps} steps");
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn gate_lr_controller_decays_when_entropy_below_collapse() -> Result<()> {
    // Phase 2A-D fix B1.3: when `entropy_ema < 0.20·log(K)` the controller
    // must back off the multiplier fast (×0.95/step). Drive
    // entropy_now = 0.1 (below 0.20·log(3) ≈ 0.220), pre-bootstrap the EMA
    // to the same value, and run 100 steps. Multiplier should decay
    // rapidly toward — and likely hit — the MIN_FLOOR of 1.0.
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let ctx = dev.cuda_context().context("controller ctx")?;
    let module = ctx
        .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec())
        .context("load rl_gate_lr_multiplier_controller cubin")?;
    let func = module
        .load_function("rl_gate_lr_multiplier_controller")
        .context("load rl_gate_lr_multiplier_controller fn")?;

    let isv = isv_zero()?;
    isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0);
    // 0.1 < 0.20·log(3) ≈ 0.220 → decay branch
    isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 0.1);
    isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 0.1);
    isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0);
    isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0);

    let mult_before = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX);
    let n_steps = 100;
    for _ in 0..n_steps {
        run_gate_lr_controller_once(&stream, &func, &isv)?;
    }
    let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX);

    // Analytical: 5.0 × 0.95^100 ≈ 5.0 × 5.92e-3 ≈ 0.0296 — pre-clamp would
    // be well below MIN_FLOOR. After clamping we expect exactly 1.0.
    assert!(
        mult_after < 1.5,
        "after {n_steps} decay steps: mult {mult_before} → {mult_after}, expected < 1.5"
    );
    assert!(
        (mult_after - 1.0).abs() < 1e-5,
        "after {n_steps} decay steps: mult = {mult_after}, expected ≈ 1.0 (clamped to MIN_FLOOR)"
    );
    eprintln!("PASS — decay branch: {mult_before} → {mult_after} (clamped to MIN_FLOOR=1.0) over {n_steps} steps");
    Ok(())
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn gate_lr_controller_clamps_at_min_floor_and_max_ceiling() -> Result<()> {
    // Phase 2A-D fix B1.3: the controller MUST clamp to [1.0, 50.0]
    // per `pearl_bootstrap_must_respect_clamp_range`. Push the multiplier
    // hard against each boundary by pre-setting an extreme value and
    // letting the kernel apply the active branch for many steps.
    let Some(dev) = build_device() else {
        return Ok(());
    };
    let stream = dev.cuda_stream()?.clone();
    let ctx = dev.cuda_context().context("controller ctx")?;
    let module = ctx
        .load_cubin(RL_GATE_LR_MULTIPLIER_CONTROLLER_CUBIN.to_vec())
        .context("load rl_gate_lr_multiplier_controller cubin")?;
    let func = module
        .load_function("rl_gate_lr_multiplier_controller")
        .context("load rl_gate_lr_multiplier_controller fn")?;

    // ── Case 1: pre-set above MAX_CEIL with escalate branch active ──
    // Pre-setting 100.0 (well above MAX_CEIL=50.0) and entropy above
    // the over-threshold should clamp DOWN to 50.0 on first kernel call.
    {
        let isv = isv_zero()?;
        isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0);
        isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 1.05);
        isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 1.05);
        isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0);
        isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 100.0);

        run_gate_lr_controller_once(&stream, &func, &isv)?;
        let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX);
        assert!(
            mult_after <= 50.0 + 1e-5,
            "above-ceil clamp: pre = 100.0 → post = {mult_after}, expected ≤ 50.0"
        );
        // 100.0 × 1.003 = 100.3, clamped to 50.0.
        assert!(
            (mult_after - 50.0).abs() < 1e-5,
            "above-ceil clamp: post = {mult_after}, expected exact 50.0"
        );
    }

    // ── Case 2: many escalate steps from inside the band shouldn't exceed MAX_CEIL ──
    {
        let isv = isv_zero()?;
        isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0);
        isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 1.05);
        isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 1.05);
        isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0);
        isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 5.0);

        // 5 × 1.003^800 ≈ 5 × 10.94 ≈ 54.7 — pre-clamp value would exceed
        // ceiling; the controller must hold it at 50.0.
        for _ in 0..800 {
            run_gate_lr_controller_once(&stream, &func, &isv)?;
        }
        let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX);
        assert!(
            mult_after <= 50.0 + 1e-5,
            "long-escalate ceiling: 800 steps from 5.0 → {mult_after}, expected ≤ 50.0"
        );
        assert!(
            (mult_after - 50.0).abs() < 1e-5,
            "long-escalate ceiling: 800 steps from 5.0 → {mult_after}, expected ≈ 50.0"
        );
    }

    // ── Case 3: pre-set below MIN_FLOOR with decay branch active ──
    {
        let isv = isv_zero()?;
        isv.write_record(RL_POLICY_NUM_HEADS_INDEX, 3.0);
        isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 0.1);
        isv.write_record(RL_POLICY_GATE_ENTROPY_EMA_INDEX, 0.1);
        isv.write_record(RL_POLICY_GATE_CONTROLLER_BOOTSTRAP_DONE_INDEX, 1.0);
        // Set BELOW MIN_FLOOR — the clamp must raise it back to 1.0.
        isv.write_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX, 0.5);

        run_gate_lr_controller_once(&stream, &func, &isv)?;
        let mult_after = isv.read_record(RL_POLICY_GATE_LR_MULTIPLIER_INDEX);
        assert!(
            mult_after >= 1.0 - 1e-5,
            "below-floor clamp: pre = 0.5 → post = {mult_after}, expected ≥ 1.0"
        );
        // 0.5 × 0.95 = 0.475, clamped UP to 1.0.
        assert!(
            (mult_after - 1.0).abs() < 1e-5,
            "below-floor clamp: post = {mult_after}, expected exact 1.0"
        );
    }
    eprintln!("PASS — multiplier clamps at both boundaries: [1.0, 50.0]");
    Ok(())
}
