//! Phase E.3 Task 23 — Composition backtest with cost sweep.
//!
//! Trains the Phase E execution-policy DQN (linear Q + Phase 1d.3
//! alpha-cache + stabilizers) on the first `--train-frac` of the
//! fxcache, then evaluates the FROZEN policy (no SGD, ε=0 greedy) over
//! the held-out remainder at multiple transaction costs. Compares
//! per-cost annualized Sharpe vs the Phase 1d.4 "always-market-when-
//! confident" baseline (+4.4 frictionless, -4.0 at half-tick).
//!
//! Goal per the plan: lift the half-tick Sharpe above 0 — i.e., let
//! the execution-policy intelligence offset the cost the
//! threshold-only baseline can't.
//!
//! ## What's swept and what's frozen
//!
//! Frozen across costs: trained Q-network weights, fill-model
//! coefficients (from `alpha_fill_coeffs.json`), alpha-logit cache
//! (from `alpha_logits_cache.bin`). One policy evaluated at multiple
//! costs.
//!
//! Swept: only `ExecutionEnvConfig.cost_per_contract`. The env's
//! mutable config is updated between cost levels (snapshots stay in
//! place, cursor re-seeded each episode).
//!
//! ## Run
//!
//! ```bash
//! cargo run -p ml --release --example alpha_compose_backtest -- \
//!   --fxcache-path /home/jgrusewski/Work/foxhunt/test_data/feature-cache/9297....fxcache \
//!   --alpha-cache  config/ml/alpha_logits_cache.bin \
//!   --fill-coeffs  config/ml/alpha_fill_coeffs.json
//! ```

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;

// Phase E.4.A.8/T12: 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 (mirror gpu_training_guard.rs MappedBuffer) ──

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

impl MappedI32 {
    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) }
    }
    /// Read all `len` entries as a slice. Caller must ensure the
    /// associated GPU kernel has issued `__threadfence_system()` before
    /// the host accesses these values (volatile reads still happen
    /// per-element under the hood via `*host_ptr.add(i)`).
    fn read_all(&self) -> Vec<i32> {
        (0..self.len)
            .map(|i| unsafe { std::ptr::read_volatile(self.host_ptr.add(i)) })
            .collect()
    }
}

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 {
    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);
        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 ml::cuda_pipeline::alpha_isv_slots::{
    RANDOM_BASELINE_MEAN_INDEX, RANDOM_BASELINE_STD_INDEX,
};
use ml::env::action_space::N_ACTIONS;
use ml::env::execution_env::{
    EpisodeState, ExecutionEnv, ExecutionEnvConfig, ReplayRng, SnapshotRow,
};

const STATE_DIM: usize = 10;
const N_WEIGHTS: usize = N_ACTIONS * STATE_DIM;
const N_BIASES: usize = N_ACTIONS;

/// Action allow-list for the pruning experiment. Same shape as the smoke
/// binary: Q-net is unchanged (9 outputs), only the selector restricts.
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];

