//! Tests for knowledge distillation functionality
//!
//! TDD: Define expected behavior for teacher-student model compression
#![cfg(feature = "disabled_tests")]

use rtx_compress::{
    Result,
    distillation::{DistillationConfig, DistillationLoss, DistillationMethod, KnowledgeDistiller},
};
use rtx_tensor::{Device, Tensor};

#[test]
fn test_distillation_config_creation() {
    let config = DistillationConfig::new(
        DistillationMethod::ResponseBased {
            temperature: 4.0,
            alpha: 0.7,
        },
        DistillationLoss::KullbackLeibler,
    );

    assert!(
        matches!(config.method, DistillationMethod::ResponseBased { temperature, alpha } if temperature == 4.0 && alpha == 0.7)
    );
    assert_eq!(config.loss_function, DistillationLoss::KullbackLeibler);
}

#[test]
fn test_basic_knowledge_distillation() {
    let device = Device::cpu();

    let config = DistillationConfig::new(
        DistillationMethod::ResponseBased {
            temperature: 3.0,
            alpha: 0.9,
        },
        DistillationLoss::KullbackLeibler,
    );

    let distiller = KnowledgeDistiller::new(config).unwrap();

    // Teacher logits (from larger model)
    let teacher_logits = Tensor::from_data(
        vec![2.0, 1.0, 0.1, -1.0], // Confident predictions
        vec![1, 4],
        &device,
    )
    .unwrap();

    // Student logits (from smaller model)
    let student_logits = Tensor::from_data(
        vec![1.5, 0.8, 0.2, -0.5], // Less confident
        vec![1, 4],
        &device,
    )
    .unwrap();

    // Compute distillation loss
    let loss = distiller
        .compute_loss(&teacher_logits, &student_logits)
        .unwrap();

    // Loss should be a scalar
    assert_eq!(loss.shape().dims(), &[1]);

    // Loss should be positive
    let loss_val = loss.to_scalar::<f32>().unwrap();
    assert!(loss_val > 0.0, "Distillation loss should be positive");
}

#[test]
fn test_temperature_scaling() {
    let device = Device::cpu();

    // Test different temperature values
    for temp in [1.0, 3.0, 5.0, 10.0] {
        let config = DistillationConfig::new(
            DistillationMethod::ResponseBased {
                temperature: temp,
                alpha: 1.0,
            },
            DistillationLoss::KullbackLeibler,
        );

        let distiller = KnowledgeDistiller::new(config).unwrap();

        let logits = Tensor::from_data(vec![10.0, 5.0, 1.0, 0.0], vec![1, 4], &device).unwrap();

        // Apply temperature scaling
        let scaled = distiller.apply_temperature(&logits).unwrap();
        let scaled_data = scaled.to_vec().unwrap();

        // Higher temperature should make distribution more uniform
        if temp > 1.0 {
            // Check that values are scaled down
            assert!(scaled_data[0] < 10.0 / temp + 0.1);
        }
    }
}

#[test]
fn test_soft_target_generation() {
    let device = Device::cpu();

    let config = DistillationConfig::default();
    let distiller = KnowledgeDistiller::new(config).unwrap();

    // Teacher logits
    let teacher_logits = Tensor::from_data(
        vec![
            3.0, 1.0, -1.0, -2.0, // Clear preference for first class
            -1.0, 2.0, 1.0, -2.0, // Preference for second class
        ],
        vec![2, 4],
        &device,
    )
    .unwrap();

    // Generate soft targets
    let soft_targets = distiller.generate_soft_targets(&teacher_logits).unwrap();

    // Check shape
    assert_eq!(soft_targets.shape().dims(), teacher_logits.shape().dims());

    // Check that soft targets sum to 1 (probability distribution)
    let soft_data = soft_targets.to_vec().unwrap();
    for batch in 0..2 {
        let start = batch * 4;
        let end = start + 4;
        let sum: f32 = soft_data[start..end].iter().sum();
        assert!((sum - 1.0).abs() < 1e-5, "Sum should be 1.0, got {}", sum);
    }
}

#[test]
fn test_attention_transfer() {
    let device = Device::cpu();

    use rtx_compress::distillation::AttentionTransferConfig;
    let config = DistillationConfig::new(
        DistillationMethod::AttentionTransfer {
            attention_config: AttentionTransferConfig::new(8, 16, false, 1.0),
            transfer_heads_individually: false,
        },
        DistillationLoss::MeanSquaredError,
    );

    let distiller = KnowledgeDistiller::new(config).unwrap();

    // Teacher attention maps [batch, heads, seq_len, seq_len]
    let teacher_attention = Tensor::randn(&[2, 8, 16, 16], &device).unwrap();

    // Student attention maps (fewer heads)
    let student_attention = Tensor::randn(&[2, 4, 16, 16], &device).unwrap();

    // Compute attention transfer loss
    let loss = distiller
        .attention_transfer_loss(&teacher_attention, &student_attention)
        .unwrap();

    // Loss should be scalar
    assert_eq!(loss.shape().dims(), &[1]);

    let loss_val = loss.to_scalar::<f32>().unwrap();
    assert!(
        loss_val >= 0.0,
        "Attention transfer loss should be non-negative"
    );
}

