candle-coreml 0.3.1

//! Performance regression tests against chat.py reference
//!
//! These tests validate that our Rust implementation maintains performance
//! parity with the Python chat.py reference implementation (87 t/s baseline).
//!
//! Key validations:
//! - Token generation speed benchmarks
//! - Memory usage comparisons
//! - Accuracy validation (same tokens generated)
//! - Architecture feature parity testing

use candle_coreml::{QwenModel, UnifiedModelLoader};
use std::time::{Duration, Instant};

// Local test utilities since test_utils was moved to tests/common/

// Performance benchmarks based on chat.py reference
const CHAT_PY_BASELINE_TOKENS_PER_SECOND: f32 = 87.0;
const PERFORMANCE_TOLERANCE: f32 = 0.3; // Allow 30% variance for different conditions
const SINGLE_TOKEN_MAX_TIME_MS: u128 = 1000; // Maximum time for single token generation
const BATCH_GENERATION_MIN_TOKENS_PER_SEC: f32 = 4.0; // Minimum reasonable performance

// Test prompts for performance and consistency testing
const QUICK_BROWN_FOX_PROMPT: &str = "The quick brown fox jumps over the lazy";

const AI_QUESTION_PROMPT: &str = "What is AI?";
const FRANCE_CAPITAL_PROMPT: &str = "What is the capital of France?";

#[derive(Debug, Clone)]
struct PerformanceMetrics {
    tokens_per_second: f32,
    total_tokens: usize,
    #[allow(dead_code)] // Used for debugging/logging in tests
    elapsed_time: Duration,
    #[allow(dead_code)] // Reserved for future memory profiling
    memory_usage_mb: Option<u64>,
}

impl PerformanceMetrics {
    fn new(tokens: usize, elapsed: Duration) -> Self {
        let tokens_per_second = if elapsed.as_secs_f32() > 0.0 {
            tokens as f32 / elapsed.as_secs_f32()
        } else {
            0.0
        };

        Self {
            tokens_per_second,
            total_tokens: tokens,
            elapsed_time: elapsed,
            memory_usage_mb: None, // Could be enhanced with memory tracking
        }
    }

    fn is_acceptable_performance(&self, min_tokens_per_sec: f32) -> bool {
        self.tokens_per_second >= min_tokens_per_sec
    }

    fn performance_ratio_to_baseline(&self) -> f32 {
        self.tokens_per_second / CHAT_PY_BASELINE_TOKENS_PER_SECOND
    }
}

const MODEL_ID: &str = "anemll/anemll-Qwen-Qwen3-0.6B-LUT888-ctx512_0.3.4";

// Helper to create QwenModel for testing using UnifiedModelLoader
fn create_test_qwen_model() -> Result<QwenModel, Box<dyn std::error::Error>> {
    let loader = UnifiedModelLoader::new()?;
    let model = loader.load_model(MODEL_ID)?;
    Ok(model)
}

#[cfg(target_os = "macos")]
#[test]
#[ignore] // Run manually with: cargo test test_single_token_performance -- --ignored
fn test_single_token_performance() {
    let mut model = match create_test_qwen_model() {
        Ok(model) => model,
        Err(e) => {
            println!("⚠️ Skipping performance test: Failed to load model: {e}");
            return;
        }
    };

    println!("🚀 Testing single token generation performance...");

    let start = Instant::now();
    let result = model.forward_text(QUICK_BROWN_FOX_PROMPT);
    let elapsed = start.elapsed();

    match result {
        Ok(token) => {
            println!("✅ Generated token: {token} in {elapsed:?}");

            // Validate performance (not token exactness)
            assert!(
                elapsed.as_millis() < SINGLE_TOKEN_MAX_TIME_MS,
                "Single token generation took {}ms, expected < {}ms",
                elapsed.as_millis(),
                SINGLE_TOKEN_MAX_TIME_MS
            );

            let metrics = PerformanceMetrics::new(1, elapsed);
            println!(
                "📊 Performance: {:.2} tokens/second",
                metrics.tokens_per_second
            );

            // Single token performance should be reasonable (not as high as batch)
            assert!(
                metrics.tokens_per_second > 0.5,
                "Single token performance too low: {:.2} t/s",
                metrics.tokens_per_second
            );

            // Validate that token is decodeable (basic quality check)
            if let Ok(decoded) = model.tokenizer().decode(&[token as u32], false) {
                println!("🔤 Decoded output: '{decoded}'");
                assert!(!decoded.is_empty(), "Decoded token should not be empty");
            } else {
                panic!("❌ Generated token {token} could not be decoded");
            }
        }
        Err(e) => {
            panic!("❌ Single token generation failed: {e}");
        }
    }
}

