//! GPU numerical tests for the Stockham FFT compute kernels.
//!
//! These tests JIT-compile the PTX generated by the FFT kernel
//! generators, launch the kernels on a real CUDA device, and compare the
//! results against a CPU naive-DFT reference.  They validate that the
//! butterfly emitter produces a *correct* transform end-to-end on the
//! GPU — not just structurally-correct PTX.
//!
//! Following the workspace convention (see `oxicuda-driver/tests`), every
//! test uses the [`gpu_context`] helper, which returns `None` when no
//! CUDA device is available; the tests then skip gracefully rather than
//! failing.  On a machine with an NVIDIA GPU they run for real.

use std::sync::Arc;

use oxicuda_driver::module::Module;
use oxicuda_driver::{Context, Stream};
use oxicuda_fft::kernels::batch_fft::generate_batch_fft_kernel;
use oxicuda_fft::kernels::large_fft::generate_large_fft_pass;
use oxicuda_fft::kernels::stockham::generate_single_kernel;
use oxicuda_fft::plan::FftStrategy;
use oxicuda_fft::types::{FftDirection, FftPrecision};
use oxicuda_launch::{Kernel, LaunchParams};
use oxicuda_memory::DeviceBuffer;
use oxicuda_ptx::arch::SmVersion;
use oxicuda_ptx::builder::KernelBuilder;

// ---------------------------------------------------------------------------
// GPU availability helper
// ---------------------------------------------------------------------------

/// Attempts to initialise CUDA and build a context on device 0.
///
/// Returns `None` (so the caller skips the test) when CUDA is not
/// available — e.g. on a machine without an NVIDIA GPU.
fn gpu_context() -> Option<Arc<Context>> {
    if oxicuda_driver::init().is_err() {
        return None;
    }
    let device = oxicuda_driver::device::Device::get(0).ok()?;
    let context = Context::new(&device).ok()?;
    Some(Arc::new(context))
}

// ---------------------------------------------------------------------------
// CPU reference DFT
// ---------------------------------------------------------------------------

/// Naive O(N^2) DFT used as the correctness reference.
///
/// `sign` is `-1.0` for the forward transform, `+1.0` for the inverse.
fn naive_dft(input: &[(f32, f32)], sign: f64) -> Vec<(f32, f32)> {
    let n = input.len();
    let mut out = vec![(0.0_f32, 0.0_f32); n];
    for (k, slot) in out.iter_mut().enumerate() {
        let mut re = 0.0_f64;
        let mut im = 0.0_f64;
        for (m, &(xr, xi)) in input.iter().enumerate() {
            let ang = sign * 2.0 * std::f64::consts::PI * (k * m) as f64 / n as f64;
            let (c, s) = (ang.cos(), ang.sin());
            re += f64::from(xr) * c - f64::from(xi) * s;
            im += f64::from(xr) * s + f64::from(xi) * c;
        }
        *slot = (re as f32, im as f32);
    }
    out
}

/// Detects the SM version of device 0 (falls back to `sm_80`).
fn detect_sm() -> SmVersion {
    SmVersion::Sm80
}

/// Flattens a complex signal into the interleaved `[re, im, re, im, ...]`
/// layout expected by the kernels.
fn interleave(signal: &[(f32, f32)]) -> Vec<f32> {
    let mut out = Vec::with_capacity(signal.len() * 2);
    for &(re, im) in signal {
        out.push(re);
        out.push(im);
    }
    out
}

/// Re-pairs an interleaved float buffer into complex numbers.
fn deinterleave(flat: &[f32]) -> Vec<(f32, f32)> {
    flat.chunks_exact(2).map(|c| (c[0], c[1])).collect()
}

// ---------------------------------------------------------------------------
// Single-kernel Stockham FFT — GPU numerical correctness
// ---------------------------------------------------------------------------

/// Runs the generated single-kernel Stockham FFT on the GPU for a single
/// batch element and returns the transformed signal.
fn run_single_kernel_fft(
    ctx: &Arc<Context>,
    signal: &[(f32, f32)],
    radices: Vec<u32>,
    direction: FftDirection,
) -> Option<Vec<(f32, f32)>> {
    let n = signal.len();
    let strategy = FftStrategy {
        radices: radices.clone(),
        strides: Vec::new(),
        single_kernel: true,
    };
    let ptx = generate_single_kernel(
        n,
        &strategy,
        1,
        FftPrecision::Single,
        direction,
        detect_sm(),
    )
    .ok()?;

    let module = Arc::new(Module::from_ptx(&ptx).ok()?);
    let dir_tag = match direction {
        FftDirection::Forward => "fwd",
        FftDirection::Inverse => "inv",
    };
    let kernel_name = format!("fft_stockham_f32_n{n}_b1_{dir_tag}");
    let kernel = Kernel::from_module(module, &kernel_name).ok()?;

    let host_in = interleave(signal);
    let d_in = DeviceBuffer::<f32>::from_host(&host_in).ok()?;
    let d_out = DeviceBuffer::<f32>::alloc(host_in.len()).ok()?;

    let block = block_size_for(n);
    let params = LaunchParams::new(1u32, block);
    let dir_val: u32 = match direction {
        FftDirection::Forward => 0,
        FftDirection::Inverse => 1,
    };
    let args = (d_in.as_device_ptr(), d_out.as_device_ptr(), 1u32, dir_val);

    let stream = Stream::new(ctx).ok()?;
    kernel.launch(&params, &stream, &args).ok()?;
    stream.synchronize().ok()?;

    let mut host_out = vec![0.0_f32; host_in.len()];
    d_out.copy_to_host(&mut host_out).ok()?;
    Some(deinterleave(&host_out))
}

