//! GPU oracle tests for cfc_step_per_branch.cu fused per-branch CfC step.
//!
//! Run with: SQLX_OFFLINE=true cargo test -p ml-alpha --test cfc_step_per_branch -- --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"), "/cfc_step_per_branch.cubin"));
const HIDDEN_DIM: usize = 128;
const N_HORIZONS: usize = 5;
const MAX_BUCKET_DIM: usize = 28;

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

/// Build the [N_HORIZONS × MAX_BUCKET_DIM] `channels_in_bucket` lookup from
/// a host-side `bucket_id_per_channel` assignment. Mirrors the device-side
/// `channels_in_bucket_kernel` (single-thread walk over channels, appending
/// to per-bucket cursors). Used by the tests below to feed the per-branch
/// kernel without depending on the transition kernel binary.
fn build_channels_in_bucket(bucket_id_per_channel: &[u8]) -> Vec<u32> {
    let mut out = vec![u32::MAX; N_HORIZONS * MAX_BUCKET_DIM];
    let mut cursors = [0usize; N_HORIZONS];
    for (c, &k) in bucket_id_per_channel.iter().enumerate() {
        let k_us = k as usize;
        debug_assert!(k_us < N_HORIZONS);
        let pos = cursors[k_us];
        debug_assert!(pos < MAX_BUCKET_DIM);
        out[k_us * MAX_BUCKET_DIM + pos] = c as u32;
        cursors[k_us] = pos + 1;
    }
    out
}

