candle-coreml 0.3.1

//! Integration tests for CoreML models
//!
//! These tests use real .mlmodelc files to validate the complete pipeline.

#![allow(clippy::needless_return)]

use candle_core::{Device, IndexOp, Tensor};
use candle_coreml::{ensure_model_downloaded, Config, CoreMLModel};
use std::path::PathBuf;

/// Helper to get the path to OpenELM test model - downloads from HuggingFace if needed
fn get_openelm_model_path() -> Option<PathBuf> {
    let model_id = "corenet-community/coreml-OpenELM-450M-Instruct";

    if std::env::var("SKIP_MODEL_DOWNLOAD").is_ok() {
        eprintln!("SKIP_MODEL_DOWNLOAD set, skipping model download");
        return None;
    }

    let cache_dir = ensure_model_downloaded(model_id, true);

    match cache_dir {
        Ok(dir) => {
            let mlpackage_path = dir.join("OpenELM-450M-Instruct-128-float32.mlpackage");
            if mlpackage_path.exists() {
                eprintln!(
                    "Found OpenELM CoreML package at: {}",
                    mlpackage_path.display()
                );
                return Some(mlpackage_path);
            }
            eprintln!("OpenELM model directory downloaded but .mlpackage not found");
            None
        }
        Err(e) => {
            eprintln!("Failed to download OpenELM model: {e}");
            None
        }
    }
}

/// Helper to get the path to test model - now downloads from HuggingFace if needed
fn get_test_model_path() -> Option<PathBuf> {
    // Try OpenELM first for backward compatibility
    get_openelm_model_path()
}

/// Test loading a real CoreML model
#[test]
fn test_load_real_model() {
    let model_path = match get_test_model_path() {
        Some(path) => path,
        None => {
            eprintln!("Skipping test: model file not found");
            return;
        }
    };

    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "OpenELM-450M-Instruct".to_string(),
    };

    let result = CoreMLModel::load_from_file(&model_path, &config);

    // On macOS, this should succeed; on other platforms, it should fail gracefully
    #[cfg(target_os = "macos")]
    {
        match result {
            Ok(model) => {
                assert_eq!(model.config().max_sequence_length, 128);
                assert_eq!(model.config().vocab_size, 32000);
                eprintln!("✅ Successfully loaded real CoreML model");
            }
            Err(err) => {
                let err_str = err.to_string();
                if err_str.contains("Compile the model") {
                    eprintln!("⚠️  Skipping test: model needs to be compiled - {err_str}");
                    return;
                } else {
                    panic!("Failed to load real CoreML model: {err}");
                }
            }
        }
    }

    #[cfg(not(target_os = "macos"))]
    {
        assert!(result.is_err(), "Should fail on non-macOS platforms");
        eprintln!("✅ Correctly failed on non-macOS platform");
    }
}

/// Test CoreML model inference with CPU tensors
#[test]
#[cfg(target_os = "macos")]
fn test_inference_cpu() {
    let model_path = match get_test_model_path() {
        Some(path) => path,
        None => {
            eprintln!("Skipping test: model file not found");
            return;
        }
    };

    let device = Device::Cpu;
    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "OpenELM-450M-Instruct".to_string(),
    };

    let model = match CoreMLModel::load_from_file(&model_path, &config) {
        Ok(model) => model,
        Err(err) => {
            let err_str = err.to_string();
            if err_str.contains("Compile the model") {
                eprintln!("Skipping test: model needs to be compiled - {err_str}");
                return;
            } else {
                panic!("Failed to load model: {err}");
            }
        }
    };

    // Create test input tensor
    let input_data = vec![1i64, 450, 4996, 17354, 1701, 29916]; // "The quick brown fox jumped over"
    let mut padded_input = input_data.clone();
    padded_input.resize(128, 0i64); // Pad to expected length

    let input_tensor = Tensor::from_slice(&padded_input, (1, 128), &device)
        .expect("Failed to create input tensor");

    let inputs = vec![&input_tensor];
    let result = model.forward(&inputs);

    match result {
        Ok(output) => {
            eprintln!("✅ CPU inference successful");
            eprintln!("Output shape: {:?}", output.dims());
            assert_eq!(output.dims(), &[1, 128, 32000]); // Expected output shape
        }
        Err(err) => {
            let err_str = err.to_string();
            if err_str.contains("Compile the model") {
                eprintln!("Skipping test: model needs to be compiled - {err_str}");
                return;
            } else {
                panic!("CPU inference failed: {err}");
            }
        }
    }
}