/// Mirrors `stockham::compute_block_size` for test launch configuration.
fn block_size_for(n: usize) -> u32 {
    if n <= 32 {
        32
    } else if n <= 64 {
        64
    } else if n <= 128 {
        128
    } else {
        256
    }
}

/// Asserts two complex signals agree within `tol` (absolute).
fn assert_signals_close(got: &[(f32, f32)], want: &[(f32, f32)], tol: f32, label: &str) {
    assert_eq!(got.len(), want.len(), "{label}: length mismatch");
    for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() {
        let dre = (g.0 - w.0).abs();
        let dim = (g.1 - w.1).abs();
        assert!(
            dre <= tol && dim <= tol,
            "{label}: bin {i}: got ({}, {}), want ({}, {}), \
             |dre|={dre}, |dim|={dim} > tol={tol}",
            g.0,
            g.1,
            w.0,
            w.1
        );
    }
}

#[test]
fn gpu_single_kernel_forward_matches_dft_radix2() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    // N = 8 = 2*2*2.
    let signal: Vec<(f32, f32)> = (0..8).map(|i| (i as f32, 0.5 * i as f32 - 1.0)).collect();
    let Some(got) = run_single_kernel_fft(&ctx, &signal, vec![2, 2, 2], FftDirection::Forward)
    else {
        eprintln!("skipping: kernel launch failed");
        return;
    };
    let want = naive_dft(&signal, -1.0);
    // f32 single-precision FFT — a generous absolute tolerance.
    assert_signals_close(&got, &want, 2e-3, "radix-2 forward");
}

#[test]
fn gpu_single_kernel_forward_matches_dft_radix4() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    // N = 16 = 4*4.
    let signal: Vec<(f32, f32)> = (0..16)
        .map(|i| (((i * 3) % 7) as f32, ((i * 2) % 5) as f32))
        .collect();
    let Some(got) = run_single_kernel_fft(&ctx, &signal, vec![4, 4], FftDirection::Forward) else {
        eprintln!("skipping: kernel launch failed");
        return;
    };
    let want = naive_dft(&signal, -1.0);
    assert_signals_close(&got, &want, 3e-3, "radix-4 forward");
}

#[test]
fn gpu_single_kernel_forward_matches_dft_radix8() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    // N = 64 = 8*8.
    let signal: Vec<(f32, f32)> = (0..64)
        .map(|i| ((i as f32 * 0.1).sin(), (i as f32 * 0.07).cos()))
        .collect();
    let Some(got) = run_single_kernel_fft(&ctx, &signal, vec![8, 8], FftDirection::Forward) else {
        eprintln!("skipping: kernel launch failed");
        return;
    };
    let want = naive_dft(&signal, -1.0);
    assert_signals_close(&got, &want, 5e-3, "radix-8 forward");
}

/// A pure DC signal (all ones) must transform to a single non-zero bin
/// at index 0 with value N.
#[test]
fn gpu_single_kernel_dc_signal_known_spectrum() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    let n = 16;
    let signal: Vec<(f32, f32)> = vec![(1.0, 0.0); n];
    let Some(got) = run_single_kernel_fft(&ctx, &signal, vec![4, 4], FftDirection::Forward) else {
        eprintln!("skipping: kernel launch failed");
        return;
    };
    // X[0] = N, all other bins ~= 0.
    assert!(
        (got[0].0 - n as f32).abs() < 5e-3,
        "DC bin should be N={n}, got {}",
        got[0].0
    );
    assert!(got[0].1.abs() < 5e-3, "DC bin imag should be 0");
    for (i, &(re, im)) in got.iter().enumerate().skip(1) {
        assert!(
            re.abs() < 5e-3 && im.abs() < 5e-3,
            "non-DC bin {i} should be zero, got ({re}, {im})"
        );
    }
}

