//! PerceptionTrainer synthetic-overfit smoke.
//!
//! Mirrors `perception_overfit.rs` but on the stacked trainer. The
//! signal is constant direction=+1, label=[1; N_HORIZONS]. Asserts the BCE loss
//! shrinks at least 40% over the training budget — proves the
//! full forward + backward (Mamba2 + CfC + heads) wires up correctly
//! end-to-end and that all 6 AdamW optimizers actually move weights.

use ml_alpha::aux_heads::N_AUX_HORIZONS;
use ml_alpha::cfc::snap_features::Mbp10RawInput;
use ml_alpha::heads::N_HORIZONS;
use ml_alpha::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig};
use ml_core::device::MlDevice;

fn test_device() -> MlDevice {
    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
}

fn synthetic_seq(
    seq_len: usize,
    mut prev_mid: f32,
    mut ts_ns: u64,
) -> (Vec<Mbp10RawInput>, Vec<[f32; N_HORIZONS]>) {
    let mut out = Vec::with_capacity(seq_len);
    for k in 0..seq_len {
        let next_mid = prev_mid + 0.25;
        let mut bid_px = [0.0f32; 10];
        let mut bid_sz = [0.0f32; 10];
        let mut ask_px = [0.0f32; 10];
        let mut ask_sz = [0.0f32; 10];
        for i in 0..10 {
            bid_px[i] = next_mid - 0.125 - 0.25 * i as f32;
            ask_px[i] = next_mid + 0.125 + 0.25 * i as f32;
            bid_sz[i] = 10.0;
            ask_sz[i] = 10.0;
        }
        let prev_ts = ts_ns;
        ts_ns += 20_000_000;
        out.push(Mbp10RawInput {
            bid_px, bid_sz, ask_px, ask_sz,
            prev_mid,
            trade_signed_vol: 1.0,
            trade_count: 1,
            ts_ns,
            prev_ts_ns: prev_ts,
            regime: [0.0; 6],
        });
        prev_mid = next_mid;
        let _ = k;
    }
    // Per-position labels: every position knows the next K snapshots all
    // move up (synthetic monotone ramp), so label = 1.0 for every horizon
    // at every position. Drives the trainer to learn "always predict 1".
    let labels = vec![[1.0; N_HORIZONS]; seq_len];
    (out, labels)
}

/// SDD-3 CB2: synthetic aux-outcome generator matching the +0.25 price
/// ramp from `synthetic_seq`. For a constant up-trend the A+B paired
/// targets are:
///   y_prof_long  = 1.0 (every long position is profitable after 2×cost)
///   y_prof_short = 0.0 (every short position loses)
///   y_size_long  ≈ +ratio (positive σ-normalized signed return)
///   y_size_short = NaN   (unprofitable side: size target masked)
/// Returns five vectors of `[N_AUX_HORIZONS]` rows + the per-(dir,horizon)
/// positive-class fraction expected by the trainer signature.
fn synthetic_aux_outcomes(
    seq_len: usize,
) -> AuxOutcomes {
    let prof_long_row = std::array::from_fn::<f32, N_AUX_HORIZONS, _>(|_| 1.0);
    let prof_short_row = std::array::from_fn::<f32, N_AUX_HORIZONS, _>(|_| 0.0);
    // Per-horizon magnitude — exact values don't matter for the smoke;
    // the invariant is sign-correctness so the dir_acc gate measures.
    let size_long_per_h: [f32; N_AUX_HORIZONS] = [0.5, 1.5, 5.0];
    let size_long_row = std::array::from_fn::<f32, N_AUX_HORIZONS, _>(|h| size_long_per_h[h]);
    // Short is unprofitable → size masked with NaN (loader convention).
    let size_short_row = std::array::from_fn::<f32, N_AUX_HORIZONS, _>(|_| f32::NAN);
    // pos_fraction is 1.0 for long (all positive) and 0.0 for short.
    let mut pos_fraction = vec![0.0_f32; 2 * N_AUX_HORIZONS];
    for h in 0..N_AUX_HORIZONS { pos_fraction[h] = 1.0; }
    AuxOutcomes {
        prof_long: vec![prof_long_row; seq_len],
        prof_short: vec![prof_short_row; seq_len],
        size_long: vec![size_long_row; seq_len],
        size_short: vec![size_short_row; seq_len],
        pos_fraction,
    }
}

