//! Phase E.1 Task 12 — H=600 DQN smoke for the kill-criteria gate.
//!
//! Trains a linear Q-network (W [9×10] + b [9], no hidden layer) on the
//! Phase E ExecutionEnv at horizon H=600. Uses ε-greedy action selection
//! with Munchausen target augmentation (alpha_munchausen_target_kernel)
//! and reports the four kill criteria at end of training:
//!
//!   ISV[539] Q_SPREAD_EMA          must be ≥ 0.05
//!   ISV[540] ACTION_ENTROPY_EMA    must be ≥ 0.5 · ln(9) ≈ 1.0986
//!   ISV[541] RETURN_VS_RANDOM_EMA  must be ≥ 0.0
//!   ISV[542] EARLY_Q_MOVEMENT_EMA  must be ≥ 0.01
//!
//! If ALL FOUR pass: H=6000 scale-up (Task 13) is viable. If ANY fail,
//! the plan calls for pivot to NoisyNet (Task 19).
//!
//! ## Architecture
//!
//! Single linear layer — no hidden layer. The 10-dim state vector
//! `[alpha_logit, alpha_confidence, spread_bps, l1_imbalance, ofi_sum_5,
//! mid_drift_5, position, step_normalized, log_tau, log_event_rate]`
//! has meaningful direct features, so linear Q captures real relations
//! like `Q[Buy] ∝ alpha_logit`. If linear can't pass the gate, no
//! architecture upgrade will save it — pivot to NoisyNet.
//!
//! All compute on GPU: forward, grad, weight update via alpha_linear_q
//! kernels; target augmentation via alpha_munchausen_target_kernel;
//! kill criteria via alpha_kill_criteria + apply_pearls_ad chain.
//! Action selection is read-only on CPU (per feedback_cpu_is_read_only,
//! reading 9 Q-values to pick argmax is not compute).
//!
//! ## Run
//!
//! ```bash
//! SQLX_OFFLINE=true cargo run -p ml --release --example alpha_dqn_h600_smoke -- \
//!   --mbp10-dir /home/jgrusewski/Work/foxhunt/test_data/futures-baseline-mbp10/ES.FUT \
//!   --fill-coeffs config/ml/alpha_fill_coeffs.json \
//!   --horizon 600 \
//!   --n-episodes 1000
//! ```

use std::fs::File;
use std::io::Write;
use std::mem::MaybeUninit;
use std::path::PathBuf;

use anyhow::{Context, Result};
use clap::Parser;
use cudarc::driver::{CudaContext, DevicePtr, DevicePtrMut};
use tracing::{info, warn};

// Phase E.4.A.8: Mamba2 temporal encoder (ml-alpha Phase 1d.1).
// Phase E.4.A.T10: Mamba2AdamW for training Mamba2 weights.
use ml_alpha::mamba2_block::{
    Mamba2Block, Mamba2BlockConfig, Mamba2AdamW, Mamba2AdamWConfig,
};
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;

// ── Mapped-pinned helpers ──────────────────────────────────────────
//
// `cuMemHostAlloc(DEVICEMAP|PORTABLE)` allocations: host and device see
// the SAME physical memory; the device pointer obtained via
// `cuMemHostGetDevicePointer_v2` lets a kernel read/write directly,
// while the host pointer lets us write inputs (per-step state) and
// read outputs (per-step action) without `memcpy_htod` / `memcpy_dtoh`.
//
// Kernels writing to mapped-pinned outputs MUST call
// `__threadfence_system()` after the write (alpha_c51.cu does this for
// the Thompson selector) so the value is PCIe-visible to the host.
//
// Mirrors the production `MappedBuffer` pattern in `gpu_training_guard.rs`
// per `feedback_no_htod_htoh_only_mapped_pinned`.

struct MappedI32 {
    host_ptr: *mut i32,
    dev_ptr: cudarc::driver::sys::CUdeviceptr,
    len: usize,
}

impl MappedI32 {
    /// # Safety
    /// Caller must ensure a CUDA context is active on the current thread.
    unsafe fn new(len: usize) -> Result<Self> {
        let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP
            | cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;
        let bytes = len * std::mem::size_of::<i32>();
        let host_ptr = cudarc::driver::result::malloc_host(bytes, flags)
            .map_err(|e| anyhow::anyhow!("mapped i32 alloc: {e}"))?
            as *mut i32;
        std::ptr::write_bytes(host_ptr, 0, len);
        let mut dev_raw = MaybeUninit::uninit();
        cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
            dev_raw.as_mut_ptr(),
            host_ptr as *mut std::ffi::c_void,
            0,
        )
        .result()
        .map_err(|e| anyhow::anyhow!("cuMemHostGetDevicePointer: {e}"))?;
        Ok(Self {
            host_ptr,
            dev_ptr: dev_raw.assume_init(),
            len,
        })
    }
    fn dev_u64(&self) -> u64 {
        self.dev_ptr as u64
    }
    fn read(&self) -> i32 {
        debug_assert!(self.len >= 1);
        unsafe { std::ptr::read_volatile(self.host_ptr) }
    }
}

impl Drop for MappedI32 {
    fn drop(&mut self) {
        unsafe {
            let _ = cudarc::driver::result::free_host(
                self.host_ptr as *mut std::ffi::c_void,
            );
        }
    }
}

struct MappedF32 {
    host_ptr: *mut f32,
    dev_ptr: cudarc::driver::sys::CUdeviceptr,
    len: usize,
}

impl MappedF32 {
    /// # Safety
    /// Caller must ensure a CUDA context is active on the current thread.
    unsafe fn new(len: usize) -> Result<Self> {
        let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP
            | cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;
        let bytes = len * std::mem::size_of::<f32>();
        let host_ptr = cudarc::driver::result::malloc_host(bytes, flags)
            .map_err(|e| anyhow::anyhow!("mapped f32 alloc({len}): {e}"))?
            as *mut f32;
        std::ptr::write_bytes(host_ptr, 0, len);
        let mut dev_raw = MaybeUninit::uninit();
        cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
            dev_raw.as_mut_ptr(),
            host_ptr as *mut std::ffi::c_void,
            0,
        )
        .result()
        .map_err(|e| anyhow::anyhow!("cuMemHostGetDevicePointer: {e}"))?;
        Ok(Self {
            host_ptr,
            dev_ptr: dev_raw.assume_init(),
            len,
        })
    }
    fn dev_u64(&self) -> u64 {
        self.dev_ptr as u64
    }
    fn write(&self, data: &[f32]) {
        debug_assert_eq!(data.len(), self.len, "MappedF32 write length mismatch");
        unsafe {
            std::ptr::copy_nonoverlapping(data.as_ptr(), self.host_ptr, self.len);
        }
    }
}

impl Drop for MappedF32 {
    fn drop(&mut self) {
        unsafe {
            let _ = cudarc::driver::result::free_host(
                self.host_ptr as *mut std::ffi::c_void,
            );
        }
    }
}

use data::providers::databento::{dbn_parser::DbnParser, mbp10::Mbp10Snapshot};
use ml::cuda_pipeline::alpha_isv_slots::{
    ACTION_ENTROPY_EMA_INDEX, ALPHA_ISV_BLOCK_LO, EARLY_Q_MOVEMENT_EMA_INDEX,
    Q_SPREAD_EMA_INDEX, RANDOM_BASELINE_MEAN_INDEX, RANDOM_BASELINE_STD_INDEX,
    RETURN_VS_RANDOM_EMA_INDEX,
};
// `COL_RAW_CLOSE` / `FEAT_DIM` are no longer used directly in this
// binary; the snapshot loader handles fxcache column layout internally.
use ml::env::action_space::N_ACTIONS;
use ml::env::execution_env::{
    EpisodeState, ExecutionEnv, ExecutionEnvConfig, SnapshotRow,
};
// `FillCoeffs` / `FillModel` are loaded via the env::loaders module
// and consumed inside ExecutionEnv — no direct references here.
use ml::trainers::dqn::collect_dbn_files_recursive;

const STATE_DIM: usize = 10;
const TICK: f32 = 0.25;
/// Number of weight floats: n_actions × state_dim = 9 × 10.
const N_WEIGHTS: usize = N_ACTIONS * STATE_DIM;
/// Number of bias floats: n_actions.
const N_BIASES: usize = N_ACTIONS;

/// Action allow-list for the pruning experiment (Phase E.3 follow-up,
/// Hypothesis A: action-variance is the alpha-extraction bottleneck).
/// Indices into the full 9-action space, kept as u8 so they drop straight
/// into the env's `step(action_id: u8, …)`. The Q-network is unchanged
/// (still 9 outputs); the selector below restricts argmax + ε-random to
/// this subset. Unused outputs get stale gradients but are never executed.
const PRUNED_ACTIONS: [u8; 4] = [0, 1, 4, 7]; // Wait, BuyMarket, SellMarket, FlatMarket
const FULL_ACTIONS: [u8; 9] = [0, 1, 2, 3, 4, 5, 6, 7, 8];

#[inline]
fn raw_price_to_f32(fixed: i64) -> f32 {
    (fixed as f64 * 1e-9) as f32
}

