//! Barlow Twins Demo
//!
//! Demonstrates Barlow Twins self-supervised learning with various configurations
//! and backbone architectures.

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

fn main() -> Result<()> {
    println!("🔥 Barlow Twins Self-Supervised Learning Demo");
    println!("===============================================");

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

    // Demo 1: Basic Barlow Twins training
    demo_basic_training(&device)?;

    // Demo 2: Custom projector configuration
    demo_custom_projector(&device)?;

    // Demo 3: Different backbone architectures
    demo_different_backbones(&device)?;

    // Demo 4: Integration with SSL framework
    demo_ssl_integration(&device)?;

    // Demo 5: Loss analysis and visualization
    demo_loss_analysis(&device)?;

    println!("✅ All Barlow Twins demos completed successfully!");

    Ok(())
}

/// Demo 1: Basic Barlow Twins training
fn demo_basic_training(device: &Device) -> Result<()> {
    println!("\n📚 Demo 1: Basic Barlow Twins Training");
    println!("--------------------------------------");

    // Create Barlow Twins config with default hyperparameters
    let config = BarlowTwinsConfig::default().with_backbone_dim(2048);

    println!("Config: {:?}", config);

    // Create backbone and trainer
    let backbone = Box::new(VisionBackbone::resnet50(device));
    let mut trainer = BarlowTwinsTrainer::new(backbone, config, device)?;

    // Create fake batch of images (batch_size=32, channels=3, height=224, width=224)
    let batch_size = 32;
    let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, device)?;

    println!("Training on batch shape: {:?}", images.shape());

    // Training step
    trainer.train();
    let result = trainer.train_step(&images, Some(42))?;

    let loss_value = result.loss.to_vec::<f32>()?[0];
    println!("Loss: {:.6}", loss_value);
    println!("Invariance loss: {:.6}", result.invariance_loss);
    println!("Redundancy loss: {:.6}", result.redundancy_loss);

    // Feature extraction
    trainer.eval();
    let features = trainer.extract_features(&images)?;
    println!("Extracted features shape: {:?}", features.shape());

    Ok(())
}

/// Demo 2: Custom projector configuration
fn demo_custom_projector(device: &Device) -> Result<()> {
    println!("\n🔧 Demo 2: Custom Projector Configuration");
    println!("------------------------------------------");

    // Custom projector with different architecture
    let custom_config = BarlowTwinsConfig::new(2048, vec![4096, 2048, 1024])
        .with_lambda_coeff(0.01) // Higher regularization
        .with_batch_norm(false) // Without batch norm
        .with_scale_loss(0.1); // Different scaling

    println!("Custom config: {:?}", custom_config);

    let backbone = Box::new(VisionBackbone::resnet50(device));
    let mut trainer = BarlowTwinsTrainer::new(backbone, custom_config, device)?;

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

    trainer.train();
    let result = trainer.train_step(&images, Some(123))?;

    let loss_value = result.loss.to_vec::<f32>()?[0];
    println!("Custom config loss: {:.6}", loss_value);

    // Analyze cross-correlation matrix
    let cross_corr = result.cross_correlation;
    let diagonal = extract_diagonal(&cross_corr)?;
    let diagonal_data = diagonal.to_vec::<f32>()?;
    let diagonal_mean = diagonal_data.iter().sum::<f32>() / diagonal_data.len() as f32;

    println!("Cross-correlation matrix shape: {:?}", cross_corr.shape());
    println!("Diagonal mean: {:.4} (should approach 1.0)", diagonal_mean);

    Ok(())
}

/// Demo 3: Different backbone architectures
fn demo_different_backbones(device: &Device) -> Result<()> {
    println!("\n🏗️  Demo 3: Different Backbone Architectures");
    println!("---------------------------------------------");

    let batch_size = 16;
    let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, device)?;

    // ResNet-50 style (2048-dim features)
    println!("Testing ResNet-50 style backbone:");
    let resnet_config = BarlowTwinsConfig::default().with_backbone_dim(2048);
    let resnet_backbone = Box::new(VisionBackbone::resnet50(device));
    let mut resnet_trainer = BarlowTwinsTrainer::new(resnet_backbone, resnet_config, device)?;

    resnet_trainer.train();
    let resnet_result = resnet_trainer.train_step(&images, Some(42))?;
    let resnet_loss = resnet_result.loss.to_vec::<f32>()?[0];
    println!("  ResNet loss: {:.6}", resnet_loss);

    // ViT-Base style (768-dim features)
    println!("Testing ViT-Base style backbone:");
    let vit_config = BarlowTwinsConfig::new(768, vec![3072, 3072, 3072]).with_lambda_coeff(0.005);
    let vit_backbone = Box::new(VisionBackbone::vit_base(device));
    let mut vit_trainer = BarlowTwinsTrainer::new(vit_backbone, vit_config, device)?;

    vit_trainer.train();
    let vit_result = vit_trainer.train_step(&images, Some(42))?;
    let vit_loss = vit_result.loss.to_vec::<f32>()?[0];
    println!("  ViT loss: {:.6}", vit_loss);

    // Feature extraction comparison
    resnet_trainer.eval();
    vit_trainer.eval();

    let resnet_features = resnet_trainer.extract_features(&images)?;
    let vit_features = vit_trainer.extract_features(&images)?;

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

    Ok(())
}