/// CB2: bundle for the A+B paired aux-outcome targets returned by
/// `synthetic_aux_outcomes` and consumed by `PerceptionTrainer::step`.
struct AuxOutcomes {
    prof_long: Vec<[f32; N_AUX_HORIZONS]>,
    prof_short: Vec<[f32; N_AUX_HORIZONS]>,
    size_long: Vec<[f32; N_AUX_HORIZONS]>,
    size_short: Vec<[f32; N_AUX_HORIZONS]>,
    pos_fraction: Vec<f32>,
}

#[test]
fn stacked_trainer_constructs_cleanly() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig::default();
    let t = PerceptionTrainer::new(&dev, &cfg).expect("init");
    drop(t);
}

#[test]
fn stacked_trainer_loss_shrinks_on_constant_signal() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16, // smaller for smoke speed; Mamba2 needs >=2
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0x4242,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 1,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let aux = synthetic_aux_outcomes(cfg.seq_len);

    // Initial loss over 8 batches.
    let mut initial_total = 0.0_f32;
    let mut ts = 1_000_000u64;
    let mut prev_mid = 5500.0_f32;
    for _ in 0..8 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        let l = trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("step warm");
        initial_total += l;
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }
    let initial_avg = initial_total / 8.0;
    eprintln!("initial_avg = {initial_avg:.4}");

    // Train 250 steps. Print every 50.
    let mut window_loss = 0.0_f32;
    let mut window_count = 0usize;
    for step_idx in 0..250 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        let l = trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("train step");
        window_loss += l;
        window_count += 1;
        if step_idx % 50 == 49 {
            eprintln!(
                "  step {}: window_avg_loss={:.4}",
                step_idx + 1,
                window_loss / window_count as f32
            );
            window_loss = 0.0;
            window_count = 0;
        }
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }

    // Final loss over 8 batches.
    let mut final_total = 0.0_f32;
    for _ in 0..8 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        let l = trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("step final");
        final_total += l;
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }
    let final_avg = final_total / 8.0;
    eprintln!("final_avg = {final_avg:.4}");

    assert!(
        final_avg < 0.6 * initial_avg || final_avg < 0.5,
        "PerceptionTrainer failed to overfit constant signal: \
         start={initial_avg:.4}, end={final_avg:.4}"
    );
}

/// Verifies the trainer still converges with a distinct seed and lr.
/// A1: decision_stride removed; dt_s is always 1.0 (event-rate).
/// The test previously verified stride=4 dt_s dynamics; those are
/// now unconditionally event-rate, so this becomes a second convergence
/// smoke with different hyperparameters.
#[test]
fn stacked_trainer_loss_shrinks_with_stride_4() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0xC4C4,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 1,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let aux = synthetic_aux_outcomes(cfg.seq_len);

    let mut initial = 0.0_f32;
    let mut ts = 1_000_000u64;
    let mut prev_mid = 5500.0_f32;
    for _ in 0..8 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        let l = trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("step");
        initial += l;
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }
    initial /= 8.0;
    eprintln!("stride=4 trainer: initial={initial:.4}");

    for _ in 0..200 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("step");
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }

    let mut final_loss = 0.0_f32;
    for _ in 0..8 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        final_loss += trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("step");
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }
    final_loss /= 8.0;
    eprintln!("stride=4 trainer: final={final_loss:.4}");
    assert!(
        final_loss < 0.6 * initial || final_loss < 0.5,
        "stride=4 trainer failed to converge: {initial:.4} → {final_loss:.4}"
    );
}

