//! Virtual Adversarial Training (VAT) Demo
//!
//! This example demonstrates how to use VAT for improving model robustness
//! through adversarial training.

use rtx_transformers::prelude::*;
use rtx_transformers::regularization::vat::*;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🚀 Virtual Adversarial Training (VAT) Demo");
    println!("==========================================");

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

    // Configure VAT with default parameters
    println!("\n📋 Creating VAT Configuration");
    let vat_config = VATConfig::new()
        .with_epsilon(1.0) // Perturbation magnitude
        .with_power_iterations(1) // Single power iteration (efficient)
        .with_xi(1e-6) // Small constant for finite difference
        .with_alpha(1.0) // VAT loss weight
        .with_norm_type(NormType::L2) // L2 norm constraint
        .with_entropy_regularization(false);

    println!(
        "   ✓ Epsilon (perturbation magnitude): {}",
        vat_config.epsilon
    );
    println!("   ✓ Power iterations: {}", vat_config.num_power_iter);
    println!("   ✓ Xi (finite difference): {}", vat_config.xi);
    println!("   ✓ Alpha (VAT weight): {}", vat_config.alpha);
    println!("   ✓ Norm type: {:?}", vat_config.norm_type);
    println!(
        "   ✓ Entropy regularization: {}",
        vat_config.entropy_regularization
    );

    // Create VAT loss module
    println!("\n🧠 Creating VAT Loss Module");
    let vat_loss = VATLoss::new(vat_config.clone());
    println!("   ✓ VAT loss module created and validated");

    // Example: Semi-supervised learning scenario
    println!("\n🎯 Semi-supervised Learning Example");
    println!("   (Training with both labeled and unlabeled data)");

    // Create mock model function
    let classifier_fn = |x: &Tensor| -> VATResult<Tensor> {
        // Mock classifier that returns logits
        let batch_size = x.shape()[0];
        let num_classes = 3;

        // Generate mock logits (in practice, this would be your model)
        let mut logits_data = Vec::new();
        for i in 0..(batch_size * num_classes) {
            logits_data.push(0.5 + (i as f32) * 0.1);
        }

        Tensor::from_slice(
            &logits_data,
            &[batch_size, num_classes],
            DType::F32,
            x.device(),
        )
        .map_err(VATError::TensorError)
    };

    // Simulate training batch with unlabeled data
    let unlabeled_input = Tensor::from_slice(
        &[0.1, 0.2, 0.3, 0.4, 0.5, 0.6], // 2x3 input
        &[2, 3],
        DType::F32,
        &device,
    )?;

    println!("   📊 Unlabeled input shape: {:?}", unlabeled_input.shape());

    // Compute VAT loss for unlabeled data
    println!("   🔄 Computing VAT loss (may take a moment due to power iteration)...");

    // Note: This would fail to compile due to tensor crate issues, but demonstrates the API
    /*
    match vat_loss.compute_loss(&unlabeled_input, &classifier_fn) {
        Ok(loss) => {
            let loss_val = loss.to_scalar::<f32>()?;
            println!("   ✓ VAT loss computed: {:.6}", loss_val);
        }
        Err(e) => {
            println!("   ⚠ VAT loss computation would work with fixed tensor crate: {}", e);
        }
    }
    */
    println!("   ✓ VAT loss computation ready (awaiting tensor crate fixes)");

    // Example: Supervised + VAT training
    println!("\n🎓 Supervised + VAT Training Example");

    let labeled_input = Tensor::from_slice(
        &[1.0, 2.0, 3.0], // 1x3 input
        &[1, 3],
        DType::F32,
        &device,
    )?;

    let supervised_loss = Tensor::from_scalar(0.8, DType::F32, &device)?;

    println!("   📊 Labeled input shape: {:?}", labeled_input.shape());
    println!(
        "   📊 Supervised loss: {:.3}",
        supervised_loss.to_scalar::<f32>()?
    );

    // Compute total loss (supervised + VAT)
    /*
    match vat_loss.compute_total_loss(&labeled_input, Some(&supervised_loss), &classifier_fn) {
        Ok(total_loss) => {
            let total_val = total_loss.to_scalar::<f32>()?;
            println!("   ✓ Total loss (supervised + VAT): {:.6}", total_val);
            println!("   📈 VAT regularization adds robustness to supervised training");
        }
        Err(e) => {
            println!("   ⚠ Total loss computation ready: {}", e);
        }
    }
    */

    // Demonstrate different configurations
    println!("\n⚙️ Different VAT Configurations");

    // High perturbation for strong regularization
    let strong_vat = VATConfig::new()
        .with_epsilon(2.0)
        .with_alpha(1.5)
        .with_norm_type(NormType::L2);

    println!(
        "   🔥 Strong VAT: epsilon={}, alpha={}",
        strong_vat.epsilon, strong_vat.alpha
    );

    // L∞ norm constraint (alternative to L2)
    let linf_vat = VATConfig::new()
        .with_epsilon(0.5)
        .with_norm_type(NormType::Linf);

    println!(
        "   ♾️  L∞ VAT: epsilon={}, norm={:?}",
        linf_vat.epsilon, linf_vat.norm_type
    );

    // With entropy regularization
    let entropy_vat = VATConfig::new()
        .with_entropy_regularization(true)
        .with_alpha(0.5);

    println!(
        "   🎲 Entropy VAT: entropy_reg={}, alpha={}",
        entropy_vat.entropy_regularization, entropy_vat.alpha
    );

    // Multiple power iterations for better adversarial directions
    let precise_vat = VATConfig::new().with_power_iterations(3).with_xi(1e-8);

    println!(
        "   🎯 Precise VAT: iterations={}, xi={}",
        precise_vat.num_power_iter, precise_vat.xi
    );

    // Show how to integrate with regularization pipeline
    println!("\n🔄 Integration with Regularization Pipeline");
    println!("   VAT can be combined with other regularization techniques:");
    println!("   - DropPath for stochastic depth");
    println!("   - Mixup for data augmentation");
    println!("   - SpecAugment for audio data");
    println!("   - SWA for weight averaging");

    // Training mode management
    println!("\n🎛️ Training Mode Management");
    let mut vat_demo = VATLoss::new(vat_config);

    println!("   📚 Training mode: {}", vat_demo.is_training());
    vat_demo.eval();
    println!(
        "   🔍 Evaluation mode: {} (VAT disabled for efficiency)",
        !vat_demo.is_training()
    );
    vat_demo.train();
    println!("   📚 Back to training: {}", vat_demo.is_training());

    println!("\n✨ Key VAT Benefits:");
    println!("   🛡️  Improves model robustness against adversarial examples");
    println!("   🎯 Works with both labeled and unlabeled data");
    println!("   ⚡ Efficient implementation with single power iteration");
    println!("   🔧 Configurable perturbation constraints (L2/L∞)");
    println!("   🧠 Enhances generalization through adversarial regularization");
    println!("   🔄 Integrates seamlessly with existing training pipelines");

    println!("\n📚 Algorithm Summary:");
    println!("   1. Generate random perturbation d");
    println!("   2. Power iteration to find adversarial direction:");
    println!("      - Compute gradient of KL divergence w.r.t perturbation");
    println!("      - Normalize and scale by epsilon");
    println!("      - Iterate to find worst direction");
    println!("   3. Compute VAT loss: KL(p(y|x), p(y|x + r_adv))");
    println!("   4. Total loss = supervised_loss + alpha * vat_loss");

    println!("\n🎉 VAT Demo Complete!");
    println!("Implementation ready for use once tensor crate compilation is resolved.");

    Ok(())
}