//! GPU oracle tests for bucket_transition_kernels.cu.
//! Run with: SQLX_OFFLINE=true cargo test -p ml-alpha --test bucket_transition_kernels -- --ignored --nocapture

use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use ml_core::device::MlDevice;

const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bucket_transition_kernels.cubin"));
const HIDDEN_DIM: usize = 128;
const N_HORIZONS: usize = 5;

fn load_kernel(stream: &std::sync::Arc<CudaStream>, fn_name: &str) -> Result<cudarc::driver::CudaFunction> {
    let ctx = stream.context();
    let module = ctx.load_cubin(KERNEL_CUBIN.to_vec()).context("load bucket_transition cubin")?;
    module.load_function(fn_name).with_context(|| format!("load {fn_name}"))
}

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

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

#[test]
#[ignore = "requires CUDA"]
fn tau_sort_kernel_sorts_tau_ascending() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "tau_sort_kernel")?;

    // Reverse-sorted input: [128, 127, ..., 1] of tau
    let tau_host: Vec<f32> = (1..=HIDDEN_DIM).rev().map(|x| x as f32).collect();
    let tau_d = upload(&stream, &tau_host)?;
    let mut sorted_indices_d = stream.alloc_zeros::<u32>(HIDDEN_DIM)?;

    let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: HIDDEN_DIM as u32 * 8 };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&tau_d).arg(&mut sorted_indices_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let sorted_indices = download(&stream, &sorted_indices_d)?;
    // After sort ascending, sorted_indices[0] should point to the smallest tau (channel 127), [127] to largest (channel 0)
    assert_eq!(sorted_indices[0], (HIDDEN_DIM - 1) as u32, "smallest tau should be at channel HIDDEN_DIM-1");
    assert_eq!(sorted_indices[HIDDEN_DIM - 1], 0, "largest tau should be at channel 0");
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn bucket_assign_kernel_assigns_quintiles() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "bucket_assign_kernel")?;

    // sorted_indices[p] = p (identity permutation; channel p has rank p)
    let sorted_indices_host: Vec<u32> = (0..HIDDEN_DIM as u32).collect();
    let sorted_indices_d = upload(&stream, &sorted_indices_host)?;
    let mut bucket_id_d = stream.alloc_zeros::<u8>(HIDDEN_DIM)?;

    let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (HIDDEN_DIM as u32,1,1), shared_mem_bytes: 0 };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&sorted_indices_d).arg(&mut bucket_id_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let bucket_id = download(&stream, &bucket_id_d)?;
    // bucket boundaries: [0,25), [25,50), [50,75), [75,100), [100,128)
    assert_eq!(bucket_id[0], 0);
    assert_eq!(bucket_id[24], 0);
    assert_eq!(bucket_id[25], 1);
    assert_eq!(bucket_id[49], 1);
    assert_eq!(bucket_id[50], 2);
    assert_eq!(bucket_id[74], 2);
    assert_eq!(bucket_id[75], 3);
    assert_eq!(bucket_id[99], 3);
    assert_eq!(bucket_id[100], 4);
    assert_eq!(bucket_id[127], 4);
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn bucket_iqr_kernel_computes_q1_q3_per_bucket() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "bucket_iqr_kernel")?;

    // tau values [1.0, 2.0, ..., 128.0] in sorted order (already sorted ascending).
    // sorted_indices[p] = p; bucket k holds taus from index k * 25 to (k+1)*25 (or 28 for k=4).
    let tau_host: Vec<f32> = (1..=HIDDEN_DIM).map(|x| x as f32).collect();
    let sorted_indices_host: Vec<u32> = (0..HIDDEN_DIM as u32).collect();
    let tau_d = upload(&stream, &tau_host)?;
    let sorted_indices_d = upload(&stream, &sorted_indices_host)?;
    let mut q1_d = stream.alloc_zeros::<f32>(N_HORIZONS)?;
    let mut q3_d = stream.alloc_zeros::<f32>(N_HORIZONS)?;

    let cfg = LaunchConfig { grid_dim: (N_HORIZONS as u32,1,1), block_dim: (32,1,1), shared_mem_bytes: 32 * 4 };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&tau_d).arg(&sorted_indices_d).arg(&mut q1_d).arg(&mut q3_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let q1 = download(&stream, &q1_d)?;
    let q3 = download(&stream, &q3_d)?;
    // Bucket 0: tau = [1..25], Q1 ≈ 7, Q3 ≈ 19 (within 1 element of true quartile)
    assert!((q1[0] - 7.0).abs() <= 1.0, "q1[0]={}", q1[0]);
    assert!((q3[0] - 19.0).abs() <= 1.0, "q3[0]={}", q3[0]);
    // Bucket 4: tau = [101..128], Q1 ≈ 108, Q3 ≈ 121
    assert!((q1[4] - 108.0).abs() <= 1.0, "q1[4]={}", q1[4]);
    assert!((q3[4] - 121.0).abs() <= 1.0, "q3[4]={}", q3[4]);
    Ok(())
}