#[derive(Debug, Parser)]
#[command(
    name = "alpha_compose_backtest",
    about = "Phase E.3 Task 23 — composition backtest with cost sweep"
)]
struct Cli {
    #[arg(long)]
    fxcache_path: PathBuf,
    #[arg(long, default_value = "config/ml/alpha_fill_coeffs.json")]
    fill_coeffs: PathBuf,
    #[arg(long)]
    alpha_cache: PathBuf,
    /// Train segment fraction. First N% of snapshots used for DQN training,
    /// the rest for evaluation.
    #[arg(long, default_value_t = 0.8)]
    train_frac: f32,
    /// Snapshots to load from the fxcache.
    #[arg(long, default_value_t = 1_500_000)]
    max_snapshots: usize,
    /// Bars to skip at the front of the fxcache before windowing. Used by
    /// walk-forward CV scripts to slide a fixed-size train+eval window
    /// across the fxcache so each fold sees a distinct chronological
    /// region.
    #[arg(long, default_value_t = 0)]
    data_start_offset: usize,
    /// Episode horizon in snapshots.
    #[arg(long, default_value_t = 600)]
    horizon: usize,
    /// DQN training episodes (on train segment).
    #[arg(long, default_value_t = 1_000)]
    n_train_episodes: usize,
    /// Frozen-policy evaluation episodes per cost level.
    #[arg(long, default_value_t = 500)]
    n_eval_episodes: usize,
    /// Comma-separated cost grid (price units per contract round-turn).
    /// Phase 1d.4 used [0.0, 0.0625, 0.125, 0.25, 0.50].
    #[arg(long, value_delimiter = ',', default_value = "0.0,0.0625,0.125,0.25,0.5")]
    cost_grid: Vec<f32>,
    /// Comma-separated alpha-confidence threshold grid for the gate at eval.
    /// Direct analogue of Phase 1d.4's `--threshold` sweep — at each
    /// threshold, the policy is forced to Wait when |sigmoid(alpha)−0.5|
    /// < threshold. Threshold 0 = no gate (original Task 23 behaviour).
    #[arg(long, value_delimiter = ',', default_value = "0.0,0.05,0.10,0.15,0.20,0.25")]
    threshold_grid: Vec<f32>,
    #[arg(long, default_value_t = 1)]
    trade_size: i32,
    #[arg(long, default_value_t = 0xCAFEBABE_u64)]
    seed: u64,
    /// Training-time cost (the policy LEARNED against this cost). Used as
    /// a fixed cost when `--train-cost-hi == train_cost`; otherwise serves
    /// as the lower bound of a uniform[lo, hi] random cost sampled per
    /// training episode (Phase E.4.A T16 cost-aware training).
    #[arg(long, default_value_t = 0.0625)]
    train_cost: f32,
    /// Upper bound on the per-episode random training cost. When > `train_cost`,
    /// each training episode samples its `cost_per_contract` uniformly from
    /// [train_cost, train_cost_hi] so the Q-network learns cost-conservative
    /// behaviour across the realistic ES futures range (quarter-tick to
    /// quad-tick). Defaults to the same value as `train_cost` (no randomization,
    /// backward-compatible).
    #[arg(long, default_value_t = 0.0625)]
    train_cost_hi: f32,
    /// Training-time alpha-confidence threshold for the gate. Forces
    /// Wait during training when `|sigmoid(alpha)−0.5| < this`, so the
    /// Q-network learns weights for the gated policy class. Default
    /// 0.39 — the equilibrium the controller stabilized to in the
    /// Phase E.2 smoke (alpha_dqn_h600_smoke at ep 200+). Set to 0.0
    /// for the original ungated-training behavior.
    #[arg(long, default_value_t = 0.39)]
    train_threshold: f32,
    /// SGD learning rate during training.
    #[arg(long, default_value_t = 1.0e-4)]
    lr: f32,
    #[arg(long, default_value_t = 0.50)]
    eps_start: f32,
    #[arg(long, default_value_t = 0.05)]
    eps_end: f32,
    #[arg(long, default_value_t = 0.99)]
    gamma: f32,
    #[arg(long, default_value_t = 0.9)]
    alpha_m: f32,
    #[arg(long, default_value_t = 0.03)]
    tau: f32,
    #[arg(long, default_value_t = -1.0)]
    log_clip_min: f32,
    #[arg(long, default_value_t = 1000.0)]
    reward_scale: f32,
    #[arg(long, default_value_t = 10)]
    target_update_every: usize,
    #[arg(long, default_value_t = 1.0)]
    grad_clip: f32,
    /// Restrict action selection to {Wait, BuyMarket, SellMarket, FlatMarket}.
    /// Applies during BOTH training and eval. **FALSIFIED 2026-05-15** —
    /// kept for reproducibility (see `pearl_action_pruning_falsified`).
    #[arg(long, default_value_t = false)]
    pruned_actions: bool,
    /// Use C51 distributional Q-network (Phase E.3 follow-up). Same
    /// semantics as the alpha_dqn_h600_smoke `--c51` flag.
    #[arg(long, default_value_t = false)]
    c51: bool,
    /// C51 atom-support lower bound (normalized reward units).
    #[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; kernel caps at 64).
    #[arg(long, default_value_t = 51)]
    c51_n_atoms: usize,
    /// Derive bid/ask from real spread_bps in fxcache instead of fixed
    /// ±0.125-tick. Phase E.3 Path 3 follow-up.
    #[arg(long, default_value_t = false)]
    real_spread: bool,
    /// Phase E.4.A: enable the temporal-encoder stack (sliding window +
    /// Mamba2 + C51).
    #[arg(long, default_value_t = false)]
    temporal: bool,
    #[arg(long, default_value_t = 16)]
    window_k: usize,
    #[arg(long, default_value_t = 32)]
    mamba2_hidden_dim: usize,
    #[arg(long, default_value_t = 16)]
    mamba2_state_dim: usize,
    /// Phase E.4.A Pillar B: fire stacker-threshold controller at the
    /// end of EVERY eval episode (not just training). Slot 543
    /// (threshold) adapts during the 2D sweep based on observed trade
    /// rate / Sharpe.
    #[arg(long, default_value_t = false)]
    isv_continual: bool,
    /// Phase E.4.A T16: enable vol-EMA regime detection + pre-emptive
    /// Kelly attenuation. When set:
    ///   - per inference step, `alpha_regime_vol_update_kernel` reads
    ///     the parallel envs' current/previous mids and maintains
    ///     ISV[549]=vol_ema (Wiener-α with 0.4 floor + Pearl A
    ///     bootstrap) and ISV[550]=vol_ref (slow tracker, β=0.005);
    ///   - per eval episode (requires `--isv-continual`), the existing
    ///     stacker controller multiplies its Sharpe-reactive Kelly
    ///     output by `clamp(vol_ref/vol_ema, 0.25, 1.0)` so vol spikes
    ///     cut Kelly *before* losses materialise.
    /// Off by default to keep the T10 baseline comparable.
    #[arg(long, default_value_t = false)]
    regime_scale: bool,
    /// Phase E.4.A.T15.batched (2026-05-15): parallel envs per training
    /// "epoch". n_train_episodes / n_train_par epochs total. Each epoch
    /// runs N_par envs in lockstep H steps + ONE batched C51 update at
    /// B = N_par * H. Default 50 keeps SGD semantics (W is updated 20
    /// times over 1000 training eps); raise to 1000 for one-epoch full-
    /// batch training.
    #[arg(long, default_value_t = 50)]
    n_train_par: usize,
    #[arg(long, default_value = "config/ml/alpha_compose_backtest.json")]
    out_path: PathBuf,
}

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)
    }
}

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
    }
}

/// Phase E.3 confidence-gated ε-greedy. If `alpha_confidence < threshold`,
/// force action=0 (Wait). Otherwise standard ε-greedy over `allowed`.
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;
    }
    epsilon_greedy(q, eps, allowed, rng)
}

