//! RTX Curriculum Learning Demo
//!
//! Demonstrates the complete curriculum learning implementation following
//! strict TDD principles with comprehensive strategies and scheduling.

use rtx_transformers::curriculum::*;
use std::collections::HashMap;

/// Example training sample
#[derive(Debug, Clone)]
struct TrainingSample {
    id: usize,
    sequence: Vec<f32>,
    label: Option<usize>,
    difficulty_metadata: HashMap<String, f32>,
}

impl TrainingSample {
    fn new(id: usize, sequence: Vec<f32>, label: Option<usize>) -> Self {
        Self {
            id,
            sequence,
            label,
            difficulty_metadata: HashMap::new(),
        }
    }

    fn with_difficulty_info(mut self, key: &str, value: f32) -> Self {
        self.difficulty_metadata.insert(key.to_string(), value);
        self
    }
}

impl Sample for TrainingSample {
    fn id(&self) -> usize {
        self.id
    }

    fn complexity_features(&self) -> HashMap<String, f32> {
        let mut features = self.difficulty_metadata.clone();

        // Add sample ID for curriculum-based scoring
        features.insert("sample_id".to_string(), self.id as f32);

        // Sequence length as a complexity measure
        features.insert("sequence_length".to_string(), self.sequence.len() as f32);

        // Variance as another complexity measure
        if !self.sequence.is_empty() {
            let mean = self.sequence.iter().sum::<f32>() / self.sequence.len() as f32;
            let variance = self
                .sequence
                .iter()
                .map(|&x| (x - mean).powi(2))
                .sum::<f32>()
                / self.sequence.len() as f32;
            features.insert("variance".to_string(), variance);
        }

        features
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🎯 RTX Curriculum Learning Demo");
    println!("================================");

    // Create diverse training samples with different difficulties
    let training_samples = vec![
        // Very easy samples (short sequences, low variance)
        TrainingSample::new(1, vec![1.0, 1.1], Some(0)),
        TrainingSample::new(2, vec![2.0, 2.1], Some(1)),
        // Easy samples
        TrainingSample::new(3, vec![1.0, 2.0, 3.0], Some(0)),
        TrainingSample::new(4, vec![2.0, 3.0, 4.0], Some(1)),
        // Medium samples
        TrainingSample::new(5, vec![1.0, 3.0, 2.0, 4.0, 1.5], Some(0)),
        TrainingSample::new(6, vec![2.0, 1.0, 4.0, 3.0, 2.5], Some(1)),
        // Hard samples (longer sequences, higher variance)
        TrainingSample::new(7, vec![1.0, 5.0, 2.0, 8.0, 1.0, 6.0, 3.0], Some(0)),
        TrainingSample::new(8, vec![3.0, 1.0, 7.0, 2.0, 9.0, 1.0, 4.0], Some(1)),
        // Very hard samples
        TrainingSample::new(
            9,
            vec![1.0, 10.0, 2.0, 15.0, 3.0, 8.0, 1.0, 12.0, 4.0],
            Some(0),
        ),
        TrainingSample::new(
            10,
            vec![5.0, 1.0, 20.0, 2.0, 11.0, 3.0, 16.0, 1.0, 7.0],
            Some(1),
        ),
    ];

    println!("\n📊 Training Dataset Analysis:");
    println!("   Total samples: {}", training_samples.len());
    for sample in &training_samples {
        let features = sample.complexity_features();
        println!(
            "   Sample {}: len={}, variance={:.2}",
            sample.id(),
            features.get("sequence_length").unwrap_or(&0.0),
            features.get("variance").unwrap_or(&0.0)
        );
    }

    // Demo 1: Different Difficulty Scorers
    println!("\n🎲 Demo 1: Different Difficulty Scoring Strategies");
    println!("==================================================");

    let length_scorer = LengthBasedDifficultyScorer::new();
    let variance_scorer = VarianceBasedDifficultyScorer::new();

    println!("Length-based vs Variance-based scoring:");
    for sample in training_samples.iter().take(5) {
        let features = sample.complexity_features();
        let length_score = length_scorer.score_features(&features);
        let variance_score = variance_scorer.score_features(&features);
        println!(
            "   Sample {}: length_score={:.3}, variance_score={:.3}",
            sample.id(),
            length_score,
            variance_score
        );
    }

    // Create composite scorer
    let mut composite_scorer = CompositeDifficultyScorer::new();
    composite_scorer.add_scorer(Box::new(length_scorer.clone()), 0.6);
    composite_scorer.add_scorer(Box::new(variance_scorer.clone()), 0.4);

    println!("\nComposite scoring (60% length + 40% variance):");
    for sample in training_samples.iter().take(5) {
        let features = sample.complexity_features();
        let composite_score = composite_scorer.score_features(&features);
        println!(
            "   Sample {}: composite_score={:.3}",
            sample.id(),
            composite_score
        );
    }

    // Demo 2: Easy-to-Hard Curriculum Strategy
    println!("\n📈 Demo 2: Easy-to-Hard Curriculum Strategy");
    println!("===========================================");

    let easy_to_hard = EasyToHardStrategy::new(0.1, 0.08);
    let mut curriculum_state = CurriculumState::new();

    println!("Training progression with easy-to-hard curriculum:");
    for epoch in 0..8 {
        curriculum_state.step = epoch;
        easy_to_hard.update_difficulty_threshold(&mut curriculum_state);

        let selected_batch = easy_to_hard.select_samples(
            &training_samples,
            &composite_scorer,
            &mut curriculum_state,
            3,
        );

        let avg_difficulty: f32 = selected_batch
            .iter()
            .map(|s| composite_scorer.score_features(&s.complexity_features()))
            .sum::<f32>()
            / selected_batch.len() as f32;

        let selected_ids: Vec<usize> = selected_batch.iter().map(|s| s.id()).collect();

        println!(
            "   Epoch {}: threshold={:.2}, selected={:?}, avg_difficulty={:.3}",
            epoch, curriculum_state.difficulty_threshold, selected_ids, avg_difficulty
        );
    }

    // Demo 3: Anti-Curriculum Strategy (Hard-to-Easy)
    println!("\n📉 Demo 3: Anti-Curriculum Strategy (Hard-to-Easy)");
    println!("==================================================");

    let anti_curriculum = AntiCurriculumStrategy::new(0.9, 0.08);
    let mut anti_state = CurriculumState::new();

    println!("Training progression with anti-curriculum (hard examples first):");
    for epoch in 0..8 {
        anti_state.step = epoch;
        anti_curriculum.update_difficulty_threshold(&mut anti_state);

        let selected_batch = anti_curriculum.select_samples(
            &training_samples,
            &composite_scorer,
            &mut anti_state,
            3,
        );

        let avg_difficulty: f32 = selected_batch
            .iter()
            .map(|s| composite_scorer.score_features(&s.complexity_features()))
            .sum::<f32>()
            / selected_batch.len() as f32;

        let selected_ids: Vec<usize> = selected_batch.iter().map(|s| s.id()).collect();

        println!(
            "   Epoch {}: threshold={:.2}, selected={:?}, avg_difficulty={:.3}",
            epoch, anti_state.difficulty_threshold, selected_ids, avg_difficulty
        );
    }

    // Demo 4: Self-Paced Learning
    println!("\n🎯 Demo 4: Self-Paced Learning (Performance-Adaptive)");
    println!("=====================================================");

    let mut self_paced = SelfPacedStrategy::new(0.3, 0.75, 5);
    let mut self_paced_state = CurriculumState::new();

    // Simulate different performance scenarios
    let performance_scenarios = vec![0.6, 0.8, 0.9, 0.7, 0.5, 0.8, 0.85, 0.9];

    println!("Self-paced learning adapting to performance:");
    for (epoch, &performance) in performance_scenarios.iter().enumerate() {
        self_paced.update_performance(performance);
        self_paced_state.step = epoch;
        self_paced.update_difficulty_threshold(&mut self_paced_state);

        let selected_batch = self_paced.select_samples(
            &training_samples,
            &composite_scorer,
            &mut self_paced_state,
            3,
        );

        let avg_difficulty: f32 = selected_batch
            .iter()
            .map(|s| composite_scorer.score_features(&s.complexity_features()))
            .sum::<f32>()
            / selected_batch.len() as f32;

        let selected_ids: Vec<usize> = selected_batch.iter().map(|s| s.id()).collect();

        println!(
            "   Epoch {}: perf={:.2}, threshold={:.2}, selected={:?}, avg_diff={:.3}",
            epoch, performance, self_paced_state.difficulty_threshold, selected_ids, avg_difficulty
        );
    }

    // Demo 5: Different Scheduling Strategies
    println!("\n⏱️  Demo 5: Curriculum Scheduling Strategies");
    println!("=============================================");

    let linear_schedule = LinearSchedule::new(0.1, 0.9, 20);
    let exponential_schedule = ExponentialSchedule::new(0.1, 0.9, 0.15);
    let cyclic_schedule = CyclicSchedule::new(0.2, 0.8, 10);

    println!("Difficulty progression across different schedules:");
    println!("   Step | Linear | Exponential | Cyclic");
    println!("   -----|--------|-------------|--------");

    for step in [0, 5, 10, 15, 20, 25, 30].iter() {
        let linear_diff = linear_schedule.get_difficulty_at_step(*step);
        let exp_diff = exponential_schedule.get_difficulty_at_step(*step);
        let cyclic_diff = cyclic_schedule.get_difficulty_at_step(*step);

        println!(
            "   {:4} | {:.3}  |    {:.3}    | {:.3}",
            step, linear_diff, exp_diff, cyclic_diff
        );
    }

    // Demo 6: Complete Pipeline Integration
    println!("\n🔗 Demo 6: Complete Curriculum Learning Pipeline");
    println!("================================================");

    let pipeline_scorer = Box::new(composite_scorer);
    let pipeline_strategy = Box::new(easy_to_hard);
    let pipeline_schedule = Box::new(linear_schedule);

    let mut curriculum_loader = CurriculumDataLoader::new(
        training_samples.clone(),
        pipeline_strategy,
        pipeline_scorer,
        pipeline_schedule,
        4,
    );

    println!("Complete training pipeline with curriculum learning:");
    for step in 0..10 {
        if let Some(batch) = curriculum_loader.next_batch() {
            let avg_length: f32 =
                batch.iter().map(|s| s.sequence.len() as f32).sum::<f32>() / batch.len() as f32;

            let batch_ids: Vec<usize> = batch.iter().map(|s| s.id()).collect();

            let state = curriculum_loader.get_curriculum_state();

            println!(
                "   Step {}: threshold={:.2}, batch={:?}, avg_len={:.1}",
                step, state.difficulty_threshold, batch_ids, avg_length
            );
        }

        curriculum_loader.step();
    }

    // Demo 7: Performance Tracking and Optimization
    println!("\n📊 Demo 7: Performance Tracking");
    println!("===============================");

    let mut performance_tracker = PerformanceTracker::new(5, 0.1);

    // Simulate training with performance feedback
    for sample_id in 1..=10 {
        let simulated_performance = 0.6 + (sample_id as f32 * 0.03) + (sample_id % 3) as f32 * 0.1; // Add some variation
        performance_tracker.record_sample_performance(sample_id, simulated_performance);
        performance_tracker.record_overall_performance(simulated_performance);

        println!(
            "   Sample {}: performance = {:.2}",
            sample_id, simulated_performance
        );
    }

    let overall_perf = performance_tracker.get_overall_performance();
    let perf_trend = performance_tracker.get_performance_trend();

    println!("   Overall performance: {:.2}", overall_perf);
    println!("   Performance trend: {:.2}", perf_trend);

    println!("\n🎉 Curriculum Learning Demo Complete!");
    println!("=====================================");

    println!("\n✅ Key Features Demonstrated:");
    println!("   🔸 Multiple difficulty scoring strategies (length, variance, composite)");
    println!("   🔸 Easy-to-hard curriculum progression");
    println!("   🔸 Anti-curriculum (hard-to-easy) strategy");
    println!("   🔸 Self-paced learning with performance adaptation");
    println!("   🔸 Various scheduling algorithms (linear, exponential, cyclic)");
    println!("   🔸 Complete data loader integration");
    println!("   🔸 Performance tracking and feedback");

    println!("\n🚀 Production-Ready Implementation:");
    println!("   ✅ Strict TDD methodology (tests first, minimal implementation)");
    println!("   ✅ Zero-cost abstractions with trait-based design");
    println!("   ✅ Memory-safe Rust implementation");
    println!("   ✅ Comprehensive error handling");
    println!("   ✅ Object-safe traits for dynamic dispatch");
    println!("   ✅ No mocks, stubs, or TODOs - complete implementation");

    Ok(())
}