use rtx_tensor::{Tensor, Device};
use pinn_mre_helmholtz::{Config, ForwardWorkspace, LffnUNet1D};
use rtx_nn::Linear;
use rtx_nn::layers::Module;

#[test]
#[cfg(feature = "cuda")]
fn test_linear_forward_out() {
    let device = Device::Cuda(0);
    
    // Create a simple linear layer  
    let in_features = 128;
    let out_features = 64;
    let batch_size = 200;
    
    let linear = Linear::new(in_features, out_features, true, &device).expect("Failed to create linear");
    
    // Create input tensor
    let input = Tensor::randn(&[batch_size, in_features], &device).expect("Failed to create input");
    println!("Input shape: {:?}", input.shape().dims());
    
    // Create output buffer
    let mut output = Tensor::zeros(&[batch_size, out_features], &device).expect("Failed to create output");
    println!("Output buffer shape: {:?}", output.shape().dims());
    
    // Test regular forward
    println!("\nTesting regular forward...");
    let result = linear.forward(&input).expect("Forward failed");
    println!("Regular forward output shape: {:?}", result.shape().dims());
    
    // Test forward_out  
    println!("\nTesting forward_out...");
    linear.forward_out(&input, &mut output).expect("forward_out failed");
    println!("forward_out completed!");
}

#[test]
#[cfg(feature = "cuda")]
fn test_fused_then_linear() {
    let device = Device::Cuda(0);
    let n_points = 200;
    
    let cfg = Config {
        n_data: n_points,
        n_pde: n_points,
        epochs: 1,
        ..Config::default()
    };
    
    // Create B weights
    let b_weights = Tensor::randn(&[1, cfg.u_ff_dim], &device).expect("Failed to create B");
    println!("B weights shape: {:?}", b_weights.shape().dims());
    
    // Create input
    let x_data: Vec<f32> = (0..n_points).map(|i| i as f32 / n_points as f32).collect();
    let x_norm = Tensor::from_slice(&x_data, &[n_points, 1], &device).expect("Failed to create input");
    println!("Input shape: {:?}", x_norm.shape().dims());
    
    // Create features buffer
    let mut features = Tensor::zeros(&[n_points, cfg.u_ff_dim * 2], &device).expect("Failed to create features");
    println!("Features buffer shape: {:?}", features.shape().dims());
    
    // Run fused fourier
    println!("\nRunning fused_fourier_features_out...");
    x_norm.fused_fourier_features_out(&b_weights, 2.0 * std::f32::consts::PI, &mut features)
        .expect("Fused failed");
    println!("Fused completed! Features shape: {:?}", features.shape().dims());
    
    // Create linear layer (128 -> 64)
    let linear = Linear::new(cfg.u_ff_dim * 2, cfg.u_hidden, true, &device).expect("Failed to create linear");
    
    // Create output buffer
    let mut output = Tensor::zeros(&[n_points, cfg.u_hidden], &device).expect("Failed to create output");
    println!("Output buffer shape: {:?}", output.shape().dims());
    
    // Test regular forward first
    println!("\nTesting regular forward with fused output...");
    let result = linear.forward(&features).expect("Forward failed");
    println!("Regular forward completed! Shape: {:?}", result.shape().dims());
    
    // Test forward_out
    println!("\nTesting forward_out with fused output...");
    linear.forward_out(&features, &mut output).expect("forward_out failed");
    println!("forward_out completed!");
}