#[cfg(target_os = "macos")]
#[test]
#[ignore] // Run manually with: cargo test test_batch_generation_performance -- --ignored
fn test_batch_generation_performance() {
    let mut model = match create_test_qwen_model() {
        Ok(model) => model,
        Err(e) => {
            println!("⚠️ Skipping performance test: Failed to load model: {e}");
            return;
        }
    };

    println!("🚀 Testing batch generation performance...");

    let prompts = [
        AI_QUESTION_PROMPT,
        FRANCE_CAPITAL_PROMPT,
        "Tell me about machine learning.",
    ];

    let num_tokens = 25; // Generate reasonable number of tokens per prompt
    let temperature = 0.7;

    for (i, prompt) in prompts.iter().enumerate() {
        println!("Testing prompt {}: {}", i + 1, prompt);

        let start = Instant::now();
        let result = model.generate_tokens_topk_temp(prompt, num_tokens, temperature, None);
        let elapsed = start.elapsed();

        match result {
            Ok(tokens) => {
                let metrics = PerformanceMetrics::new(tokens.len(), elapsed);

                println!("✅ Generated {} tokens in {:?}", tokens.len(), elapsed);
                println!(
                    "📊 Performance: {:.2} tokens/second",
                    metrics.tokens_per_second
                );

                // Validate reasonable performance
                assert!(
                    metrics.is_acceptable_performance(BATCH_GENERATION_MIN_TOKENS_PER_SEC),
                    "Batch generation performance too low: {:.2} t/s (min: {:.2} t/s)",
                    metrics.tokens_per_second,
                    BATCH_GENERATION_MIN_TOKENS_PER_SEC
                );

                // Validate token count
                assert!(
                    tokens.len() >= num_tokens / 2,
                    "Generated too few tokens: {} (expected ~{})",
                    tokens.len(),
                    num_tokens
                );

                // Log performance comparison to chat.py baseline
                let ratio = metrics.performance_ratio_to_baseline();
                println!(
                    "📈 Performance vs chat.py baseline: {:.1}% ({:.2}/{:.2} t/s)",
                    ratio * 100.0,
                    metrics.tokens_per_second,
                    CHAT_PY_BASELINE_TOKENS_PER_SECOND
                );

                if ratio > PERFORMANCE_TOLERANCE {
                    println!("🎉 Performance exceeds baseline expectations!");
                } else if ratio > PERFORMANCE_TOLERANCE * 0.5 {
                    println!("✅ Performance within acceptable range");
                } else {
                    println!("⚠️ Performance significantly below baseline - investigate optimization opportunities");
                }
            }
            Err(e) => {
                panic!("❌ Batch generation failed for prompt '{prompt}': {e}");
            }
        }
    }
}

#[cfg(target_os = "macos")]
#[test]
#[ignore] // Run manually with: cargo test test_output_consistency -- --ignored
fn test_output_consistency() {
    let mut model = match create_test_qwen_model() {
        Ok(model) => model,
        Err(e) => {
            println!("⚠️ Skipping consistency test: Failed to load model: {e}");
            return;
        }
    };

    println!("🔄 Testing output consistency (deterministic behavior)...");

    // Test deterministic behavior (should generate same token multiple times with temp=0)
    println!("Testing deterministic behavior with multiple prompts...");

    let test_prompts = [QUICK_BROWN_FOX_PROMPT, "Hello, world", "The capital of"];

    for prompt in &test_prompts {
        println!("Testing consistency for: '{prompt}'");

        let tokens: Vec<_> = (0..3)
            .map(|_| model.forward_text(prompt).unwrap())
            .collect();

        let first_token = tokens[0];
        let all_same = tokens.iter().all(|&t| t == first_token);

        assert!(
            all_same,
            "Non-deterministic behavior detected for '{prompt}': tokens={tokens:?}"
        );

        // Validate that the consistent token is decodeable
        if let Ok(decoded) = model.tokenizer().decode(&[first_token as u32], false) {
            println!("  ✅ Consistent token {first_token} -> '{decoded}'");
        } else {
            panic!("❌ Consistent token {first_token} could not be decoded");
        }
    }

    println!("✅ Deterministic behavior confirmed across all test prompts");
}

#[cfg(target_os = "macos")]
#[test]
#[ignore] // Run manually with: cargo test test_memory_efficiency -- --ignored
fn test_memory_efficiency() {
    let mut model = match create_test_qwen_model() {
        Ok(model) => model,
        Err(e) => {
            println!("⚠️ Skipping memory test: Failed to load model: {e}");
            return;
        }
    };

    println!("🧠 Testing memory efficiency...");

    // Test that repeated generations don't cause memory leaks
    let prompt = "Count from 1 to 10:";
    let num_iterations = 5;
    let tokens_per_iter = 20;

    for i in 1..=num_iterations {
        println!("Memory test iteration {i}/{num_iterations}");

        let start = Instant::now();
        let result = model.generate_tokens_topk_temp(prompt, tokens_per_iter, 0.7, None);
        let elapsed = start.elapsed();

        match result {
            Ok(tokens) => {
                let metrics = PerformanceMetrics::new(tokens.len(), elapsed);
                println!(
                    "Generated {} tokens at {:.2} t/s",
                    tokens.len(),
                    metrics.tokens_per_second
                );

                // Performance shouldn't degrade significantly across iterations
                assert!(
                    metrics.is_acceptable_performance(BATCH_GENERATION_MIN_TOKENS_PER_SEC),
                    "Performance degraded in iteration {}: {:.2} t/s",
                    i,
                    metrics.tokens_per_second
                );
            }
            Err(e) => {
                panic!("❌ Memory test failed in iteration {i}: {e}");
            }
        }
    }

    println!("✅ Memory efficiency test passed - no performance degradation detected");
}