/// Verifies the trainer converges at n_batch=32 — the FIRST test that
/// exercises the cross-batch reducer code path. Existing
/// `stacked_trainer_loss_shrinks_*` tests all use n_batch=1 so the new
/// per-batch scratch + reducer logic was never previously hit.
#[test]
fn stacked_trainer_loss_shrinks_at_batch_32() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0xB32B,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 32,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");

    let mut ts_base = 1_000_000u64;
    let mut prev_mid = 5500.0_f32;

    let make_batch = |prev_mid: &mut f32, ts_base: &mut u64,
                      cfg: &PerceptionTrainerConfig|
     -> (Vec<Vec<Mbp10RawInput>>, Vec<Vec<[f32; N_HORIZONS]>>) {
        let mut seqs: Vec<Vec<Mbp10RawInput>> = Vec::with_capacity(cfg.n_batch);
        let mut labels: Vec<Vec<[f32; N_HORIZONS]>> = Vec::with_capacity(cfg.n_batch);
        for _ in 0..cfg.n_batch {
            let (seq, lbl) = synthetic_seq(cfg.seq_len, *prev_mid, *ts_base);
            *prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
            *ts_base = seq.last().unwrap().ts_ns;
            seqs.push(seq);
            labels.push(lbl);
        }
        (seqs, labels)
    };
    // SDD-3 CB2: synthetic aux outcomes are seq_len-determined and
    // direction-asymmetric; the shared per-sample rows for both prof-binary
    // and size targets are broadcast across all B samples per step.
    let aux_seq = synthetic_aux_outcomes(cfg.seq_len);
    let prof_long_batch: Vec<Vec<[f32; N_AUX_HORIZONS]>> =
        (0..cfg.n_batch).map(|_| aux_seq.prof_long.clone()).collect();
    let prof_short_batch: Vec<Vec<[f32; N_AUX_HORIZONS]>> =
        (0..cfg.n_batch).map(|_| aux_seq.prof_short.clone()).collect();
    let size_long_batch: Vec<Vec<[f32; N_AUX_HORIZONS]>> =
        (0..cfg.n_batch).map(|_| aux_seq.size_long.clone()).collect();
    let size_short_batch: Vec<Vec<[f32; N_AUX_HORIZONS]>> =
        (0..cfg.n_batch).map(|_| aux_seq.size_short.clone()).collect();
    let prof_long_refs: Vec<&[[f32; N_AUX_HORIZONS]]> =
        prof_long_batch.iter().map(|v| v.as_slice()).collect();
    let prof_short_refs: Vec<&[[f32; N_AUX_HORIZONS]]> =
        prof_short_batch.iter().map(|v| v.as_slice()).collect();
    let size_long_refs: Vec<&[[f32; N_AUX_HORIZONS]]> =
        size_long_batch.iter().map(|v| v.as_slice()).collect();
    let size_short_refs: Vec<&[[f32; N_AUX_HORIZONS]]> =
        size_short_batch.iter().map(|v| v.as_slice()).collect();

    let mut initial = 0.0_f32;
    for warmup in 0..4 {
        let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
        let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
        let lbl_refs: Vec<&[[f32; N_HORIZONS]]> = labels.iter().map(|l| l.as_slice()).collect();
        let l = trainer
            .step_batched(
                &seq_refs,
                &lbl_refs,
                &prof_long_refs,
                &prof_short_refs,
                &size_long_refs,
                &size_short_refs,
                &aux_seq.pos_fraction,
            )
            .expect("warm step");
        if warmup >= 2 {
            initial += l;
        }
    }
    initial /= 2.0;
    eprintln!("B=32 trainer: initial={initial:.4}");

    for _ in 0..200 {
        let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
        let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
        let lbl_refs: Vec<&[[f32; N_HORIZONS]]> = labels.iter().map(|l| l.as_slice()).collect();
        trainer
            .step_batched(
                &seq_refs,
                &lbl_refs,
                &prof_long_refs,
                &prof_short_refs,
                &size_long_refs,
                &size_short_refs,
                &aux_seq.pos_fraction,
            )
            .expect("train step");
    }

    let mut final_loss = 0.0_f32;
    for _ in 0..4 {
        let (seqs, labels) = make_batch(&mut prev_mid, &mut ts_base, &cfg);
        let seq_refs: Vec<&[Mbp10RawInput]> = seqs.iter().map(|s| s.as_slice()).collect();
        let lbl_refs: Vec<&[[f32; N_HORIZONS]]> = labels.iter().map(|l| l.as_slice()).collect();
        final_loss += trainer
            .step_batched(
                &seq_refs,
                &lbl_refs,
                &prof_long_refs,
                &prof_short_refs,
                &size_long_refs,
                &size_short_refs,
                &aux_seq.pos_fraction,
            )
            .expect("final step");
    }
    final_loss /= 4.0;
    eprintln!("B=32 trainer: final={final_loss:.4}");
    assert!(
        final_loss < 0.6 * initial || final_loss < 0.1,
        "B=32 trainer failed to converge: {initial:.4} → {final_loss:.4}"
    );
}