/// Test CoreML model inference with Metal tensors (if available)
#[test]
#[cfg(target_os = "macos")]
fn test_inference_metal() {
    let model_path = match get_test_model_path() {
        Some(path) => path,
        None => {
            eprintln!("Skipping test: model file not found");
            return;
        }
    };

    let device = match Device::new_metal(0) {
        Ok(device) => device,
        Err(_) => {
            eprintln!("Metal not available, skipping Metal inference test");
            return;
        }
    };

    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "OpenELM-450M-Instruct".to_string(),
    };

    let model = match CoreMLModel::load_from_file(&model_path, &config) {
        Ok(model) => model,
        Err(err) => {
            let err_str = err.to_string();
            if err_str.contains("Compile the model") {
                eprintln!("Skipping test: model needs to be compiled - {err_str}");
                return;
            } else {
                panic!("Failed to load model: {err}");
            }
        }
    };

    // Create test input tensor on Metal
    let input_data = vec![1i64, 450, 4996, 17354, 1701, 29916];
    let mut padded_input = input_data.clone();
    padded_input.resize(128, 0i64);

    let input_tensor = Tensor::from_slice(&padded_input, (1, 128), &device)
        .expect("Failed to create Metal input tensor");

    let inputs = vec![&input_tensor];
    let result = model.forward(&inputs);

    match result {
        Ok(output) => {
            eprintln!("✅ Metal inference successful");
            eprintln!("Output shape: {:?}", output.dims());
            assert_eq!(output.dims(), &[1, 128, 32000]);
        }
        Err(err) => {
            let err_str = err.to_string();
            if err_str.contains("Compile the model") {
                eprintln!("Skipping test: model needs to be compiled - {err_str}");
                return;
            } else {
                panic!("Metal inference failed: {err}");
            }
        }
    }
}

/// Test device validation - CUDA tensors should be rejected
#[test]
#[cfg(target_os = "macos")]
fn test_device_validation_cuda_rejection() {
    let model_path = match get_test_model_path() {
        Some(path) => path,
        None => {
            eprintln!("Skipping test: model file not found");
            return;
        }
    };

    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "OpenELM-450M-Instruct".to_string(),
    };

    let model = match CoreMLModel::load_from_file(&model_path, &config) {
        Ok(model) => model,
        Err(err) => {
            let err_str = err.to_string();
            if err_str.contains("Compile the model") {
                eprintln!("Skipping test: model needs to be compiled - {err_str}");
                return;
            } else {
                panic!("Failed to load model: {err}");
            }
        }
    };

    // Try to create a CUDA tensor (this should fail on macOS, but we're testing the error handling)
    let cpu_device = Device::Cpu;
    let input_data = vec![1i64, 450, 4996, 17354, 1701, 29916];
    let mut padded_input = input_data.clone();
    padded_input.resize(128, 0i64);

    let input_tensor = Tensor::from_slice(&padded_input, (1, 128), &cpu_device)
        .expect("Failed to create CPU input tensor");

    // Note: On macOS, CUDA tensors cannot be created, so we test with a CPU tensor
    // The actual validation happens in the CoreML integration layer
    let inputs = vec![&input_tensor];
    let result = model.forward(&inputs);

    // This should succeed since we're using CPU tensors
    assert!(result.is_ok(), "CPU tensors should be accepted");
    eprintln!("✅ Device validation test passed");
}

/// Test tensor round-trip conversion
#[test]
fn test_tensor_roundtrip() {
    let device = Device::Cpu;
    let original_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
    let tensor =
        Tensor::from_slice(&original_data, (6,), &device).expect("Failed to create tensor");

    // Test that tensor data can be extracted
    let extracted: Vec<f32> = tensor.to_vec1().expect("Failed to extract tensor data");
    assert_eq!(original_data, extracted);
    eprintln!("✅ Tensor round-trip conversion successful");
}

/// Test configuration validation
#[test]
fn test_config_validation() {
    // Test valid config
    let valid_config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "test-model".to_string(),
    };

    // Basic validation - ensure required fields are present
    assert!(!valid_config.input_names.is_empty());
    assert!(!valid_config.output_name.is_empty());
    assert!(valid_config.max_sequence_length > 0);
    assert!(valid_config.vocab_size > 0);

    eprintln!("✅ Configuration validation test passed");
}

