/// Integration test: verify Rust inference matches PyTorch reference outputs.
///
/// Requires:
///   data/model_base.safetensors  (converted weights)
///   data/config.json
///   data/reference_outputs.json  (from scripts/dump_reference.py)
///
/// Run: cargo test --test reference_test -- --nocapture

use burn::backend::NdArray as B;
use burn::prelude::*;
use burn::tensor::TensorData;
use std::path::Path;

#[test]
fn test_model_loads_and_runs() {
    let config_path = Path::new("data/config.json");
    let weights_path = Path::new("data/model_base.safetensors");

    if !config_path.exists() || !weights_path.exists() {
        eprintln!("Skipping test: weights not found at data/model_base.safetensors");
        return;
    }

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let (encoder, ms) = sleepfm::SleepFmEncoder::<B>::load(
        config_path, weights_path, device.clone(),
    ).expect("Failed to load model");

    eprintln!("Model loaded in {ms:.1}ms");
    eprintln!("{}", encoder.describe());

    // Test all 4 modalities with zero input
    for n_channels in [2, 4, 7, 10] {
        let t_samples = encoder.model_cfg.samples_per_chunk();
        let signal = Tensor::<B, 3>::zeros([1, n_channels, t_samples], &device);
        let mask_data: Vec<bool> = vec![false; n_channels];
        let mask = Tensor::<B, 2, Bool>::from_data(
            TensorData::new(mask_data, vec![1, n_channels]),
            &device,
        );

        let result = encoder.run(signal, mask).expect("Inference failed");
        assert_eq!(result.pooled.len(), 128, "Expected 128-dim pooled output");
        assert_eq!(result.patches_shape, vec![60, 128], "Expected [60, 128] patches");
        eprintln!("  {n_channels} channels: OK (pooled[0] = {:.6})", result.pooled[0]);
    }
}

#[test]
fn test_embedding_dimensions() {
    let config_path = Path::new("data/config.json");
    let weights_path = Path::new("data/model_base.safetensors");

    if !config_path.exists() || !weights_path.exists() {
        eprintln!("Skipping test: weights not found");
        return;
    }

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let (encoder, _) = sleepfm::SleepFmEncoder::<B>::load(
        config_path, weights_path, device.clone(),
    ).unwrap();

    // Verify config
    assert_eq!(encoder.model_cfg.embed_dim, 128);
    assert_eq!(encoder.model_cfg.num_heads, 8);
    assert_eq!(encoder.model_cfg.num_layers, 6);
    assert_eq!(encoder.model_cfg.patch_size, 640);
    assert_eq!(encoder.model_cfg.samples_per_chunk(), 38400);
    assert_eq!(encoder.model_cfg.num_patches(), 60);
}

#[test]
fn test_with_channel_masking() {
    let config_path = Path::new("data/config.json");
    let weights_path = Path::new("data/model_base.safetensors");

    if !config_path.exists() || !weights_path.exists() {
        eprintln!("Skipping test: weights not found");
        return;
    }

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let (encoder, _) = sleepfm::SleepFmEncoder::<B>::load(
        config_path, weights_path, device.clone(),
    ).unwrap();

    let n_channels = 10;
    let t_samples = encoder.model_cfg.samples_per_chunk();

    // 10 channels but mask last 3 as padding
    let signal = Tensor::<B, 3>::zeros([1, n_channels, t_samples], &device);
    let mut mask_vec = vec![false; n_channels];
    for i in 7..10 {
        mask_vec[i] = true;  // last 3 are padding
    }
    let mask = Tensor::<B, 2, Bool>::from_data(
        TensorData::new(mask_vec, vec![1, n_channels]),
        &device,
    );

    let result = encoder.run(signal, mask).expect("Masked inference failed");
    assert_eq!(result.pooled.len(), 128);
    eprintln!("  Masked (7/10 valid): pooled[0] = {:.6}", result.pooled[0]);
}
