// GPU integration tests for the tensor abstraction. Both require nvcc + a GPU,
// so they are gated behind `not(no_cuda)`. On a GPU-less machine build.rs sets
// the `no_cuda` cfg and these compile out, keeping host `cargo check` green.
#![cfg(not(no_cuda))]

use xtrain_cuda::device;
use xtrain_tensor::{Device, Tensor};

/// (a) Host → device → host roundtrip preserves the data exactly.
#[test]
fn host_device_roundtrip() {
    assert!(
        device::device_count().expect("device count") > 0,
        "no CUDA device"
    );
    device::set_device(0).unwrap();

    let host: Vec<f32> = (0..1024).map(|i| i as f32 * 0.5).collect();
    let cpu = Tensor::from_slice(&host, &[1024]);

    let gpu = cpu.to_device(Device::Cuda(0));
    assert_eq!(gpu.device(), Device::Cuda(0));
    assert_eq!(gpu.shape(), &[1024]);

    let back = gpu.to_device(Device::Cpu);
    assert_eq!(back.device(), Device::Cpu);
    assert_eq!(back.as_slice::<f32>(), host.as_slice());
    println!("roundtrip OK: {} elems preserved", host.len());
}

/// (b) The elementwise `scale` kernel produces correct results.
#[test]
fn elementwise_scale_kernel() {
    assert!(
        device::device_count().expect("device count") > 0,
        "no CUDA device"
    );
    device::set_device(0).unwrap();

    let host: Vec<f32> = (0..2048).map(|i| i as f32).collect();
    let alpha = 3.0f32;
    let expected: Vec<f32> = host.iter().map(|x| x * alpha).collect();

    let gpu = Tensor::from_slice(&host, &[2048]).to_device(Device::Cuda(0));
    let scaled = gpu.scale(alpha);
    let result = scaled.to_device(Device::Cpu);

    assert_eq!(result.shape(), &[2048]);
    assert_eq!(result.as_slice::<f32>(), expected.as_slice());
    let r = result.as_slice::<f32>();
    println!(
        "scale OK (alpha={alpha}): first={} mid={} last={} ({} elems)",
        r[0],
        r[r.len() / 2],
        r[r.len() - 1],
        r.len()
    );
}