//! Basic CUDA functionality test
//! Tests if basic tensor operations work on CUDA device without hanging

use anyhow::Result;
use candle_core::{Device, Tensor};

fn main() -> Result<()> {
    println!("╔═══════════════════════════════════════════════════════════╗");
    println!("║              Basic CUDA Test                              ║");
    println!("╚═══════════════════════════════════════════════════════════╝");

    // Test 1: Device creation
    println!("\n[TEST 1] Creating CUDA device...");
    let device = Device::new_cuda(0)?;
    println!("✓ CUDA device created successfully");

    // Test 2: Simple tensor creation on CUDA
    println!("\n[TEST 2] Creating tensor on CUDA...");
    let tensor_cuda = Tensor::zeros((32, 54), candle_core::DType::F64, &device)?;
    println!("✓ Tensor created on CUDA: {:?}", tensor_cuda.dims());

    // Test 3: CPU tensor creation and device transfer
    println!("\n[TEST 3] Creating CPU tensor and transferring to CUDA...");
    let data: Vec<f64> = (0..100).map(|i| i as f64).collect();
    let tensor_cpu = Tensor::new(&data[..], &Device::Cpu)?;
    println!("  CPU tensor shape: {:?}", tensor_cpu.dims());

    let tensor_moved = tensor_cpu.to_device(&device)?;
    println!("✓ Tensor moved to CUDA: {:?}", tensor_moved.dims());

    // Test 4: Basic arithmetic on CUDA
    println!("\n[TEST 4] Testing arithmetic operations on CUDA...");
    let a = Tensor::ones((10, 10), candle_core::DType::F64, &device)?;
    let b = Tensor::ones((10, 10), candle_core::DType::F64, &device)?;
    let c = (&a + &b)?;
    println!("✓ Addition completed: result shape {:?}", c.dims());

    // Test 5: Matrix multiplication on CUDA
    println!("\n[TEST 5] Testing matmul on CUDA...");
    let x = Tensor::randn(0.0, 1.0, (32, 54), &device)?;
    let y = Tensor::randn(0.0, 1.0, (54, 16), &device)?;
    let z = x.matmul(&y)?;
    println!("✓ Matmul completed: result shape {:?}", z.dims());

    // Test 6: Concatenation on CUDA
    println!("\n[TEST 6] Testing concatenation on CUDA...");
    let t1 = Tensor::ones((1, 60, 54), candle_core::DType::F64, &device)?;
    let t2 = Tensor::ones((1, 60, 54), candle_core::DType::F64, &device)?;
    let t_cat = Tensor::cat(&[&t1, &t2], 0)?;
    println!("✓ Concatenation completed: result shape {:?}", t_cat.dims());

    // Test 7: Device transfer after concatenation
    println!("\n[TEST 7] Testing device transfer after concatenation...");
    let t1_cpu = Tensor::ones((1, 60, 54), candle_core::DType::F64, &Device::Cpu)?;
    let t2_cpu = Tensor::ones((1, 60, 54), candle_core::DType::F64, &Device::Cpu)?;
    let t_cat_cpu = Tensor::cat(&[&t1_cpu, &t2_cpu], 0)?;
    println!("  CPU concatenated tensor: {:?}", t_cat_cpu.dims());

    let t_cat_cuda = t_cat_cpu.to_device(&device)?;
    println!("✓ Device transfer completed: {:?}", t_cat_cuda.dims());

    println!("\n╔═══════════════════════════════════════════════════════════╗");
    println!("║          All CUDA Tests PASSED                            ║");
    println!("╚═══════════════════════════════════════════════════════════╝");

    Ok(())
}