//! GPU oracle tests for `aux_heads.cu` + `aux_loss.cu` (SDD-3 CB3+CB4).
//!
//! Run with:
//!   SQLX_OFFLINE=true cargo test -p ml-alpha --test aux_heads \
//!     -- --ignored --nocapture
//!
//! These tests are `#[ignore]` because they require a CUDA device. They
//! drive the cubins directly through the host wrappers in
//! `crates/ml-alpha/src/aux_heads.rs` and compare against a deterministic
//! CPU naive reference (per `feedback_no_cpu_test_fallbacks.md` the
//! production path stays GPU-only; the CPU implementation here exists
//! ONLY as a unit-test oracle).
//!
//! Covers the A+B paired-head architecture:
//!   * 4 outputs per direction (prof_long_logit, size_long_pred,
//!     prof_short_logit, size_short_pred)
//!   * Class-weighted BCE on the prof heads (with per-horizon
//!     `pos_weight` scaling positive-class gradient)
//!   * NaN-masked Huber on the size heads (conditional-Huber via the
//!     loader's `y_size = NaN` at `y_prof = 0`)

use anyhow::Result;
use ml_alpha::aux_heads::{
    aux_bce_loss_gpu, aux_heads_bwd_gpu, aux_heads_fwd_gpu, aux_huber_masked_loss_gpu, AuxBceLoss,
    AuxHeads, AuxHeadsConfig, AuxMaskedHuberLoss, DEFAULT_HUBER_DELTA, N_AUX_HORIZONS,
};
use ml_alpha::cfc::AUX_HIDDEN;
use ml_core::device::MlDevice;

fn upload(
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    host: &[f32],
) -> Result<cudarc::driver::CudaSlice<f32>> {
    let mut d = stream.alloc_zeros::<f32>(host.len())?;
    stream.memcpy_htod(host, &mut d)?;
    Ok(d)
}

fn upload_i32(
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    host: &[i32],
) -> Result<cudarc::driver::CudaSlice<i32>> {
    let mut d = stream.alloc_zeros::<i32>(host.len())?;
    stream.memcpy_htod(host, &mut d)?;
    Ok(d)
}

fn download(
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    d: &cudarc::driver::CudaSlice<f32>,
) -> Result<Vec<f32>> {
    let mut h = vec![0.0_f32; d.len()];
    stream.memcpy_dtoh(d, &mut h)?;
    Ok(h)
}

fn download_i32(
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    d: &cudarc::driver::CudaSlice<i32>,
) -> Result<Vec<i32>> {
    let mut h = vec![0_i32; d.len()];
    stream.memcpy_dtoh(d, &mut h)?;
    Ok(h)
}