#[derive(Debug, Parser)]
#[command(
    name = "alpha_dqn_h600_smoke",
    about = "Phase E.1 Task 12 — H=600 DQN smoke for kill-criteria gate"
)]
struct Cli {
    /// Snapshot source 1: raw MBP-10 directory (bid/ask from real LOB,
    /// alpha_logit=0 placeholder unless --alpha-cache also set).
    #[arg(long)]
    mbp10_dir: Option<PathBuf>,
    /// Snapshot source 2: precomputed fxcache (bid/ask synthesized from
    /// mid at fixed half-tick, 81-dim feature row available). Required when
    /// --alpha-cache is set since the cache indexing is fxcache-aligned.
    #[arg(long)]
    fxcache_path: Option<PathBuf>,
    /// Optional alpha-logits cache produced by `alpha_train_stacker
    /// --alpha-cache-out`. Binary file: [u32 n] [f32 logits[n]]. When set,
    /// SnapshotRow.alpha_logit is populated from the cache instead of 0.0.
    /// Requires --fxcache-path (cache indices align to fxcache bars).
    #[arg(long)]
    alpha_cache: Option<PathBuf>,
    #[arg(long, default_value = "config/ml/alpha_fill_coeffs.json")]
    fill_coeffs: PathBuf,
    #[arg(long, default_value_t = 600)]
    horizon: usize,
    #[arg(long, default_value_t = 1_000)]
    n_episodes: usize,
    #[arg(long, default_value_t = 1)]
    trade_size: i32,
    #[arg(long, default_value_t = 0.0625)]
    cost_per_contract: f32,
    #[arg(long, default_value_t = 0xCAFEBABE_u64)]
    seed: u64,
    #[arg(long, default_value_t = 50)]
    snapshot_interval: usize,
    #[arg(long, default_value_t = 500_000)]
    max_snapshots: usize,
    /// SGD learning rate. With the stabilizer trio (--reward-scale,
    /// --target-update-every, --grad-clip), 1e-4 is the right operating
    /// point — gradients are O(1) in normalized units, target net
    /// breaks the V_soft chase-tail loop, grad clip catches bursts.
    /// Without the stabilizers, lr=1e-4 diverges to NaN within 20
    /// episodes; lr=1e-6 stays finite but undertrains.
    #[arg(long, default_value_t = 1.0e-4)]
    lr: f32,
    /// ε-greedy: ε at start of training.
    #[arg(long, default_value_t = 0.50)]
    eps_start: f32,
    /// ε-greedy: ε at end of training.
    #[arg(long, default_value_t = 0.05)]
    eps_end: f32,
    /// DQN discount factor.
    #[arg(long, default_value_t = 0.99)]
    gamma: f32,
    /// Munchausen scale (Vieillard 2020 default 0.9).
    #[arg(long, default_value_t = 0.9)]
    alpha_m: f32,
    /// Boltzmann temperature for Munchausen (Vieillard default 0.03).
    #[arg(long, default_value_t = 0.03)]
    tau: f32,
    /// Lower clip for τ·log π(a|s) (Vieillard default -1.0).
    #[arg(long, default_value_t = -1.0)]
    log_clip_min: f32,
    /// Episodes between kill-criteria pipeline launches.
    #[arg(long, default_value_t = 50)]
    kill_criteria_every: usize,
    /// Phase E.2 stacker-threshold controller — target trade rate (ISV[544]).
    /// Set once at training start, never reset (TrainingPersist anchor).
    #[arg(long, default_value_t = 0.08)]
    trade_rate_target: f32,
    /// Stacker-threshold controller P-gain on rate error.
    #[arg(long, default_value_t = 0.01)]
    k_threshold: f32,
    /// Kelly-attenuation controller P-gain on Sharpe error.
    #[arg(long, default_value_t = 0.005)]
    k_atten: f32,
    /// Target rollout Sharpe for the Kelly-attenuation controller.
    #[arg(long, default_value_t = 0.5)]
    target_sharpe: f32,
    /// Pearl D Wiener-α floor on the observed-rate EMA (slot 545). 0.4 per
    /// `pearl_wiener_alpha_floor_for_nonstationary` — controller co-adapts
    /// with the policy, so the EMA needs to stay responsive.
    #[arg(long, default_value_t = 0.4)]
    wiener_alpha_floor: f32,
    /// Meta-α for the Wiener variance updates inside the controller.
    #[arg(long, default_value_t = 0.1)]
    ctl_alpha_meta: f32,
    /// Reward normalization scale. Rewards are divided by this before
    /// being passed to the Munchausen target — the network learns
    /// normalized Q-values. Default 1000 ≈ |median reward| at H=600
    /// (≈-2800 median random-baseline reward / sane order). Action
    /// selection and rollout-R reporting use ORIGINAL rewards.
    #[arg(long, default_value_t = 1000.0)]
    reward_scale: f32,
    /// Episodes between hard updates of the target network (W_target,
    /// b_target ← W_online, b_online). Decouples the V_soft(s')
    /// bootstrap target from the policy's per-step drift — essential
    /// for Munchausen DQN stability.
    #[arg(long, default_value_t = 10)]
    target_update_every: usize,
    /// Per-element gradient clip bound. After alpha_linear_q_grad and
    /// before the SGD step, |dW[i]|, |db[i]| are clamped to ±this value.
    /// Default 1.0 — combined with lr=1e-4 caps per-step weight delta
    /// at 1e-4 per element, which is safe.
    #[arg(long, default_value_t = 1.0)]
    grad_clip: f32,
    /// Restrict action selection to {Wait, BuyMarket, SellMarket, FlatMarket}.
    /// Q-network still outputs 9 actions; selector + ε-greedy filter to the
    /// 4-subset. Tests Hypothesis A from the Phase E.3 close-out memo
    /// (action variance is the alpha-extraction bottleneck, not capacity
    /// or fill economics). **FALSIFIED 2026-05-15** — kept for reproducibility.
    #[arg(long, default_value_t = false)]
    pruned_actions: bool,
    /// Use C51 distributional Q-network instead of scalar linear Q.
    /// Output is 9 × n_atoms probability distributions per state; action
    /// selection uses Thompson sampling (inverse-CDF) on GPU; target is
    /// the categorical Bellman projection with Huber negative-tail
    /// compression. Tests the calibration hypothesis (Phase E.3 close-out).
    /// Replaces Munchausen target — C51 has its own implicit entropy bonus.
    #[arg(long, default_value_t = false)]
    c51: bool,
    /// C51 atom-support lower bound (in normalized reward units, i.e.
    /// reward / `reward_scale`). Default -10 covers ~2σ of the random
    /// baseline's normalized terminal-reward distribution.
    #[arg(long, default_value_t = -10.0)]
    c51_vmin: f32,
    /// C51 atom-support upper bound.
    #[arg(long, default_value_t = 10.0)]
    c51_vmax: f32,
    /// C51 atom count. Canonical 51; the kernel caps at 64. delta_z is
    /// derived: `(vmax - vmin) / (n_atoms - 1)`.
    #[arg(long, default_value_t = 51)]
    c51_n_atoms: usize,
    /// Path 3 (Phase E.3 follow-up, 2026-05-15): derive bid/ask from
    /// the fxcache's real `spread_bps` feature instead of fixed
    /// ±0.125-tick synthesis. L2/L3 still synthesized at ±tick from
    /// real L1. Tests whether the variable-spread env improves the
    /// half-tick gap to the Phase 1d.4 baseline.
    #[arg(long, default_value_t = false)]
    real_spread: bool,
    /// Phase E.4.A: enable the temporal-encoder stack (sliding window
    /// buffer + Mamba2 SSM + C51 head). When false, runs the existing
    /// C51-flat path on the current state vector only. The buffer is
    /// maintained but unused until Task 8 wires Mamba2 over it.
    #[arg(long, default_value_t = false)]
    temporal: bool,
    /// Sliding-window length K (in snapshots). 16 starting point; the
    /// kernel max is `MAMBA2_KERNEL_SEQ_MAX = 32`. Sweep 16 / 32 in
    /// later iterations.
    #[arg(long, default_value_t = 16)]
    window_k: usize,
    /// Mamba2 hidden dimension (= sh2 in kernel). 32 starting point.
    /// Also the input dim to the C51 head when --temporal is set
    /// (C51 W shape: [N_ACTIONS * n_atoms, mamba2_hidden_dim]).
    #[arg(long, default_value_t = 32)]
    mamba2_hidden_dim: usize,
    /// Mamba2 SSM state dimension. Kernel max is
    /// `MAMBA2_KERNEL_STATE_MAX = 16`.
    #[arg(long, default_value_t = 16)]
    mamba2_state_dim: usize,
    /// Output JSON path for final verdict + ISV readings.
    #[arg(long, default_value = "config/ml/alpha_dqn_h600_smoke.json")]
    out_path: PathBuf,
}

/// Build the env::SnapshotRow stream from MBP-10 files. Same loader pattern
/// as `alpha_random_baseline.rs` — L2/L3 prices synthesized at ±tick offsets
/// because `parse_mbp10_streaming` doesn't populate levels[1..10].
fn load_snapshots(
    parser: &DbnParser,
    mbp10_dir: &std::path::Path,
    snapshot_interval: usize,
    max_snapshots: usize,
) -> Result<Vec<SnapshotRow>> {
    let files = collect_dbn_files_recursive(mbp10_dir);
    if files.is_empty() {
        anyhow::bail!("no MBP-10 .dbn[.zst] files in {:?}", mbp10_dir);
    }
    info!("Found {} MBP-10 file(s)", files.len());
    let mut rows: Vec<SnapshotRow> = Vec::with_capacity(max_snapshots);
    'files: for file in &files {
        let mut hit_limit = false;
        let result = parser.parse_mbp10_streaming(file, snapshot_interval, |snap: &Mbp10Snapshot| {
            if rows.len() >= max_snapshots {
                hit_limit = true;
                return;
            }
            if snap.levels.is_empty() {
                return;
            }
            let l1 = &snap.levels[0];
            let bid_l1 = raw_price_to_f32(l1.bid_px);
            let ask_l1 = raw_price_to_f32(l1.ask_px);
            if bid_l1 <= 0.0 || ask_l1 <= 0.0 || bid_l1 >= ask_l1 {
                return;
            }
            let mid = 0.5 * (bid_l1 + ask_l1);
            let spread_bps = 10_000.0 * (ask_l1 - bid_l1) / mid;
            let bid_sz = l1.bid_sz as f32;
            let ask_sz = l1.ask_sz as f32;
            let l1_imbalance = if bid_sz + ask_sz > 0.0 {
                bid_sz / (bid_sz + ask_sz)
            } else {
                0.5
            };
            let ofi_sum_5 = bid_sz - ask_sz;
            // Phase E.4.A.3: L1 from real MBP-10; L2-L10 synthesized.
            // Real L2-L10 lookup (per snapshot) is also available from
            // `snap.levels[1..10]` after the parser fix 5c0bcb1fd —
            // wire that in Task 5 follow-on.
            let mut bid_l = [0.0_f32; 10];
            let mut ask_l = [0.0_f32; 10];
            for k in 0..10 {
                bid_l[k] = bid_l1 - (k as f32) * TICK;
                ask_l[k] = ask_l1 + (k as f32) * TICK;
            }
            rows.push(SnapshotRow {
                mid_price: mid,
                bid_l,
                ask_l,
                alpha_logit: 0.0,
                alpha_confidence: 0.5,
                spread_bps,
                l1_imbalance,
                ofi_sum_5,
                mid_drift_5: 0.0,
                time_since_trade_s: 0.0,
                book_event_rate: 5.0,
            });
        });
        match result {
            Ok(c) => info!(
                "  {} → {} streamed snapshots",
                file.file_name().unwrap_or_default().to_string_lossy(),
                c
            ),
            Err(e) => warn!("parser error on {}: {}", file.display(), e),
        }
        if hit_limit {
            break 'files;
        }
    }
    Ok(rows)
}