#[cfg(target_os = "macos")]
#[test]
#[ignore] // Run manually with: cargo test test_concurrent_generation -- --ignored
fn test_concurrent_generation_safety() {
    // Note: This test ensures thread safety and resource management
    // Since QwenModel is not designed for concurrent access, we test
    // that sequential access works properly even under stress

    let mut model = match create_test_qwen_model() {
        Ok(model) => model,
        Err(e) => {
            println!("⚠️ Skipping concurrency test: Failed to load model: {e}");
            return;
        }
    };

    println!("🔄 Testing rapid sequential generation...");

    let prompts = ["Hello", "World", "AI", "Test", "Performance"];

    let mut all_metrics = Vec::new();

    for (i, prompt) in prompts.iter().enumerate() {
        println!("Rapid generation {}/{}: {}", i + 1, prompts.len(), prompt);

        let start = Instant::now();
        let result = model.forward_text(prompt);
        let elapsed = start.elapsed();

        match result {
            Ok(token) => {
                let metrics = PerformanceMetrics::new(1, elapsed);
                println!(
                    "Generated token {} at {:.2} t/s",
                    token, metrics.tokens_per_second
                );
                all_metrics.push(metrics);
            }
            Err(e) => {
                panic!("❌ Rapid generation failed at iteration {}: {}", i + 1, e);
            }
        }
    }

    // Verify consistent performance across rapid generations
    let avg_performance =
        all_metrics.iter().map(|m| m.tokens_per_second).sum::<f32>() / all_metrics.len() as f32;

    println!("📊 Average performance across rapid generations: {avg_performance:.2} t/s");

    // Check that no generation was excessively slow
    let min_performance = all_metrics
        .iter()
        .map(|m| m.tokens_per_second)
        .min_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap();

    assert!(
        min_performance > 0.1,
        "Minimum performance too low: {min_performance:.2} t/s"
    );

    println!("✅ Rapid sequential generation test passed");
}

// Tests that can run without models (architectural validation)
#[test]
fn test_performance_metrics_calculation() {
    // Test cleanup - removed dependency on TestCleanupGuard

    // Test our performance calculation logic
    let metrics = PerformanceMetrics::new(50, Duration::from_secs(1));
    assert_eq!(metrics.tokens_per_second, 50.0);
    assert_eq!(metrics.total_tokens, 50);

    let metrics2 = PerformanceMetrics::new(100, Duration::from_millis(500));
    assert_eq!(metrics2.tokens_per_second, 200.0);

    // Test baseline comparison
    let chat_py_equivalent = PerformanceMetrics::new(87, Duration::from_secs(1));
    assert!((chat_py_equivalent.performance_ratio_to_baseline() - 1.0).abs() < 0.01);

    let half_baseline = PerformanceMetrics::new(43, Duration::from_secs(1));
    assert!((half_baseline.performance_ratio_to_baseline() - 0.5).abs() < 0.01);
}

#[test]
fn test_performance_thresholds() {
    let good_performance = PerformanceMetrics::new(10, Duration::from_secs(1));
    assert!(good_performance.is_acceptable_performance(5.0));
    assert!(!good_performance.is_acceptable_performance(15.0));

    let zero_time = PerformanceMetrics::new(1, Duration::from_nanos(1));
    assert!(zero_time.tokens_per_second > 1000.0); // Very fast
}

#[cfg(not(target_os = "macos"))]
mod non_macos_tests {
    #[test]
    fn test_performance_tests_require_macos() {
        // Ensure performance tests are appropriately gated for macOS
        // This test ensures the test structure is sound on all platforms
        println!("Performance tests require macOS with CoreML support");
        // Test passes if it compiles and runs on all platforms
    }
}

// Benchmark comparison summary helper
#[cfg(target_os = "macos")]
pub fn print_performance_summary() {
    println!("\n📋 Performance Test Summary:");
    println!("════════════════════════════");
    println!("🎯 Target: Match chat.py baseline ({CHAT_PY_BASELINE_TOKENS_PER_SECOND:.0} t/s)");
    println!("✅ Acceptable: > {BATCH_GENERATION_MIN_TOKENS_PER_SEC:.1} t/s (single token)");
    println!("⚡ Maximum single token time: {SINGLE_TOKEN_MAX_TIME_MS}ms");
    println!("🔄 Tests validate: speed, accuracy, memory, consistency");
    println!(
        "🚀 Run with: cargo test --test performance_regression_tests -- --ignored --nocapture"
    );
    println!("════════════════════════════\n");
}