// NOTE on scratch-clears testing:
//
// A direct "scratch is zero between steps" test is structurally hard to
// write against this trainer. The lifecycle is:
//   step 1: uncaptured warmup execute      → scratch ends with grad
//   step 2: begin_capture/end_capture       → records kernels but does
//                                              NOT execute them, scratch
//                                              unchanged from step 1
//   step 3+: graph.launch_replay            → executes the captured graph
//
// A post-step snapshot read on step 2 returns step 1's residual,
// which is bit-identical between any two captures regardless of input
// data. Two replay steps would work but require either three+ trainer
// calls or capture-internal scratch readback, both of which require
// substantial test infrastructure.
//
// Instead we rely on the `stacked_trainer_loss_shrinks_at_batch_32`
// smoke as the implicit scratch-zero validation: if the K-loop scratch
// wasn't being zeroed at step start, gradients from step N would
// pollute step N+1, training would diverge after a handful of steps,
// and the convergence assertion would fail. The B=32 smoke explicitly
// runs 200 training steps + measures final loss, so the absence of
// divergence IS the scratch-zero guarantee.


/// Eval alone must work — proves the eval path is fine WITHOUT any
/// prior captured-graph training. If this passes but
/// `evaluate_works_after_captured_training_step` fails, the bug is in
/// the train→eval transition (captured graph leaving state hostile to
/// direct kernel launches).
#[test]
fn evaluate_alone_succeeds() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0x6262,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 1,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let ts = 1_000_000u64;
    let prev_mid = 5500.0_f32;
    let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
    let (loss, probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval alone");
    assert!(loss.is_finite(), "eval loss must be finite, got {loss}");
    assert_eq!(probs.len(), cfg.seq_len * N_HORIZONS);
}

/// Regression test for z2w9w cluster run: training step (which captures
/// a CUDA Graph) must NOT break the subsequent `evaluate()` call.
/// z2w9w hit CUDA_ERROR_INVALID_VALUE at "eval snap_batched fwd" the
/// very first time eval ran after training; the captured graph or some
/// of its launch state left the stream in a state hostile to the
/// non-captured eval path.
#[test]
fn evaluate_works_after_captured_training_step() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0x5151,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 1,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let aux = synthetic_aux_outcomes(cfg.seq_len);

    // Drive enough training steps to exercise warmup → capture → replay.
    let mut ts = 1_000_000u64;
    let mut prev_mid = 5500.0_f32;
    for _ in 0..5 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("train step");
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }

    // Now call evaluate — must NOT fail with CUDA_ERROR_INVALID_VALUE.
    let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
    let (loss, probs) = trainer.evaluate(&seq, labels.as_slice()).expect("evaluate after train");
    assert!(loss.is_finite(), "eval loss must be finite, got {loss}");
    assert_eq!(probs.len(), cfg.seq_len * N_HORIZONS, "eval probs must be [K, N_HORIZONS] flat");
    assert!(probs.iter().all(|p| p.is_finite()), "eval probs must be finite");
}

/// 2 steps = warmup + capture (no replay yet). Does eval fail right
/// after capture but BEFORE the first graph.launch?
#[test]
fn evaluate_works_after_capture_no_replay() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0x8181,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 1,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let aux = synthetic_aux_outcomes(cfg.seq_len);
    let ts = 1_000_000u64;
    let prev_mid = 5500.0_f32;
    let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
    trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("warmup step");
    trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("capture step");
    // Eval right after capture, no replay.
    let (loss, _probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval after capture");
    assert!(loss.is_finite(), "eval loss must be finite, got {loss}");
}

