//! GPU-oracle tests for the IQN head and the `rl_sample_tau` device-side
//! PRNG kernel.
//!
//! Three analytical invariants, zero CPU reference implementations:
//!
//!   1. `rl_sample_tau_produces_uniform_01` — xorshift32 kernel outputs
//!      lie in [0, 1), and a second call with the same PRNG state buffer
//!      produces different values (seed advances).
//!   2. `iqn_forward_output_shape` — `IqnHead::forward` produces finite
//!      Q-values of shape `[B, N_TAU, N_ACTIONS]` given valid h_t + tau.
//!   3. `iqn_expected_q_is_mean_over_tau` — `IqnHead::expected_q` output
//!      equals the arithmetic mean of `q_values` over the tau dimension,
//!      within 1e-4. This is a definitional identity.
//!
//! Per `feedback_no_cpu_test_fallbacks`: GPU oracle only, no CPU
//! reference implementations.
//! Per `feedback_no_atomicadd`: no atomics.
//!
//! Run with:
//!   `cargo test -p ml-alpha --test iqn_oracle -- --ignored --nocapture`

use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use ml_alpha::heads::HIDDEN_DIM;
use ml_alpha::pinned_mem::MappedF32Buffer;
use ml_alpha::rl::common::N_ACTIONS;
use ml_alpha::rl::iqn::{IqnHead, IqnHeadConfig};
use ml_core::device::MlDevice;

const B: usize = 2;
const N_TAU: usize = 32;

const RL_SAMPLE_TAU_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/rl_sample_tau.cubin"));

fn get_device() -> Option<(MlDevice, Arc<CudaStream>)> {
    let dev = match MlDevice::cuda(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            return None;
        }
    };
    let stream = dev.cuda_stream().expect("cuda_stream").clone();
    Some((dev, stream))
}

fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> CudaSlice<f32> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }.expect("staging alloc");
    staging.write_from_slice(host);
    let mut d = stream.alloc_zeros::<f32>(n).expect("alloc f32");
    unsafe {
        let (dst, _g) = d.device_ptr_mut(stream);
        cudarc::driver::result::memcpy_dtod_async(
            dst, staging.dev_ptr, n * 4, stream.cu_stream(),
        ).expect("upload DtoD");
    }
    stream.synchronize().expect("upload sync");
    d
}

fn download_f32(stream: &Arc<CudaStream>, d: &CudaSlice<f32>, n: usize) -> Vec<f32> {
    let staging = unsafe { MappedF32Buffer::new(n) }.expect("staging alloc");
    unsafe {
        let (src, _g) = d.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(
            staging.dev_ptr, src, n * 4, stream.cu_stream(),
        ).expect("download DtoD");
    }
    stream.synchronize().expect("download sync");
    let mut out = vec![0.0_f32; n];
    for i in 0..n {
        out[i] = unsafe { std::ptr::read_volatile(staging.host_ptr.add(i)) };
    }
    out
}

// ---------------------------------------------------------------------------
// Test 1: rl_sample_tau kernel produces U(0,1) and advances PRNG
// ---------------------------------------------------------------------------

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn rl_sample_tau_produces_uniform_01() {
    let Some((dev, stream)) = get_device() else { return };
    let ctx = dev.cuda_context().expect("cuda_context");

    let module = ctx
        .load_cubin(RL_SAMPLE_TAU_CUBIN.to_vec())
        .expect("load rl_sample_tau cubin");
    let sample_tau_fn = module
        .load_function("rl_sample_tau")
        .expect("load rl_sample_tau fn");

    // alloc_zeros self-seeds the xorshift32 on first call (seed==0 →
    // Knuth golden-ratio hash from batch index).
    let mut prng_state_d = stream
        .alloc_zeros::<u32>(B)
        .expect("alloc prng_state_d");
    let mut tau_d = stream
        .alloc_zeros::<f32>(B * N_TAU)
        .expect("alloc tau_d");

    let b_i = B as i32;
    let n_tau_i = N_TAU as i32;
    let tau_cfg = LaunchConfig {
        grid_dim: (B as u32, 1, 1),
        block_dim: (N_TAU as u32, 1, 1),
        shared_mem_bytes: 0,
    };

    // --- First launch ---
    unsafe {
        stream
            .launch_builder(&sample_tau_fn)
            .arg(&mut prng_state_d)
            .arg(&mut tau_d)
            .arg(&b_i)
            .arg(&n_tau_i)
            .launch(tau_cfg)
            .expect("rl_sample_tau launch 1");
    }
    stream.synchronize().expect("sync 1");

    let tau_host_1 = download_f32(&stream, &tau_d, B * N_TAU);

    for (i, &t) in tau_host_1.iter().enumerate() {
        assert!(
            t >= 0.0 && t < 1.0,
            "tau[{i}] = {t} outside [0, 1) on first call"
        );
    }

    // --- Second launch (PRNG advances) ---
    unsafe {
        stream
            .launch_builder(&sample_tau_fn)
            .arg(&mut prng_state_d)
            .arg(&mut tau_d)
            .arg(&b_i)
            .arg(&n_tau_i)
            .launch(tau_cfg)
            .expect("rl_sample_tau launch 2");
    }
    stream.synchronize().expect("sync 2");

    let tau_host_2 = download_f32(&stream, &tau_d, B * N_TAU);

    for (i, &t) in tau_host_2.iter().enumerate() {
        assert!(
            t >= 0.0 && t < 1.0,
            "tau[{i}] = {t} outside [0, 1) on second call"
        );
    }

    // At least one value must differ between calls — PRNG must advance.
    let any_differ = tau_host_1
        .iter()
        .zip(tau_host_2.iter())
        .any(|(a, b)| (a - b).abs() > f32::EPSILON);
    assert!(
        any_differ,
        "Second rl_sample_tau call produced identical values — PRNG did not advance"
    );

    eprintln!(
        "IQN.1 OK — rl_sample_tau: all {} values in [0,1), second call differs",
        B * N_TAU
    );
}