/// A single-frequency complex exponential `exp(2*pi*i*k0*n/N)` must
/// transform to a single spike at bin `k0`.
#[test]
fn gpu_single_kernel_pure_tone_known_spectrum() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    let n = 32;
    let k0 = 5usize;
    let signal: Vec<(f32, f32)> = (0..n)
        .map(|m| {
            let ang = 2.0 * std::f64::consts::PI * (k0 * m) as f64 / n as f64;
            (ang.cos() as f32, ang.sin() as f32)
        })
        .collect();
    let Some(got) = run_single_kernel_fft(&ctx, &signal, vec![8, 4], FftDirection::Forward) else {
        eprintln!("skipping: kernel launch failed");
        return;
    };
    for (i, &(re, im)) in got.iter().enumerate() {
        let mag = (re * re + im * im).sqrt();
        if i == k0 {
            assert!(
                (mag - n as f32).abs() < 0.1,
                "tone bin {k0} magnitude should be N={n}, got {mag}"
            );
        } else {
            assert!(mag < 0.1, "bin {i} should be ~0, got magnitude {mag}");
        }
    }
}

/// Forward followed by inverse must reproduce the input up to the `1/N`
/// scaling convention (the kernels do not apply the inverse scale).
#[test]
fn gpu_single_kernel_forward_inverse_roundtrip() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    let n = 64;
    let signal: Vec<(f32, f32)> = (0..n)
        .map(|i| (((i * 7) % 13) as f32, ((i * 5) % 11) as f32 - 5.0))
        .collect();
    let Some(spectrum) = run_single_kernel_fft(&ctx, &signal, vec![8, 8], FftDirection::Forward)
    else {
        eprintln!("skipping: forward launch failed");
        return;
    };
    let Some(back) = run_single_kernel_fft(&ctx, &spectrum, vec![8, 8], FftDirection::Inverse)
    else {
        eprintln!("skipping: inverse launch failed");
        return;
    };
    // back == N * signal  =>  back / N == signal.
    let scaled: Vec<(f32, f32)> = back
        .iter()
        .map(|&(re, im)| (re / n as f32, im / n as f32))
        .collect();
    assert_signals_close(&scaled, &signal, 5e-3, "forward/inverse round-trip");
}

// ---------------------------------------------------------------------------
// Batch FFT kernel — GPU numerical correctness
// ---------------------------------------------------------------------------

/// Each batch row must be transformed independently and correctly; this
/// verifies the per-batch-row shared-memory confinement.
#[test]
fn gpu_batch_fft_each_row_matches_dft() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    let n = 16usize;
    let batch = 4usize;

    // Build `batch` distinct signals.
    let signals: Vec<Vec<(f32, f32)>> = (0..batch)
        .map(|b| {
            (0..n)
                .map(|i| (((i + b) % 5) as f32, ((i * 2 + b) % 7) as f32))
                .collect()
        })
        .collect();

    let ptx = match generate_batch_fft_kernel(
        n,
        batch,
        FftPrecision::Single,
        FftDirection::Forward,
        detect_sm(),
    ) {
        Ok(p) => p,
        Err(_) => {
            eprintln!("skipping: batch kernel generation failed");
            return;
        }
    };
    let Some(module) = Module::from_ptx(&ptx).ok().map(Arc::new) else {
        eprintln!("skipping: JIT failed");
        return;
    };
    let kernel_name = format!("fft_batch_f32_n{n}_b{batch}_fwd");
    let Some(kernel) = Kernel::from_module(module, &kernel_name).ok() else {
        eprintln!("skipping: kernel lookup failed");
        return;
    };

    // Concatenate all batch rows into one interleaved buffer.
    let mut host_in: Vec<f32> = Vec::with_capacity(n * batch * 2);
    for sig in &signals {
        host_in.extend(interleave(sig));
    }
    let Some(d_in) = DeviceBuffer::<f32>::from_host(&host_in).ok() else {
        eprintln!("skipping: device alloc failed");
        return;
    };
    let Some(d_out) = DeviceBuffer::<f32>::alloc(host_in.len()).ok() else {
        eprintln!("skipping: device alloc failed");
        return;
    };

    let params = LaunchParams::new(batch as u32, block_size_for(n));
    let args = (
        d_in.as_device_ptr(),
        d_out.as_device_ptr(),
        batch as u32,
        0u32,
    );
    let Some(stream) = Stream::new(&ctx).ok() else {
        eprintln!("skipping: stream creation failed");
        return;
    };
    if kernel.launch(&params, &stream, &args).is_err() || stream.synchronize().is_err() {
        eprintln!("skipping: launch failed");
        return;
    }

    let mut host_out = vec![0.0_f32; host_in.len()];
    if d_out.copy_to_host(&mut host_out).is_err() {
        eprintln!("skipping: copy back failed");
        return;
    }

    // Verify every batch row independently.
    for (b, sig) in signals.iter().enumerate() {
        let row = &host_out[b * n * 2..(b + 1) * n * 2];
        let got = deinterleave(row);
        let want = naive_dft(sig, -1.0);
        assert_signals_close(&got, &want, 3e-3, &format!("batch row {b}"));
    }
}