/// Forward shape correctness + value match against a naive CPU oracle.
///
/// Validates:
///   * Output shapes ([B × N_AUX_HORIZONS] per (direction, head)).
///   * Per-element value within tolerance (5e-4) of the naive reference
///     for ALL FOUR outputs (prof_long, size_long, prof_short, size_short).
///   * The four heads use INDEPENDENT parameter sets — flipping a single
///     prof_long weight does not perturb a size_long output nor any
///     short-direction output (verified implicitly via the closed-form
///     match).
#[test]
#[ignore = "requires CUDA"]
fn fwd_matches_naive_reference() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let heads = AuxHeads::new(&dev, AuxHeadsConfig { seed: 0xAE_AB_BEEF })?;

    let b_sz: usize = 4;
    let h_aux: Vec<f32> = (0..b_sz * AUX_HIDDEN)
        .map(|i| ((i as f32) * 0.013).sin() * 0.6)
        .collect();
    let h_aux_d = upload(&stream, &h_aux)?;
    let mut prof_long_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
    let mut size_long_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
    let mut prof_short_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;
    let mut size_short_d = stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS)?;

    aux_heads_fwd_gpu(
        &stream,
        &heads.fwd_fn,
        &heads.w_prof_long_d,
        &heads.b_prof_long_d,
        &heads.w_size_long_d,
        &heads.b_size_long_d,
        &heads.w_prof_short_d,
        &heads.b_prof_short_d,
        &heads.w_size_short_d,
        &heads.b_size_short_d,
        &h_aux_d,
        b_sz as i32,
        &mut prof_long_d,
        &mut size_long_d,
        &mut prof_short_d,
        &mut size_short_d,
    )?;
    stream.synchronize()?;

    let prof_long = download(&stream, &prof_long_d)?;
    let size_long = download(&stream, &size_long_d)?;
    let prof_short = download(&stream, &prof_short_d)?;
    let size_short = download(&stream, &size_short_d)?;
    let w = heads.download_weights()?;

    for batch in 0..b_sz {
        for k in 0..N_AUX_HORIZONS {
            // Closed-form expected outputs for all four heads.
            let mut e_pl = w.b_prof_long[k];
            let mut e_sl = w.b_size_long[k];
            let mut e_ps = w.b_prof_short[k];
            let mut e_ss = w.b_size_short[k];
            for c in 0..AUX_HIDDEN {
                let hc = h_aux[batch * AUX_HIDDEN + c];
                e_pl += w.w_prof_long [k * AUX_HIDDEN + c] * hc;
                e_sl += w.w_size_long [k * AUX_HIDDEN + c] * hc;
                e_ps += w.w_prof_short[k * AUX_HIDDEN + c] * hc;
                e_ss += w.w_size_short[k * AUX_HIDDEN + c] * hc;
            }
            let g_pl = prof_long [batch * N_AUX_HORIZONS + k];
            let g_sl = size_long [batch * N_AUX_HORIZONS + k];
            let g_ps = prof_short[batch * N_AUX_HORIZONS + k];
            let g_ss = size_short[batch * N_AUX_HORIZONS + k];
            for (name, got, expected) in [
                ("prof_long",  g_pl, e_pl),
                ("size_long",  g_sl, e_sl),
                ("prof_short", g_ps, e_ps),
                ("size_short", g_ss, e_ss),
            ] {
                let err = (got - expected).abs();
                assert!(
                    err < 5e-4,
                    "{name}[{batch},{k}] = {got} vs expected {expected} (err={err})",
                );
            }
        }
    }

    Ok(())
}