// ---------------------------------------------------------------------------
// Test 2: IqnHead::forward produces finite Q-values of correct shape
// ---------------------------------------------------------------------------

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn iqn_forward_output_shape() {
    let Some((dev, stream)) = get_device() else { return };

    let cfg = IqnHeadConfig {
        hidden_dim: HIDDEN_DIM,
        n_tau: N_TAU,
        ..IqnHeadConfig::default()
    };
    let iqn = IqnHead::new(&dev, cfg).expect("IqnHead::new");

    // h_t [B, HIDDEN_DIM] — small deterministic values
    let h_t_host: Vec<f32> = (0..B * HIDDEN_DIM)
        .map(|i| ((i as f32) * 0.001 - 0.064).sin() * 0.1)
        .collect();
    let h_t_d = upload_f32(&stream, &h_t_host);

    // tau [B, N_TAU] — evenly spaced quantile fractions
    let tau_host: Vec<f32> = (0..B * N_TAU)
        .map(|i| {
            let within_row = i % N_TAU;
            (within_row as f32 + 0.5) / N_TAU as f32
        })
        .collect();
    let mut tau_d = upload_f32(&stream, &tau_host);
    let mut prng_state_d = stream.alloc_zeros::<u32>(B).expect("alloc prng_state");

    let mut q_values_d = stream
        .alloc_zeros::<f32>(B * N_TAU * N_ACTIONS)
        .expect("alloc q_values");

    iqn.forward(&mut prng_state_d, &h_t_d, &mut tau_d, B, N_TAU, &mut q_values_d)
        .expect("IqnHead::forward");
    stream.synchronize().expect("sync forward");

    let q_host = download_f32(&stream, &q_values_d, B * N_TAU * N_ACTIONS);

    for (i, &q) in q_host.iter().enumerate() {
        assert!(
            q.is_finite(),
            "q_values[{i}] = {q} is not finite"
        );
    }

    eprintln!(
        "IQN.2 OK — IqnHead::forward: all {} values finite (B={B}, N_TAU={N_TAU}, N_ACTIONS={N_ACTIONS})",
        q_host.len()
    );
}

// ---------------------------------------------------------------------------
// Test 3: expected_q == mean over tau dimension (analytical invariant)
// ---------------------------------------------------------------------------

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn iqn_expected_q_is_mean_over_tau() {
    let Some((dev, stream)) = get_device() else { return };

    let cfg = IqnHeadConfig {
        hidden_dim: HIDDEN_DIM,
        n_tau: N_TAU,
        ..IqnHeadConfig::default()
    };
    let iqn = IqnHead::new(&dev, cfg).expect("IqnHead::new");

    // h_t [B, HIDDEN_DIM] — varied values to avoid degenerate Q output
    let h_t_host: Vec<f32> = (0..B * HIDDEN_DIM)
        .map(|i| ((i as f32) * 0.0073 + 1.0).sin() * 0.2)
        .collect();
    let h_t_d = upload_f32(&stream, &h_t_host);

    // tau [B, N_TAU] — spread across the unit interval
    let tau_host: Vec<f32> = (0..B * N_TAU)
        .map(|i| {
            let within_row = i % N_TAU;
            (within_row as f32 + 0.5) / N_TAU as f32
        })
        .collect();
    let mut tau_d = upload_f32(&stream, &tau_host);
    let mut prng_state_d = stream.alloc_zeros::<u32>(B).expect("alloc prng_state");

    // Forward pass
    let mut q_values_d = stream
        .alloc_zeros::<f32>(B * N_TAU * N_ACTIONS)
        .expect("alloc q_values");
    iqn.forward(&mut prng_state_d, &h_t_d, &mut tau_d, B, N_TAU, &mut q_values_d)
        .expect("IqnHead::forward");

    // Expected Q via the kernel
    let mut expected_q_d = stream
        .alloc_zeros::<f32>(B * N_ACTIONS)
        .expect("alloc expected_q");
    iqn.expected_q(&q_values_d, B, N_TAU, &mut expected_q_d)
        .expect("IqnHead::expected_q");
    stream.synchronize().expect("sync expected_q");

    // Read back both
    let q_host = download_f32(&stream, &q_values_d, B * N_TAU * N_ACTIONS);
    let eq_host = download_f32(&stream, &expected_q_d, B * N_ACTIONS);

    // Verify: expected_q[b, a] == mean_tau(q_values[b, :, a])
    // q_values layout: [B, N_TAU, N_ACTIONS] row-major
    let mut max_err: f32 = 0.0;
    for b in 0..B {
        for a in 0..N_ACTIONS {
            let mut sum = 0.0_f32;
            for t in 0..N_TAU {
                sum += q_host[b * N_TAU * N_ACTIONS + t * N_ACTIONS + a];
            }
            let mean = sum / N_TAU as f32;
            let kernel_val = eq_host[b * N_ACTIONS + a];
            let err = (mean - kernel_val).abs();
            if err > max_err {
                max_err = err;
            }
            assert!(
                err < 1e-4,
                "expected_q[{b},{a}] mismatch: kernel={kernel_val}, \
                 mean_tau={mean}, err={err}"
            );
        }
    }

    eprintln!(
        "IQN.3 OK — expected_q matches mean over tau for all (b,a) pairs \
         (max_err={max_err:.2e}, B={B}, N_TAU={N_TAU}, N_ACTIONS={N_ACTIONS})"
    );
}
