//! Mean Teacher Demo
//!
//! Demonstrates Mean Teacher semi-supervised learning implementation.
//! This demo shows how to train a model using both labeled and unlabeled data
//! with consistency regularization.

use rtx_transformers::prelude::*;
use rtx_transformers::ssl::*;

fn main() -> Result<()> {
    println!("🎓 Mean Teacher Semi-Supervised Learning Demo");
    println!("==============================================");

    let device = Device::cuda(0).unwrap_or(Device::default());

    // Configure Mean Teacher with default hyperparameters
    let config = MeanTeacherConfig::default()
        .with_ema_decay(0.999)
        .with_consistency_weight(100.0)
        .with_consistency_rampup(5)
        .with_noise_level(0.15)
        .with_augmentation_strategy(AugmentationStrategy::Gaussian);

    println!("📋 Configuration:");
    println!("   EMA Decay: {}", config.ema_decay);
    println!("   Consistency Weight: {}", config.consistency_weight);
    println!(
        "   Consistency Ramp-up: {} epochs",
        config.consistency_rampup
    );
    println!("   Noise Level: {}", config.noise_level);
    println!(
        "   Augmentation Strategy: {:?}",
        config.augmentation_strategy
    );

    // Create Mean Teacher trainer
    let input_dim = 784; // MNIST-like input
    let hidden_dim = 256;
    let output_dim = 10; // 10 classes

    let mut trainer = MeanTeacherTrainer::new(config, input_dim, hidden_dim, output_dim, &device)?;

    println!("\n🏗️  Model Architecture:");
    println!("   Input: {} features", input_dim);
    println!("   Hidden: {} units", hidden_dim);
    println!("   Output: {} classes", output_dim);

    // Generate synthetic training data
    let labeled_batch_size = 16;
    let unlabeled_batch_size = 32;

    println!("\n📊 Training Data:");
    println!("   Labeled batch size: {}", labeled_batch_size);
    println!("   Unlabeled batch size: {}", unlabeled_batch_size);

    trainer.train();

    println!("\n🚀 Training Loop:");
    println!("================");

    for epoch in 0..10 {
        trainer.set_epoch(epoch);

        // Generate labeled data (images + labels)
        let labeled_images =
            Tensor::randn(vec![labeled_batch_size, input_dim], DType::F32, &device)?;
        let labels = Tensor::randint(
            0,
            output_dim as i64,
            vec![labeled_batch_size],
            DType::I64,
            &device,
        )?;

        // Generate unlabeled data (images only)
        let unlabeled_images =
            Tensor::randn(vec![unlabeled_batch_size, input_dim], DType::F32, &device)?;

        // Mixed training step (labeled + unlabeled)
        let result = trainer.mixed_step(&labeled_images, &labels, &unlabeled_images, epoch)?;

        println!(
            "Epoch {:2}: Supervised Loss = {:.4}, Consistency Loss = {:.4}, Total Loss = {:.4}",
            epoch, result.supervised_loss, result.consistency_loss, result.total_loss
        );

        // Show ramp-up progress
        if epoch < 6 {
            let rampup = ConsistencyRampUp::new(5, 100.0);
            let weight = rampup.get_weight(epoch);
            println!("           Consistency Weight: {:.2}", weight);
        }
    }

    println!("\n🔍 Component Demonstrations:");
    println!("============================");

    // Demonstrate EMA updater
    println!("\n1. EMA Updater:");
    let mut ema = EMAUpdater::new(0.9, &device)?;
    let teacher_param = Tensor::zeros(vec![5, 5], DType::F32, &device)?;
    let student_param = Tensor::ones(vec![5, 5], DType::F32, &device)?;
    let updated = ema.update(&teacher_param, &student_param)?;
    println!("   Updated parameter shape: {:?}", updated.shape());
    println!("   EMA steps: {}", ema.step());

    // Demonstrate noise augmenter
    println!("\n2. Noise Augmenter:");
    let augmenter = NoiseAugmenter::new(AugmentationStrategy::Gaussian, 0.1, &device)?;
    let clean_input = Tensor::zeros(vec![4, 10], DType::F32, &device)?;
    let noisy_input = augmenter.apply_noise(&clean_input, Some(42))?;
    println!("   Clean input shape: {:?}", clean_input.shape());
    println!("   Noisy input shape: {:?}", noisy_input.shape());

    // Demonstrate consistency loss
    println!("\n3. Consistency Loss:");
    let pred1 = Tensor::randn(vec![8, 10], DType::F32, &device)?;
    let pred2 = Tensor::randn(vec![8, 10], DType::F32, &device)?;
    let loss = compute_consistency_loss(&pred1, &pred2)?;
    println!("   Student prediction shape: {:?}", pred1.shape());
    println!("   Teacher prediction shape: {:?}", pred2.shape());
    println!("   Consistency loss: {:.6}", loss.to_vec::<f32>()?[0]);

    // Demonstrate ramp-up scheduler
    println!("\n4. Consistency Ramp-up:");
    let scheduler = ConsistencyRampUp::new(5, 100.0);
    for epoch in 0..8 {
        let weight = scheduler.get_weight(epoch);
        println!("   Epoch {}: Weight = {:.2}", epoch, weight);
    }

    // Switch to evaluation mode and extract features
    trainer.eval();
    let test_images = Tensor::randn(vec![5, input_dim], DType::F32, &device)?;
    let features = trainer.extract_features(&test_images)?;

    println!("\n📈 Evaluation:");
    println!("==============");
    println!("   Test images shape: {:?}", test_images.shape());
    println!("   Extracted features shape: {:?}", features.shape());
    println!("   Model in training mode: {}", trainer.is_training());

    println!("\n✅ Mean Teacher demo completed successfully!");
    println!("\n📚 Key Concepts Demonstrated:");
    println!("   • Student-teacher architecture with EMA updates");
    println!("   • Consistency regularization on unlabeled data");
    println!("   • Noise augmentation for robustness");
    println!("   • Linear ramp-up scheduling for consistency weight");
    println!("   • Mixed supervised and unsupervised training");

    Ok(())
}