/// Backward grad shapes finite, no NaN/Inf, and bias-grad finite-difference
/// matches within tolerance against the implicit loss
/// `L = Σ Σ grad_out[head] * out[head]` (linear in each output so all
/// dL/dW and dL/db are closed-form and stable under FD).
///
/// Also asserts the `grad_h_aux` per-batch slot is non-zero and finite,
/// which is the load-bearing signal that flows back into `aux_trunk_bwd`.
#[test]
#[ignore = "requires CUDA"]
fn bwd_finite_diff_matches_bias_sample() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let heads = AuxHeads::new(&dev, AuxHeadsConfig { seed: 0xBEEF_FEED })?;

    let b_sz: usize = 3;
    let h_aux: Vec<f32> = (0..b_sz * AUX_HIDDEN)
        .map(|i| ((i as f32) * 0.011).cos() * 0.4)
        .collect();
    // Structured per-element grad (non-zero everywhere) so every grad
    // path exercises a non-trivial accumulation. Distinct seeds across
    // the four heads catch any cross-head wiring mistake.
    let g_pl: Vec<f32> = (0..b_sz * N_AUX_HORIZONS)
        .map(|i| ((i as f32) * 0.17).sin() * 0.5 + 0.10)
        .collect();
    let g_sl: Vec<f32> = (0..b_sz * N_AUX_HORIZONS)
        .map(|i| ((i as f32) * 0.21).cos() * 0.4 + 0.05)
        .collect();
    let g_ps: Vec<f32> = (0..b_sz * N_AUX_HORIZONS)
        .map(|i| ((i as f32) * 0.23).cos() * 0.4 - 0.05)
        .collect();
    let g_ss: Vec<f32> = (0..b_sz * N_AUX_HORIZONS)
        .map(|i| ((i as f32) * 0.29).sin() * 0.3 - 0.02)
        .collect();

    let h_aux_d = upload(&stream, &h_aux)?;
    let g_pl_d = upload(&stream, &g_pl)?;
    let g_sl_d = upload(&stream, &g_sl)?;
    let g_ps_d = upload(&stream, &g_ps)?;
    let g_ss_d = upload(&stream, &g_ss)?;

    let alloc_w = || stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS * AUX_HIDDEN);
    let alloc_b = || stream.alloc_zeros::<f32>(b_sz * N_AUX_HORIZONS);
    let mut gw_pl = alloc_w()?;
    let mut gb_pl = alloc_b()?;
    let mut gw_sl = alloc_w()?;
    let mut gb_sl = alloc_b()?;
    let mut gw_ps = alloc_w()?;
    let mut gb_ps = alloc_b()?;
    let mut gw_ss = alloc_w()?;
    let mut gb_ss = alloc_b()?;
    let mut grad_h_aux_d = stream.alloc_zeros::<f32>(b_sz * AUX_HIDDEN)?;

    aux_heads_bwd_gpu(
        &stream,
        &heads.bwd_fn,
        &heads.w_prof_long_d,
        &heads.w_size_long_d,
        &heads.w_prof_short_d,
        &heads.w_size_short_d,
        &h_aux_d,
        &g_pl_d,
        &g_sl_d,
        &g_ps_d,
        &g_ss_d,
        b_sz as i32,
        &mut gw_pl,
        &mut gb_pl,
        &mut gw_sl,
        &mut gb_sl,
        &mut gw_ps,
        &mut gb_ps,
        &mut gw_ss,
        &mut gb_ss,
        &mut grad_h_aux_d,
    )?;
    stream.synchronize()?;

    let gwl_pl = download(&stream, &gw_pl)?;
    let gbl_pl = download(&stream, &gb_pl)?;
    let gwl_sl = download(&stream, &gw_sl)?;
    let gbl_sl = download(&stream, &gb_sl)?;
    let gwl_ps = download(&stream, &gw_ps)?;
    let gbl_ps = download(&stream, &gb_ps)?;
    let gwl_ss = download(&stream, &gw_ss)?;
    let gbl_ss = download(&stream, &gb_ss)?;
    let gha = download(&stream, &grad_h_aux_d)?;

    for (name, v) in [
        ("grad_w_prof_long",  &gwl_pl),
        ("grad_b_prof_long",  &gbl_pl),
        ("grad_w_size_long",  &gwl_sl),
        ("grad_b_size_long",  &gbl_sl),
        ("grad_w_prof_short", &gwl_ps),
        ("grad_b_prof_short", &gbl_ps),
        ("grad_w_size_short", &gwl_ss),
        ("grad_b_size_short", &gbl_ss),
        ("grad_h_aux",        &gha),
    ] {
        assert!(
            v.iter().all(|x| x.is_finite()),
            "{name} contains NaN/Inf",
        );
    }
    assert!(
        gha.iter().any(|x| x.abs() > 0.0),
        "grad_h_aux is all-zero — kernel didn't write the trunk feedback?",
    );

    // Closed-form bias-grad: dL/db[head][k] = Σ_batch grad_out[head][batch, k]
    let w = heads.download_weights()?;
    for k in 0..N_AUX_HORIZONS {
        // Per-head closed-form bias sum across batch.
        let mut got = [0.0_f32; 4];
        let mut expected = [0.0_f32; 4];
        for batch in 0..b_sz {
            let idx = batch * N_AUX_HORIZONS + k;
            got[0] += gbl_pl[idx];
            got[1] += gbl_sl[idx];
            got[2] += gbl_ps[idx];
            got[3] += gbl_ss[idx];
            expected[0] += g_pl[idx];
            expected[1] += g_sl[idx];
            expected[2] += g_ps[idx];
            expected[3] += g_ss[idx];
        }
        for (name, gv, ev) in [
            ("grad_b_prof_long",  got[0], expected[0]),
            ("grad_b_size_long",  got[1], expected[1]),
            ("grad_b_prof_short", got[2], expected[2]),
            ("grad_b_size_short", got[3], expected[3]),
        ] {
            let err = (gv - ev).abs();
            assert!(
                err < 1e-5,
                "{name}[{k}] mismatch: got={gv} expected={ev}",
            );
        }
    }

    // Spot-check weight entries against the closed form. Per head:
    //   dL/dW[head][k, c] = Σ_batch grad_out[head][batch, k] * h_aux[batch, c]
    for &(k, c) in &[(0_usize, 0_usize), (1, 7), (2, 33), (0, 63)] {
        let mut got = [0.0_f32; 4];
        let mut expected = [0.0_f32; 4];
        for batch in 0..b_sz {
            let off = batch * N_AUX_HORIZONS * AUX_HIDDEN + k * AUX_HIDDEN + c;
            got[0] += gwl_pl[off];
            got[1] += gwl_sl[off];
            got[2] += gwl_ps[off];
            got[3] += gwl_ss[off];
            let g_idx = batch * N_AUX_HORIZONS + k;
            let h_c = h_aux[batch * AUX_HIDDEN + c];
            expected[0] += g_pl[g_idx] * h_c;
            expected[1] += g_sl[g_idx] * h_c;
            expected[2] += g_ps[g_idx] * h_c;
            expected[3] += g_ss[g_idx] * h_c;
        }
        for (name, gv, ev) in [
            ("grad_w_prof_long",  got[0], expected[0]),
            ("grad_w_size_long",  got[1], expected[1]),
            ("grad_w_prof_short", got[2], expected[2]),
            ("grad_w_size_short", got[3], expected[3]),
        ] {
            let err = (gv - ev).abs();
            let tol = (5e-4_f32 * ev.abs().max(1e-3)).max(5e-5);
            assert!(
                err < tol,
                "{name}[{k},{c}] mismatch: got={gv} expected={ev} (err={err})",
            );
        }
    }

    // grad_h_aux is per-batch — closed-form per (batch, c). The
    // expected value sums across all FOUR heads + N_AUX_HORIZONS.
    for &batch in &[0_usize, b_sz - 1] {
        for &c in &[0_usize, 17, AUX_HIDDEN - 1] {
            let mut expected = 0.0_f32;
            for k in 0..N_AUX_HORIZONS {
                let g_idx = batch * N_AUX_HORIZONS + k;
                expected += g_pl[g_idx] * w.w_prof_long [k * AUX_HIDDEN + c];
                expected += g_sl[g_idx] * w.w_size_long [k * AUX_HIDDEN + c];
                expected += g_ps[g_idx] * w.w_prof_short[k * AUX_HIDDEN + c];
                expected += g_ss[g_idx] * w.w_size_short[k * AUX_HIDDEN + c];
            }
            let got = gha[batch * AUX_HIDDEN + c];
            let err = (got - expected).abs();
            let tol = (5e-4_f32 * expected.abs().max(1e-3)).max(5e-4);
            assert!(
                err < tol,
                "grad_h_aux[{batch},{c}] mismatch: got={got} expected={expected} (err={err})",
            );
        }
    }

    Ok(())
}

