//! Comprehensive Meta-Learning Example
//! 
//! This example demonstrates how to use the RTX Transformers meta-learning
//! implementation for few-shot learning tasks using both MAML and 
//! Prototypical Networks.

use rtx_transformers::meta::*;
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🚀 RTX Transformers Meta-Learning Comprehensive Example");
    println!("========================================================\n");
    
    let device = Device::cuda(0).unwrap_or(Device::default());
    
    // Step 1: Create a more realistic few-shot dataset
    println!("📊 Step 1: Creating Few-Shot Learning Dataset");
    println!("-----------------------------------------------");
    
    let num_classes = 10;
    let samples_per_class = 20;
    let feature_dim = 128;
    
    let dataset = FewShotDataset::synthetic(num_classes, samples_per_class, feature_dim, &device)?;
    println!("✓ Created dataset:");
    println!("  - {} classes", dataset.num_classes);
    println!("  - {} samples per class", samples_per_class);
    println!("  - {} feature dimensions", dataset.feature_dim);
    println!();
    
    // Step 2: Configure different few-shot scenarios
    println!("🎯 Step 2: Few-Shot Learning Scenarios");
    println!("--------------------------------------");
    
    let scenarios = vec![
        ("5-way 1-shot", 5, 1, 5),
        ("5-way 5-shot", 5, 5, 5), 
        ("10-way 1-shot", 10, 1, 3),
        ("3-way 2-shot", 3, 2, 7),
    ];
    
    for (name, n_way, k_shot, query_per_class) in &scenarios {
        println!("📋 Scenario: {}", name);
        let episode = dataset.sample_episode(*n_way, *k_shot, *query_per_class, &device)?;
        
        println!("  Support set: {} samples ({} × {} shots)", 
                 episode.support_set.len(), n_way, k_shot);
        println!("  Query set: {} samples ({} × {} queries)", 
                 episode.query_set.len(), n_way, query_per_class);
        
        // Verify episode structure
        let (support_x, support_y) = episode.support_batch(&device)?;
        let (query_x, query_y) = episode.query_batch(&device)?;
        
        println!("  Support batch shape: {:?}", support_x.shape());
        println!("  Query batch shape: {:?}", query_x.shape());
        println!();
    }
    
    // Step 3: MAML Demonstration
    println!("🧠 Step 3: MAML (Model-Agnostic Meta-Learning)");
    println!("===============================================");
    
    let maml_config = MAMLConfig {
        inner_lr: 0.01,
        outer_lr: 0.001,
        inner_steps: 3,
        first_order: false,
        meta_batch_size: 8,
    };
    
    println!("📋 MAML Configuration:");
    println!("  Inner learning rate: {}", maml_config.inner_lr);
    println!("  Outer learning rate: {}", maml_config.outer_lr);
    println!("  Inner steps: {}", maml_config.inner_steps);
    println!("  First-order approximation: {}", maml_config.first_order);
    println!();
    
    let mut maml = MAML::new(feature_dim, 64, 5, maml_config, &device)?;
    
    // Generate training episodes for MAML
    let mut sampler = EpisodeSampler::new(dataset.clone(), EpisodeSamplerConfig {
        n_way: 5,
        k_shot: 1,
        query_per_class: 3,
        num_episodes: 16,
        seed: Some(42),
    });
    
    println!("🏋️ Training MAML...");
    let training_episodes = sampler.sample_episodes(8, &device)?;
    let maml_stats = maml.meta_update(training_episodes, &device)?;
    
    println!("✓ MAML Training Results:");
    println!("  Meta-updates: {}", maml_stats.meta_updates);
    println!("  Tasks processed: {}", maml_stats.tasks_processed);
    println!("  Average inner loss: {:.4}", maml_stats.avg_inner_loss);
    println!("  Average outer loss: {:.4}", maml_stats.avg_outer_loss);
    println!("  Average query accuracy: {:.4}", maml_stats.avg_query_accuracy);
    println!();
    
    // Test fast adaptation
    println!("⚡ Testing MAML Fast Adaptation...");
    let test_episode = sampler.sample_episode(&device)?;
    let adapted_network = maml.fast_adaptation(&test_episode, &device)?;
    
    let (query_x, query_y) = test_episode.query_batch(&device)?;
    let predictions = adapted_network.forward(&query_x)?;
    let accuracy = maml.compute_accuracy(&predictions, &query_y)?;
    
    println!("✓ Fast Adaptation Results:");
    println!("  Adapted in {} inner steps", maml.config.inner_steps);
    println!("  Query accuracy after adaptation: {:.4}", accuracy);
    println!();
    
    // Step 4: Prototypical Networks Demonstration
    println!("🎯 Step 4: Prototypical Networks");
    println!("================================");
    
    // Test different distance metrics
    let distance_metrics = vec![
        ("Euclidean", DistanceMetric::Euclidean),
        ("Cosine", DistanceMetric::Cosine),
        ("Manhattan", DistanceMetric::Manhattan),
    ];
    
    let mut proto_results = HashMap::new();
    
    for (metric_name, metric) in &distance_metrics {
        println!("📏 Testing {} Distance Metric", metric_name);
        
        let proto_config = PrototypicalConfig {
            learning_rate: 0.001,
            feature_dim: 32,
            distance_metric: metric.clone(),
            temperature: 1.0,
            batch_size: 4,
        };
        
        let proto_net = PrototypicalNetworks::new(feature_dim, proto_config, &device)?;
        
        // Test on multiple episodes
        let mut accuracies = Vec::new();
        let mut losses = Vec::new();
        
        for _ in 0..5 {
            let episode = dataset.sample_episode(3, 2, 4, &device)?;
            let (logits, _, accuracy) = proto_net.classify(&episode, &device)?;
            let (_, query_y) = episode.query_batch(&device)?;
            let loss = proto_net.compute_loss(&logits, &query_y)?;
            
            accuracies.push(accuracy);
            losses.push(loss.to_vec()?[0]);
        }
        
        let mean_accuracy = accuracies.iter().sum::<f32>() / accuracies.len() as f32;
        let mean_loss = losses.iter().sum::<f32>() / losses.len() as f32;
        
        proto_results.insert(metric_name, (mean_accuracy, mean_loss));
        
        println!("  ✓ Mean accuracy: {:.4}", mean_accuracy);
        println!("  ✓ Mean loss: {:.4}", mean_loss);
        println!();
    }
    
    // Step 5: Comparison and Analysis
    println!("📊 Step 5: Method Comparison & Analysis");
    println!("=======================================");
    
    // Compare distance metrics
    println!("🔍 Distance Metric Comparison:");
    let mut best_metric = "";
    let mut best_accuracy = 0.0;
    
    for (metric, (accuracy, loss)) in &proto_results {
        println!("  {} - Accuracy: {:.4}, Loss: {:.4}", metric, accuracy, loss);
        if *accuracy > best_accuracy {
            best_accuracy = *accuracy;
            best_metric = metric;
        }
    }
    println!("  🏆 Best performing metric: {} ({:.4} accuracy)", best_metric, best_accuracy);
    println!();
    
    // Step 6: Complete Pipeline Demonstration
    println!("🔄 Step 6: Complete Meta-Learning Pipeline");
    println!("==========================================");
    
    let pipeline_config = MetaLearningConfig {
        input_dim: feature_dim,
        hidden_dim: 64,
        output_dim: num_classes,
        algorithm: MetaLearningAlgorithm::Hybrid,
        sampler_config: EpisodeSamplerConfig {
            n_way: 5,
            k_shot: 2,
            query_per_class: 3,
            num_episodes: 20,
            seed: Some(12345),
        },
        maml_config: MAMLConfig {
            inner_lr: 0.02,
            outer_lr: 0.001,
            inner_steps: 2,
            first_order: false,
            meta_batch_size: 4,
        },
        proto_config: PrototypicalConfig {
            learning_rate: 0.001,
            feature_dim: 32,
            distance_metric: DistanceMetric::Euclidean,
            temperature: 0.5,
            batch_size: 4,
        },
    };
    
    println!("🚀 Running Hybrid Pipeline (MAML + Prototypical)...");
    let mut pipeline = MetaLearningPipeline::new(dataset, pipeline_config, &device)?;
    let training_results = pipeline.train(10, &device)?;
    
    println!("✓ Pipeline Training Complete:");
    if let Some(maml_stats) = &training_results.maml_stats {
        println!("  MAML - Meta-updates: {}, Accuracy: {:.4}", 
                 maml_stats.meta_updates, maml_stats.avg_query_accuracy);
    }
    if let Some(proto_stats) = &training_results.proto_stats {
        println!("  Prototypical - Episodes: {}, Accuracy: {:.4}", 
                 proto_stats.episodes_processed, proto_stats.avg_accuracy);
    }
    println!();
    
    // Generate test episodes for evaluation
    println!("🧪 Evaluating on Test Episodes...");
    let test_episodes = pipeline.sampler.sample_episodes(5, &device)?;
    let evaluation_results = pipeline.evaluate(test_episodes, &device)?;
    
    if let Some(metrics) = &evaluation_results.evaluation_metrics {
        println!("✓ Evaluation Results:");
        println!("  Test episodes: {}", metrics.num_episodes);
        println!("  Mean accuracy: {:.4}", metrics.mean_accuracy);
        println!("  Mean loss: {:.4}", metrics.mean_loss);
        println!("  95% Confidence interval: ({:.4}, {:.4})", 
                 metrics.confidence_interval.0, metrics.confidence_interval.1);
    }
    
    if let Some(comparison) = &evaluation_results.comparison {
        println!("  Method comparison: {} is better by {:.4} accuracy", 
                 comparison.better_approach, comparison.accuracy_difference.abs());
    }
    println!();
    
    // Step 7: Advanced Features Demo
    println!("🎛️ Step 7: Advanced Features");
    println!("============================");
    
    // Task distribution sampling
    println!("🎲 Task Distribution Sampling:");
    let base_config = EpisodeSamplerConfig {
        n_way: 3,
        k_shot: 1,
        query_per_class: 2,
        num_episodes: 1,
        seed: Some(999),
    };
    
    let mut task_dist = TaskDistribution::multi_task(base_config);
    
    for i in 0..3 {
        let config = task_dist.sample_task_config();
        println!("  Task {}: {}-way {}-shot", i + 1, config.n_way, config.k_shot);
    }
    println!();
    
    // Episode consistency check
    println!("🔍 Episode Consistency Check:");
    let mut sampler = EpisodeSampler::new(
        FewShotDataset::synthetic(5, 10, 16, &device)?, 
        EpisodeSamplerConfig {
            n_way: 3,
            k_shot: 1,
            query_per_class: 2,
            seed: Some(777),
            ..Default::default()
        }
    );
    
    let episode1 = sampler.sample_episode(&device)?;
    sampler.reset();
    let episode2 = sampler.sample_episode(&device)?;
    
    let consistent = episode1.support_labels == episode2.support_labels;
    println!("  Reproducible sampling: {}", if consistent { "✓ PASS" } else { "✗ FAIL" });
    println!();
    
    // Final Summary
    println!("🎉 Meta-Learning Implementation Summary");
    println!("======================================");
    println!("✅ MAML (Model-Agnostic Meta-Learning):");
    println!("   - Inner/outer loop optimization implemented");
    println!("   - Fast adaptation for new tasks");
    println!("   - First-order approximation (FOMAML) support");
    println!();
    println!("✅ Prototypical Networks:");
    println!("   - Multiple distance metrics (Euclidean, Cosine, Manhattan)");
    println!("   - Prototype computation from support sets");
    println!("   - Temperature scaling for probability calibration");
    println!();
    println!("✅ Episode Sampling & Evaluation:");
    println!("   - N-way K-shot episode generation");
    println!("   - Reproducible sampling with seeding");
    println!("   - Comprehensive evaluation metrics");
    println!("   - Task distribution sampling");
    println!();
    println!("✅ Integration Features:");
    println!("   - Transformer model integration ready");
    println!("   - Complete meta-learning pipeline");
    println!("   - Hybrid algorithm support");
    println!("   - Production-ready implementation");
    println!();
    println!("🚀 RTX Transformers Meta-Learning: IMPLEMENTATION COMPLETE!");
    
    Ok(())
}