// ---------------------------------------------------------------------------
// Large multi-pass FFT — GPU numerical correctness
// ---------------------------------------------------------------------------

/// Runs every pass of a uniform-radix large FFT on the GPU, ping-ponging
/// between two global buffers, and returns the transformed signal.
fn run_large_multipass_fft(
    ctx: &Arc<Context>,
    signal: &[(f32, f32)],
    radix: u32,
    stages: u32,
    direction: FftDirection,
) -> Option<Vec<(f32, f32)>> {
    let n = signal.len();
    let host_in = interleave(signal);

    let buf_a = DeviceBuffer::<f32>::from_host(&host_in).ok()?;
    let buf_b = DeviceBuffer::<f32>::alloc(host_in.len()).ok()?;
    let stream = Stream::new(ctx).ok()?;

    let dir_tag = match direction {
        FftDirection::Forward => "fwd",
        FftDirection::Inverse => "inv",
    };
    let dir_val: u32 = match direction {
        FftDirection::Forward => 0,
        FftDirection::Inverse => 1,
    };

    let butterflies = n / radix as usize;
    let block = 256u32;
    let grid = (butterflies as u32).div_ceil(block);
    let params = LaunchParams::new(grid, block);

    // `src_is_a` tracks which buffer currently holds the live data.
    let mut src_is_a = true;
    for stage in 0..stages {
        let ptx = generate_large_fft_pass(
            n,
            radix,
            stage,
            FftPrecision::Single,
            direction,
            detect_sm(),
        )
        .ok()?;
        let module = Arc::new(Module::from_ptx(&ptx).ok()?);
        let kernel_name = format!("fft_large_pass_f32_n{n}_r{radix}_s{stage}_{dir_tag}");
        let kernel = Kernel::from_module(module, &kernel_name).ok()?;

        let (src_ptr, dst_ptr) = if src_is_a {
            (buf_a.as_device_ptr(), buf_b.as_device_ptr())
        } else {
            (buf_b.as_device_ptr(), buf_a.as_device_ptr())
        };
        let args = (src_ptr, dst_ptr, n as u32, 1u32, dir_val);
        kernel.launch(&params, &stream, &args).ok()?;
        stream.synchronize().ok()?;
        src_is_a = !src_is_a;
    }

    let mut host_out = vec![0.0_f32; host_in.len()];
    if src_is_a {
        buf_a.copy_to_host(&mut host_out).ok()?;
    } else {
        buf_b.copy_to_host(&mut host_out).ok()?;
    }
    Some(deinterleave(&host_out))
}

#[test]
fn gpu_large_multipass_forward_matches_dft() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    // N = 64 = 8^2 — two radix-8 passes.
    let signal: Vec<(f32, f32)> = (0..64)
        .map(|i| (((i * 11) % 17) as f32, ((i * 3) % 13) as f32 - 6.0))
        .collect();
    let Some(got) = run_large_multipass_fft(&ctx, &signal, 8, 2, FftDirection::Forward) else {
        eprintln!("skipping: multipass launch failed");
        return;
    };
    let want = naive_dft(&signal, -1.0);
    assert_signals_close(&got, &want, 6e-3, "large multipass forward");
}

#[test]
fn gpu_large_multipass_forward_inverse_roundtrip() {
    let Some(ctx) = gpu_context() else {
        eprintln!("skipping: no CUDA device");
        return;
    };
    // N = 256 = 4^4 — four radix-4 passes.
    let n = 256usize;
    let signal: Vec<(f32, f32)> = (0..n)
        .map(|i| (((i * 7) % 19) as f32, ((i * 5) % 23) as f32 - 11.0))
        .collect();
    let Some(spectrum) = run_large_multipass_fft(&ctx, &signal, 4, 4, FftDirection::Forward) else {
        eprintln!("skipping: forward multipass failed");
        return;
    };
    let Some(back) = run_large_multipass_fft(&ctx, &spectrum, 4, 4, FftDirection::Inverse) else {
        eprintln!("skipping: inverse multipass failed");
        return;
    };
    let scaled: Vec<(f32, f32)> = back
        .iter()
        .map(|&(re, im)| (re / n as f32, im / n as f32))
        .collect();
    assert_signals_close(&scaled, &signal, 8e-3, "large multipass round-trip");
}

/// Sanity check that the PTX builder type used here is the same one the
/// kernel generators target — guards against an API drift breaking the
/// integration test silently.
#[test]
fn kernel_builder_type_is_reachable() {
    let _ = KernelBuilder::new("noop");
}