/// Class-weighted BCE forward + backward correctness on a structured
/// target tensor that exercises BOTH y=0 and y=1 paths and verifies the
/// `pos_weight` parameter scales the positive-class gradient correctly.
///
/// Validates:
///   * Sigmoid+BCE math matches naive CPU oracle within 5e-4.
///   * `dL/dz` matches `pos_weight*(p-1)` for y=1 and `p` for y=0.
///   * `pos_weight=10` produces positive-class gradient scaled by ~10×
///     relative to `pos_weight=1` for the same logit/label.
#[test]
#[ignore = "requires CUDA"]
fn aux_bce_loss_matches_naive_reference() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let bce = AuxBceLoss::new(&dev)?;

    let b_sz: usize = 8;
    let n_total = b_sz * N_AUX_HORIZONS;
    // Mix of positive (y=1) and negative (y=0) labels across horizons
    // so the per-horizon `pos_weight` routing exercises every slot.
    let y_prof: Vec<f32> = (0..n_total)
        .map(|i| if (i % 3) != 0 { 1.0 } else { 0.0 })
        .collect();
    // Spread logits across the saturation regions so we hit the clamp
    // (p → 0 or 1) and the linear regime.
    let logits: Vec<f32> = (0..n_total)
        .map(|i| (i as f32 - (n_total as f32 / 2.0)) * 0.3)
        .collect();
    // Per-horizon positive-class weights — distinct values catch a
    // mis-indexing bug where the kernel reads the wrong horizon's weight.
    let pos_weight: [f32; N_AUX_HORIZONS] = [1.0, 5.0, 12.0];

    let logits_d = upload(&stream, &logits)?;
    let y_prof_d = upload(&stream, &y_prof)?;
    let pw_d = upload(&stream, &pos_weight)?;
    let mut loss_d = stream.alloc_zeros::<f32>(1)?;
    let mut grad_d = stream.alloc_zeros::<f32>(n_total)?;
    let mut valid_d = upload_i32(&stream, &[0_i32])?;

    aux_bce_loss_gpu(
        &stream,
        &bce.fn_,
        &logits_d,
        &y_prof_d,
        &pw_d,
        n_total as i32,
        &mut loss_d,
        &mut grad_d,
        &mut valid_d,
    )?;
    stream.synchronize()?;

    let loss_v = download(&stream, &loss_d)?[0];
    let grad_v = download(&stream, &grad_d)?;
    let valid_v = download_i32(&stream, &valid_d)?[0];
    assert_eq!(valid_v as usize, n_total, "valid count mismatch");

    let mut expected_loss = 0.0_f32;
    let mut saw_pos = false;
    let mut saw_neg = false;
    for i in 0..n_total {
        let z = logits[i];
        let y = y_prof[i];
        let h = i % N_AUX_HORIZONS;
        let pw = pos_weight[h];
        let p = 1.0_f32 / (1.0 + (-z).exp());
        let p_c = p.max(1e-7);
        let one_m_p = (1.0 - p).max(1e-7);
        let (l_e, g_e) = if y > 0.5 {
            saw_pos = true;
            (-pw * p_c.ln(), pw * (p - 1.0))
        } else {
            saw_neg = true;
            (-one_m_p.ln(), p)
        };
        expected_loss += l_e;
        let err = (grad_v[i] - g_e).abs();
        let tol = (5e-5_f32 * g_e.abs().max(1.0)).max(5e-5);
        assert!(
            err < tol,
            "grad_prof_logit[{i}] = {} vs expected {} (z={}, y={}, pw={})",
            grad_v[i], g_e, z, y, pw,
        );
    }
    assert!(saw_pos, "test fixture missed any y=1 sample");
    assert!(saw_neg, "test fixture missed any y=0 sample");
    let err = (loss_v - expected_loss).abs();
    assert!(
        err < 5e-4,
        "total BCE loss mismatch: got={loss_v} expected={expected_loss}",
    );

    Ok(())
}