#[derive(Debug, Clone, serde::Serialize)]
struct CostBin {
    cost: f32,
    threshold: f32,
    n_episodes: usize,
    mean_reward: f32,
    std_reward: f32,
    sharpe_per_episode: f32,
    /// Sharpe scaled by sqrt(episodes_per_year). Time span derived from
    /// the eval window's mid_price index span × bar_seconds.
    sharpe_annualised: f32,
    win_rate: f32,
    avg_n_trades: f32,
    p05: f32,
    p50: f32,
    p95: f32,
}

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.3 Task 23 — composition backtest starting");
    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
    };

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

    // --- Load cubins ---
    let lq_module = ctx
        .load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_LINEAR_Q_CUBIN.to_vec())
        .context("alpha_linear_q cubin")?;
    let lq_fwd = lq_module.load_function("alpha_linear_q_forward_kernel")?;
    let lq_grad = lq_module.load_function("alpha_linear_q_grad_kernel")?;
    let lq_sgd = lq_module.load_function("alpha_linear_q_sgd_step_kernel")?;
    let lq_clip = lq_module.load_function("alpha_clip_inplace_kernel")?;
    let munch_cubin: Vec<u8> = std::fs::read(concat!(
        env!("OUT_DIR"),
        "/alpha_munchausen_target.cubin"
    ))?;
    let munch_module = ctx.load_cubin(munch_cubin)?;
    let munch_kernel =
        munch_module.load_function("alpha_munchausen_target_kernel")?;

    // Phase E.3 follow-up: C51 kernels (always loaded; only used when --c51).
    let c51_module = ctx
        .load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_C51_CUBIN.to_vec())
        .context("alpha_c51 cubin")?;
    let c51_fwd_kernel = c51_module.load_function("alpha_c51_forward_kernel")?;
    let c51_project_kernel = c51_module.load_function("alpha_c51_project_kernel")?;
    let c51_grad_kernel = c51_module.load_function("alpha_c51_grad_kernel")?;
    let c51_thompson_kernel = c51_module.load_function("alpha_c51_thompson_select_kernel")?;
    // Phase E.4.A.6: window push kernel (used in --temporal path).
    let push_module = ctx
        .load_cubin(ml::cuda_pipeline::alpha_kernels::ALPHA_WINDOW_PUSH_CUBIN.to_vec())
        .context("alpha_window_push cubin")?;
    let push_kernel = push_module.load_function("alpha_window_push_kernel")?;
    let h_store_kernel = push_module.load_function("alpha_h_enriched_store_kernel")?;
    let push_batched_kernel = push_module.load_function("alpha_window_push_batched_kernel")?;
    let h_store_batched_kernel = push_module.load_function("alpha_h_enriched_store_batched_kernel")?;
    // Phase E.4.A.7 controller (used in --isv-continual eval path).
    let ctl_module = ctx
        .load_cubin(ml::cuda_pipeline::alpha_kernels::STACKER_THRESHOLD_CONTROLLER_CUBIN.to_vec())
        .context("stacker controller cubin")?;
    let ctl_kernel = ctl_module
        .load_function("stacker_threshold_controller_update")?;
    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 ({} atoms, [{:.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]");
        }
        if c51_v_max <= c51_v_min {
            anyhow::bail!("--c51-vmax must exceed --c51-vmin");
        }
    } 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 fill model");
    let alpha_cache = ml::env::loaders::load_alpha_cache(&cli.alpha_cache)?;
    info!("Loaded alpha cache: {} entries", alpha_cache.len());
    let rows = ml::env::loaders::load_snapshots_from_fxcache_at(
        &cli.fxcache_path,
        cli.data_start_offset,
        cli.max_snapshots,
        Some(&alpha_cache),
        cli.real_spread,
        None,  // E.4.A T4: mbp10_dir wiring + --use-real-depth flag lands in T5
    )?;
    let n_total = rows.len();
    info!("Loaded {} snapshots", n_total);

    let n_train = ((n_total as f32) * cli.train_frac) as usize;
    let n_eval = n_total - n_train;
    if n_train <= cli.horizon || n_eval <= cli.horizon {
        anyhow::bail!(
            "insufficient snapshots: train={}, eval={}, horizon={}",
            n_train,
            n_eval,
            cli.horizon
        );
    }
    info!(
        "Split: train={} bars (cursor 0..{}), eval={} bars (cursor {}..{})",
        n_train, n_train, n_eval, n_train, n_total
    );

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

    // --- Initialize Q-network ---
    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 mut rng = SmokeRng::new(cli.seed.wrapping_add(0xDEAD_BEEF));
    let xavier_scale = (2.0_f32 / c51_input_dim as f32).sqrt();
    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 mut w_dev = stream.clone_htod(&w_init)?;
    let mut b_dev = stream.clone_htod(&b_init)?;
    let mut w_target_dev = stream.clone_htod(&w_init)?;
    let mut b_target_dev = stream.clone_htod(&b_init)?;
    let mut dw_dev = stream.alloc_zeros::<f32>(n_weights_eff)?;
    let mut db_dev = stream.alloc_zeros::<f32>(n_biases_eff)?;

    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)?;
    let mut next_states_dev = stream.alloc_zeros::<f32>(cli.horizon * STATE_DIM)?;
    let mut actions_dev = stream.alloc_zeros::<i32>(cli.horizon)?;
    let mut rewards_dev = stream.alloc_zeros::<f32>(cli.horizon)?;
    let mut dones_dev = stream.alloc_zeros::<f32>(cli.horizon)?;
    let mut q_current_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS)?;
    let mut q_next_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS)?;
    let mut target_dev = stream.alloc_zeros::<f32>(cli.horizon)?;
    let mut single_state_dev = stream.alloc_zeros::<f32>(STATE_DIM)?;
    let mut single_q_dev = stream.alloc_zeros::<f32>(N_ACTIONS)?;
    // C51 device buffers (always allocated; cheap).
    let mut probs_current_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS * c51_n_atoms)?;
    let mut probs_next_dev = stream.alloc_zeros::<f32>(cli.horizon * N_ACTIONS * c51_n_atoms)?;
    let mut m_dev = stream.alloc_zeros::<f32>(cli.horizon * c51_n_atoms)?;
    let mut single_probs_dev = stream.alloc_zeros::<f32>(N_ACTIONS * c51_n_atoms)?;
    // Mapped-pinned per-step inference buffers (C51 path).
    let state_pinned = unsafe { MappedF32::new(STATE_DIM)? };
    let action_pinned = unsafe { MappedI32::new(1)? };

    // Phase E.4.A.8 (T12): temporal buffers, allocated unconditionally.
    let mut window_tensor = GpuTensor::zeros(
        &[1, cli.window_k, STATE_DIM], &stream
    ).map_err(|e| anyhow::anyhow!("alloc window: {e}"))?;
    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")?;
    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 (only 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
    };
    // Load the C51 grad-input kernel (for backward chain into Mamba2).
    let c51_grad_input_kernel = c51_module
        .load_function("alpha_c51_grad_input_kernel")
        .context("c51 grad_input load")?;
    // T10: train-time window store (captures windows during inference for
    // the end-of-epoch batched Mamba2 forward).
    let train_window_store_kernel = push_module
        .load_function("alpha_train_window_store_batched_kernel")
        .context("train_window_store kernel load")?;
    let n_atoms_i = c51_n_atoms as i32;
    // ISV buffer + Wiener state for the eval-time controller path.
    let mut isv_host: Vec<f32> = vec![0.0; 552];
    isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_MEAN_INDEX] = -5191.53;
    isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_STD_INDEX] = 4963.62;
    isv_host[ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX] = 0.08;
    let mut isv_dev = stream.clone_htod(&isv_host).context("upload isv")?;
    let mut ctl_wiener_dev = stream.alloc_zeros::<f32>(3).context("alloc ctl wiener")?;

    // ============================================================
    // Phase 1: TRAIN DQN — BATCHED PARALLEL ENVS (Phase E.4.A.T15)
    // ============================================================
    let train_n_par = cli.n_train_par;
    let train_n_epochs = (cli.n_train_episodes + train_n_par - 1) / train_n_par;
    info!("=== Training phase: {} epochs × N_par={} parallel envs × H={} steps (BATCHED) ===",
          train_n_epochs, train_n_par, cli.horizon);
    let train_max_start = (n_train.saturating_sub(cli.horizon + 1)).max(1);
    let train_h = cli.horizon;
    let train_batch = train_n_par * train_h;
    let mut states_train_dev = stream.alloc_zeros::<f32>(train_batch * STATE_DIM)?;
    let mut next_states_train_dev = stream.alloc_zeros::<f32>(train_batch * STATE_DIM)?;
    let mut actions_train_dev = stream.alloc_zeros::<i32>(train_batch)?;
    let mut rewards_train_dev = stream.alloc_zeros::<f32>(train_batch)?;
    let mut dones_train_dev = stream.alloc_zeros::<f32>(train_batch)?;
    let mut probs_curr_train_dev = stream.alloc_zeros::<f32>(train_batch * N_ACTIONS * c51_n_atoms)?;
    let mut probs_next_train_dev = stream.alloc_zeros::<f32>(train_batch * N_ACTIONS * c51_n_atoms)?;
    let mut m_train_dev = stream.alloc_zeros::<f32>(train_batch * c51_n_atoms)?;
    let h_enriched_train_capacity = (train_h + 1) * train_n_par * cli.mamba2_hidden_dim;
    let mut h_enriched_train_dev = stream.alloc_zeros::<f32>(h_enriched_train_capacity)?;
    let batched_state_pinned_train = unsafe { MappedF32::new(train_n_par * STATE_DIM)? };
    let batched_action_pinned_train = unsafe { MappedI32::new(train_n_par)? };
    let mut batched_window_tensor_train = GpuTensor::zeros(
        &[train_n_par, cli.window_k, STATE_DIM], &stream
    ).map_err(|e| anyhow::anyhow!("alloc batched window train: {e}"))?;
    // T10 buffers: all-step train windows and d_h_enriched for the backward chain.
    let mut train_windows_tensor_train = GpuTensor::zeros(
        &[train_batch, cli.window_k, STATE_DIM], &stream
    ).map_err(|e| anyhow::anyhow!("alloc train windows train: {e}"))?;
    let mut d_h_enriched_tensor_train = GpuTensor::zeros(
        &[train_batch, cli.mamba2_hidden_dim], &stream
    ).map_err(|e| anyhow::anyhow!("alloc d_h_enriched train: {e}"))?;
    let mut batched_probs_inference_train = stream
        .alloc_zeros::<f32>(train_n_par * N_ACTIONS * c51_n_atoms)?;
    let snapshots_arc_train = env.snapshots_arc();
    let fill_model_for_train = env.fill_model.clone();
    let env_config_template_train = env.config.clone();
    let mut episode_rng = SmokeRng::new(cli.seed.wrapping_add(0xFEED));

    for epoch in 0..train_n_epochs {
        let _eps = cli.eps_start + (cli.eps_end - cli.eps_start)
            * (epoch as f32 / train_n_epochs.max(1) as f32);
        // T16 cost-aware training: when --train-cost-hi > --train-cost,
        // sample the per-episode cost uniformly from that range, so the
        // policy learns to behave robustly across the realistic ES cost
        // spectrum (quarter-tick to quad-tick) instead of memorizing a
        // single low-cost world.
        let cost_for_epoch: f32 = if cli.train_cost_hi > cli.train_cost {
            let u = (episode_rng.next_u64() as f32) / (u64::MAX as f32);
            cli.train_cost + u * (cli.train_cost_hi - cli.train_cost)
        } else {
            cli.train_cost
        };
        let mut par_envs: Vec<ExecutionEnv> = (0..train_n_par).map(|_| {
            let mut cfg = env_config_template_train.clone();
            cfg.cost_per_contract = cost_for_epoch;
            ExecutionEnv::new_arc(
                cfg, fill_model_for_train.clone(),
                std::sync::Arc::clone(&snapshots_arc_train),
                0,
            )
        }).collect();
        for env_i in par_envs.iter_mut() {
            let start = (episode_rng.next_u64() as usize) % train_max_start;
            let seed = episode_rng.next_u64();
            env_i.reset_at(seed, start);
        }
        let mut par_states: Vec<EpisodeState> = vec![EpisodeState::new(); train_n_par];
        let mut done_flags = vec![false; train_n_par];
        if cli.temporal {
            stream.memset_zeros(batched_window_tensor_train.data_mut())?;
            stream.memset_zeros(train_windows_tensor_train.data_mut())?;
            stream.memset_zeros(d_h_enriched_tensor_train.data_mut())?;
        }
        stream.memset_zeros(&mut h_enriched_train_dev)?;
        let mut states_host_train: Vec<f32> = vec![0.0; train_batch * STATE_DIM];
        let mut next_states_host_train: Vec<f32> = vec![0.0; train_batch * STATE_DIM];
        let mut actions_host_train: Vec<i32> = vec![0; train_batch];
        let mut rewards_host_train: Vec<f32> = vec![0.0; train_batch];
        let mut dones_host_train: Vec<f32> = vec![0.0; train_batch];

        for step in 0..train_h {
            let mut batched_states_host = vec![0.0_f32; train_n_par * STATE_DIM];
            for i in 0..train_n_par {
                if !done_flags[i] {
                    let s = par_envs[i].state(&par_states[i]);
                    batched_states_host[i * STATE_DIM..(i + 1) * STATE_DIM].copy_from_slice(&s);
                }
            }
            batched_state_pinned_train.write(&batched_states_host);

            if cli.temporal {
                {
                    let (w_ptr, _g) = batched_window_tensor_train.data_mut().device_ptr_mut(&stream);
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push_batched(
                            &stream, &push_batched_kernel,
                            batched_state_pinned_train.dev_u64(), w_ptr,
                            train_n_par as i32, cli.window_k as i32, state_dim_i,
                        )?;
                    }
                }
                // T10: capture post-push windows into [H*N_par, K, in_dim] buffer
                // at step row offset `step * N_par`.
                {
                    let (src_ptr, _g_src) = batched_window_tensor_train.data().device_ptr(&stream);
                    let (dst_ptr, _g_dst) = train_windows_tensor_train.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, (step * train_n_par) as i32,
                            train_n_par as i32, cli.window_k as i32, state_dim_i,
                        )?;
                    }
                }
                let block = mamba2_block.as_ref().expect("Mamba2Block missing");
                let (_logit, cache) = block.forward_train(&batched_window_tensor_train)
                    .map_err(|e| anyhow::anyhow!("train mamba2: {e}"))?;
                {
                    let (src_p, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream);
                    let (buf_p, _g_buf) = h_enriched_train_dev.device_ptr_mut(&stream);
                    let step_row_offset = (step * train_n_par) as i32;
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store_batched(
                            &stream, &h_store_batched_kernel,
                            src_p, buf_p, step_row_offset,
                            train_n_par as i32, cli.mamba2_hidden_dim as i32,
                        )?;
                    }
                }
                let (h_ptr, _g_h) = cache.h_enriched.cuda_data().device_ptr(&stream);
                let (w_ptr, _g0) = w_dev.device_ptr(&stream);
                let (b_ptr, _g1) = b_dev.device_ptr(&stream);
                let (p_ptr, _g3) = batched_probs_inference_train.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
                        &stream, &c51_fwd_kernel,
                        w_ptr, b_ptr, h_ptr, p_ptr,
                        train_n_par as i32, 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, _g3) = batched_probs_inference_train.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
                        &stream, &c51_fwd_kernel,
                        w_ptr, b_ptr, batched_state_pinned_train.dev_u64(), p_ptr,
                        train_n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
                    )?;
                }
            }
            let step_seed = episode_rng.next_u64() as u32;
            {
                let (p_ptr, _g0) = batched_probs_inference_train.device_ptr(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
                        &stream, &c51_thompson_kernel,
                        p_ptr,
                        batched_state_pinned_train.dev_u64(),
                        cli.train_threshold,
                        1, state_dim_i, c51_v_min, c51_delta_z, step_seed,
                        batched_action_pinned_train.dev_u64(),
                        train_n_par as i32, n_act_i, n_atoms_i,
                    )?;
                }
            }
            stream.synchronize()?;
            let actions = batched_action_pinned_train.read_all();

            for i in 0..train_n_par {
                if done_flags[i] { continue; }
                let s = &batched_states_host[i * STATE_DIM..(i + 1) * STATE_DIM];
                let row = step * train_n_par + i;
                states_host_train[row * STATE_DIM..(row + 1) * STATE_DIM].copy_from_slice(s);
                actions_host_train[row] = actions[i];
                let action = actions[i] as u8;
                let (next_state_arr, reward, done) = par_envs[i].step(action, &mut par_states[i])
                    .ok_or_else(|| anyhow::anyhow!("env step None"))?;
                next_states_host_train[row * STATE_DIM..(row + 1) * STATE_DIM]
                    .copy_from_slice(&next_state_arr);
                rewards_host_train[row] = reward / cli.reward_scale;
                dones_host_train[row] = if done { 1.0 } else { 0.0 };
                if done {
                    done_flags[i] = true;
                }
            }
        }

        stream.memcpy_htod(&states_host_train, &mut states_train_dev)?;
        stream.memcpy_htod(&next_states_host_train, &mut next_states_train_dev)?;
        stream.memcpy_htod(&actions_host_train, &mut actions_train_dev)?;
        stream.memcpy_htod(&rewards_host_train, &mut rewards_train_dev)?;
        stream.memcpy_htod(&dones_host_train, &mut dones_train_dev)?;

        let bt = train_batch as i32;
        // T10: re-forward Mamba2 on collected train windows to get a cache
        // that backward_from_h_enriched can consume. cache.h_enriched
        // replaces h_enriched_train_dev curr offset (bit-identical since
        // Mamba2 weights haven't been updated yet this epoch).
        let cache_train: Option<ml_alpha::mamba2_block::Mamba2ForwardCache> = if cli.temporal {
            let block = mamba2_block.as_ref().expect("Mamba2Block missing");
            let (_logit, cache) = block.forward_train(&train_windows_tensor_train)
                .map_err(|e| anyhow::anyhow!("train mamba2 (epoch 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_curr, g) = cache.h_enriched.cuda_data().device_ptr(&stream);
            let (h_base_next, g_next) = h_enriched_train_dev.device_ptr(&stream);
            curr_input_ptr_guard = g;
            next_input_ptr_guard = g_next;
            let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
            (
                h_curr,
                h_base_next + (train_n_par as u64) * (cli.mamba2_hidden_dim as u64) * 4u64,
            )
        } else {
            let (s_ptr, g_c) = states_train_dev.device_ptr(&stream);
            let (n_ptr, g_n) = next_states_train_dev.device_ptr(&stream);
            curr_input_ptr_guard = g_c;
            next_input_ptr_guard = g_n;
            let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
            (s_ptr, n_ptr)
        };

        {
            let (w_ptr, _g0) = w_dev.device_ptr(&stream);
            let (b_ptr, _g1) = b_dev.device_ptr(&stream);
            let (p_ptr, _g3) = probs_curr_train_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,
                    bt, c51_input_dim as i32, n_act_i, n_atoms_i,
                )?;
            }
        }
        {
            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_train_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,
                    bt, c51_input_dim as i32, n_act_i, n_atoms_i,
                )?;
            }
        }
        {
            let (pn_ptr, _g0) = probs_next_train_dev.device_ptr(&stream);
            let (r_ptr, _g1) = rewards_train_dev.device_ptr(&stream);
            let (d_ptr, _g2) = dones_train_dev.device_ptr(&stream);
            let (m_ptr, _g3) = m_train_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, bt, n_act_i, n_atoms_i,
                )?;
            }
        }
        {
            let (p_ptr, _g0) = probs_curr_train_dev.device_ptr(&stream);
            let (m_ptr, _g1) = m_train_dev.device_ptr(&stream);
            let (a_ptr, _g2) = actions_train_dev.device_ptr(&stream);
            let (s_ptr, _g3) = states_train_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,
                    bt, c51_input_dim as i32, n_act_i, n_atoms_i,
                    1.0 / bt as f32,
                )?;
            }
        }
        // T10 backward chain: C51 grad_input → Mamba2 backward → AdamW step.
        if cli.temporal {
            {
                let (p_ptr, _g0) = probs_curr_train_dev.device_ptr(&stream);
                let (m_ptr, _g1) = m_train_dev.device_ptr(&stream);
                let (a_ptr, _g2) = actions_train_dev.device_ptr(&stream);
                let (w_ptr, _g3) = w_dev.device_ptr(&stream);
                let (dh_ptr, _g4) = d_h_enriched_tensor_train.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,
                        bt, c51_input_dim as i32, n_act_i, n_atoms_i,
                        1.0 / bt 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_train)
                .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}"))?;
        }
        {
            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,
                )?;
            }
        }
        {
            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,
                )?;
            }
        }
        {
            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,
                )?;
            }
        }
        if (epoch + 1) % cli.target_update_every == 0 {
            stream.synchronize()?;
            let w_now = stream.clone_dtoh(&w_dev)?;
            let b_now = stream.clone_dtoh(&b_dev)?;
            stream.memcpy_htod(&w_now, &mut w_target_dev)?;
            stream.memcpy_htod(&b_now, &mut b_target_dev)?;
        }
        if (epoch + 1) % 5 == 0 || epoch + 1 == train_n_epochs {
            info!("  train epoch {}/{}", epoch + 1, train_n_epochs);
        }
    }
    info!("Training complete. Frozen policy ready for eval.");

    // ---------------------------------------------------------------
    // Phase 2: 2D SWEEP — frozen-policy greedy eval per (threshold, cost).
    // Direct analogue of Phase 1d.4's threshold × cost table.
    //
    // Phase E.4.A.T14.batched (2026-05-15): batched eval — N=n_eval_episodes
    // environments run in LOCKSTEP per cell, with ONE GPU forward + ONE
    // sync per step (instead of N forwards × N syncs). Snapshots shared
    // across envs via Arc to avoid 50MB × N duplication.
    // ---------------------------------------------------------------
    info!("=== Eval phase: {} episodes × {} thresholds × {} costs (batched) ===",
          cli.n_eval_episodes, cli.threshold_grid.len(), cli.cost_grid.len());
    let eval_max_start = (n_total - n_train).saturating_sub(cli.horizon + 1).max(1);
    let n_par = cli.n_eval_episodes;
    // Batched-eval GPU buffers (allocated once, reused across cells).
    let batched_state_pinned = unsafe { MappedF32::new(n_par * STATE_DIM)? };
    let batched_action_pinned = unsafe { MappedI32::new(n_par)? };
    let mut batched_window_tensor = GpuTensor::zeros(
        &[n_par, cli.window_k, STATE_DIM], &stream
    ).map_err(|e| anyhow::anyhow!("alloc batched window: {e}"))?;
    let mut batched_probs_dev = stream
        .alloc_zeros::<f32>(n_par * N_ACTIONS * c51_n_atoms)
        .context("alloc batched probs")?;
    // T16 regime kernel state. mids buffers are mapped-pinned so the host
    // gather (one f32 per env per step) writes directly into device-mapped
    // memory with no htod copy.
    let mids_curr_pinned = unsafe { MappedF32::new(n_par)? };
    let mids_prev_pinned = unsafe { MappedF32::new(n_par)? };
    let regime_module = if cli.regime_scale {
        Some(ctx.load_cubin(
            ml::cuda_pipeline::alpha_kernels::ALPHA_REGIME_VOL_UPDATE_CUBIN.to_vec()
        ).context("alpha_regime_vol_update cubin load")?)
    } else { None };
    let regime_kernel = regime_module.as_ref().map(|m|
        m.load_function("alpha_regime_vol_update_kernel")
        .expect("alpha_regime_vol_update_kernel load")
    );
    let snapshots_arc = env.snapshots_arc();
    let fill_model_for_par = env.fill_model.clone();
    let env_config_template = env.config.clone();
    let mut bins: Vec<CostBin> =
        Vec::with_capacity(cli.cost_grid.len() * cli.threshold_grid.len());
    for &threshold in &cli.threshold_grid {
    for &cost in &cli.cost_grid {
        env.config.cost_per_contract = cost;
        // Build N parallel envs sharing snapshots; reset each at a
        // random eval-start cursor + seed.
        let mut par_envs: Vec<ExecutionEnv> = (0..n_par).map(|_| {
            let mut cfg = env_config_template.clone();
            cfg.cost_per_contract = cost;
            ExecutionEnv::new_arc(
                cfg, fill_model_for_par.clone(),
                std::sync::Arc::clone(&snapshots_arc),
                0,  // seed re-set below
            )
        }).collect();
        for env_i in par_envs.iter_mut() {
            let start = n_train + (episode_rng.next_u64() as usize) % eval_max_start;
            let seed = episode_rng.next_u64();
            env_i.reset_at(seed, start);
        }
        let mut par_states: Vec<EpisodeState> = vec![EpisodeState::new(); n_par];
        let mut terminal_rs = vec![0.0_f32; n_par];
        let mut ep_n_trades_par = vec![0_u32; n_par];
        let mut done_flags = vec![false; n_par];
        // Zero batched windows.
        if cli.temporal {
            stream.memset_zeros(batched_window_tensor.data_mut())
                .context("zero batched windows")?;
        }
        // T16: reset prev_mids so the first-step Pearl-A bootstrap fires
        // cleanly (kernel guards against curr<=0 || prev<=0 internally).
        if cli.regime_scale {
            let zeros = vec![0.0_f32; n_par];
            mids_curr_pinned.write(&zeros);
            mids_prev_pinned.write(&zeros);
        }
        // Lockstep step loop.
        for _step in 0..cli.horizon {
            // Gather current states (CPU; <100μs for N=500).
            let mut batched_states_host = vec![0.0_f32; n_par * STATE_DIM];
            for i in 0..n_par {
                if !done_flags[i] {
                    let s = par_envs[i].state(&par_states[i]);
                    batched_states_host[i * STATE_DIM..(i + 1) * STATE_DIM].copy_from_slice(&s);
                }
            }
            batched_state_pinned.write(&batched_states_host);
            // T16: capture current mids + fire regime kernel BEFORE the
            // env step advances. Regime signal therefore reflects vol
            // observed up to and including this step, which is what the
            // controller's per-episode Kelly update will read.
            if let Some(rk) = regime_kernel.as_ref() {
                let mut mids_host = vec![0.0_f32; n_par];
                for i in 0..n_par {
                    if !done_flags[i] {
                        mids_host[i] = par_envs[i].current_mid();
                    }
                }
                mids_curr_pinned.write(&mids_host);
                unsafe {
                    let (isv_ptr, _g) = isv_dev.device_ptr_mut(&stream);
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_regime_vol_update(
                        &stream, rk,
                        mids_curr_pinned.dev_u64(),
                        mids_prev_pinned.dev_u64(),
                        isv_ptr,
                        n_par as i32,
                        ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_EMA_INDEX as i32,
                        ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_INDEX as i32,
                        0.005,  // β: 200-step horizon for vol_ref tracking
                    )?;
                }
                // Swap prev ← curr for the next step. Mapped-pinned write
                // is a CPU-side fill, no device-to-device copy needed.
                mids_prev_pinned.write(&mids_host);
            }
            if cli.temporal {
                // Batched window push.
                {
                    let (w_ptr, _g) = batched_window_tensor.data_mut().device_ptr_mut(&stream);
                    unsafe {
                        ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push_batched(
                            &stream, &push_batched_kernel,
                            batched_state_pinned.dev_u64(), w_ptr,
                            n_par as i32, cli.window_k as i32, state_dim_i,
                        )?;
                    }
                }
                // Batched Mamba2 forward → cache.h_enriched [N, hidden].
                let block = mamba2_block.as_ref().expect("Mamba2Block missing");
                let (_logit, cache) = block.forward_train(&batched_window_tensor)
                    .map_err(|e| anyhow::anyhow!("batched mamba2 forward (eval): {e}"))?;
                // Batched C51 forward.
                let (h_ptr, _g_h) = cache.h_enriched.cuda_data().device_ptr(&stream);
                let (w_ptr, _g0) = w_dev.device_ptr(&stream);
                let (b_ptr, _g1) = b_dev.device_ptr(&stream);
                let (p_ptr, _g3) = batched_probs_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
                        &stream, &c51_fwd_kernel,
                        w_ptr, b_ptr, h_ptr, p_ptr,
                        n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
                    )?;
                }
            } else {
                // C51 reads directly from batched_state_pinned (B=N).
                let (w_ptr, _g0) = w_dev.device_ptr(&stream);
                let (b_ptr, _g1) = b_dev.device_ptr(&stream);
                let (p_ptr, _g3) = batched_probs_dev.device_ptr_mut(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
                        &stream, &c51_fwd_kernel,
                        w_ptr, b_ptr, batched_state_pinned.dev_u64(), p_ptr,
                        n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
                    )?;
                }
            }
            // Batched Thompson select.
            let step_seed = episode_rng.next_u64() as u32;
            {
                let (p_ptr, _g0) = batched_probs_dev.device_ptr(&stream);
                unsafe {
                    ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
                        &stream, &c51_thompson_kernel,
                        p_ptr,
                        batched_state_pinned.dev_u64(),
                        threshold,
                        1,
                        state_dim_i,
                        c51_v_min, c51_delta_z,
                        step_seed,
                        batched_action_pinned.dev_u64(),
                        n_par as i32, n_act_i, n_atoms_i,
                    )?;
                }
            }
            // ONE sync per step (instead of N).
            stream.synchronize()?;
            let actions = batched_action_pinned.read_all();
            // Step all envs on CPU.
            for i in 0..n_par {
                if done_flags[i] { continue; }
                let action = actions[i] as u8;
                if action != 0 { ep_n_trades_par[i] += 1; }
                if let Some((_, r, done)) = par_envs[i].step(action, &mut par_states[i]) {
                    if done {
                        terminal_rs[i] = r;
                        done_flags[i] = true;
                    }
                }
            }
        }
        // Aggregate stats. Declared mut to keep the dead-code legacy
        // block (under `if false`) borrow-check clean.
        #[allow(unused_mut)]
        let mut rewards: Vec<f32> = terminal_rs;
        #[allow(unused_mut)]
        let mut trade_count_total: u64 = ep_n_trades_par.iter().map(|&t| t as u64).sum();
        #[allow(unused_mut)]
        let mut win_count: usize = rewards.iter().filter(|&&r| r > 0.0).count();
        // ISV-continual at cell-level (single fire with aggregate stats).
        if cli.isv_continual {
            let mean_terminal_r = rewards.iter().sum::<f32>() / (n_par as f32);
            let trade_count_avg = trade_count_total as f32 / (n_par as f32);
            let decisions_avg = cli.horizon as f32;
            let baseline_std = isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_STD_INDEX].max(1e-6);
            let rollout_sharpe = mean_terminal_r / baseline_std;
            unsafe {
                let (isv_ptr, _gi) = isv_dev.device_ptr_mut(&stream);
                let (wv_ptr, _gw) = ctl_wiener_dev.device_ptr_mut(&stream);
                ml::cuda_pipeline::alpha_kernels::launch_stacker_threshold_controller(
                    &stream, &ctl_kernel,
                    trade_count_avg, decisions_avg, rollout_sharpe,
                    0.5, 0.01, 0.005, 0.4, 0.1,
                    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. Active when --isv-continual is on;
                    // the regime kernel writes slots 549/550 per-step during
                    // inference. Floor 0.25 = max defensive 4× Kelly cut.
                    if cli.regime_scale {
                        ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_EMA_INDEX as i32
                    } else { -1 },
                    if cli.regime_scale {
                        ml::cuda_pipeline::alpha_isv_slots::REGIME_VOL_REF_INDEX as i32
                    } else { -1 },
                    0.25,
                    isv_ptr, wv_ptr,
                )?;
            }
        }
        let n = rewards.len() as f64;
        let mean = rewards.iter().map(|r| *r as f64).sum::<f64>() / n;
        let var = rewards
            .iter()
            .map(|r| (*r as f64 - mean).powi(2))
            .sum::<f64>()
            / n;
        let std = var.sqrt().max(1e-6);
        let sharpe_per = (mean / std) as f32;
        let episodes_per_year = 252.0 * 6.5 * 3600.0 / (cli.horizon as f64 * 12.0);
        let sharpe_ann = (sharpe_per as f64 * episodes_per_year.sqrt()) as f32;
        let win_rate = win_count as f32 / cli.n_eval_episodes as f32;
        let avg_n_trades = trade_count_total as f32 / cli.n_eval_episodes as f32;
        let mut sorted = rewards.clone();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let pick = |q: f64| -> f32 {
            let idx = (q * (n - 1.0)).round() as usize;
            sorted[idx.min(sorted.len() - 1)]
        };
        let bin = CostBin {
            cost,
            threshold,
            n_episodes: cli.n_eval_episodes,
            mean_reward: mean as f32,
            std_reward: std as f32,
            sharpe_per_episode: sharpe_per,
            sharpe_annualised: sharpe_ann,
            win_rate,
            avg_n_trades,
            p05: pick(0.05),
            p50: pick(0.50),
            p95: pick(0.95),
        };
        info!(
            "  τ={:.2} cost={:>6.4}  mean={:>+9.2}  Sharpe/ep={:+.3}  Sharpe_ann={:+.3}  win={:.3}  trades/ep={:.1}",
            threshold, cost, mean, sharpe_per, sharpe_ann, win_rate, avg_n_trades
        );
        bins.push(bin);
    }
    }

    // --- Print table ---
    info!("");
    info!("=== Phase E.3 2D sweep (threshold × cost, vs Phase 1d.4 baseline) ===");
    info!(
        "  {:>6}  {:>6}  {:>9}  {:>9}  {:>10}  {:>11}  {:>8}  {:>10}",
        "τ", "cost", "mean_R", "std_R", "Sharpe/ep", "Sharpe_ann", "win_rate", "trades/ep"
    );
    for b in &bins {
        info!(
            "  {:>6.3}  {:>6.4}  {:>9.2}  {:>9.2}  {:>10.4}  {:>11.4}  {:>8.3}  {:>10.2}",
            b.threshold, b.cost, b.mean_reward, b.std_reward,
            b.sharpe_per_episode, b.sharpe_annualised, b.win_rate, b.avg_n_trades
        );
    }
    info!("");
    info!("BEST per-cost (max Sharpe_ann across all τ at each cost):");
    for &cost in &cli.cost_grid {
        let best = bins
            .iter()
            .filter(|b| (b.cost - cost).abs() < 1e-6)
            .max_by(|a, b| {
                a.sharpe_annualised
                    .partial_cmp(&b.sharpe_annualised)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
        if let Some(b) = best {
            info!(
                "  cost={:>6.4}  best τ={:.3}  Sharpe_ann={:+.3}  win={:.3}  trades/ep={:.1}",
                b.cost, b.threshold, b.sharpe_annualised, b.win_rate, b.avg_n_trades
            );
        }
    }
    info!("");
    info!("Phase 1d.4 baseline for comparison: +4.4 annualised at cost=0,");
    info!("                                    -4.0 annualised at cost=0.125.");

    // --- Save JSON ---
    let json = serde_json::json!({
        "phase": "E.3 Task 23 (2D sweep)",
        "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 },
        "temporal": cli.temporal,
        "window_k": if cli.temporal { cli.window_k } else { 0 },
        "mamba2_hidden_dim": if cli.temporal { cli.mamba2_hidden_dim } else { 0 },
        "mamba2_state_dim": if cli.temporal { cli.mamba2_state_dim } else { 0 },
        "isv_continual": cli.isv_continual,
        "horizon": cli.horizon,
        "train_frac": cli.train_frac,
        "n_train_episodes": cli.n_train_episodes,
        "n_eval_episodes": cli.n_eval_episodes,
        "train_cost": cli.train_cost,
        "cost_grid": cli.cost_grid,
        "threshold_grid": cli.threshold_grid,
        "bins": bins,
    });
    let mut f = File::create(&cli.out_path)?;
    write!(f, "{}", serde_json::to_string_pretty(&json)?)?;
    info!("Wrote table to {}", cli.out_path.display());

    // Touch unused ISV slot constants so they're imported for future expansion.
    let _ = (RANDOM_BASELINE_MEAN_INDEX, RANDOM_BASELINE_STD_INDEX);
    let _: Option<ReplayRng> = None;
    let _ = SnapshotRow {
        mid_price: 0.0,
        bid_l: [0.0; 10],
        ask_l: [0.0; 10],
        alpha_logit: 0.0,
        alpha_confidence: 0.0,
        spread_bps: 0.0,
        l1_imbalance: 0.0,
        ofi_sum_5: 0.0,
        mid_drift_5: 0.0,
        time_since_trade_s: 0.0,
        book_event_rate: 0.0,
    };

    Ok(())
}