/// Compute env-state observation as Vec<f32> for direct upload.
fn state_as_vec(env: &ExecutionEnv, ep: &EpisodeState) -> Vec<f32> {
    env.state(ep).to_vec()
}

/// L2 (Frobenius) norm of W concatenated with b — diagnostic anchor for
/// early-Q-movement. Pure read-side reduction; computed CPU-side after
/// dtoh (not on the hot loop, only at episode boundaries).
fn weight_norm(w: &[f32], b: &[f32]) -> f32 {
    let s: f32 = w.iter().map(|x| x * x).sum::<f32>() + b.iter().map(|x| x * x).sum::<f32>();
    s.sqrt()
}

/// L2 distance from init: `||W_now − W_init||_F + ||b_now − b_init||_F` (combined).
/// Captures *direction change* in weight space, not just growth — the kill-
/// criteria kernel computes `|q_early − q_init| / |q_init|`, so the call
/// site sets `q_early = q_init + weight_distance_from_init(...)`, making
/// the kernel's ratio = `distance_from_init / ||W_init||_F`. A pure
/// rotation that keeps `||W||` constant still produces a non-zero distance.
fn weight_distance_from_init(w_now: &[f32], b_now: &[f32], w_init: &[f32], b_init: &[f32]) -> f32 {
    let mut s: f32 = 0.0;
    for (a, b) in w_now.iter().zip(w_init.iter()) {
        let d = a - b;
        s += d * d;
    }
    for (a, b) in b_now.iter().zip(b_init.iter()) {
        let d = a - b;
        s += d * d;
    }
    s.sqrt()
}

/// Snapshot of all trainable Mamba2 parameters at init time. Concatenated
/// flat Vec<f32> in deterministic order so a later snapshot subtracts
/// element-wise. Used by the EARLY_Q_MOVEMENT_EMA diagnostic so the
/// kill criterion captures encoder movement (not just C51 head), which
/// would otherwise under-report system learning when --temporal is on.
fn mamba2_snapshot(
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    block: &Mamba2Block,
) -> anyhow::Result<Vec<f32>> {
    let mut out = Vec::new();
    out.extend(stream.clone_dtoh(&block.w_in.weight).context("dtoh w_in")?);
    out.extend(stream.clone_dtoh(&block.w_in.bias).context("dtoh w_in.b")?);
    out.extend(stream.clone_dtoh(&block.w_a.weight).context("dtoh w_a")?);
    out.extend(stream.clone_dtoh(&block.w_a.bias).context("dtoh w_a.b")?);
    out.extend(stream.clone_dtoh(&block.w_b.weight).context("dtoh w_b")?);
    out.extend(stream.clone_dtoh(&block.w_b.bias).context("dtoh w_b.b")?);
    out.extend(stream.clone_dtoh(&block.w_c).context("dtoh w_c")?);
    out.extend(stream.clone_dtoh(&block.w_out.weight).context("dtoh w_out")?);
    out.extend(stream.clone_dtoh(&block.w_out.bias).context("dtoh w_out.b")?);
    Ok(out)
}

/// L2 distance ||M_now − M_init||_F across the same concatenated layout
/// produced by `mamba2_snapshot`. Returns 0.0 if init has different
/// length (shouldn't happen unless model topology changed).
fn mamba2_weight_distance(now: &[f32], init: &[f32]) -> f32 {
    if now.len() != init.len() {
        return 0.0;
    }
    let mut s: f32 = 0.0;
    for (a, b) in now.iter().zip(init.iter()) {
        let d = a - b;
        s += d * d;
    }
    s.sqrt()
}

/// SplitMix64 RNG — same as ExecutionEnv's ReplayRng so seeds compose
/// cleanly when the smoke is reproduced.
struct SmokeRng {
    state: u64,
}

impl SmokeRng {
    fn new(seed: u64) -> Self {
        Self { state: seed }
    }
    fn next_u64(&mut self) -> u64 {
        self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.state;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }
    fn next_f32(&mut self) -> f32 {
        ((self.next_u64() >> 40) as f32) / ((1u64 << 24) as f32)
    }
}

/// Phase E.3 confidence-gated ε-greedy. If `alpha_confidence < threshold`,
/// force action=0 (Wait) regardless of Q. Otherwise standard ε-greedy
/// restricted to `allowed` (the action allow-list — full 9 or pruned 4).
fn epsilon_greedy_gated(
    q: &[f32],
    alpha_confidence: f32,
    threshold: f32,
    eps: f32,
    allowed: &[u8],
    rng: &mut SmokeRng,
) -> u8 {
    if alpha_confidence < threshold {
        return 0; // Wait
    }
    epsilon_greedy(q, eps, allowed, rng)
}

/// Picks an action via ε-greedy over the `allowed` subset. With probability
/// `eps` uniform-random from `allowed`; otherwise argmax of `q[a]` for
/// `a ∈ allowed`. Returns the raw u8 action index (in the full 9-action
/// numbering, so it drops straight into `env.step`).
fn epsilon_greedy(q: &[f32], eps: f32, allowed: &[u8], rng: &mut SmokeRng) -> u8 {
    if rng.next_f32() < eps {
        allowed[(rng.next_u64() as usize) % allowed.len()]
    } else {
        let mut best_a: u8 = allowed[0];
        let mut best_v: f32 = q[allowed[0] as usize];
        for &a in &allowed[1..] {
            let v = q[a as usize];
            if v > best_v {
                best_v = v;
                best_a = a;
            }
        }
        best_a
    }
}

fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();
    let cli = Cli::parse();
    info!("Phase E.1 Task 12 — H=600 DQN smoke starting");
    info!("  horizon={}, n_episodes={}, lr={}, eps={:.2}→{:.2}",
          cli.horizon, cli.n_episodes, cli.lr, cli.eps_start, cli.eps_end);
    let allowed_actions: &[u8] = if cli.pruned_actions {
        info!("  ACTION SET: pruned ({} actions: Wait, BuyMarket, SellMarket, FlatMarket)", PRUNED_ACTIONS.len());
        &PRUNED_ACTIONS
    } else {
        info!("  ACTION SET: full ({} actions)", FULL_ACTIONS.len());
        &FULL_ACTIONS
    };

    // --- CUDA + cubins ---
    let ctx = CudaContext::new(0).context("CUDA context init")?;
    let stream = ctx.default_stream();

    // Linear Q kernels
    let lq_module = ctx
        .load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_LINEAR_Q_CUBIN.to_vec())
        .context("alpha_linear_q cubin load")?;
    let lq_fwd = lq_module.load_function("alpha_linear_q_forward_kernel")
        .context("forward load")?;
    let lq_grad = lq_module.load_function("alpha_linear_q_grad_kernel")
        .context("grad load")?;
    let lq_sgd = lq_module.load_function("alpha_linear_q_sgd_step_kernel")
        .context("sgd load")?;
    let lq_clip = lq_module.load_function("alpha_clip_inplace_kernel")
        .context("clip load")?;

    // Munchausen target
    let munch_cubin: Vec<u8> = std::fs::read(
        concat!(env!("OUT_DIR"), "/alpha_munchausen_target.cubin")
    ).map_err(|e| anyhow::anyhow!("read munch cubin: {e}"))?;
    let munch_module = ctx.load_cubin(munch_cubin).context("munch cubin load")?;
    let munch_kernel = munch_module.load_function("alpha_munchausen_target_kernel")
        .context("munch load")?;

    // Kill criteria + apply_pearls
    let kc_cubin: Vec<u8> = std::fs::read(
        concat!(env!("OUT_DIR"), "/alpha_kill_criteria.cubin")
    ).map_err(|e| anyhow::anyhow!("read kc cubin: {e}"))?;
    let kc_module = ctx.load_cubin(kc_cubin).context("kc cubin load")?;
    let kc_kernel = kc_module.load_function("alpha_kill_criteria_compute_kernel")
        .context("kc kernel load")?;

    let pearls_cubin: Vec<u8> = std::fs::read(
        concat!(env!("OUT_DIR"), "/apply_pearls_kernel.cubin")
    ).map_err(|e| anyhow::anyhow!("read pearls cubin: {e}"))?;
    let pearls_module = ctx.load_cubin(pearls_cubin).context("pearls cubin load")?;
    let pearls_kernel = pearls_module.load_function("apply_pearls_ad_kernel")
        .context("pearls load")?;

    // Phase E.2 Task 16: stacker-threshold + Kelly-attenuation controller
    let ctl_module = ctx
        .load_cubin(ml::cuda_pipeline::alpha_kernels::STACKER_THRESHOLD_CONTROLLER_CUBIN.to_vec())
        .context("controller cubin load")?;
    let ctl_kernel = ctl_module
        .load_function("stacker_threshold_controller_update")
        .context("controller kernel load")?;

    // Phase E.3 follow-up: C51 distributional Q kernels. Loaded
    // unconditionally — cubin is small and the linear-Q path simply
    // ignores it. When --c51 is set, the training loop swaps in the
    // C51 forward / project / grad / Thompson chain.
    let c51_module = ctx
        .load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_C51_CUBIN.to_vec())
        .context("alpha_c51 cubin load")?;
    // Phase E.4.A.6: circular-buffer push kernel for sliding window.
    let push_module = ctx
        .load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_WINDOW_PUSH_CUBIN.to_vec())
        .context("alpha_window_push cubin load")?;
    let push_kernel = push_module
        .load_function("alpha_window_push_kernel")
        .context("window_push kernel load")?;
    // Phase E.4.A.8.fix: GPU-side h_enriched slot copy (replaces
    // dtoh/htod that was violating feedback_cpu_is_read_only).
    let h_store_kernel = push_module
        .load_function("alpha_h_enriched_store_kernel")
        .context("h_enriched_store kernel load")?;
    if cli.temporal {
        info!(
            "  TEMPORAL: sliding window K={} × state_dim={} = {} floats",
            cli.window_k, STATE_DIM, cli.window_k * STATE_DIM,
        );
    }
    let c51_fwd_kernel = c51_module
        .load_function("alpha_c51_forward_kernel")
        .context("c51 forward load")?;
    let c51_project_kernel = c51_module
        .load_function("alpha_c51_project_kernel")
        .context("c51 project load")?;
    let c51_grad_kernel = c51_module
        .load_function("alpha_c51_grad_kernel")
        .context("c51 grad load")?;
    let c51_expq_kernel = c51_module
        .load_function("alpha_c51_expected_q_kernel")
        .context("c51 expected_q load")?;
    let c51_thompson_kernel = c51_module
        .load_function("alpha_c51_thompson_select_kernel")
        .context("c51 thompson load")?;

    // C51 atom-support derived (constant across episodes; ISV-adaptive
    // would lift this to a controller — out of scope for the minimal
    // borrow per the Phase E.3 wisdom rationale).
    let c51_n_atoms = cli.c51_n_atoms;
    let c51_v_min = cli.c51_vmin;
    let c51_v_max = cli.c51_vmax;
    let c51_delta_z = (c51_v_max - c51_v_min) / (c51_n_atoms.saturating_sub(1).max(1) as f32);
    if cli.c51 {
        info!(
            "  Q-network: C51 distributional ({} atoms, support [{:.2}, {:.2}], Δz={:.4})",
            c51_n_atoms, c51_v_min, c51_v_max, c51_delta_z
        );
        if c51_n_atoms == 0 || c51_n_atoms > 64 {
            anyhow::bail!("--c51-n-atoms must be in (0, 64], got {}", c51_n_atoms);
        }
        if c51_v_max <= c51_v_min {
            anyhow::bail!("--c51-vmax ({}) must exceed --c51-vmin ({})", c51_v_max, c51_v_min);
        }
    } else {
        info!("  Q-network: linear scalar (9 outputs)");
    }

    // --- Load env data ---
    let fill_model = ml::env::loaders::load_fill_model_from_json(&cli.fill_coeffs)?;
    info!("Loaded FillModel from {}", cli.fill_coeffs.display());

    // Load alpha cache first (if any) so we can pass it to the fxcache loader.
    let alpha_cache_vec: Option<Vec<f32>> = if let Some(p) = cli.alpha_cache.as_ref() {
        let cache = ml::env::loaders::load_alpha_cache(p)?;
        info!("Loaded alpha cache: {} entries from {}", cache.len(), p.display());
        Some(cache)
    } else {
        None
    };

    // Choose snapshot source. Exactly one of --fxcache-path / --mbp10-dir must
    // be set (or both, in which case fxcache wins because alpha_cache needs
    // fxcache-aligned indices).
    let rows: Vec<SnapshotRow> = match (cli.fxcache_path.as_ref(), cli.mbp10_dir.as_ref()) {
        (Some(fxc), _) => {
            info!("Loading snapshots from fxcache: {}", fxc.display());
            ml::env::loaders::load_snapshots_from_fxcache(
                fxc,
                cli.max_snapshots,
                alpha_cache_vec.as_deref(),
                cli.real_spread,
                None,  // E.4.A T4: mbp10_dir wiring + --use-real-depth flag lands in T5
            )?
        }
        (None, Some(mbp10)) => {
            if cli.alpha_cache.is_some() {
                anyhow::bail!(
                    "--alpha-cache requires --fxcache-path (cache indices align to fxcache bars)"
                );
            }
            info!("Loading snapshots from MBP-10 dir: {}", mbp10.display());
            let parser = DbnParser::new().context("DbnParser::new")?;
            load_snapshots(&parser, mbp10, cli.snapshot_interval, cli.max_snapshots)?
        }
        (None, None) => {
            anyhow::bail!("must set either --fxcache-path or --mbp10-dir");
        }
    };
    info!("Loaded {} snapshots into env", rows.len());
    if rows.len() <= cli.horizon {
        anyhow::bail!("not enough snapshots ({}) for horizon ({})", rows.len(), cli.horizon);
    }
    let n_rows = rows.len();

    let mut env = ExecutionEnv::new(
        ExecutionEnvConfig {
            horizon_snapshots: cli.horizon,
            trade_size_contracts: cli.trade_size,
            cost_per_contract: cli.cost_per_contract,
        },
        fill_model,
        rows,
        cli.seed,
    );

    // --- Initialize Q-network: Xavier ---
    // C51 lifts the output dimension to N_ACTIONS × n_atoms (logits per
    // atom, softmax-normalized across atoms). Linear-Q stays at N_ACTIONS.
    // --temporal additionally swaps the C51 input feature dim from
    // STATE_DIM=10 to mamba2_hidden_dim (typically 32) — Q-net consumes
    // Mamba2's `h_enriched` instead of raw state.
    let c51_input_dim: usize = if cli.c51 && cli.temporal {
        cli.mamba2_hidden_dim
    } else {
        STATE_DIM
    };
    let n_weights_eff: usize = if cli.c51 {
        N_ACTIONS * c51_n_atoms * c51_input_dim
    } else {
        N_WEIGHTS
    };
    let n_biases_eff: usize = if cli.c51 {
        N_ACTIONS * c51_n_atoms
    } else {
        N_BIASES
    };
    let xavier_scale = (2.0_f32 / c51_input_dim as f32).sqrt();
    let mut rng = SmokeRng::new(cli.seed.wrapping_add(0xDEAD_BEEF));
    let w_init: Vec<f32> = (0..n_weights_eff)
        .map(|_| xavier_scale * 2.0 * (rng.next_f32() - 0.5))
        .collect();
    let b_init: Vec<f32> = vec![0.0; n_biases_eff];
    let q_init_norm = weight_norm(&w_init, &b_init);
    info!(
        "Q-net init: ||W||₂ = {:.4}  (n_weights={}, n_biases={})",
        q_init_norm, n_weights_eff, n_biases_eff
    );

    let mut w_dev = stream.clone_htod(&w_init).context("upload W")?;
    let mut b_dev = stream.clone_htod(&b_init).context("upload b")?;
    // Target network: identical to online at init; periodic hard-update.
    // V_soft / projection bootstrap reads from this.
    let mut w_target_dev = stream.clone_htod(&w_init).context("upload W_target")?;
    let mut b_target_dev = stream.clone_htod(&b_init).context("upload b_target")?;
    let mut dw_dev = stream.alloc_zeros::<f32>(n_weights_eff).context("alloc dW")?;
    let mut db_dev = stream.alloc_zeros::<f32>(n_biases_eff).context("alloc db")?;

    // --- ISV buffer (552 floats) with TrainingPersist anchors set ---
    let mut isv_host: Vec<f32> = vec![0.0; 552];
    // Updated 2026-05-15 with the new FillModel (L2-regularized fit + MBP-10
    // parser full-levels-copy fix). Negligible drift from the pre-fix run
    // (-5185.13 / 4952.85) — confirms the smoke verdict is robust to these
    // bug fixes. See alpha_random_baseline.json for the source numbers.
    isv_host[RANDOM_BASELINE_MEAN_INDEX] = -5191.53;
    isv_host[RANDOM_BASELINE_STD_INDEX] = 4963.62;
    // Phase E.2 Task 17: trade-rate target anchor for the stacker-threshold
    // controller. Set at training start, never reset across folds
    // (TrainingPersist). Default 0.08 = ~8% per-step trade rate.
    isv_host[ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX] =
        cli.trade_rate_target;
    let mut isv_dev = stream.clone_htod(&isv_host).context("upload isv")?;

    // Wiener state for the 4 kill-criteria slots: 4 × [sample_var, diff_var, x_lag] = 12 floats.
    let mut wiener_dev = stream.alloc_zeros::<f32>(12).context("alloc wiener")?;
    // Scratch buffer for kill-criteria producer output: 4 floats.
    let mut kc_scratch_dev = stream.alloc_zeros::<f32>(4).context("alloc kc_scratch")?;
    // Phase E.2: Wiener state for slot 545 (observed-rate EMA): 3 floats.
    let mut ctl_wiener_dev = stream.alloc_zeros::<f32>(3).context("alloc ctl wiener")?;

    // --- Allocate per-episode batch buffers (sized to horizon, reused) ---
    let _h = cli.horizon as i32; // unused but kept for log diagnostics
    let state_dim_i = STATE_DIM as i32;
    let n_act_i = N_ACTIONS as i32;
    let n_atoms_i = c51_n_atoms as i32;
    let mut states_dev = stream.alloc_zeros::<f32>(cli.horizon * STATE_DIM).context("alloc states")?;
    let mut next_states_dev = stream.alloc_zeros::<f32>(cli.horizon * STATE_DIM).context("alloc next_states")?;
    let mut actions_dev = stream.alloc_zeros::<i32>(cli.horizon).context("alloc actions")?;
    let mut rewards_dev = stream.alloc_zeros::<f32>(cli.horizon).context("alloc rewards")?;
    let mut dones_dev = stream.alloc_zeros::<f32>(cli.horizon).context("alloc dones")?;
    // Linear-Q scalar Q outputs (used by linear-Q path and as the
    // input to kill-criteria in BOTH paths — in C51 the expected_q
    // kernel writes E[Z] into the same buffer shape).
    let mut q_current_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS).context("alloc q_current")?;
    let mut q_next_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS).context("alloc q_next")?;
    // Linear-Q TD target (scalar per sample). Unused in C51 path.
    let mut target_dev = stream.alloc_zeros::<f32>(cli.horizon).context("alloc target")?;
    let mut single_state_dev = stream.alloc_zeros::<f32>(STATE_DIM).context("alloc single_state")?;
    let mut single_q_dev = stream.alloc_zeros::<f32>(N_ACTIONS).context("alloc single_q")?;

    // C51 distributional buffers — allocated unconditionally; small
    // overhead when --c51 is off.
    let probs_size_full = cli.horizon * N_ACTIONS * c51_n_atoms;
    let mut probs_current_dev = stream.alloc_zeros::<f32>(probs_size_full)
        .context("alloc probs_current (C51)")?;
    let mut probs_next_dev = stream.alloc_zeros::<f32>(probs_size_full)
        .context("alloc probs_next (C51)")?;
    let mut m_dev = stream.alloc_zeros::<f32>(cli.horizon * c51_n_atoms)
        .context("alloc m (C51 projection target)")?;
    let mut single_probs_dev = stream.alloc_zeros::<f32>(N_ACTIONS * c51_n_atoms)
        .context("alloc single_probs (C51)")?;
    // Mapped-pinned per-step inference buffers (C51 path). Host writes
    // state, kernel reads via dev_ptr; kernel writes action with
    // __threadfence_system(), host reads via volatile.
    let state_pinned = unsafe { MappedF32::new(STATE_DIM)? };
    let action_pinned = unsafe { MappedI32::new(1)? };

    // Phase E.4.A.7: GPU-resident sliding-window state buffer for the
    // temporal encoder. Allocated unconditionally; only populated when
    // cli.temporal is set. Capacity: window_k × STATE_DIM floats.
    // Phase E.4.A.6+A.8: shift+insert chronological window, allocated
    // as a GpuTensor of shape [1, K, in_dim] so Mamba2Block can consume
    // it directly via `forward_train(&window_tensor)`. The push kernel
    // writes via the underlying CudaSlice (no reshape needed — same
    // 160 floats either way).
    let mut window_tensor = GpuTensor::zeros(
        &[1, cli.window_k, STATE_DIM], &stream
    ).map_err(|e| anyhow::anyhow!("alloc window tensor: {e}"))?;
    // Phase E.4.A.8 batched training: store each step's h_enriched
    // (Mamba2 output) so the batched C51 forward at episode end reads
    // pre-computed temporal representations instead of recomputing
    // Mamba2 over each step's window. Size = (horizon + 1) so the
    // terminal next-state h_enriched fits at the end.
    let h_enriched_buf_capacity = (cli.horizon + 1) * cli.mamba2_hidden_dim;
    let mut h_enriched_buf_dev = stream
        .alloc_zeros::<f32>(h_enriched_buf_capacity)
        .context("alloc h_enriched buffer")?;

    // Phase E.4.A.8: construct the Mamba2 temporal encoder when
    // --temporal is set. Hidden_dim = c51_input_dim (so the encoder's
    // output flows directly into the C51 head without an additional
    // projection). State_dim ≤ 16 (kernel max).
    let mut mamba2_block: Option<Mamba2Block> = if cli.temporal {
        let cfg = Mamba2BlockConfig {
            in_dim: STATE_DIM,
            hidden_dim: cli.mamba2_hidden_dim,
            state_dim: cli.mamba2_state_dim,
            seq_len: cli.window_k,
        };
        info!(
            "  TEMPORAL Mamba2: in={} hidden={} state={} K={}",
            cfg.in_dim, cfg.hidden_dim, cfg.state_dim, cfg.seq_len
        );
        Some(Mamba2Block::new(cfg, stream.clone())
            .map_err(|e| anyhow::anyhow!("Mamba2Block init: {e}"))?)
    } else {
        None
    };
    // Phase E.4.A.T10: AdamW for Mamba2 weights (active when --temporal).
    let mut mamba2_adamw: Option<Mamba2AdamW> = if let Some(block) = mamba2_block.as_ref() {
        Some(Mamba2AdamW::new(block, Mamba2AdamWConfig::default())
            .map_err(|e| anyhow::anyhow!("Mamba2AdamW init: {e}"))?)
    } else {
        None
    };
    // Snapshot Mamba2 weights at init so EARLY_Q_MOVEMENT can include
    // encoder movement (not just C51 head). Without this, the diagnostic
    // under-reports system learning when --temporal: the encoder absorbs
    // the gradient signal, the head moves little, and the KC reads ~0.
    let mamba2_init_snap: Option<Vec<f32>> = if let Some(block) = mamba2_block.as_ref() {
        Some(mamba2_snapshot(&stream, block)?)
    } else {
        None
    };
    // d_h_enriched as GpuTensor so it can feed `backward_from_h_enriched`
    // after the C51 grad-input kernel writes per-row gradients into it.
    let mut d_h_enriched_tensor_smoke = GpuTensor::zeros(
        &[cli.horizon, cli.mamba2_hidden_dim], &stream
    ).map_err(|e| anyhow::anyhow!("alloc d_h_enriched smoke: {e}"))?;
    // Per-step train-time windows, written by the new
    // `alpha_train_window_store_batched_kernel` during inference and
    // re-forwarded as one big batched Mamba2 call at end-of-episode.
    let mut train_windows_tensor_smoke = GpuTensor::zeros(
        &[cli.horizon, cli.window_k, STATE_DIM], &stream
    ).map_err(|e| anyhow::anyhow!("alloc train windows smoke: {e}"))?;
    // Load the C51 grad-input kernel and the train-window store kernel.
    let c51_grad_input_kernel = c51_module
        .load_function("alpha_c51_grad_input_kernel")
        .context("c51 grad_input load")?;
    let train_window_store_kernel = push_module
        .load_function("alpha_train_window_store_batched_kernel")
        .context("train_window_store kernel load")?;

    // Kill-criteria inputs: action_counts (i32 × n_actions),
    // scalar_inputs (f32 × 3: rollout_R_mean, q_init_norm, q_early_norm).
    let mut kc_action_counts_dev = stream.alloc_zeros::<i32>(N_ACTIONS).context("alloc kc actions")?;
    let mut kc_scalar_dev = stream.alloc_zeros::<f32>(3).context("alloc kc scalar")?;

    // --- Training loop ---
    let mut episode_rng = SmokeRng::new(cli.seed.wrapping_add(0xFEED));
    let mut action_counts_host = vec![0_i32; N_ACTIONS];
    // Phase E.3: cached threshold from controller (slot 543). Updated after
    // each rollout-end controller invocation; used in action selection for
    // the NEXT episode. Initialised to 0 (no gate) so the first episode's
    // exploration is unrestricted, letting the controller accumulate enough
    // observations to drive a meaningful threshold.
    let mut current_threshold: f32 = 0.0;
    let mut recent_returns: Vec<f32> = Vec::new();
    let max_start = n_rows.saturating_sub(cli.horizon + 1).max(1);

    let mut kc_logs: Vec<(usize, [f32; 4])> = Vec::new();

    for ep in 0..cli.n_episodes {
        let eps = cli.eps_start
            + (cli.eps_end - cli.eps_start) * (ep as f32 / cli.n_episodes.max(1) as f32);
        let start_cursor = (episode_rng.next_u64() as usize) % max_start;
        let env_seed = episode_rng.next_u64();
        env.reset_at(env_seed, start_cursor);
        let mut state = EpisodeState::new();

        // Phase E.4.A.7: clear the temporal window on episode start so
        // Mamba2 sees zero-history for the first window_k-1 steps.
        // Phase E.4.A.T10: also zero the train-time window buffer and the
        // d_h_enriched buffer so terminal slots (past ep_len) contribute
        // zero gradient through the Mamba2 backward.
        if cli.temporal {
            stream
                .memset_zeros(window_tensor.data_mut())
                .context("zero window on episode reset")?;
            stream
                .memset_zeros(train_windows_tensor_smoke.data_mut())
                .context("zero train windows on episode reset")?;
            stream
                .memset_zeros(d_h_enriched_tensor_smoke.data_mut())
                .context("zero d_h_enriched on episode reset")?;
        }

        let mut states_host: Vec<f32> = Vec::with_capacity(cli.horizon * STATE_DIM);
        let mut next_states_host: Vec<f32> = Vec::with_capacity(cli.horizon * STATE_DIM);
        let mut actions_host: Vec<i32> = Vec::with_capacity(cli.horizon);
        let mut rewards_host: Vec<f32> = Vec::with_capacity(cli.horizon);
        let mut dones_host: Vec<f32> = Vec::with_capacity(cli.horizon);

        loop {
            // Per-step forward pass (batch=1). C51 path uses mapped-pinned
            // state input + GPU Thompson selector + mapped-pinned action
            // output — zero memcpy_htod / memcpy_dtoh per step. Linear-Q
            // path retains the legacy CPU-readback + ε-greedy selector.
            let s_vec = state_as_vec(&env, &state);
            let action: u8 = if cli.c51 {
                state_pinned.write(&s_vec);
                // Phase E.4.A.7: shift-and-insert push of current state.
                if cli.temporal {
                    {
                        let (w_ptr, _g1) = window_tensor.data_mut().device_ptr_mut(&stream);
                        unsafe {
                            ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push(
                                &stream, &push_kernel,
                                state_pinned.dev_u64(), w_ptr,
                                cli.window_k as i32, state_dim_i,
                            )?;
                        }
                    }
                    // T10: capture post-push window for end-of-episode re-forward.
                    {
                        let (src_ptr, _g_src) = window_tensor.data().device_ptr(&stream);
                        let (dst_ptr, _g_dst) = train_windows_tensor_smoke.data_mut().device_ptr_mut(&stream);
                        unsafe {
                            ml::cuda_pipeline::alpha_kernels::launch_alpha_train_window_store_batched(
                                &stream, &train_window_store_kernel,
                                src_ptr, dst_ptr, state.step as i32,
                                1, cli.window_k as i32, state_dim_i,
                            )?;
                        }
                    }
                }
                // Phase E.4.A.8: Mamba2 forward over the window → h_enriched
                // (`cache.h_enriched [1, hidden_dim]`). The C51 head reads
                // this device buffer in place of `state_pinned`.
                //
                // The cache is dropped at the end of this scope —
                // inference doesn't need backward (we only train via
                // the batched update at end-of-episode). Backward
                // wiring with cache retention lands in T10.
                // Phase E.4.A.8 + .8.fix: per-step Mamba2 forward; store
                // h_enriched into per-step buffer slot via GPU kernel
                // (no CPU roundtrip); feed h_enriched to C51 forward.
                if cli.temporal {
                    let block = mamba2_block.as_ref()
                        .expect("temporal flag set but Mamba2Block missing");
                    let (_logit, cache) = block.forward_train(&window_tensor)
                        .map_err(|e| anyhow::anyhow!("mamba2 forward: {e}"))?;
                    // GPU-side slot copy: cache.h_enriched → h_enriched_buf_dev[slot..]
                    {
                        let slot_offset = (state.step * cli.mamba2_hidden_dim) as i32;
                        let (src_ptr, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream);
                        let (buf_ptr, _g_buf) = h_enriched_buf_dev.device_ptr_mut(&stream);
                        unsafe {
                            ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store(
                                &stream, &h_store_kernel,
                                src_ptr, buf_ptr, slot_offset,
                                cli.mamba2_hidden_dim as i32,
                            )?;
                        }
                    }
                    let (w_ptr, _g0) = w_dev.device_ptr(&stream);
                    let (b_ptr, _g1) = b_dev.device_ptr(&stream);
                    let (p_ptr, _g2) = single_probs_dev.device_ptr_mut(&stream);
                    let (h_ptr, _g3) = cache.h_enriched.cuda_data().device_ptr(&stream);
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
                            &stream, &c51_fwd_kernel,
                            w_ptr, b_ptr, h_ptr, p_ptr,
                            1, c51_input_dim as i32, n_act_i, n_atoms_i,
                        )?;
                    }
                } else {
                    let (w_ptr, _g0) = w_dev.device_ptr(&stream);
                    let (b_ptr, _g1) = b_dev.device_ptr(&stream);
                    let (p_ptr, _g2) = single_probs_dev.device_ptr_mut(&stream);
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
                            &stream, &c51_fwd_kernel,
                            w_ptr, b_ptr, state_pinned.dev_u64(), p_ptr,
                            1, c51_input_dim as i32, n_act_i, n_atoms_i,
                        )?;
                    }
                }
                let step_seed = episode_rng.next_u64() as u32;
                {
                    let (p_ptr, _g0) = single_probs_dev.device_ptr(&stream);
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
                            &stream, &c51_thompson_kernel,
                            p_ptr,
                            state_pinned.dev_u64(),
                            current_threshold,
                            1, // alpha_confidence is state index 1
                            state_dim_i,
                            c51_v_min, c51_delta_z,
                            step_seed,
                            action_pinned.dev_u64(),
                            1, n_act_i, n_atoms_i,
                        )?;
                    }
                }
                stream.synchronize().context("sync after C51 inference")?;
                action_pinned.read() as u8
            } else {
                stream.memcpy_htod(&s_vec, &mut single_state_dev)
                    .context("htod single state")?;
                {
                    let (w_ptr, _g0) = w_dev.device_ptr(&stream);
                    let (b_ptr, _g1) = b_dev.device_ptr(&stream);
                    let (s_ptr, _g2) = single_state_dev.device_ptr(&stream);
                    let (q_ptr, _g3) = single_q_dev.device_ptr_mut(&stream);
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_forward(
                            &stream, &lq_fwd,
                            w_ptr, b_ptr, s_ptr, q_ptr,
                            1, state_dim_i, n_act_i,
                        )?;
                    }
                }
                stream.synchronize().context("sync after per-step fwd")?;
                let q_host = stream.clone_dtoh(&single_q_dev).context("dtoh Q")?;
                epsilon_greedy_gated(
                    &q_host,
                    s_vec[1],
                    current_threshold,
                    eps,
                    allowed_actions,
                    &mut episode_rng,
                )
            };

            let (_next_state_arr, reward, done) = env.step(action, &mut state)
                .ok_or_else(|| anyhow::anyhow!("step returned None"))?;

            // Buffer transition for batched update
            states_host.extend_from_slice(&s_vec);
            // next-state observation: get the env state AT THE NEW CURSOR.
            // env.state(&state) gives current state; after env.step, cursor
            // and state are updated to s'.
            let s_next_vec = state_as_vec(&env, &state);
            next_states_host.extend_from_slice(&s_next_vec);
            actions_host.push(action as i32);
            rewards_host.push(reward);
            dones_host.push(if done { 1.0 } else { 0.0 });
            action_counts_host[action as usize] += 1;

            if done {
                recent_returns.push(reward);
                break;
            }
        }

        let ep_len = actions_host.len() as i32;
        if ep_len < 2 {
            continue;
        }

        // --- Batched update on this episode's transitions ---
        // Upload
        stream.memcpy_htod(&states_host, &mut states_dev)
            .context("htod states")?;
        stream.memcpy_htod(&next_states_host, &mut next_states_dev)
            .context("htod next_states")?;
        stream.memcpy_htod(&actions_host, &mut actions_dev)
            .context("htod actions")?;
        // Reward normalization: divide by cli.reward_scale so the
        // Munchausen target produces normalized targets. Original
        // rewards stay in rewards_host for reporting / kill-criteria.
        let rewards_normalized: Vec<f32> = rewards_host
            .iter()
            .map(|r| r / cli.reward_scale)
            .collect();
        stream.memcpy_htod(&rewards_normalized, &mut rewards_dev)
            .context("htod rewards")?;
        stream.memcpy_htod(&dones_host, &mut dones_dev)
            .context("htod dones")?;

        if cli.c51 {
            // ── C51 batched compute ─────────────────────────────────
            // Phase E.4.A.8 + .8.fix: in temporal mode, read pre-stored
            // h_enriched from h_enriched_buf_dev (populated GPU-side
            // during per-step inference via alpha_h_enriched_store).
            // Run ONE extra Mamba2 forward on the terminal window to
            // fill slot ep_len (the "next-state" representation for
            // the last training transition).
            //
            // Phase E.4.A.T10: also run a second Mamba2 forward on the
            // collected train_windows_tensor_smoke to get a cache that
            // covers all ep_len visited windows. We use the cache's
            // h_enriched as the curr_input for C51 (so the C51 grad can
            // flow back into Mamba2 via backward_from_h_enriched), while
            // keeping h_enriched_buf_dev offset for next_input (frozen
            // target side, no backward).
            let cache_train: Option<ml_alpha::mamba2_block::Mamba2ForwardCache> = if cli.temporal {
                let block = mamba2_block.as_ref().expect("Mamba2Block missing");
                // Terminal forward fills h_enriched_buf_dev[ep_len].
                {
                    let (_logit, cache) = block.forward_train(&window_tensor)
                        .map_err(|e| anyhow::anyhow!("mamba2 terminal forward: {e}"))?;
                    let term_offset = ((ep_len as usize) * cli.mamba2_hidden_dim) as i32;
                    let (src_ptr, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream);
                    let (buf_ptr, _g_buf) = h_enriched_buf_dev.device_ptr_mut(&stream);
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store(
                            &stream, &h_store_kernel,
                            src_ptr, buf_ptr, term_offset,
                            cli.mamba2_hidden_dim as i32,
                        )?;
                    }
                }
                // Training-time forward on the full [horizon, K, in_dim]
                // batch (slots ep_len..horizon are zero, contribute zero
                // gradient since d_h_enriched stays zero there).
                let (_logit, cache) = block.forward_train(&train_windows_tensor_smoke)
                    .map_err(|e| anyhow::anyhow!("mamba2 train forward: {e}"))?;
                Some(cache)
            } else {
                None
            };
            let (curr_input_ptr_guard, next_input_ptr_guard);
            let (curr_input_ptr, next_input_ptr): (u64, u64) = if cli.temporal {
                let cache = cache_train.as_ref().expect("cache_train missing");
                let (h_ptr_curr, g_curr) = cache.h_enriched.cuda_data().device_ptr(&stream);
                let (h_ptr_next_base, g_next) = h_enriched_buf_dev.device_ptr(&stream);
                curr_input_ptr_guard = g_curr;
                next_input_ptr_guard = g_next;
                let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
                (
                    h_ptr_curr,
                    h_ptr_next_base + (cli.mamba2_hidden_dim as u64) * 4u64,
                )
            } else {
                let (s_ptr, g_curr) = states_dev.device_ptr(&stream);
                let (n_ptr, g_next) = next_states_dev.device_ptr(&stream);
                curr_input_ptr_guard = g_curr;
                next_input_ptr_guard = g_next;
                let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
                (s_ptr, n_ptr)
            };
            // Forward online → probs_current
            {
                let (w_ptr, _g0) = w_dev.device_ptr(&stream);
                let (b_ptr, _g1) = b_dev.device_ptr(&stream);
                let (p_ptr, _g3) = probs_current_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
                        &stream, &c51_fwd_kernel,
                        w_ptr, b_ptr, curr_input_ptr, p_ptr,
                        ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
                    )?;
                }
            }
            // Forward target net → probs_next
            {
                let (w_ptr, _g0) = w_target_dev.device_ptr(&stream);
                let (b_ptr, _g1) = b_target_dev.device_ptr(&stream);
                let (p_ptr, _g3) = probs_next_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
                        &stream, &c51_fwd_kernel,
                        w_ptr, b_ptr, next_input_ptr, p_ptr,
                        ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
                    )?;
                }
            }
            // Bellman categorical projection → m
            {
                let (pn_ptr, _g0) = probs_next_dev.device_ptr(&stream);
                let (r_ptr, _g1) = rewards_dev.device_ptr(&stream);
                let (d_ptr, _g2) = dones_dev.device_ptr(&stream);
                let (m_ptr, _g3) = m_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_project(
                        &stream, &c51_project_kernel,
                        pn_ptr, r_ptr, d_ptr,
                        c51_v_min, c51_v_max, cli.gamma, c51_delta_z,
                        m_ptr, ep_len, n_act_i, n_atoms_i,
                    )?;
                }
            }
            // CE gradient on (probs_current, m, actions). With --temporal
            // the input to the C51 layer is h_enriched (hidden_dim);
            // gradient w.r.t. W computes dW = (p - m) ⊗ input^T, so
            // input pointer must match the forward's input.
            {
                let (p_ptr, _g0) = probs_current_dev.device_ptr(&stream);
                let (m_ptr, _g1) = m_dev.device_ptr(&stream);
                let (a_ptr, _g2) = actions_dev.device_ptr(&stream);
                let (s_ptr, _g3) = states_dev.device_ptr(&stream);
                let (dw_ptr, _g4) = dw_dev.device_ptr_mut(&stream);
                let (db_ptr, _g5) = db_dev.device_ptr_mut(&stream);
                let grad_input_ptr = if cli.temporal { curr_input_ptr } else { s_ptr };
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad(
                        &stream, &c51_grad_kernel,
                        p_ptr, m_ptr, a_ptr, grad_input_ptr,
                        dw_ptr, db_ptr,
                        ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
                        1.0 / ep_len as f32,
                    )?;
                }
            }
            // T10 backward chain: C51 grad_input → Mamba2 backward → AdamW.
            // The C51 grad_input kernel writes dL/d_h_enriched into the
            // first ep_len rows of d_h_enriched_tensor_smoke; remaining
            // rows stay zero (zeroed on episode reset).
            if cli.temporal {
                {
                    let (p_ptr, _g0) = probs_current_dev.device_ptr(&stream);
                    let (m_ptr, _g1) = m_dev.device_ptr(&stream);
                    let (a_ptr, _g2) = actions_dev.device_ptr(&stream);
                    let (w_ptr, _g3) = w_dev.device_ptr(&stream);
                    let (dh_ptr, _g4) = d_h_enriched_tensor_smoke.data_mut().device_ptr_mut(&stream);
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad_input(
                            &stream, &c51_grad_input_kernel,
                            p_ptr, m_ptr, a_ptr, w_ptr,
                            dh_ptr,
                            ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
                            1.0 / ep_len as f32,
                        )?;
                    }
                }
                let cache = cache_train.as_ref().expect("cache_train missing");
                let block_ref = mamba2_block.as_ref().expect("Mamba2Block missing");
                let grads = block_ref
                    .backward_from_h_enriched(cache, &d_h_enriched_tensor_smoke)
                    .map_err(|e| anyhow::anyhow!("Mamba2 backward: {e}"))?;
                let _ = cache_train;
                let block_mut = mamba2_block.as_mut().expect("Mamba2Block missing");
                let adamw = mamba2_adamw.as_mut().expect("Mamba2AdamW missing");
                adamw
                    .step(block_mut, &grads)
                    .map_err(|e| anyhow::anyhow!("Mamba2AdamW step: {e}"))?;
            }
        } else {
            // ── Linear-Q batched compute (Munchausen target) ────────
            // Forward Q_current on states
            {
                let (w_ptr, _g0) = w_dev.device_ptr(&stream);
                let (b_ptr, _g1) = b_dev.device_ptr(&stream);
                let (s_ptr, _g2) = states_dev.device_ptr(&stream);
                let (q_ptr, _g3) = q_current_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_forward(
                        &stream, &lq_fwd,
                        w_ptr, b_ptr, s_ptr, q_ptr,
                        ep_len, state_dim_i, n_act_i,
                    )?;
                }
            }
            // Forward Q_next on next_states USING TARGET WEIGHTS.
            {
                let (w_ptr, _g0) = w_target_dev.device_ptr(&stream);
                let (b_ptr, _g1) = b_target_dev.device_ptr(&stream);
                let (s_ptr, _g2) = next_states_dev.device_ptr(&stream);
                let (q_ptr, _g3) = q_next_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_forward(
                        &stream, &lq_fwd,
                        w_ptr, b_ptr, s_ptr, q_ptr,
                        ep_len, state_dim_i, n_act_i,
                    )?;
                }
            }
            // Munchausen target → target_dev[0..ep_len]
            {
                let (qn_ptr, _g0) = q_next_dev.device_ptr(&stream);
                let (qc_ptr, _g1) = q_current_dev.device_ptr(&stream);
                let (a_ptr, _g2) = actions_dev.device_ptr(&stream);
                let (r_ptr, _g3) = rewards_dev.device_ptr(&stream);
                let (d_ptr, _g4) = dones_dev.device_ptr(&stream);
                let (t_ptr, _g5) = target_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_munchausen_target(
                        &stream, &munch_kernel,
                        qn_ptr, qc_ptr, a_ptr, r_ptr, d_ptr,
                        cli.gamma, cli.alpha_m, cli.tau, cli.log_clip_min,
                        t_ptr, ep_len, n_act_i,
                    )?;
                }
            }
            // Linear-Q gradient
            {
                let (qc_ptr, _g0) = q_current_dev.device_ptr(&stream);
                let (t_ptr, _g1) = target_dev.device_ptr(&stream);
                let (a_ptr, _g2) = actions_dev.device_ptr(&stream);
                let (s_ptr, _g3) = states_dev.device_ptr(&stream);
                let (dw_ptr, _g4) = dw_dev.device_ptr_mut(&stream);
                let (db_ptr, _g5) = db_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_grad(
                        &stream, &lq_grad,
                        qc_ptr, t_ptr, a_ptr, s_ptr,
                        dw_ptr, db_ptr,
                        ep_len, state_dim_i, n_act_i,
                        1.0 / ep_len as f32,
                    )?;
                }
            }
        }
        // ── Gradient clip + SGD (shared kernels; size depends on mode) ──
        {
            let (dw_ptr, _g0) = dw_dev.device_ptr_mut(&stream);
            unsafe {
                ml::cuda_pipeline::alpha_kernels::launch_alpha_clip_inplace(
                    &stream, &lq_clip,
                    dw_ptr, cli.grad_clip, n_weights_eff as i32,
                )?;
            }
        }
        {
            let (db_ptr, _g0) = db_dev.device_ptr_mut(&stream);
            unsafe {
                ml::cuda_pipeline::alpha_kernels::launch_alpha_clip_inplace(
                    &stream, &lq_clip,
                    db_ptr, cli.grad_clip, n_biases_eff as i32,
                )?;
            }
        }
        // SGD step on W
        {
            let (w_ptr, _g0) = w_dev.device_ptr_mut(&stream);
            let (dw_ptr, _g1) = dw_dev.device_ptr(&stream);
            unsafe {
                ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_sgd_step(
                    &stream, &lq_sgd,
                    w_ptr, dw_ptr, cli.lr, n_weights_eff as i32,
                )?;
            }
        }
        // SGD step on b
        {
            let (b_ptr, _g0) = b_dev.device_ptr_mut(&stream);
            let (db_ptr, _g1) = db_dev.device_ptr(&stream);
            unsafe {
                ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_sgd_step(
                    &stream, &lq_sgd,
                    b_ptr, db_ptr, cli.lr, n_biases_eff as i32,
                )?;
            }
        }

        // Phase E.2 Task 17: Stacker-threshold controller at rollout end.
        // Trade rate = (count of non-Wait actions) / total decisions in the
        // episode. Realized Sharpe ≈ mean(returns) / std(returns) over the
        // window — for a single rollout this collapses to the rollout's
        // terminal reward divided by the random baseline std (a per-rollout
        // analog of the rvr metric, lets the controller decide whether to
        // tighten or loosen Kelly).
        let trade_count_ep: f32 = actions_host
            .iter()
            .filter(|&&a| a != 0) // 0 == Wait
            .count() as f32;
        let decisions_ep = actions_host.len() as f32;
        // Single-episode Sharpe proxy: terminal reward / baseline std.
        // Note: terminal reward is the LAST element of rewards_host
        // (sum of step rewards from env.step into ep.cumulative_pnl).
        let ep_terminal_r = *rewards_host.last().unwrap_or(&0.0);
        let baseline_std = isv_host[RANDOM_BASELINE_STD_INDEX].max(1e-6);
        let rollout_sharpe = ep_terminal_r / baseline_std;
        unsafe {
            let (isv_ptr, _g0) = isv_dev.device_ptr_mut(&stream);
            let (w_ptr, _g1) = ctl_wiener_dev.device_ptr_mut(&stream);
            ml::cuda_pipeline::alpha_kernels::launch_stacker_threshold_controller(
                &stream,
                &ctl_kernel,
                trade_count_ep,
                decisions_ep,
                rollout_sharpe,
                cli.target_sharpe,
                cli.k_threshold,
                cli.k_atten,
                cli.wiener_alpha_floor,
                cli.ctl_alpha_meta,
                ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX as i32,
                ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX as i32,
                ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX as i32,
                ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX as i32,
                // T16 regime hookup. Smoke doesn't currently run the vol
                // update kernel so we pass -1/-1 to disable the regime
                // scale here (backward compat).
                -1,
                -1,
                0.25,
                // T16 floor controller — also disabled in smoke.
                -1,
                -1,
                0.5,
                0.1,
                1.0e10,
                isv_ptr,
                w_ptr,
            )?;
        }
        // Phase E.3: refresh CPU-side threshold cache from device after the
        // controller update, for the NEXT episode's gating decisions.
        stream.synchronize().context("sync after controller")?;
        let isv_now = stream.clone_dtoh(&isv_dev).context("dtoh isv post-controller")?;
        current_threshold =
            isv_now[ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX];

        // Target network hard update every K episodes: w_target ← w_online.
        // Standard DQN stabilizer — V_soft(s') bootstrap reads w_target,
        // so updating it slowly breaks the chase-its-own-tail divergence
        // of online-only Munchausen.
        if (ep + 1) % cli.target_update_every == 0 {
            stream.synchronize().context("sync before target update")?;
            let w_now = stream.clone_dtoh(&w_dev).context("dtoh W for target update")?;
            let b_now = stream.clone_dtoh(&b_dev).context("dtoh b for target update")?;
            stream.memcpy_htod(&w_now, &mut w_target_dev)
                .context("htod W_target hard update")?;
            stream.memcpy_htod(&b_now, &mut b_target_dev)
                .context("htod b_target hard update")?;
        }

        // Periodic kill-criteria pipeline launch
        if (ep + 1) % cli.kill_criteria_every == 0 {
            stream.synchronize().context("sync before kc")?;
            // Snapshot weights for q_early_norm
            let w_now = stream.clone_dtoh(&w_dev).context("dtoh W")?;
            let b_now = stream.clone_dtoh(&b_dev).context("dtoh b")?;
            // q_early = q_init + ||W_now − W_init||_F so the kernel's
            // |q_early − q_init| / |q_init| ratio yields the relative
            // weight-space distance from init (direction-sensitive).
            //
            // When --temporal is on, include Mamba2 weight movement too:
            // the encoder absorbs most of the gradient signal so the C51
            // head moves little on its own. Without this addition, the
            // diagnostic systematically under-reports system learning.
            let mut dist = weight_distance_from_init(&w_now, &b_now, &w_init, &b_init);
            if let (Some(block), Some(init_snap)) =
                (mamba2_block.as_ref(), mamba2_init_snap.as_ref())
            {
                let now_snap = mamba2_snapshot(&stream, block)?;
                dist += mamba2_weight_distance(&now_snap, init_snap);
            }
            let q_early_norm = q_init_norm + dist;
            let rollout_r_mean: f32 = if recent_returns.is_empty() {
                0.0
            } else {
                let n = recent_returns.len();
                let take = n.min(cli.kill_criteria_every);
                let slice = &recent_returns[n - take..];
                slice.iter().sum::<f32>() / take as f32
            };
            let scalar_inputs = vec![rollout_r_mean, q_init_norm, q_early_norm];
            stream.memcpy_htod(&scalar_inputs, &mut kc_scalar_dev)
                .context("htod kc scalars")?;
            stream.memcpy_htod(&action_counts_host, &mut kc_action_counts_dev)
                .context("htod kc action_counts")?;

            // C51 path: convert distributions to scalar E[Z(s,a)] so the
            // existing kill-criteria kernel (which expects scalar Q[B, A])
            // can consume them. Reads probs_current_dev (filled by the
            // most recent batched forward), writes q_current_dev.
            if cli.c51 {
                let (p_ptr, _g0) = probs_current_dev.device_ptr(&stream);
                let (q_ptr, _g1) = q_current_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_expected_q(
                        &stream, &c51_expq_kernel,
                        p_ptr, c51_v_min, c51_delta_z, q_ptr,
                        ep_len, n_act_i, n_atoms_i,
                    )?;
                }
            }

            // Launch kill-criteria producer → kc_scratch[0..4]
            {
                let (q_ptr, _g0) = q_current_dev.device_ptr(&stream);
                let (ac_ptr, _g1) = kc_action_counts_dev.device_ptr(&stream);
                let (sc_ptr, _g2) = kc_scalar_dev.device_ptr(&stream);
                let (isv_ptr, _g3) = isv_dev.device_ptr(&stream);
                let (out_ptr, _g4) = kc_scratch_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_kill_criteria(
                        &stream, &kc_kernel,
                        q_ptr, ac_ptr, sc_ptr, isv_ptr,
                        ep_len, n_act_i, out_ptr,
                    )?;
                }
            }
            // Chained Pearls A+D applicator → ISV[539..543]
            {
                let (sc_ptr, _g0) = kc_scratch_dev.device_ptr(&stream);
                let (isv_ptr, _g1) = isv_dev.device_ptr_mut(&stream);
                let (w_ptr, _g2) = wiener_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::sp4_wiener_ema::launch_apply_pearls(
                        &stream, &pearls_kernel,
                        sc_ptr, 0,
                        isv_ptr, ALPHA_ISV_BLOCK_LO as i32,
                        w_ptr, 0,
                        4,
                        ml::cuda_pipeline::sp4_wiener_ema::ALPHA_META,
                    )?;
                }
            }
            stream.synchronize().context("sync after kc")?;

            let isv = stream.clone_dtoh(&isv_dev).context("dtoh isv")?;
            let kc = [
                isv[Q_SPREAD_EMA_INDEX],
                isv[ACTION_ENTROPY_EMA_INDEX],
                isv[RETURN_VS_RANDOM_EMA_INDEX],
                isv[EARLY_Q_MOVEMENT_EMA_INDEX],
            ];
            // Phase E.2: controller-driven slots for visibility.
            let ctl_threshold = isv[ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX];
            let ctl_obs_rate = isv[ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX];
            let ctl_atten = isv[ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX];
            info!(
                "ep {:>4} | ε={:.3} | R_mean={:>+9.1} | KC q={:.3} H={:.3} rvr={:+.3} ΔQ={:.3} | CTL thresh={:.4} obs={:.3} atten={:.3}",
                ep + 1, eps, rollout_r_mean, kc[0], kc[1], kc[2], kc[3],
                ctl_threshold, ctl_obs_rate, ctl_atten,
            );
            kc_logs.push((ep + 1, kc));
        }
    }

    stream.synchronize().context("final sync")?;
    let isv_final = stream.clone_dtoh(&isv_dev).context("dtoh isv final")?;
    let kc_final = [
        isv_final[Q_SPREAD_EMA_INDEX],
        isv_final[ACTION_ENTROPY_EMA_INDEX],
        isv_final[RETURN_VS_RANDOM_EMA_INDEX],
        isv_final[EARLY_Q_MOVEMENT_EMA_INDEX],
    ];

    // --- Verdict ---
    let entropy_threshold = 0.5 * (N_ACTIONS as f32).ln();
    let pass_q_spread = kc_final[0] >= 0.05;
    let pass_entropy = kc_final[1] >= entropy_threshold;
    let pass_rvr = kc_final[2] >= 0.0;
    let pass_early = kc_final[3] >= 0.01;
    let all_pass = pass_q_spread && pass_entropy && pass_rvr && pass_early;

    info!("=== Final kill-criteria verdict ===");
    info!(
        "  Q_SPREAD_EMA          = {:.4}   threshold ≥ 0.05    [{}]",
        kc_final[0], if pass_q_spread { "PASS" } else { "FAIL" }
    );
    info!(
        "  ACTION_ENTROPY_EMA    = {:.4}   threshold ≥ {:.4}  [{}]",
        kc_final[1], entropy_threshold,
        if pass_entropy { "PASS" } else { "FAIL" }
    );
    info!(
        "  RETURN_VS_RANDOM_EMA  = {:+.4}  threshold ≥ 0.0     [{}]",
        kc_final[2], if pass_rvr { "PASS" } else { "FAIL" }
    );
    info!(
        "  EARLY_Q_MOVEMENT_EMA  = {:.4}   threshold ≥ 0.01    [{}]",
        kc_final[3], if pass_early { "PASS" } else { "FAIL" }
    );
    info!(
        "  Overall: {} (H=6000 scale-up {})",
        if all_pass { "PASS" } else { "FAIL" },
        if all_pass { "VIABLE" } else { "BLOCKED → pivot to NoisyNet (Task 19)" }
    );

    // --- Save JSON ---
    let json = serde_json::json!({
        "phase": "E.1 Task 12",
        "pruned_actions": cli.pruned_actions,
        "n_allowed_actions": allowed_actions.len(),
        "c51": cli.c51,
        "c51_n_atoms": if cli.c51 { c51_n_atoms } else { 0 },
        "c51_vmin": if cli.c51 { c51_v_min } else { 0.0 },
        "c51_vmax": if cli.c51 { c51_v_max } else { 0.0 },
        "horizon": cli.horizon,
        "n_episodes": cli.n_episodes,
        "lr": cli.lr,
        "eps_start": cli.eps_start,
        "eps_end": cli.eps_end,
        "gamma": cli.gamma,
        "alpha_m": cli.alpha_m,
        "tau": cli.tau,
        "reward_scale": cli.reward_scale,
        "target_update_every": cli.target_update_every,
        "grad_clip": cli.grad_clip,
        "q_init_norm": q_init_norm,
        "final_stacker_threshold": isv_final[ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX],
        "final_trade_rate_observed_ema": isv_final[ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX],
        "final_stacker_kelly_attenuation": isv_final[ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX],
        "trade_rate_target": cli.trade_rate_target,
        "q_spread_ema": kc_final[0],
        "action_entropy_ema": kc_final[1],
        "return_vs_random_ema": kc_final[2],
        "early_q_movement_ema": kc_final[3],
        "pass_q_spread": pass_q_spread,
        "pass_entropy": pass_entropy,
        "pass_rvr": pass_rvr,
        "pass_early": pass_early,
        "all_pass": all_pass,
        "kc_log": kc_logs.iter().map(|(ep, kc)| {
            serde_json::json!({
                "episode": ep,
                "q_spread": kc[0],
                "entropy": kc[1],
                "rvr": kc[2],
                "early_mvmt": kc[3],
            })
        }).collect::<Vec<_>>(),
    });
    let mut f = File::create(&cli.out_path).context("create out")?;
    write!(f, "{}", serde_json::to_string_pretty(&json)?)?;
    info!("Wrote verdict + KC trajectory to {}", cli.out_path.display());

    Ok(())
}