/// `pos_weight = 10` scales the positive-class gradient by exactly 10×
/// relative to `pos_weight = 1` for the same logit/label. This is the
/// load-bearing class-balance signal — if the kernel ignores the
/// pos_weight buffer, the test catches it immediately.
#[test]
#[ignore = "requires CUDA"]
fn aux_bce_pos_weight_scales_positive_gradient() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let bce = AuxBceLoss::new(&dev)?;

    // Single batch, all-positive labels at the same logit so the gradient
    // is purely the positive-class term `pw * (p - 1)`. Three horizons
    // with pos_weights [1, 10, 1] — only horizon 1's gradient should be
    // 10× the horizon 0 gradient.
    let b_sz: usize = 1;
    let n_total = b_sz * N_AUX_HORIZONS;
    let logits: Vec<f32> = vec![0.5; n_total]; // identical logit per horizon
    let y_prof: Vec<f32> = vec![1.0; n_total]; // all positives
    let pos_weight: [f32; N_AUX_HORIZONS] = [1.0, 10.0, 1.0];

    let logits_d = upload(&stream, &logits)?;
    let y_prof_d = upload(&stream, &y_prof)?;
    let pw_d = upload(&stream, &pos_weight)?;
    let mut loss_d = stream.alloc_zeros::<f32>(1)?;
    let mut grad_d = stream.alloc_zeros::<f32>(n_total)?;
    let mut valid_d = upload_i32(&stream, &[0_i32])?;

    aux_bce_loss_gpu(
        &stream,
        &bce.fn_,
        &logits_d,
        &y_prof_d,
        &pw_d,
        n_total as i32,
        &mut loss_d,
        &mut grad_d,
        &mut valid_d,
    )?;
    stream.synchronize()?;

    let g = download(&stream, &grad_d)?;
    // grad_h0 with pos_weight=1, grad_h1 with pos_weight=10, grad_h2 with pos_weight=1.
    // Expected ratio grad_h1 / grad_h0 == 10.
    let ratio = g[1] / g[0];
    assert!(
        (ratio - 10.0).abs() < 1e-4,
        "pos_weight=10 expected to scale gradient by 10× — got ratio={ratio} \
         (g[0]={}, g[1]={}, g[2]={})",
        g[0], g[1], g[2],
    );
    // Sanity: h0 and h2 share pos_weight=1 → identical gradient at
    // identical (logit, label).
    assert!(
        (g[0] - g[2]).abs() < 1e-6,
        "same pos_weight should give same gradient: g[0]={}, g[2]={}",
        g[0], g[2],
    );
    Ok(())
}