/// Test error handling for missing files
#[test]
fn test_missing_file_error() {
    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "test-model".to_string(),
    };

    let result = CoreMLModel::load_from_file("/nonexistent/path.mlmodelc", &config);
    assert!(result.is_err(), "Should fail for nonexistent file");
    eprintln!("✅ Missing file error handling test passed");
}

/// Test basic state creation functionality
#[test]
#[cfg(target_os = "macos")]
fn test_state_creation() {
    let model_path = match get_test_model_path() {
        Some(path) => path,
        None => {
            eprintln!("Skipping test: model file not found");
            return;
        }
    };

    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "OpenELM-450M-Instruct".to_string(),
    };

    let _model = match CoreMLModel::load_from_file(&model_path, &config) {
        Ok(model) => model,
        Err(err) => {
            let err_str = err.to_string();
            if err_str.contains("Compile the model") {
                eprintln!("Skipping test: model needs to be compiled - {err_str}");
                return;
            } else {
                panic!("Failed to load model: {err}");
            }
        }
    };

    // Test state creation - not all models support stateful predictions
    eprintln!("ℹ️  State creation test - not all CoreML models support stateful predictions");
    eprintln!("✅ State creation test passed (OpenELM typically doesn't support states)");
}

/// Test stateful prediction functionality
#[test]
#[cfg(target_os = "macos")]
fn test_stateful_prediction() {
    let model_path = match get_test_model_path() {
        Some(path) => path,
        None => {
            eprintln!("Skipping test: model file not found");
            return;
        }
    };

    let _device = Device::Cpu;
    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "OpenELM-450M-Instruct".to_string(),
    };

    let _model = match CoreMLModel::load_from_file(&model_path, &config) {
        Ok(model) => model,
        Err(err) => {
            let err_str = err.to_string();
            if err_str.contains("Compile the model") {
                eprintln!("Skipping test: model needs to be compiled - {err_str}");
                return;
            } else {
                panic!("Failed to load model: {err}");
            }
        }
    };

    // Most OpenELM models don't support stateful predictions
    eprintln!("ℹ️  Stateful prediction test - OpenELM models typically don't support states");
    eprintln!("✅ Stateful prediction test passed (skipped for OpenELM compatibility)");
}

/// Test stateful prediction with multiple calls (persistence check)
#[test]
#[cfg(target_os = "macos")]
fn test_stateful_prediction_persistence() {
    let model_path = match get_test_model_path() {
        Some(path) => path,
        None => {
            eprintln!("Skipping test: model file not found");
            return;
        }
    };

    let _device = Device::Cpu;
    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "OpenELM-450M-Instruct".to_string(),
    };

    let _model = match CoreMLModel::load_from_file(&model_path, &config) {
        Ok(model) => model,
        Err(err) => {
            let err_str = err.to_string();
            if err_str.contains("Compile the model") {
                eprintln!("Skipping test: model needs to be compiled - {err_str}");
                return;
            } else {
                panic!("Failed to load model: {err}");
            }
        }
    };

    // OpenELM models don't support stateful predictions with persistent state
    eprintln!(
        "ℹ️  Stateful persistence test - OpenELM models typically don't support persistent states"
    );
    eprintln!("✅ Stateful persistence test passed (skipped for OpenELM compatibility)");
}

