//! SSL Demo
//!
//! Demonstration of BYOL and MAE self-supervised learning frameworks.

use rtx_transformers::prelude::*;

fn main() -> Result<()> {
    println!("🦀 RTX Transformers SSL Demo");
    println!("==============================");

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

    // BYOL Demo
    println!("\n🔄 BYOL (Bootstrap Your Own Latent) Demo");
    byol_demo(&device)?;

    // MAE Demo
    println!("\n🎭 MAE (Masked Autoencoder) Demo");
    mae_demo(&device)?;

    // Unified SSL Training Demo
    println!("\n🚀 Unified SSL Training Demo");
    unified_ssl_demo(&device)?;

    println!("\n✅ SSL Demo completed successfully!");

    Ok(())
}

fn byol_demo(device: &Device) -> Result<()> {
    println!("Creating BYOL configuration...");

    let byol_config = BYOLConfig::new(256, 128, 64)
        .with_tau(0.996)
        .with_temperature(0.1);

    let ssl_config = SSLTrainingConfig {
        method: SSLMethod::BYOL(byol_config),
        learning_rate: 1e-3,
        batch_size: 8,
        epochs: 5,
        ..Default::default()
    };

    println!("Creating vision backbones...");
    let online_backbone = VisionBackbone::resnet50(device);
    let target_backbone = VisionBackbone::resnet50(device);

    println!("Initializing BYOL trainer...");
    let mut trainer = SSLTrainer::new(online_backbone, Some(target_backbone), ssl_config, device)?;

    println!("Running BYOL training steps...");
    for epoch in 0..3 {
        trainer.set_epoch(epoch);
        trainer.train();

        // Create synthetic batch
        let images = Tensor::randn(vec![8, 3, 224, 224], DType::F32, device)?;
        let metrics = trainer.train_step(&images, Some(epoch as u64))?;

        println!("  Epoch {}: Loss = {:.4}", epoch, metrics.train_loss);

        if let MethodMetrics::BYOL { similarity } = metrics.method_metrics {
            println!("           Similarity = {:.4}", similarity);
        }
    }

    println!("Testing feature extraction...");
    trainer.eval();
    let test_images = Tensor::randn(vec![4, 3, 224, 224], DType::F32, device)?;
    let features = trainer.extract_features(&test_images)?;
    println!("  Extracted features shape: {:?}", features.shape());

    Ok(())
}

fn mae_demo(device: &Device) -> Result<()> {
    println!("Creating MAE configuration...");

    let mae_config = MAEConfig::new(
        384,  // encoder_dim
        256,  // decoder_dim
        6,    // encoder_layers
        4,    // decoder_layers
        6,    // num_heads
        16,   // patch_size
        0.75, // mask_ratio
    );

    let ssl_config = SSLTrainingConfig {
        method: SSLMethod::MAE(mae_config),
        learning_rate: 1e-3,
        batch_size: 16,
        epochs: 10,
        ..Default::default()
    };

    println!("Initializing MAE trainer...");
    let backbone = VisionBackbone::vit_base(device);
    let mut trainer = SSLTrainer::new(
        backbone, None, // MAE doesn't need separate target backbone
        ssl_config, device,
    )?;

    println!("Running MAE training steps...");
    for epoch in 0..3 {
        trainer.set_epoch(epoch);
        trainer.train();

        // Create synthetic batch
        let images = Tensor::randn(vec![16, 3, 224, 224], DType::F32, device)?;
        let metrics = trainer.train_step(&images, Some(epoch as u64))?;

        println!("  Epoch {}: Loss = {:.4}", epoch, metrics.train_loss);

        if let MethodMetrics::MAE {
            reconstruction_acc,
            mask_ratio,
        } = metrics.method_metrics
        {
            println!(
                "           Reconstruction Acc = {:.3}, Mask Ratio = {:.2}",
                reconstruction_acc, mask_ratio
            );
        }
    }

    println!("Testing feature extraction...");
    trainer.eval();
    let test_images = Tensor::randn(vec![4, 3, 224, 224], DType::F32, device)?;
    let features = trainer.extract_features(&test_images)?;
    println!("  Extracted features shape: {:?}", features.shape());

    Ok(())
}

fn unified_ssl_demo(device: &Device) -> Result<()> {
    println!("Demonstrating data augmentation pipeline...");

    let aug_config = AugmentationConfig {
        flip_prob: 0.5,
        color_jitter: 0.4,
        crop_scale: (0.2, 1.0),
        blur_prob: 0.1,
        normalize_mean: vec![0.485, 0.456, 0.406],
        normalize_std: vec![0.229, 0.224, 0.225],
    };

    let pipeline = AugmentationPipeline::new(aug_config, device);
    let image = Tensor::randn(vec![1, 3, 224, 224], DType::F32, device)?;

    let (view1, view2) = pipeline.augment_pair(&image, Some(42))?;
    println!("  Original shape: {:?}", image.shape());
    println!("  Augmented view1 shape: {:?}", view1.shape());
    println!("  Augmented view2 shape: {:?}", view2.shape());

    println!("Demonstrating SSL evaluation...");
    let evaluator = SSLEvaluator::new(device);

    // Create simple BYOL trainer for evaluation
    let byol_config = BYOLConfig::new(128, 64, 32);
    let eval_ssl_config = SSLTrainingConfig {
        method: SSLMethod::BYOL(byol_config),
        ..Default::default()
    };

    let backbone = VisionBackbone::new(128, device);
    let mut eval_trainer =
        SSLTrainer::new(backbone.clone(), Some(backbone), eval_ssl_config, device)?;

    // Create synthetic evaluation data
    let train_images = Tensor::randn(vec![50, 3, 32, 32], DType::F32, device)?;
    let train_labels = Tensor::randint(0, 10, vec![50], DType::F32, device)?;
    let val_images = Tensor::randn(vec![20, 3, 32, 32], DType::F32, device)?;
    let val_labels = Tensor::randint(0, 10, vec![20], DType::F32, device)?;

    println!("Running linear probe evaluation...");
    let linear_acc = evaluator.linear_probe(
        &mut eval_trainer,
        &train_images,
        &train_labels,
        &val_images,
        &val_labels,
    )?;
    println!("  Linear probing accuracy: {:.3}", linear_acc);

    println!("Running k-NN evaluation...");
    let knn_acc = evaluator.knn_evaluation(
        &mut eval_trainer,
        &train_images,
        &train_labels,
        &val_images,
        &val_labels,
        5,
    )?;
    println!("  k-NN accuracy (k=5): {:.3}", knn_acc);

    println!("Testing different vision backbones...");
    let resnet50 = VisionBackbone::resnet50(device);
    let vit_base = VisionBackbone::vit_base(device);

    let test_input = Tensor::randn(vec![2, 3, 224, 224], DType::F32, device)?;

    let resnet_features = resnet50.forward(&test_input)?;
    let vit_features = vit_base.forward(&test_input)?;

    println!("  ResNet-50 features: {:?}", resnet_features.shape());
    println!("  ViT-Base features: {:?}", vit_features.shape());

    Ok(())
}