/// ISV-driven per-horizon EMA + lambda — after a few training steps,
/// `loss_ema` should contain positive BCE values for every horizon
/// (proves the BCE kernel writes the per-horizon channel correctly),
/// and `lambda` should land inside the ASYMMETRIC clamp [1.0, 2.0]
/// (boost-only — never demote below uniform, per
/// `pearl_audit_unboundedness_for_implicit_asymmetry.md`).
#[test]
fn horizon_ema_and_lambda_track_after_training() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0x9292,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 1,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let aux = synthetic_aux_outcomes(cfg.seq_len);

    // EMA / lambda before any training step.
    let ema0 = trainer.loss_ema_snapshot().expect("ema initial");
    let lam0 = trainer.lambda_snapshot().expect("lambda initial");
    assert!(ema0.iter().all(|v| *v == 0.0), "ema should be zero-init sentinel, got {ema0:?}");
    assert!(lam0.iter().all(|v| *v == 0.0), "lambda should be zero before first kernel run, got {lam0:?}");

    let mut ts = 1_000_000u64;
    let mut prev_mid = 5500.0_f32;
    for _ in 0..5 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("train step");
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }

    let ema = trainer.loss_ema_snapshot().expect("ema after train");
    let lam = trainer.lambda_snapshot().expect("lambda after train");
    eprintln!("loss_ema = {ema:?}");
    eprintln!("lambda   = {lam:?}");
    // Each horizon must have a positive finite EMA (BCE > 0 unless model is perfect).
    for (h, &v) in ema.iter().enumerate() {
        assert!(v.is_finite(), "ema[{h}] not finite: {v}");
        assert!(v > 0.0, "ema[{h}] should be positive, got {v}");
    }
    // Lambda asymmetric clamp: every entry in [1.0, 2.0] (boost-only).
    for (h, &v) in lam.iter().enumerate() {
        assert!(v.is_finite(), "lambda[{h}] not finite: {v}");
        assert!(v >= 1.0 - 1e-6 && v <= 2.0 + 1e-6,
            "lambda[{h}] = {v} outside asymmetric clamp [1.0, 2.0]");
    }
    // Mean of lambda must be >= 1.0 (boost-only, never demote).
    // Upper-bounded by 2.0 (ceiling clamp). On synthetic-constant data
    // the per-horizon EMAs converge, ratios approach 1.0, mean → 1.0.
    let lam_mean: f32 = lam.iter().sum::<f32>() / lam.len() as f32;
    assert!(lam_mean >= 1.0 - 1e-6 && lam_mean <= 2.0 + 1e-6,
        "lambda mean {lam_mean} outside [1.0, 2.0] envelope");
}