/// Baseline test: OpenELM "quick brown fox" completion
///
/// This test serves as our "known truth" baseline to verify CoreML infrastructure works perfectly.
/// It downloads the OpenELM model and tests the classic "The quick brown fox jumped over the lazy" -> "dog" completion.
///
/// This test is ignored by default to avoid downloading large models in CI.
///
/// To run manually:
/// ```bash
/// cargo test test_openelm_baseline_text_completion -- --ignored --nocapture
/// ```
///
/// Expected behavior:
/// - Downloads OpenELM-450M-Instruct model (~1.7GB)
/// - Tests "quick brown fox" completion
/// - MUST predict "dog" (token 11203) as top prediction
/// - Confidence score should be > 10.0
/// - Test will FAIL if our CoreML infrastructure is broken
#[test]
#[ignore = "downloads large model - run manually to verify baseline"]
fn test_openelm_baseline_text_completion() {
    println!("🦙 BASELINE TEST: OpenELM 'quick brown fox' completion");
    println!("This test verifies our CoreML infrastructure works perfectly");

    let device = Device::Cpu;

    // Download and load OpenELM model
    let model_path = get_openelm_model_path().expect("Failed to download OpenELM model");

    let config = Config {
        input_names: vec!["input_ids".to_string()],
        output_name: "logits".to_string(),
        max_sequence_length: 128,
        vocab_size: 32000,
        model_type: "OpenELM-450M-Instruct".to_string(),
    };

    let model = CoreMLModel::load_from_file(&model_path, &config)
        .expect("Should be able to load OpenELM model");

    println!("✅ Model loaded successfully");

    // Test the classic "quick brown fox" completion
    // Simulated tokens for: "The quick brown fox jumped over the lazy"
    let test_tokens = vec![1i64, 450, 4996, 17354, 1701, 29916, 12500, 975, 278, 17366];
    let mut padded_tokens = test_tokens.clone();
    padded_tokens.resize(128, 0i64);

    println!("📝 Testing completion for: 'The quick brown fox jumped over the lazy'");
    println!("🔤 Token sequence: {:?}...", &test_tokens);

    let input_tensor = Tensor::from_slice(&padded_tokens, (1, 128), &device)
        .expect("Failed to create input tensor");

    let inputs = vec![&input_tensor];
    let output = model.forward(&inputs).expect("Model inference failed");

    println!("✅ Inference successful, output shape: {:?}", output.dims());

    // Get logits for the last meaningful position (position 9, after "lazy")
    let logits = output
        .i((0, 9))
        .expect("Failed to extract logits at position 9");
    let logits_vec: Vec<f32> = logits.to_vec1().expect("Failed to convert logits to vec");

    // Find top prediction
    let (top_token_id, top_token_score) = logits_vec
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
        .map(|(i, &score)| (i, score))
        .expect("Failed to find top prediction");

    println!("🏆 Top prediction: Token {top_token_id} with score {top_token_score:.3}");

    // Get top 5 predictions for debugging
    let mut indexed_logits: Vec<(usize, f32)> = logits_vec
        .iter()
        .enumerate()
        .map(|(i, &v)| (i, v))
        .collect();
    indexed_logits.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

    println!("📊 Top 5 predictions:");
    for (rank, &(token_id, score)) in indexed_logits.iter().take(5).enumerate() {
        println!("  {}. Token {}: {:.3}", rank + 1, token_id, score);
    }

    // Expected: Token 11203 should be "dog" with high confidence
    assert_eq!(
        top_token_id, 11203,
        "Expected 'dog' (token 11203) as top prediction, got token {top_token_id}"
    );

    assert!(
        top_token_score > 10.0,
        "Expected high confidence (>10.0), got {top_token_score:.3}"
    );

    let logit_range = indexed_logits[0].1 - indexed_logits.last().unwrap().1;
    println!(
        "📈 Logit range: {:.3} (max: {:.3}, min: {:.3})",
        logit_range,
        indexed_logits[0].1,
        indexed_logits.last().unwrap().1
    );

    println!("🎉 BASELINE TEST PASSED!");
    println!("  ✅ OpenELM correctly predicts 'dog' for 'quick brown fox' completion");
    println!("  ✅ CoreML infrastructure working perfectly");
    println!("  ✅ Confidence: {top_token_score:.3}, Range: {logit_range:.3}");
    println!("  ✅ This confirms our implementation is solid");
}

/// Test state parameter validation
#[test]
fn test_stateful_prediction_validation() {
    let config = Config::bert_config("logits", 128, 30522);

    // Test that config creation works
    assert_eq!(config.max_sequence_length, 128);
    assert_eq!(config.vocab_size, 30522);
    assert_eq!(config.output_name, "logits");

    eprintln!("✅ State parameter validation test passed");
}

/// Test device compatibility for stateful predictions
#[test]
fn test_stateful_device_validation() {
    let device = Device::Cpu;

    // Test basic tensor creation and device handling
    let test_data = vec![1.0f32, 2.0, 3.0, 4.0];
    let tensor =
        Tensor::from_slice(&test_data, (2, 2), &device).expect("Failed to create test tensor");

    // Device comparison - just check that tensor was created successfully
    assert_eq!(tensor.dims(), &[2, 2]);

    eprintln!("✅ Device validation test passed");
}