/// Conditional-Huber: `y_size = NaN` at masked slots must produce ZERO
/// gradient and contribute nothing to the loss accumulator. This is the
/// load-bearing semantic — without the mask a NaN label would poison
/// the trunk gradient and the optimiser update would explode to NaN.
#[test]
#[ignore = "requires CUDA"]
fn aux_huber_masked_does_not_propagate() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let huber = AuxMaskedHuberLoss::new(&dev)?;

    let n: usize = 12;
    let y_hat: Vec<f32> = (0..n).map(|i| (i as f32) * 0.1 - 0.5).collect();
    // Half the labels are NaN-masked (simulating `y_size = NaN @
    // y_prof = 0`); the other half are 0.0. The kernel must skip the
    // masked slots and write a finite loss covering only the unmasked.
    let y_size_true: Vec<f32> = (0..n)
        .map(|i| if i % 2 == 0 { f32::NAN } else { 0.0 })
        .collect();

    let y_hat_d = upload(&stream, &y_hat)?;
    let y_true_d = upload(&stream, &y_size_true)?;
    let mut loss_d = stream.alloc_zeros::<f32>(1)?;
    let mut grad_d = stream.alloc_zeros::<f32>(n)?;
    // Pre-seed with sentinels so we can verify the kernel actually wrote.
    let mut valid_d = upload_i32(&stream, &[-1_i32])?;

    aux_huber_masked_loss_gpu(
        &stream,
        &huber.fn_,
        &y_hat_d,
        &y_true_d,
        DEFAULT_HUBER_DELTA,
        n as i32,
        &mut loss_d,
        &mut grad_d,
        &mut valid_d,
    )?;
    stream.synchronize()?;

    let loss_v = download(&stream, &loss_d)?[0];
    let grad_v = download(&stream, &grad_d)?;
    let valid_v = download_i32(&stream, &valid_d)?[0];

    assert_eq!(valid_v, (n / 2) as i32, "NaN-masked valid count wrong");
    assert!(loss_v.is_finite(), "loss is NaN/Inf");

    let mut expected_loss = 0.0_f32;
    let mut hit_quadratic = false;
    let mut hit_linear = false;
    for i in 0..n {
        if !y_size_true[i].is_finite() {
            // NaN label: grad must be exactly 0, no loss contribution.
            assert_eq!(
                grad_v[i], 0.0,
                "NaN-masked grad_y_size_hat[{i}] = {} (must be 0)",
                grad_v[i],
            );
        } else {
            let r = y_hat[i] - y_size_true[i];
            let ar = r.abs();
            let (l, g) = if ar <= DEFAULT_HUBER_DELTA {
                hit_quadratic = true;
                (0.5 * r * r, r)
            } else {
                hit_linear = true;
                (
                    DEFAULT_HUBER_DELTA * (ar - 0.5 * DEFAULT_HUBER_DELTA),
                    DEFAULT_HUBER_DELTA * r.signum(),
                )
            };
            expected_loss += l;
            let err = (grad_v[i] - g).abs();
            assert!(
                err < 5e-5,
                "grad_y_size_hat[{i}] = {} vs expected {} (r={}, |r|={})",
                grad_v[i], g, r, ar,
            );
        }
    }
    // Confirm the test fixture exercised both branches at least once.
    assert!(hit_quadratic, "test fixture missed the quadratic branch");
    // Linear branch may not fire on this small fixture; widen y_hat
    // range to also test it.
    let _ = hit_linear; // signal-only; explicit linear-branch test below.

    let err = (loss_v - expected_loss).abs();
    assert!(
        err < 5e-4,
        "NaN-masked Huber loss mismatch: got={loss_v} expected={expected_loss}",
    );

    Ok(())
}

