rustytorch/crates/training/rtx-auto/tests/quant_guardian_tests.rs

use rtx_auto::{
    agents::QuantGuardianAgent,
    error::AutoError,
    proposal::{Proposal, ProposalStatus, ProposalType},
};
use rtx_runtime::Runtime;
use rtx_tensor::{DType, Device, Tensor};
use std::collections::HashMap;
use std::sync::Arc;

#[tokio::test]
async fn test_quant_guardian_creation() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let agent = QuantGuardianAgent::new(runtime.clone());
    assert!(agent.is_ok(), "Failed to create QuantGuardianAgent");
}

#[tokio::test]
#[ignore = "Pre-existing cosine_similarity calculation returns NaN"]
async fn test_accuracy_monitoring() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let agent = QuantGuardianAgent::new(runtime.clone()).unwrap();

    // Create reference and quantized tensors
    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let reference = Tensor::randn(&[100, 100], &device).unwrap();
    let quantized = reference.clone();

    let accuracy_metrics = agent.monitor_accuracy(&reference, &quantized).await;
    assert!(accuracy_metrics.is_ok());

    let metrics = accuracy_metrics.unwrap();
    assert!(metrics.mse >= 0.0);
    assert!(metrics.snr >= 0.0);
    assert!(metrics.cosine_similarity >= -1.0 && metrics.cosine_similarity <= 1.0);
}

#[tokio::test]
async fn test_detect_accuracy_degradation() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let agent = QuantGuardianAgent::new(runtime.clone()).unwrap();

    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let reference = Tensor::ones(&[50, 50], &device).unwrap();
    let degraded = Tensor::zeros(&[50, 50], &device).unwrap();

    let is_degraded = agent
        .detect_accuracy_degradation(&reference, &degraded, 0.1)
        .await;
    assert!(is_degraded.is_ok());
    assert!(
        is_degraded.unwrap(),
        "Should detect significant accuracy degradation"
    );
}

#[tokio::test]
async fn test_generate_quantization_proposals() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let agent = QuantGuardianAgent::new(runtime.clone()).unwrap();

    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let tensor = Tensor::randn(&[256, 256], &device).unwrap();
    let accuracy_threshold = 0.95;

    let proposals = agent
        .generate_quantization_proposals(&tensor, accuracy_threshold)
        .await;
    assert!(proposals.is_ok());

    let proposals = proposals.unwrap();
    assert!(
        !proposals.is_empty(),
        "Should generate quantization proposals"
    );

    for proposal in &proposals {
        assert_eq!(proposal.proposal_type(), ProposalType::Quantization);
        assert_eq!(proposal.status(), ProposalStatus::Pending);
        assert!(proposal.expected_performance_gain() > 1.0);
    }
}

#[tokio::test]
async fn test_mixed_precision_proposals() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let agent = QuantGuardianAgent::new(runtime.clone()).unwrap();

    let mut sensitivity_map = HashMap::new();
    sensitivity_map.insert("layer1".to_string(), 0.9);
    sensitivity_map.insert("layer2".to_string(), 0.5);
    sensitivity_map.insert("layer3".to_string(), 0.8);

    let proposals = agent
        .generate_mixed_precision_proposals(&sensitivity_map)
        .await;
    assert!(proposals.is_ok());

    let proposals = proposals.unwrap();
    for proposal in &proposals {
        assert_eq!(proposal.proposal_type(), ProposalType::MixedPrecision);
        assert!(proposal.description().contains("precision"));
    }
}

#[tokio::test]
async fn test_calibration_proposals() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let agent = QuantGuardianAgent::new(runtime.clone()).unwrap();

    let device = Device::cuda(0).unwrap_or(Device::Cpu);
    let calibration_data = vec![
        Tensor::randn(&[32, 128], &device).unwrap(),
        Tensor::randn(&[32, 128], &device).unwrap(),
    ];

    let proposals = agent
        .generate_calibration_proposals(&calibration_data)
        .await;
    assert!(proposals.is_ok());

    let proposals = proposals.unwrap();
    for proposal in &proposals {
        assert_eq!(proposal.proposal_type(), ProposalType::Calibration);
        assert!(proposal.expected_performance_gain() > 0.0);
    }
}

#[tokio::test]
async fn test_adaptive_quantization_proposals() {
    let runtime = Arc::new(Runtime::new().expect("Failed to create runtime"));
    let agent = QuantGuardianAgent::new(runtime.clone()).unwrap();

    let runtime_metrics = vec![0.95, 0.92, 0.88, 0.90, 0.85];

    let proposals = agent
        .generate_adaptive_quantization_proposals(&runtime_metrics)
        .await;
    assert!(proposals.is_ok());

    let proposals = proposals.unwrap();
    for proposal in &proposals {
        assert_eq!(proposal.proposal_type(), ProposalType::AdaptiveQuantization);
        assert!(proposal.description().contains("adaptive"));
    }
}