/// Demo 4: Integration with SSL framework
fn demo_ssl_integration(device: &Device) -> Result<()> {
    println!("\n🔄 Demo 4: SSL Framework Integration");
    println!("------------------------------------");

    // Create Barlow Twins SSL config
    let barlow_config = BarlowTwinsConfig::new(2048, vec![8192, 8192, 8192])
        .with_lambda_coeff(0.005)
        .with_batch_norm(true);

    let ssl_config = SSLTrainingConfig {
        method: SSLMethod::BarlowTwins(barlow_config),
        learning_rate: 1e-3,
        batch_size: 32,
        epochs: 10,
        warmup_epochs: 2,
        weight_decay: 1e-4,
        eval_freq: 5,
        augmentation: AugmentationConfig {
            flip_prob: 0.5,
            color_jitter: 0.4,
            crop_scale: (0.08, 1.0),
            blur_prob: 0.1,
            normalize_mean: vec![0.485, 0.456, 0.406],
            normalize_std: vec![0.229, 0.224, 0.225],
        },
    };

    println!("SSL config created for Barlow Twins");

    // Create SSL trainer
    let backbone = VisionBackbone::resnet50(device);
    let mut trainer = SSLTrainer::new(
        backbone, None, // Barlow Twins doesn't need separate target backbone
        ssl_config, device,
    )?;

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

    // Training loop simulation
    for epoch in 0..5 {
        trainer.set_epoch(epoch);
        trainer.train();

        let metrics = trainer.train_step(&images, Some(epoch as u64))?;

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

        if let MethodMetrics::BarlowTwins {
            invariance_loss,
            redundancy_loss,
            diagonal_mean,
            off_diagonal_rms,
        } = metrics.method_metrics
        {
            println!(
                "  Invariance: {:.4}, Redundancy: {:.4}",
                invariance_loss, redundancy_loss
            );
            println!(
                "  Diagonal mean: {:.4}, Off-diag RMS: {:.4}",
                diagonal_mean, off_diagonal_rms
            );
        }
    }

    // Evaluation
    trainer.eval();
    let features = trainer.extract_features(&images)?;
    println!("Final features shape: {:?}", features.shape());

    Ok(())
}

/// Demo 5: Loss analysis and visualization
fn demo_loss_analysis(device: &Device) -> Result<()> {
    println!("\n📊 Demo 5: Loss Analysis");
    println!("------------------------");

    let batch_size = 64;
    let feature_dim = 1024;
    let lambda_coeff = 0.005;
    let scale_loss = 1.0 / 32.0;

    // Test with random embeddings
    println!("Analyzing loss with random embeddings:");
    let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, device)?;
    let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, device)?;

    let loss_result = compute_barlow_twins_loss(&y1, &y2, lambda_coeff, scale_loss)?;
    let loss_value = loss_result.loss.to_vec::<f32>()?[0];

    println!("  Total loss: {:.6}", loss_value);
    println!("  Invariance loss: {:.6}", loss_result.invariance_loss);
    println!("  Redundancy loss: {:.6}", loss_result.redundancy_loss);

    // Test with identical embeddings (should give low loss)
    println!("\nAnalyzing loss with identical embeddings:");
    let y_identical = normalize_features(&y1)?;
    let identical_loss =
        compute_barlow_twins_loss(&y_identical, &y_identical, lambda_coeff, scale_loss)?;
    let identical_loss_value = identical_loss.loss.to_vec::<f32>()?[0];

    println!("  Total loss: {:.6}", identical_loss_value);
    println!("  Invariance loss: {:.6}", identical_loss.invariance_loss);
    println!("  Redundancy loss: {:.6}", identical_loss.redundancy_loss);

    // Test symmetry
    println!("\nTesting loss symmetry:");
    let loss_12 = compute_barlow_twins_loss(&y1, &y2, lambda_coeff, scale_loss)?;
    let loss_21 = compute_barlow_twins_loss(&y2, &y1, lambda_coeff, scale_loss)?;

    let loss_12_val = loss_12.loss.to_vec::<f32>()?[0];
    let loss_21_val = loss_21.loss.to_vec::<f32>()?[0];
    let symmetry_diff = (loss_12_val - loss_21_val).abs();

    println!("  Loss(y1, y2): {:.6}", loss_12_val);
    println!("  Loss(y2, y1): {:.6}", loss_21_val);
    println!("  Symmetry difference: {:.8} (should be ~0)", symmetry_diff);

    // Cross-correlation matrix analysis
    println!("\nCross-correlation matrix analysis:");
    let cross_corr = compute_cross_correlation_matrix(&y1, &y2)?;
    let diagonal = extract_diagonal(&cross_corr)?;
    let diagonal_data = diagonal.to_vec::<f32>()?;

    let diagonal_mean = diagonal_data.iter().sum::<f32>() / diagonal_data.len() as f32;
    let diagonal_std = {
        let variance = diagonal_data
            .iter()
            .map(|&x| (x - diagonal_mean).powi(2))
            .sum::<f32>()
            / diagonal_data.len() as f32;
        variance.sqrt()
    };

    println!("  Matrix shape: {:?}", cross_corr.shape());
    println!(
        "  Diagonal mean: {:.4} ± {:.4}",
        diagonal_mean, diagonal_std
    );

    Ok(())
}

/// Simple test backbone for demonstrations
#[derive(Debug, Clone)]
struct TestBackbone {
    output_dim: usize,
}

impl TestBackbone {
    #[allow(dead_code)]
    fn new(output_dim: usize) -> Self {
        Self { output_dim }
    }
}

impl Backbone for TestBackbone {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        let batch_size = input.shape()[0];
        let device = input.device();

        // Simple projection to output_dim
        Tensor::randn(vec![batch_size, self.output_dim], DType::F32, device)
    }

    fn output_dim(&self) -> usize {
        self.output_dim
    }
}