#[test]
fn test_feature_matching() {
    let device = Device::cpu();

    use rtx_compress::distillation::FeatureMatchingConfig;
    let config = DistillationConfig::new(
        DistillationMethod::FeatureBased {
            feature_config: FeatureMatchingConfig::new(
                vec![256, 512, 1024],
                vec![0.3, 0.3, 0.4],
                "mse".to_string(),
            ),
            intermediate_matching: true,
        },
        DistillationLoss::MeanSquaredError,
    );

    let distiller = KnowledgeDistiller::new(config).unwrap();

    // Teacher features from intermediate layers
    let teacher_features = vec![
        Tensor::randn(&[2, 256, 32, 32], &device).unwrap(), // Layer 1
        Tensor::randn(&[2, 512, 16, 16], &device).unwrap(), // Layer 2
        Tensor::randn(&[2, 1024, 8, 8], &device).unwrap(),  // Layer 3
    ];

    // Student features (potentially different dimensions)
    let student_features = vec![
        Tensor::randn(&[2, 128, 32, 32], &device).unwrap(), // Layer 1
        Tensor::randn(&[2, 256, 16, 16], &device).unwrap(), // Layer 2
        Tensor::randn(&[2, 512, 8, 8], &device).unwrap(),   // Layer 3
    ];

    // Compute feature matching loss
    let loss = distiller
        .feature_matching_loss(&teacher_features, &student_features)
        .unwrap();

    assert_eq!(loss.shape().dims(), &[1]);
    let loss_val = loss.to_scalar::<f32>().unwrap();
    assert!(loss_val >= 0.0);
}

#[test]
fn test_distillation_with_labels() {
    let device = Device::cpu();

    let config = DistillationConfig::new(
        DistillationMethod::ResponseBased {
            temperature: 4.0,
            alpha: 0.7,
        },
        DistillationLoss::KullbackLeibler,
    );

    let distiller = KnowledgeDistiller::new(config).unwrap();

    // Logits
    let teacher_logits = Tensor::randn(&[4, 10], &device).unwrap();
    let student_logits = Tensor::randn(&[4, 10], &device).unwrap();

    // Ground truth labels
    let labels = Tensor::from_data(vec![0.0, 1.0, 2.0, 3.0], vec![4], &device).unwrap();

    // Compute combined loss (distillation + hard targets)
    let loss = distiller
        .compute_combined_loss(&teacher_logits, &student_logits, &labels)
        .unwrap();

    assert_eq!(loss.shape().dims(), &[1]);
    let loss_val = loss.to_scalar::<f32>().unwrap();
    assert!(loss_val > 0.0);
}

#[test]
fn test_student_compression_ratio() {
    let device = Device::cpu();

    let config = DistillationConfig::default();
    let distiller = KnowledgeDistiller::new(config).unwrap();

    // Teacher model size (parameters)
    let teacher_params = 100_000_000; // 100M parameters

    // Student model size
    let student_params = 10_000_000; // 10M parameters

    let compression_ratio = distiller.compute_compression_ratio(teacher_params, student_params);

    assert!(
        (compression_ratio - 10.0).abs() < 1e-6,
        "Expected 10.0, got {}",
        compression_ratio
    );
}

#[test]
fn test_gradual_distillation() {
    let device = Device::cpu();

    // Test temperature annealing during training
    let config = DistillationConfig::new_with_adaptive_temperature(
        DistillationMethod::ResponseBased {
            temperature: 4.0,
            alpha: 0.7,
        },
        DistillationLoss::KullbackLeibler,
        true,  // adaptive
        0.001, // rate
    );
    let mut distiller = KnowledgeDistiller::new(config).unwrap();

    // Check temperature at different steps
    let temp_0 = distiller.get_temperature_at_step(0);
    assert!((temp_0 - 4.0).abs() < 1e-6, "Expected 4.0, got {}", temp_0);

    let temp_500 = distiller.get_temperature_at_step(500);
    assert!(temp_500 < 4.0 && temp_500 > 1.0);

    let temp_999 = distiller.get_temperature_at_step(999);
    assert!(
        (temp_999 - 1.0).abs() < 0.1,
        "Expected ~1.0, got {}",
        temp_999
    );
}

#[test]
fn test_distillation_metrics() {
    let device = Device::cpu();

    let config = DistillationConfig::default();
    let distiller = KnowledgeDistiller::new(config).unwrap();

    // Simulate teacher and student predictions
    let teacher_preds = Tensor::from_data(
        vec![0.0, 1.0, 2.0, 1.0], // Teacher predictions
        vec![4],
        &device,
    )
    .unwrap();

    let student_preds = Tensor::from_data(
        vec![0.0, 1.0, 1.0, 1.0], // Student predictions (one mistake)
        vec![4],
        &device,
    )
    .unwrap();

    let labels = Tensor::from_data(
        vec![0.0, 1.0, 2.0, 1.0], // Ground truth
        vec![4],
        &device,
    )
    .unwrap();

    let metrics = distiller
        .compute_metrics(&teacher_preds, &student_preds, &labels)
        .unwrap();

    assert_eq!(metrics.teacher_accuracy, 1.0); // Teacher is perfect
    assert_eq!(metrics.student_accuracy, 0.75); // Student has 75% accuracy
    assert_eq!(metrics.agreement_rate, 0.75); // They agree on 3/4 samples
}