// ALPHA fix 2026-05-21: `tau_reorder_kernel` removed (it caused systematic
// indexing mismatch between the reordered τ layout and the downstream
// kernels' original-channel reads). Replaced by `channels_in_bucket_kernel`
// + `zero_off_bucket_kernel`; tests below.

#[test]
#[ignore = "requires CUDA"]
fn channels_in_bucket_kernel_populates_lookup_in_original_channel_order() -> Result<()> {
    // For each bucket k, channels_in_bucket[k * MAX_BUCKET_DIM .. k * MAX_BUCKET_DIM + bdim[k]]
    // should hold the original channel indices in ASCENDING order (kernel walks
    // channels in c = 0..HIDDEN_DIM and appends to per-bucket cursors). Slots
    // beyond `bucket_dim_k[k]` hold the sentinel 0xFFFFFFFF.
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "channels_in_bucket_kernel")?;

    const MAX_BUCKET_DIM: usize = 28;
    // Quintile assignment: channel c → bucket c/25 capped at 4.
    let bucket_id_host: Vec<u8> = (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let bucket_dim_k_host: Vec<u32> = vec![25, 25, 25, 25, 28];
    let bucket_id_d = upload(&stream, &bucket_id_host)?;
    let bucket_dim_k_d = upload(&stream, &bucket_dim_k_host)?;
    let mut channels_in_bucket_d =
        stream.alloc_zeros::<u32>(N_HORIZONS * MAX_BUCKET_DIM)?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&bucket_id_d)
        .arg(&bucket_dim_k_d)
        .arg(&mut channels_in_bucket_d);
    unsafe {
        launch.launch(cfg)?;
    }
    stream.synchronize()?;

    let channels_in_bucket = download(&stream, &channels_in_bucket_d)?;
    // Bucket 0 = channels [0..25), bucket 1 = [25..50), … bucket 4 = [100..128).
    for k in 0..N_HORIZONS {
        let bdim = bucket_dim_k_host[k] as usize;
        let bucket_start = k * 25;
        for i in 0..bdim {
            let expected_c = bucket_start + i;
            let got = channels_in_bucket[k * MAX_BUCKET_DIM + i];
            assert_eq!(
                got, expected_c as u32,
                "channels_in_bucket[{}][{}]={} expected {}",
                k, i, got, expected_c
            );
        }
        // Sentinel positions: slots [bdim, MAX_BUCKET_DIM) should be 0xFFFFFFFF.
        for i in bdim..MAX_BUCKET_DIM {
            let got = channels_in_bucket[k * MAX_BUCKET_DIM + i];
            assert_eq!(
                got, 0xFFFFFFFFu32,
                "channels_in_bucket[{}][{}]={} expected sentinel (0xFFFFFFFF)",
                k, i, got
            );
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn channels_in_bucket_kernel_handles_non_contiguous_assignment() -> Result<()> {
    // Non-trivial bucket assignment: round-robin (c % N_HORIZONS). Every
    // bucket gets a non-contiguous set of channels. The kernel must still
    // produce the correct lookup table (channels listed in ascending
    // original-channel-index order within each bucket).
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "channels_in_bucket_kernel")?;

    const MAX_BUCKET_DIM: usize = 28;
    let bucket_id_host: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| (c % N_HORIZONS) as u8).collect();
    // With round-robin assignment over 128 channels and 5 buckets:
    //   bucket 0: c = 0, 5, 10, ..., 125 → 26 channels
    //   bucket 1: c = 1, 6, ..., 126 → 26 channels
    //   bucket 2: c = 2, 7, ..., 127 → 26 channels
    //   bucket 3: c = 3, 8, ..., 123 → 25 channels
    //   bucket 4: c = 4, 9, ..., 124 → 25 channels
    let mut bucket_counts = [0u32; N_HORIZONS];
    for c in 0..HIDDEN_DIM {
        bucket_counts[c % N_HORIZONS] += 1;
    }
    let bucket_dim_k_host: Vec<u32> = bucket_counts.to_vec();
    let bucket_id_d = upload(&stream, &bucket_id_host)?;
    let bucket_dim_k_d = upload(&stream, &bucket_dim_k_host)?;
    let mut channels_in_bucket_d =
        stream.alloc_zeros::<u32>(N_HORIZONS * MAX_BUCKET_DIM)?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&bucket_id_d)
        .arg(&bucket_dim_k_d)
        .arg(&mut channels_in_bucket_d);
    unsafe {
        launch.launch(cfg)?;
    }
    stream.synchronize()?;

    let channels_in_bucket = download(&stream, &channels_in_bucket_d)?;
    // Host reference: same walk.
    let mut expected: Vec<Vec<u32>> = vec![Vec::new(); N_HORIZONS];
    for c in 0..HIDDEN_DIM {
        let k = c % N_HORIZONS;
        expected[k].push(c as u32);
    }
    for k in 0..N_HORIZONS {
        let bdim = bucket_dim_k_host[k] as usize;
        for i in 0..bdim {
            let got = channels_in_bucket[k * MAX_BUCKET_DIM + i];
            assert_eq!(got, expected[k][i], "k={} i={}", k, i);
        }
        for i in bdim..MAX_BUCKET_DIM {
            let got = channels_in_bucket[k * MAX_BUCKET_DIM + i];
            assert_eq!(got, 0xFFFFFFFFu32, "sentinel k={} i={}", k, i);
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn zero_off_bucket_kernel_zeros_off_diagonal_positions() -> Result<()> {
    // Tensor [N_HORIZONS × HIDDEN_DIM] initialised to 1.0. After applying
    // zero_off_bucket, position (h, c) should be:
    //   1.0 if bucket_id_per_channel[c] == h (in-bucket)
    //   0.0 otherwise (off-bucket)
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "zero_off_bucket_kernel")?;

    let bucket_id_host: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let tensor_host: Vec<f32> = vec![1.0; N_HORIZONS * HIDDEN_DIM];

    let bucket_id_d = upload(&stream, &bucket_id_host)?;
    let mut tensor_d = upload(&stream, &tensor_host)?;

    let cfg = LaunchConfig {
        grid_dim: (N_HORIZONS as u32, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&bucket_id_d).arg(&mut tensor_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let tensor_out = download(&stream, &tensor_d)?;
    for h in 0..N_HORIZONS {
        for c in 0..HIDDEN_DIM {
            let idx = h * HIDDEN_DIM + c;
            let in_bucket = (c / 25).min(4) == h;
            let expected = if in_bucket { 1.0 } else { 0.0 };
            assert!(
                (tensor_out[idx] - expected).abs() < 1e-6,
                "h={} c={} got={} expected={}",
                h, c, tensor_out[idx], expected
            );
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn zero_off_bucket_kernel_keeps_off_bucket_zero_after_many_mock_adam_steps() -> Result<()> {
    // ALPHA fix invariant test: simulate the post-Adam projection wired in
    // perception.rs. Start with heads_w_skip = all-ones at in-bucket positions,
    // zero elsewhere (post-mask_init state). Run N "Adam steps" of the form
    //   heads_w_skip += eps  (uniform tiny push that would breach the
    //                         block-diagonal invariant without projection)
    // and verify that after each post-step zero_off_bucket call, off-bucket
    // entries return to exactly 0.0, while in-bucket entries accumulate the
    // push.
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func_zero = load_kernel(&stream, "zero_off_bucket_kernel")?;
    let func_mask_init = load_kernel(&stream, "heads_w_skip_mask_init_kernel")?;

    let bucket_id_host: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let bucket_id_d = upload(&stream, &bucket_id_host)?;
    let mut heads_w_skip_d =
        upload(&stream, &vec![1.0_f32; N_HORIZONS * HIDDEN_DIM])?;
    let mut mask_d = stream.alloc_zeros::<f32>(N_HORIZONS * HIDDEN_DIM)?;

    // Initialise: zeros off-bucket weights + populates mask.
    let cfg_mask_init = LaunchConfig {
        grid_dim: (N_HORIZONS as u32, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    {
        let mut launch = stream.launch_builder(&func_mask_init);
        launch
            .arg(&bucket_id_d)
            .arg(&mut heads_w_skip_d)
            .arg(&mut mask_d);
        unsafe {
            launch.launch(cfg_mask_init)?;
        }
    }
    stream.synchronize()?;

    // Run N mock steps: add a tiny uniform value, then re-project.
    let n_steps = 50;
    let eps_push: f32 = 1e-3;
    let push_host: Vec<f32> = vec![eps_push; N_HORIZONS * HIDDEN_DIM];
    let push_d = upload(&stream, &push_host)?;

    // Simple add via the existing "use a tiny kernel" pattern would inflate
    // the test scope; instead we read back, add on host, re-upload, and
    // project. This validates the kernel's projection behaviour against
    // arbitrary intermediate states.
    for step in 0..n_steps {
        let mut current = download(&stream, &heads_w_skip_d)?;
        let push_back = download(&stream, &push_d)?;
        for i in 0..current.len() {
            current[i] += push_back[i];
        }
        stream.memcpy_htod(&current, &mut heads_w_skip_d)?;
        // Apply the post-Adam projection.
        let cfg_zero = LaunchConfig {
            grid_dim: (N_HORIZONS as u32, 1, 1),
            block_dim: (HIDDEN_DIM as u32, 1, 1),
            shared_mem_bytes: 0,
        };
        let mut launch = stream.launch_builder(&func_zero);
        launch.arg(&bucket_id_d).arg(&mut heads_w_skip_d);
        unsafe {
            launch.launch(cfg_zero)?;
        }
        stream.synchronize()?;
        let post = download(&stream, &heads_w_skip_d)?;
        for h in 0..N_HORIZONS {
            for c in 0..HIDDEN_DIM {
                let idx = h * HIDDEN_DIM + c;
                let in_bucket = (c / 25).min(4) == h;
                if !in_bucket {
                    assert_eq!(
                        post[idx], 0.0,
                        "off-bucket [{}, {}] = {} after step {} — projection failed",
                        h, c, post[idx], step
                    );
                } else {
                    // In-bucket positions accumulate the push:
                    //   starts at 1.0; after (step + 1) pushes, value = 1.0 + (step+1) * eps_push.
                    let expected = 1.0_f32 + ((step + 1) as f32) * eps_push;
                    assert!(
                        (post[idx] - expected).abs() < 1e-5,
                        "in-bucket [{}, {}] = {} after step {} expected {}",
                        h, c, post[idx], step, expected
                    );
                }
            }
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn heads_compact_kernel_reorders_w_skip() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "heads_compact_kernel")?;

    // Original heads_w_skip is [N_HORIZONS × HIDDEN_DIM] = 640 floats.
    // For test: w_skip[h, c] = h * 1000 + c.
    let w_skip_orig_host: Vec<f32> = (0..N_HORIZONS).flat_map(|h| {
        (0..HIDDEN_DIM).map(move |c| (h * 1000 + c) as f32)
    }).collect();
    let bucket_id_host: Vec<u8> = (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let bucket_offset_host: Vec<u32> = vec![0, 25, 50, 75, 100, 128];
    let w_skip_orig_d = upload(&stream, &w_skip_orig_host)?;
    let bucket_id_d = upload(&stream, &bucket_id_host)?;
    let bucket_offset_d = upload(&stream, &bucket_offset_host)?;
    let mut w_skip_compact_d = stream.alloc_zeros::<f32>(HIDDEN_DIM)?;

    let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (HIDDEN_DIM as u32,1,1), shared_mem_bytes: (N_HORIZONS as u32 + 1) * 4 };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&w_skip_orig_d).arg(&bucket_id_d).arg(&bucket_offset_d).arg(&mut w_skip_compact_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let w_skip_compact = download(&stream, &w_skip_compact_d)?;
    // For each channel c in bucket k, compact[c] = orig[k * HIDDEN_DIM + c]
    for c in 0..HIDDEN_DIM {
        let k = (c / 25).min(4);
        let expected = (k * 1000 + c) as f32;
        assert!((w_skip_compact[c] - expected).abs() < 1e-6, "compact[{}]={} expected {}", c, w_skip_compact[c], expected);
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn tau_change_frobenius_kernel_sums_squared_diff() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "tau_change_frobenius_kernel")?;

    let tau_t: Vec<f32> = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.1).collect();
    let tau_prev: Vec<f32> = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.1 + 0.5).collect();
    let tau_t_d = upload(&stream, &tau_t)?;
    let tau_prev_d = upload(&stream, &tau_prev)?;
    let mut out_d = stream.alloc_zeros::<f32>(1)?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: HIDDEN_DIM as u32 * 4,
    };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&tau_t_d).arg(&tau_prev_d).arg(&mut out_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let out = download(&stream, &out_d)?;
    // Each diff = -0.5; squared = 0.25; sum over 128 channels = 32.0
    assert!((out[0] - 32.0).abs() < 1e-3, "out={} expected ~32.0", out[0]);
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn tau_clamp_kernel_bounds_tau_within_bucket_iqr() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "tau_clamp_kernel")?;

    // All channels start with τ = 100 + c * 0.1 (way above any bucket's hi).
    // Bucket assignment: channel c → bucket c/25 capped at 4.
    // IQR bounds: bucket k has [k*10, k*10 + 5].
    // Expected outcome: every channel clamps to its bucket's hi value.
    let tau_host: Vec<f32> =
        (0..HIDDEN_DIM).map(|c| 100.0 + (c as f32) * 0.1).collect();
    let bucket_id_host: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let iqr_lo_host: Vec<f32> =
        (0..N_HORIZONS).map(|k| (k as f32) * 10.0).collect();
    let iqr_hi_host: Vec<f32> =
        (0..N_HORIZONS).map(|k| (k as f32) * 10.0 + 5.0).collect();

    let mut tau_d = upload(&stream, &tau_host)?;
    let bucket_id_d = upload(&stream, &bucket_id_host)?;
    let iqr_lo_d = upload(&stream, &iqr_lo_host)?;
    let iqr_hi_d = upload(&stream, &iqr_hi_host)?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&mut tau_d)
        .arg(&bucket_id_d)
        .arg(&iqr_lo_d)
        .arg(&iqr_hi_d);
    unsafe {
        launch.launch(cfg)?;
    }
    stream.synchronize()?;

    let tau_out = download(&stream, &tau_d)?;
    for c in 0..HIDDEN_DIM {
        let k = (c / 25).min(4);
        let expected = (k as f32) * 10.0 + 5.0; // upper bound of bucket
        assert!(
            (tau_out[c] - expected).abs() < 1e-6,
            "c={} k={} got={} expected={}",
            c,
            k,
            tau_out[c],
            expected
        );
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn heads_lr_multiplier_scale_kernel_rescales_delta_per_horizon() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "heads_lr_multiplier_scale_kernel")?;

    // Per-horizon stride 8: param tensor shape [N_HORIZONS * 8] = 40 floats.
    let per_horizon_stride: i32 = 8;
    let total: i32 = N_HORIZONS as i32 * per_horizon_stride;
    let param_pre: Vec<f32> =
        (0..total).map(|i| (i as f32) * 0.1).collect();
    // Uniform delta of 1.0 applied "by Adam" — simulated post-Adam value.
    let param_post: Vec<f32> = param_pre.iter().map(|x| x + 1.0).collect();
    // Distinct multipliers per horizon, including 0 (full attenuation) and 1
    // (no change). After scaling: delta'[k] = lr_mult[k] = 0, 0.25, 0.5, 0.75, 1.0.
    let lr_mult: Vec<f32> = vec![0.0, 0.25, 0.5, 0.75, 1.0];

    let mut param_d = upload(&stream, &param_post)?;
    let param_pre_d = upload(&stream, &param_pre)?;
    let lr_mult_d = upload(&stream, &lr_mult)?;

    let block: u32 = 64;
    let grid: u32 = ((total as u32) + block - 1) / block;
    let cfg = LaunchConfig {
        grid_dim: (grid, 1, 1),
        block_dim: (block, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&mut param_d)
        .arg(&param_pre_d)
        .arg(&lr_mult_d)
        .arg(&per_horizon_stride)
        .arg(&total);
    unsafe {
        launch.launch(cfg)?;
    }
    stream.synchronize()?;

    let param_out = download(&stream, &param_d)?;
    for i in 0..total as usize {
        let k = i / per_horizon_stride as usize;
        let expected = param_pre[i] + lr_mult[k];
        assert!(
            (param_out[i] - expected).abs() < 1e-6,
            "i={} k={} got={} expected={}",
            i,
            k,
            param_out[i],
            expected
        );
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn h_mag_per_bucket_kernel_computes_mean_abs_per_bucket() -> Result<()> {
    // ALPHA fix 2026-05-21: kernel now reads `channels_in_bucket` lookup
    // (bucket → original-channel) instead of `bucket_channel_offset` —
    // h_state stays in ORIGINAL channel layout post-ALPHA.
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "h_mag_per_bucket_kernel")?;
    let func_lookup = load_kernel(&stream, "channels_in_bucket_kernel")?;

    const MAX_BUCKET_DIM: usize = 28;
    // Synthetic h_state: every channel = (channel_index + 1) * 0.1 across
    // B=4 samples (same value per sample). Sign is intentionally negated
    // for odd channels to confirm the kernel takes absolute value.
    let b_sz: usize = 4;
    let bucket_dim_k_host: Vec<u32> = vec![25, 25, 25, 25, 28];
    let bucket_id_host: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let h_state_host: Vec<f32> = (0..b_sz * HIDDEN_DIM)
        .map(|i| {
            let c = i % HIDDEN_DIM;
            let val = (c as f32 + 1.0) * 0.1;
            if c % 2 == 1 { -val } else { val }
        })
        .collect();

    let h_state_d = upload(&stream, &h_state_host)?;
    let bucket_id_d = upload(&stream, &bucket_id_host)?;
    let bucket_dim_k_d = upload(&stream, &bucket_dim_k_host)?;
    let mut channels_in_bucket_d =
        stream.alloc_zeros::<u32>(N_HORIZONS * MAX_BUCKET_DIM)?;
    let mut h_mag_d = stream.alloc_zeros::<f32>(N_HORIZONS)?;

    // Build the bucket → original-channel lookup using the same kernel
    // the trainer uses at transition.
    {
        let cfg = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (HIDDEN_DIM as u32, 1, 1),
            shared_mem_bytes: 0,
        };
        let mut launch = stream.launch_builder(&func_lookup);
        launch
            .arg(&bucket_id_d)
            .arg(&bucket_dim_k_d)
            .arg(&mut channels_in_bucket_d);
        unsafe { launch.launch(cfg)?; }
    }

    let b_i32: i32 = b_sz as i32;
    let cfg = LaunchConfig {
        grid_dim: (N_HORIZONS as u32, 1, 1),
        block_dim: (32, 1, 1),
        shared_mem_bytes: 32 * 4,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&h_state_d)
        .arg(&channels_in_bucket_d)
        .arg(&bucket_dim_k_d)
        .arg(&b_i32)
        .arg(&mut h_mag_d);
    unsafe {
        launch.launch(cfg)?;
    }
    stream.synchronize()?;

    let h_mag = download(&stream, &h_mag_d)?;
    // Host reference: same reduction using the quintile assignment.
    for k in 0..N_HORIZONS {
        let start = k * 25;
        let end = if k == 4 { HIDDEN_DIM } else { (k + 1) * 25 };
        let bdim = end - start;
        let mut sum_abs = 0.0_f32;
        for _b in 0..b_sz {
            for c in start..end {
                let val = (c as f32 + 1.0) * 0.1;
                sum_abs += val; // signs cancel under fabsf
            }
        }
        let expected = sum_abs / (bdim as f32 * b_sz as f32);
        assert!(
            (h_mag[k] - expected).abs() < 1e-4,
            "h_mag[{}]={} expected {}",
            k,
            h_mag[k],
            expected
        );
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn bucket_iqr_double_widening_kernel_doubles_one_bucket_iqr() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "bucket_iqr_double_widening_kernel")?;

    // Pre-widen all 5 buckets to known values; widening kernel should only
    // affect the targeted bucket index, others untouched.
    let iqr_lo_host: Vec<f32> = vec![1.0, 2.0, 4.0, 8.0, 16.0];
    let iqr_hi_host: Vec<f32> = vec![10.0, 20.0, 40.0, 80.0, 160.0];
    let mut iqr_lo_d = upload(&stream, &iqr_lo_host)?;
    let mut iqr_hi_d = upload(&stream, &iqr_hi_host)?;

    let target_bucket: i32 = 2;
    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (1, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&mut iqr_lo_d)
        .arg(&mut iqr_hi_d)
        .arg(&target_bucket);
    unsafe {
        launch.launch(cfg)?;
    }
    stream.synchronize()?;

    let lo = download(&stream, &iqr_lo_d)?;
    let hi = download(&stream, &iqr_hi_d)?;
    for k in 0..N_HORIZONS {
        if k as i32 == target_bucket {
            // bucket 2 was {4, 40}; expect halved+doubled = {2, 80}.
            assert!(
                (lo[k] - iqr_lo_host[k] * 0.5).abs() < 1e-6,
                "lo[{}]={} expected {}",
                k,
                lo[k],
                iqr_lo_host[k] * 0.5
            );
            assert!(
                (hi[k] - iqr_hi_host[k] * 2.0).abs() < 1e-6,
                "hi[{}]={} expected {}",
                k,
                hi[k],
                iqr_hi_host[k] * 2.0
            );
        } else {
            assert!(
                (lo[k] - iqr_lo_host[k]).abs() < 1e-6,
                "lo[{}] was modified (untouched bucket)",
                k
            );
            assert!(
                (hi[k] - iqr_hi_host[k]).abs() < 1e-6,
                "hi[{}] was modified (untouched bucket)",
                k
            );
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn heads_w_skip_mask_init_zeros_off_bucket() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "heads_w_skip_mask_init_kernel")?;

    // Bucket assignment: channel c → bucket c/25 capped at 4 (quintile
    // boundaries used elsewhere in this file: bucket_dim_k = [25, 25, 25,
    // 25, 28]). heads_w_skip starts all 1.0; mask + zero-out result is
    // 1.0 at in-bucket positions, 0.0 elsewhere.
    let bucket_id_host: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let heads_w_skip_host: Vec<f32> = vec![1.0; N_HORIZONS * HIDDEN_DIM];

    let bucket_id_d = upload(&stream, &bucket_id_host)?;
    let mut heads_w_skip_d = upload(&stream, &heads_w_skip_host)?;
    let mut mask_d = stream.alloc_zeros::<f32>(N_HORIZONS * HIDDEN_DIM)?;

    let cfg = LaunchConfig {
        grid_dim: (N_HORIZONS as u32, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&bucket_id_d)
        .arg(&mut heads_w_skip_d)
        .arg(&mut mask_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let mask = download(&stream, &mask_d)?;
    let heads_w_skip_out = download(&stream, &heads_w_skip_d)?;
    for h in 0..N_HORIZONS {
        for c in 0..HIDDEN_DIM {
            let idx = h * HIDDEN_DIM + c;
            let in_bucket = (c / 25).min(4) == h;
            let expected = if in_bucket { 1.0 } else { 0.0 };
            assert!(
                (mask[idx] - expected).abs() < 1e-6,
                "mask[{},{}]={} expected {}",
                h,
                c,
                mask[idx],
                expected
            );
            assert!(
                (heads_w_skip_out[idx] - expected).abs() < 1e-6,
                "heads_w_skip[{},{}]={} expected {}",
                h,
                c,
                heads_w_skip_out[idx],
                expected
            );
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn heads_w_skip_grad_mask_apply_zeros_off_bucket_grad() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "heads_w_skip_grad_mask_apply_kernel")?;

    // Build a mask manually for the same quintile layout used above; feed
    // an all-ones grad tensor and expect off-bucket entries to be zeroed.
    let mask_host: Vec<f32> = (0..N_HORIZONS)
        .flat_map(|h| {
            (0..HIDDEN_DIM).map(move |c| {
                if (c / 25).min(4) == h { 1.0_f32 } else { 0.0_f32 }
            })
        })
        .collect();
    let grad_host: Vec<f32> = vec![1.0; N_HORIZONS * HIDDEN_DIM];

    let mask_d = upload(&stream, &mask_host)?;
    let mut grad_d = upload(&stream, &grad_host)?;

    let cfg = LaunchConfig {
        grid_dim: (N_HORIZONS as u32, 1, 1),
        block_dim: (HIDDEN_DIM as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&mask_d).arg(&mut grad_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let grad_out = download(&stream, &grad_d)?;
    for h in 0..N_HORIZONS {
        for c in 0..HIDDEN_DIM {
            let idx = h * HIDDEN_DIM + c;
            let in_bucket = (c / 25).min(4) == h;
            let expected = if in_bucket { 1.0 } else { 0.0 };
            assert!(
                (grad_out[idx] - expected).abs() < 1e-6,
                "grad[{},{}]={} expected {}",
                h,
                c,
                grad_out[idx],
                expected
            );
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn slack_factor_apply_kernel_widens_iqr_geometrically() -> Result<()> {
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "slack_factor_apply_kernel")?;

    // Q1 = [1, 4, 9, 16, 25], Q3 = [4, 16, 36, 64, 100]
    // → slack = sqrt(Q3/Q1) = [2, 2, 2, 2, 2]
    // → out: lo = [0.5, 2, 4.5, 8, 12.5], hi = [8, 32, 72, 128, 200]
    let q1: Vec<f32> = vec![1.0, 4.0, 9.0, 16.0, 25.0];
    let q3: Vec<f32> = vec![4.0, 16.0, 36.0, 64.0, 100.0];
    let mut lo_d = upload(&stream, &q1)?;
    let mut hi_d = upload(&stream, &q3)?;

    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (N_HORIZONS as u32, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch.arg(&mut lo_d).arg(&mut hi_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let lo = download(&stream, &lo_d)?;
    let hi = download(&stream, &hi_d)?;
    let expected_lo = [0.5, 2.0, 4.5, 8.0, 12.5];
    let expected_hi = [8.0, 32.0, 72.0, 128.0, 200.0];
    for k in 0..N_HORIZONS {
        assert!((lo[k] - expected_lo[k]).abs() < 1e-3, "lo[{}]={} expected {}", k, lo[k], expected_lo[k]);
        assert!((hi[k] - expected_hi[k]).abs() < 1e-3, "hi[{}]={} expected {}", k, hi[k], expected_hi[k]);
    }
    Ok(())
}