/// SDD-3 Layer B5: verifies the aux supervision Huber loss converges on
/// the same synthetic constant-direction up-ramp the BCE smoke uses.
/// Asserts:
///   * Aux per-horizon Huber EMA is finite and bounded after training.
///   * Aux per-horizon directional accuracy EMA is high (the synthetic
///     long > short separation is unambiguous).
///   * Asymmetric stop-grad to encoder either stays masked (Phase 1
///     auditor behaviour) OR lifts after the gates are satisfied —
///     either is a valid outcome on the small synthetic budget.
#[test]
fn stacked_trainer_aux_supervision_converges_on_constant_signal() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0xA1B5,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 1,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let aux = synthetic_aux_outcomes(cfg.seq_len);

    let mut ts = 1_000_000u64;
    let mut prev_mid = 5500.0_f32;
    // Train enough steps for the aux Huber + dir_acc EMAs to settle.
    for _ in 0..400 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        trainer
            .step(
                &seq,
                labels.as_slice(),
                aux.prof_long.as_slice(),
                aux.prof_short.as_slice(),
                aux.size_long.as_slice(),
                aux.size_short.as_slice(),
                &aux.pos_fraction,
            )
            .expect("train step");
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }

    // After ≥400 steps (CB5):
    //   * Aux prof BCE EMA must be finite + non-negative on every horizon
    //     and STRICTLY BELOW the random-baseline ln(2) ≈ 0.693 — the
    //     constant-up signal is perfectly classifiable so even a small
    //     budget should beat random. (The CB5 kernel returns a joint
    //     scalar that we broadcast to every horizon, so the same value
    //     appears at h=0..N — but that's still load-bearing as an
    //     invariant check.)
    //   * Aux size Huber EMA must be finite + non-negative (no upper
    //     bound assertion because the σ-normalised size target can be
    //     anywhere in ℝ and we're not testing the regression accuracy,
    //     only that the loss is well-defined and the kernel ran).
    //   * Per-horizon dir_acc EMA must exceed chance (0.5) on every
    //     horizon — the synthetic long/short separation is unambiguous,
    //     so a well-wired aux head should clear it easily.
    const RANDOM_BCE: f32 = std::f32::consts::LN_2;
    for (h, &v) in trainer.aux_prof_bce_ema_per_h.iter().enumerate() {
        assert!(v.is_finite(), "aux_prof_bce_ema_per_h[{h}] not finite: {v}");
        assert!(v >= 0.0, "aux_prof_bce_ema_per_h[{h}] negative: {v}");
        assert!(
            v < RANDOM_BCE,
            "aux_prof_bce_ema_per_h[{h}] = {v} >= random baseline ln(2)={RANDOM_BCE}; \
             aux prof head should beat random on unambiguous signal"
        );
    }
    for (h, &v) in trainer.aux_size_huber_ema_per_h.iter().enumerate() {
        assert!(v.is_finite(), "aux_size_huber_ema_per_h[{h}] not finite: {v}");
        assert!(v >= 0.0, "aux_size_huber_ema_per_h[{h}] negative: {v}");
    }
    for (h, &v) in trainer.aux_dir_acc_ema_per_h.iter().enumerate() {
        assert!(v.is_finite(), "aux_dir_acc_ema_per_h[{h}] not finite: {v}");
        assert!(
            (0.0..=1.0).contains(&v),
            "aux_dir_acc_ema_per_h[{h}] out of [0, 1]: {v}"
        );
        assert!(
            v >= 0.5,
            "aux_dir_acc_ema_per_h[{h}] = {v} below chance on unambiguous signal"
        );
    }
    eprintln!(
        "aux_prof_bce_ema_per_h   = {:?}",
        trainer.aux_prof_bce_ema_per_h
    );
    eprintln!(
        "aux_size_huber_ema_per_h = {:?}",
        trainer.aux_size_huber_ema_per_h
    );
    eprintln!(
        "aux_dir_acc_ema_per_h    = {:?}",
        trainer.aux_dir_acc_ema_per_h
    );
    eprintln!(
        "stop_grad_aux_to_encoder = {}",
        trainer.stop_grad_aux_to_encoder
    );
}

/// Does the bug surface after just ONE step (warmup only, no capture)?
/// If yes, the warmup dispatch path itself breaks subsequent eval.
/// If no, the captured graph instantiation / launch is what breaks eval.
#[test]
fn evaluate_works_after_warmup_only() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0x7171,
        horizon_weights: [1.0; N_HORIZONS],
        n_batch: 1,
        smoothness_base_lambda: 0.0,
        kernel_step_trace_path: None,
        bucket_warmup_cap_override: None,
        lr_aux: 3e-3,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let aux = synthetic_aux_outcomes(cfg.seq_len);
    let ts = 1_000_000u64;
    let prev_mid = 5500.0_f32;
    let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
    // ONLY one step — warmup, no capture yet.
    trainer.step(
            &seq,
            labels.as_slice(),
            aux.prof_long.as_slice(),
            aux.prof_short.as_slice(),
            aux.size_long.as_slice(),
            aux.size_short.as_slice(),
            &aux.pos_fraction,
        ).expect("warmup step");
    // Now eval.
    let (loss, _probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval after warmup");
    assert!(loss.is_finite(), "eval loss must be finite, got {loss}");
}