#[test]
#[ignore = "requires CUDA"]
fn fwd_matches_naive_per_branch() -> Result<()> {
    // ALPHA fix 2026-05-21: kernel reads `channels_in_bucket` lookup
    // (bucket → original-channel) instead of a `bucket_channel_offset`
    // contiguous reorder. All weight tensors stay in ORIGINAL channel
    // layout.
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "cfc_step_per_branch_fwd")?;

    let b_sz: usize = 2;
    // Inputs: x [B, HIDDEN_DIM], h_old [B, HIDDEN_DIM]
    // W_in, W_rec shared; b [HIDDEN_DIM]; tau_all [HIDDEN_DIM] (original channel layout)
    let x: Vec<f32> = (0..b_sz * HIDDEN_DIM).map(|i| (i as f32) * 0.001).collect();
    let h_old: Vec<f32> = (0..b_sz * HIDDEN_DIM)
        .map(|i| ((i as f32) * 0.003).sin())
        .collect();
    let w_in: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
        .map(|i| ((i % 17) as f32) * 0.01)
        .collect();
    let w_rec: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
        .map(|i| ((i % 13) as f32) * 0.005)
        .collect();
    let b_vec: Vec<f32> = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.002).collect();
    let tau_all: Vec<f32> = (0..HIDDEN_DIM)
        .map(|c| ((c / 25 + 1) as f32) * 10.0)
        .collect();
    // Bucket assignment: quintiles by channel (matches the contiguous case
    // from the prior reorder-based test, so the naive reference stays the
    // same; the new lookup table just maps tid → (bucket_start + tid)).
    let bucket_id_per_channel: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let channels_in_bucket = build_channels_in_bucket(&bucket_id_per_channel);
    let bucket_dim_k: Vec<u32> = vec![25, 25, 25, 25, 28];

    // Upload
    let x_d = upload(&stream, &x)?;
    let h_old_d = upload(&stream, &h_old)?;
    let w_in_d = upload(&stream, &w_in)?;
    let w_rec_d = upload(&stream, &w_rec)?;
    let b_d = upload(&stream, &b_vec)?;
    let tau_d = upload(&stream, &tau_all)?;
    let cib_d = upload(&stream, &channels_in_bucket)?;
    let bdk_d = upload(&stream, &bucket_dim_k)?;
    let mut h_new_d = stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM)?;

    let dt: f32 = 1.0;
    let b_i32: i32 = b_sz as i32;

    // shared_mem_bytes = 2 * HIDDEN_DIM * sizeof(float) = x_local + h_old_local
    let cfg = LaunchConfig {
        grid_dim: (b_sz as u32, N_HORIZONS as u32, 1),
        block_dim: (MAX_BUCKET_DIM as u32, 1, 1),
        shared_mem_bytes: (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&w_in_d)
        .arg(&w_rec_d)
        .arg(&b_d)
        .arg(&tau_d)
        .arg(&x_d)
        .arg(&h_old_d)
        .arg(&dt)
        .arg(&b_i32)
        .arg(&cib_d)
        .arg(&bdk_d)
        .arg(&mut h_new_d);
    unsafe {
        launch.launch(cfg)?;
    }
    stream.synchronize()?;

    let h_new = download(&stream, &h_new_d)?;

    // Naive reference: same math per-channel (original layout).
    for batch in 0..b_sz {
        for c in 0..HIDDEN_DIM {
            let mut pre = b_vec[c];
            for k in 0..HIDDEN_DIM {
                pre += w_in[c * HIDDEN_DIM + k] * x[batch * HIDDEN_DIM + k];
            }
            for k in 0..HIDDEN_DIM {
                pre += w_rec[c * HIDDEN_DIM + k] * h_old[batch * HIDDEN_DIM + k];
            }
            let decay = (-dt / tau_all[c].max(1e-6)).exp();
            let expected =
                h_old[batch * HIDDEN_DIM + c] * decay + (1.0 - decay) * pre.tanh();
            let got = h_new[batch * HIDDEN_DIM + c];
            assert!(
                (got - expected).abs() < 1e-4,
                "h_new[{},{}]={} expected {}",
                batch,
                c,
                got,
                expected
            );
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn fwd_writes_only_in_bucket_channels_under_non_contiguous_assignment() -> Result<()> {
    // ALPHA invariant: with arbitrary (non-contiguous) bucket assignment, the
    // kernel should still write h_new for every channel exactly once and the
    // result should equal the original-channel-layout formula. This proves
    // the kernel's indexing no longer depends on a contiguous bucket-grouped
    // layout (the bug fixed in ALPHA).
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func = load_kernel(&stream, "cfc_step_per_branch_fwd")?;

    let b_sz: usize = 1;
    let x: Vec<f32> = (0..b_sz * HIDDEN_DIM).map(|i| (i as f32) * 0.002).collect();
    let h_old: Vec<f32> = (0..b_sz * HIDDEN_DIM)
        .map(|i| ((i as f32) * 0.007).cos())
        .collect();
    let w_in: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
        .map(|i| ((i % 11) as f32) * 0.013)
        .collect();
    let w_rec: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
        .map(|i| ((i % 9) as f32) * 0.007)
        .collect();
    let b_vec: Vec<f32> = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.003).collect();
    let tau_all: Vec<f32> = (0..HIDDEN_DIM)
        .map(|c| 5.0 + (c as f32) * 0.5)
        .collect();
    // Round-robin assignment: bucket k owns channels {k, k+5, k+10, ...}.
    // Every bucket gets ~26 non-contiguous channels.
    let bucket_id_per_channel: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| (c % N_HORIZONS) as u8).collect();
    let channels_in_bucket = build_channels_in_bucket(&bucket_id_per_channel);
    let mut bucket_dim_k = vec![0u32; N_HORIZONS];
    for &k in &bucket_id_per_channel {
        bucket_dim_k[k as usize] += 1;
    }

    let x_d = upload(&stream, &x)?;
    let h_old_d = upload(&stream, &h_old)?;
    let w_in_d = upload(&stream, &w_in)?;
    let w_rec_d = upload(&stream, &w_rec)?;
    let b_d = upload(&stream, &b_vec)?;
    let tau_d = upload(&stream, &tau_all)?;
    let cib_d = upload(&stream, &channels_in_bucket)?;
    let bdk_d = upload(&stream, &bucket_dim_k)?;
    // Initialise h_new with a sentinel so we can detect untouched slots.
    let sentinel = f32::NAN;
    let h_new_init: Vec<f32> = vec![sentinel; b_sz * HIDDEN_DIM];
    let mut h_new_d = upload(&stream, &h_new_init)?;

    let dt: f32 = 1.0;
    let b_i32: i32 = b_sz as i32;
    let cfg = LaunchConfig {
        grid_dim: (b_sz as u32, N_HORIZONS as u32, 1),
        block_dim: (MAX_BUCKET_DIM as u32, 1, 1),
        shared_mem_bytes: (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&w_in_d).arg(&w_rec_d).arg(&b_d).arg(&tau_d)
        .arg(&x_d).arg(&h_old_d).arg(&dt).arg(&b_i32)
        .arg(&cib_d).arg(&bdk_d)
        .arg(&mut h_new_d);
    unsafe { launch.launch(cfg)?; }
    stream.synchronize()?;

    let h_new = download(&stream, &h_new_d)?;
    // Every channel must have been written (no NaN sentinels remain).
    for batch in 0..b_sz {
        for c in 0..HIDDEN_DIM {
            let got = h_new[batch * HIDDEN_DIM + c];
            assert!(
                got.is_finite(),
                "h_new[{}, {}] = {} — channel was not written (channel layout mismatch)",
                batch, c, got
            );

            let mut pre = b_vec[c];
            for k in 0..HIDDEN_DIM {
                pre += w_in[c * HIDDEN_DIM + k] * x[batch * HIDDEN_DIM + k];
            }
            for k in 0..HIDDEN_DIM {
                pre += w_rec[c * HIDDEN_DIM + k] * h_old[batch * HIDDEN_DIM + k];
            }
            let decay = (-dt / tau_all[c].max(1e-6)).exp();
            let expected =
                h_old[batch * HIDDEN_DIM + c] * decay + (1.0 - decay) * pre.tanh();
            assert!(
                (got - expected).abs() < 1e-4,
                "h_new[{},{}]={} expected {}",
                batch, c, got, expected
            );
        }
    }
    Ok(())
}

#[test]
#[ignore = "requires CUDA"]
fn bwd_finite_diff_matches() -> Result<()> {
    // Backward smoke for Task 4: kernel loads cleanly + writes non-NaN gradients.
    // Full numeric finite-diff validation deferred to Task 14 local validation.
    let dev = MlDevice::cuda(0)?;
    let stream = dev.cuda_stream()?.clone();
    let func_bwd = load_kernel(&stream, "cfc_step_per_branch_bwd")?;

    let b_sz: usize = 2;
    let x: Vec<f32> = (0..b_sz * HIDDEN_DIM).map(|i| (i as f32) * 0.001).collect();
    let h_old: Vec<f32> = (0..b_sz * HIDDEN_DIM)
        .map(|i| ((i as f32) * 0.003).sin())
        .collect();
    let w_in: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
        .map(|i| ((i % 17) as f32) * 0.01)
        .collect();
    let w_rec: Vec<f32> = (0..HIDDEN_DIM * HIDDEN_DIM)
        .map(|i| ((i % 13) as f32) * 0.005)
        .collect();
    let b_vec: Vec<f32> = (0..HIDDEN_DIM).map(|i| (i as f32) * 0.002).collect();
    let tau_all: Vec<f32> = (0..HIDDEN_DIM)
        .map(|c| ((c / 25 + 1) as f32) * 10.0)
        .collect();
    let grad_h_new: Vec<f32> = (0..b_sz * HIDDEN_DIM)
        .map(|i| ((i as f32) * 0.007).cos() * 0.1)
        .collect();
    // ALPHA fix 2026-05-21: bwd kernel reads `channels_in_bucket` lookup.
    let bucket_id_per_channel: Vec<u8> =
        (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
    let channels_in_bucket = build_channels_in_bucket(&bucket_id_per_channel);
    let bucket_dim_k: Vec<u32> = vec![25, 25, 25, 25, 28];

    let x_d = upload(&stream, &x)?;
    let h_old_d = upload(&stream, &h_old)?;
    let w_in_d = upload(&stream, &w_in)?;
    let w_rec_d = upload(&stream, &w_rec)?;
    let b_d = upload(&stream, &b_vec)?;
    let tau_d = upload(&stream, &tau_all)?;
    let grad_h_new_d = upload(&stream, &grad_h_new)?;
    let cib_d = upload(&stream, &channels_in_bucket)?;
    let bdk_d = upload(&stream, &bucket_dim_k)?;

    // Per-batch gradient outputs (caller reduces across batch).
    let mut grad_w_in_d =
        stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM * HIDDEN_DIM)?;
    let mut grad_w_rec_d =
        stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM * HIDDEN_DIM)?;
    let mut grad_b_d = stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM)?;
    let mut grad_tau_d = stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM)?;
    let mut grad_h_old_d = stream.alloc_zeros::<f32>(b_sz * HIDDEN_DIM)?;

    let dt: f32 = 1.0;
    let b_i32: i32 = b_sz as i32;

    let cfg = LaunchConfig {
        grid_dim: (b_sz as u32, N_HORIZONS as u32, 1),
        block_dim: (MAX_BUCKET_DIM as u32, 1, 1),
        shared_mem_bytes: (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
    };
    let mut launch = stream.launch_builder(&func_bwd);
    launch
        .arg(&w_in_d)
        .arg(&w_rec_d)
        .arg(&b_d)
        .arg(&tau_d)
        .arg(&x_d)
        .arg(&h_old_d)
        .arg(&grad_h_new_d)
        .arg(&dt)
        .arg(&b_i32)
        .arg(&cib_d)
        .arg(&bdk_d)
        .arg(&mut grad_w_in_d)
        .arg(&mut grad_w_rec_d)
        .arg(&mut grad_b_d)
        .arg(&mut grad_tau_d)
        .arg(&mut grad_h_old_d);
    unsafe {
        launch.launch(cfg)?;
    }
    stream.synchronize()?;

    // Smoke: no NaNs in any output buffer.
    let g_w_in = download(&stream, &grad_w_in_d)?;
    let g_w_rec = download(&stream, &grad_w_rec_d)?;
    let g_b = download(&stream, &grad_b_d)?;
    let g_tau = download(&stream, &grad_tau_d)?;
    let g_h_old = download(&stream, &grad_h_old_d)?;
    assert!(g_w_in.iter().all(|v| v.is_finite()), "grad_w_in has NaN/Inf");
    assert!(
        g_w_rec.iter().all(|v| v.is_finite()),
        "grad_w_rec has NaN/Inf"
    );
    assert!(g_b.iter().all(|v| v.is_finite()), "grad_b has NaN/Inf");
    assert!(g_tau.iter().all(|v| v.is_finite()), "grad_tau has NaN/Inf");
    assert!(
        g_h_old.iter().all(|v| v.is_finite()),
        "grad_h_old has NaN/Inf"
    );
    // At least one non-zero gradient (kernel wrote something).
    assert!(
        g_b.iter().any(|v| v.abs() > 0.0),
        "grad_b is all zero — kernel didn't write?"
    );
    assert!(
        g_h_old.iter().any(|v| v.abs() > 0.0),
        "grad_h_old is all zero — kernel didn't write?"
    );
    Ok(())
}