/// Huber math correctness across BOTH branches (|r| ≤ δ quadratic and
/// |r| > δ linear). Mirrors the original B4 Huber test but for the new
/// `aux_huber_masked_fwd_bwd` kernel signature.
#[test]
#[ignore = "requires CUDA"]
fn aux_huber_masked_covers_both_branches() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let huber = AuxMaskedHuberLoss::new(&dev)?;

    // Construct y_hat and y_true so that some residuals land in each
    // branch. With δ = 1.0: residuals in [-0.5, 0.5] are quadratic,
    // residuals in [-2.4, -1.0] ∪ [1.0, 2.2] are linear.
    let n: usize = 24;
    let y_hat: Vec<f32> = (0..n)
        .map(|i| (i as f32 - 12.0) * 0.2) // -2.4 .. 2.2 step 0.2
        .collect();
    let y_true: Vec<f32> = vec![0.0; n];

    let delta = DEFAULT_HUBER_DELTA;
    let y_hat_d = upload(&stream, &y_hat)?;
    let y_true_d = upload(&stream, &y_true)?;
    let mut loss_d = stream.alloc_zeros::<f32>(1)?;
    let mut grad_d = stream.alloc_zeros::<f32>(n)?;
    let mut valid_d = upload_i32(&stream, &[0_i32])?;

    aux_huber_masked_loss_gpu(
        &stream,
        &huber.fn_,
        &y_hat_d,
        &y_true_d,
        delta,
        n as i32,
        &mut loss_d,
        &mut grad_d,
        &mut valid_d,
    )?;
    stream.synchronize()?;

    let loss_v = download(&stream, &loss_d)?[0];
    let grad_v = download(&stream, &grad_d)?;
    let valid_v = download_i32(&stream, &valid_d)?[0];
    assert_eq!(valid_v as usize, n, "valid count mismatch");

    let mut expected_loss = 0.0_f32;
    let mut hit_quadratic = false;
    let mut hit_linear = false;
    for i in 0..n {
        let r = y_hat[i] - y_true[i];
        let ar = r.abs();
        let (l, g) = if ar <= delta {
            hit_quadratic = true;
            (0.5 * r * r, r)
        } else {
            hit_linear = true;
            (delta * (ar - 0.5 * delta), delta * r.signum())
        };
        expected_loss += l;
        let err = (grad_v[i] - g).abs();
        assert!(
            err < 5e-5,
            "grad_y_size_hat[{i}] = {} vs expected {} (r={}, |r|={})",
            grad_v[i], g, r, ar,
        );
    }
    assert!(hit_quadratic, "test fixture missed the quadratic branch");
    assert!(hit_linear, "test fixture missed the linear branch");
    let err = (loss_v - expected_loss).abs();
    assert!(
        err < 5e-4,
        "Huber total loss mismatch: got={loss_v} expected={expected_loss}",
    );

    Ok(())
}

/// NaN-masked BCE: `y_prof = NaN` must produce ZERO gradient and skip
/// the loss accumulator — same defence as the Huber mask path.
#[test]
#[ignore = "requires CUDA"]
fn aux_bce_nan_mask_does_not_propagate() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let bce = AuxBceLoss::new(&dev)?;

    let b_sz: usize = 4;
    let n_total = b_sz * N_AUX_HORIZONS;
    let logits: Vec<f32> = (0..n_total).map(|i| (i as f32) * 0.05 - 0.2).collect();
    // Mark even-indexed slots as NaN, odd as 0/1 alternating.
    let y_prof: Vec<f32> = (0..n_total)
        .map(|i| {
            if i % 2 == 0 {
                f32::NAN
            } else if (i / 2) % 2 == 0 {
                0.0
            } else {
                1.0
            }
        })
        .collect();
    let pos_weight: [f32; N_AUX_HORIZONS] = [3.0, 3.0, 3.0];

    let logits_d = upload(&stream, &logits)?;
    let y_prof_d = upload(&stream, &y_prof)?;
    let pw_d = upload(&stream, &pos_weight)?;
    let mut loss_d = stream.alloc_zeros::<f32>(1)?;
    let mut grad_d = stream.alloc_zeros::<f32>(n_total)?;
    let mut valid_d = upload_i32(&stream, &[-1_i32])?;

    aux_bce_loss_gpu(
        &stream,
        &bce.fn_,
        &logits_d,
        &y_prof_d,
        &pw_d,
        n_total as i32,
        &mut loss_d,
        &mut grad_d,
        &mut valid_d,
    )?;
    stream.synchronize()?;

    let loss_v = download(&stream, &loss_d)?[0];
    let grad_v = download(&stream, &grad_d)?;
    let valid_v = download_i32(&stream, &valid_d)?[0];

    let expected_valid = y_prof.iter().filter(|y| y.is_finite()).count() as i32;
    assert_eq!(valid_v, expected_valid, "BCE NaN-masked valid count wrong");
    assert!(loss_v.is_finite(), "BCE loss is NaN/Inf");
    for i in 0..n_total {
        if !y_prof[i].is_finite() {
            assert_eq!(
                grad_v[i], 0.0,
                "NaN-masked grad_prof_logit[{i}] = {} (must be 0)",
                grad_v[i],
            );
        } else {
            assert!(grad_v[i].is_finite(), "non-masked grad NaN at i={i}");
        }
    }
    Ok(())
}
