#!/usr/bin/env cargo
//! Pseudo-labeling (Self-training) Demonstration
//!
//! This example shows how to use the pseudo-labeling framework for semi-supervised learning.
//! It demonstrates the complete workflow of generating pseudo-labels from unlabeled data
//! and iteratively improving the model.

use std::collections::HashMap;

/// Simple demonstration of pseudo-labeling components
fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("RTX Transformers - Pseudo-labeling (Self-training) Demo");
    println!("======================================================\n");

    // 1. Configuration with default hyperparameters
    println!("1. Configuration Setup");
    let config = PseudoLabelingConfig::default();
    println!("  - Confidence threshold: {}", config.confidence_threshold);
    println!("  - Threshold schedule: {}", config.threshold_schedule);
    println!("  - Label smoothing: {}", config.label_smoothing);
    println!("  - Pseudo weight: {}", config.pseudo_weight);
    println!("  - Class balance: {}", config.class_balance);
    println!("  - Soft labels: {}", config.soft_labels);
    println!();

    // 2. Demonstrate confidence scoring
    println!("2. Confidence Scoring");
    let predictions = vec![
        vec![0.1, 0.2, 0.7], // High confidence for class 2
        vec![0.8, 0.1, 0.1], // High confidence for class 0
        vec![0.3, 0.3, 0.4], // Low confidence (uncertain)
    ];

    let scorer = ConfidenceScorer::new(ConfidenceMetric::MaxProbability);
    let confidences = scorer.compute_confidence(&predictions);
    println!("  - Max probability confidences: {:?}", confidences);

    let entropy_scorer = ConfidenceScorer::new(ConfidenceMetric::Entropy);
    let entropy_confidences = entropy_scorer.compute_confidence(&predictions);
    println!("  - Entropy-based confidences: {:?}", entropy_confidences);
    println!();

    // 3. Demonstrate threshold scheduling
    println!("3. Threshold Scheduling");
    let linear_scheduler = ThresholdScheduler::new(ThresholdSchedule::Linear {
        start: 0.95,
        end: 0.7,
    });

    for iteration in 0..=4 {
        let threshold = linear_scheduler.get_threshold(iteration, 4);
        println!("  - Iteration {}: threshold = {:.3}", iteration, threshold);
    }
    println!();

    // 4. Demonstrate pseudo-label generation
    println!("4. Pseudo-label Generation");
    let hard_config = PseudoLabelingConfig {
        soft_labels: false,
        ..Default::default()
    };

    let soft_config = PseudoLabelingConfig {
        soft_labels: true,
        label_smoothing: 0.1,
        ..Default::default()
    };

    let hard_generator = PseudoLabelGenerator::new(hard_config);
    let soft_generator = PseudoLabelGenerator::new(soft_config);

    let hard_result = hard_generator.generate_labels(&predictions);
    let soft_result = soft_generator.generate_labels(&predictions);

    println!("  - Hard labels: {:?}", hard_result.labels);
    println!("  - Soft labels: {:?}", soft_result.labels);
    println!("  - Confidences: {:?}", hard_result.confidences);
    println!();

    // 5. Demonstrate class balancing
    println!("5. Class Balancing");
    let balancer = ClassBalancer::new();
    let selected_indices = balancer.balance_selection(&predictions, &confidences, 0.6, Some(1));
    println!(
        "  - Selected indices for balanced pseudo-labels: {:?}",
        selected_indices
    );
    println!();

    // 6. Complete training iteration simulation
    println!("6. Training Iteration Simulation");
    let mut stats = PseudoLabelingStats::new();

    for iteration in 1..=5 {
        // Simulate training data
        let labeled_data = simulate_labeled_data(50, 3);
        let unlabeled_data = simulate_unlabeled_data(200, 3);

        // Simulate model predictions (improving over iterations)
        let confidence_boost = iteration as f32 * 0.05;
        let pseudo_results =
            simulate_training_iteration(&labeled_data, &unlabeled_data, &config, confidence_boost);

        stats.record_iteration(iteration, &pseudo_results);

        println!(
            "  - Iteration {}: {} pseudo-labels generated, avg confidence: {:.3}",
            iteration, pseudo_results.num_pseudo_labels, pseudo_results.avg_pseudo_confidence
        );
    }

    println!();
    println!("7. Final Statistics");
    println!("{}", stats.summary());

    println!("\n✅ Pseudo-labeling demo completed successfully!");
    println!(
        "📊 The algorithm demonstrates iterative self-training with confidence-based selection."
    );
    println!(
        "🎯 Key features: threshold scheduling, class balancing, and both hard/soft labeling."
    );

    Ok(())
}

// Simplified implementations for demonstration

#[derive(Debug, Clone)]
struct PseudoLabelingConfig {
    confidence_threshold: f32,
    threshold_schedule: String,
    label_smoothing: f32,
    pseudo_weight: f32,
    class_balance: bool,
    soft_labels: bool,
    max_pseudo_per_class: Option<usize>,
    num_iterations: usize,
}

impl Default for PseudoLabelingConfig {
    fn default() -> Self {
        Self {
            confidence_threshold: 0.95,
            threshold_schedule: "linear".to_string(),
            label_smoothing: 0.1,
            pseudo_weight: 1.0,
            class_balance: true,
            soft_labels: false,
            max_pseudo_per_class: None,
            num_iterations: 5,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum ConfidenceMetric {
    MaxProbability,
    Entropy,
    Margin,
}

struct ConfidenceScorer {
    metric: ConfidenceMetric,
}

impl ConfidenceScorer {
    fn new(metric: ConfidenceMetric) -> Self {
        Self { metric }
    }

    fn compute_confidence(&self, predictions: &[Vec<f32>]) -> Vec<f32> {
        match self.metric {
            ConfidenceMetric::MaxProbability => predictions
                .iter()
                .map(|pred| pred.iter().fold(0.0f32, |a, &b| a.max(b)))
                .collect(),
            ConfidenceMetric::Entropy => predictions
                .iter()
                .map(|pred| {
                    let entropy: f32 = pred
                        .iter()
                        .filter(|&&p| p > 0.0)
                        .map(|&p| -p * p.ln())
                        .sum();
                    let max_entropy = (pred.len() as f32).ln();
                    1.0 - (entropy / max_entropy).min(1.0)
                })
                .collect(),
            ConfidenceMetric::Margin => predictions
                .iter()
                .map(|pred| {
                    let mut sorted = pred.clone();
                    sorted.sort_by(|a, b| b.total_cmp(a));
                    if sorted.len() >= 2 {
                        sorted[0] - sorted[1]
                    } else {
                        sorted[0]
                    }
                })
                .collect(),
        }
    }
}

#[derive(Debug, Clone)]
enum ThresholdSchedule {
    Fixed(f32),
    Linear { start: f32, end: f32 },
}

struct ThresholdScheduler {
    schedule: ThresholdSchedule,
}

impl ThresholdScheduler {
    fn new(schedule: ThresholdSchedule) -> Self {
        Self { schedule }
    }

    fn get_threshold(&self, iteration: usize, total_iterations: usize) -> f32 {
        match &self.schedule {
            ThresholdSchedule::Fixed(threshold) => *threshold,
            ThresholdSchedule::Linear { start, end } => {
                if total_iterations == 0 {
                    *start
                } else {
                    let progress = iteration as f32 / total_iterations as f32;
                    start + progress * (end - start)
                }
            }
        }
    }
}

#[derive(Debug)]
struct PseudoLabelResult {
    labels: Vec<i32>,
    confidences: Vec<f32>,
}

struct PseudoLabelGenerator {
    config: PseudoLabelingConfig,
}

impl PseudoLabelGenerator {
    fn new(config: PseudoLabelingConfig) -> Self {
        Self { config }
    }

    fn generate_labels(&self, predictions: &[Vec<f32>]) -> PseudoLabelResult {
        let scorer = ConfidenceScorer::new(ConfidenceMetric::MaxProbability);
        let confidences = scorer.compute_confidence(predictions);

        let labels: Vec<i32> = if self.config.soft_labels {
            // For soft labels, we still return the argmax but apply smoothing in practice
            predictions
                .iter()
                .map(|pred| {
                    pred.iter()
                        .enumerate()
                        .fold(
                            (0, 0.0f32),
                            |acc, (i, &val)| {
                                if val > acc.1 { (i, val) } else { acc }
                            },
                        )
                        .0 as i32
                })
                .collect()
        } else {
            predictions
                .iter()
                .map(|pred| {
                    pred.iter()
                        .enumerate()
                        .fold(
                            (0, 0.0f32),
                            |acc, (i, &val)| {
                                if val > acc.1 { (i, val) } else { acc }
                            },
                        )
                        .0 as i32
                })
                .collect()
        };

        PseudoLabelResult {
            labels,
            confidences,
        }
    }
}

struct ClassBalancer;

impl ClassBalancer {
    fn new() -> Self {
        Self
    }

    fn balance_selection(
        &self,
        predictions: &[Vec<f32>],
        confidences: &[f32],
        threshold: f32,
        max_per_class: Option<usize>,
    ) -> Vec<usize> {
        let mut class_samples: HashMap<i32, Vec<(usize, f32)>> = HashMap::new();

        // Group samples by predicted class
        for (idx, (pred, &conf)) in predictions.iter().zip(confidences.iter()).enumerate() {
            if conf >= threshold {
                let class = pred
                    .iter()
                    .enumerate()
                    .fold(
                        (0, 0.0f32),
                        |acc, (i, &val)| {
                            if val > acc.1 { (i, val) } else { acc }
                        },
                    )
                    .0 as i32;

                class_samples
                    .entry(class)
                    .or_insert_with(Vec::new)
                    .push((idx, conf));
            }
        }

        // Balance selection across classes
        let mut selected = Vec::new();

        for (_, mut samples) in class_samples {
            // Sort by confidence (highest first)
            samples.sort_by(|a, b| b.1.total_cmp(&a.1));

            let take_count = if let Some(max) = max_per_class {
                std::cmp::min(samples.len(), max)
            } else {
                samples.len()
            };

            selected.extend(samples.into_iter().take(take_count).map(|(idx, _)| idx));
        }

        selected
    }
}

#[derive(Debug, Clone)]
struct PseudoLabelingResult {
    supervised_loss: f32,
    pseudo_loss: f32,
    total_loss: f32,
    num_pseudo_labels: usize,
    avg_pseudo_confidence: f32,
}

struct PseudoLabelingStats {
    iterations: Vec<PseudoLabelingResult>,
}

impl PseudoLabelingStats {
    fn new() -> Self {
        Self {
            iterations: Vec::new(),
        }
    }

    fn record_iteration(&mut self, _iteration: usize, result: &PseudoLabelingResult) {
        self.iterations.push(result.clone());
    }

    fn summary(&self) -> String {
        let total_pseudo = self
            .iterations
            .iter()
            .map(|r| r.num_pseudo_labels)
            .sum::<usize>();
        let avg_confidence: f32 = self
            .iterations
            .iter()
            .map(|r| r.avg_pseudo_confidence)
            .sum::<f32>()
            / self.iterations.len() as f32;

        format!(
            "  - Total pseudo-labels generated: {}\n  - Average confidence: {:.3}\n  - Iterations completed: {}",
            total_pseudo,
            avg_confidence,
            self.iterations.len()
        )
    }
}

// Simulation functions
fn simulate_labeled_data(num_samples: usize, num_classes: usize) -> Vec<(Vec<f32>, i32)> {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};

    (0..num_samples)
        .map(|i| {
            let mut hasher = DefaultHasher::new();
            i.hash(&mut hasher);
            let seed = hasher.finish();

            let class = (seed % num_classes as u64) as i32;
            let features: Vec<f32> = (0..10)
                .map(|j| {
                    let mut hasher = DefaultHasher::new();
                    (i * 10 + j).hash(&mut hasher);
                    (hasher.finish() % 1000) as f32 / 1000.0
                })
                .collect();

            (features, class)
        })
        .collect()
}

fn simulate_unlabeled_data(num_samples: usize, num_classes: usize) -> Vec<Vec<f32>> {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};

    (0..num_samples)
        .map(|i| {
            (0..10)
                .map(|j| {
                    let mut hasher = DefaultHasher::new();
                    (i * 10 + j + 12345).hash(&mut hasher);
                    (hasher.finish() % 1000) as f32 / 1000.0
                })
                .collect()
        })
        .collect()
}

fn simulate_training_iteration(
    _labeled_data: &[(Vec<f32>, i32)],
    unlabeled_data: &[Vec<f32>],
    config: &PseudoLabelingConfig,
    confidence_boost: f32,
) -> PseudoLabelingResult {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};

    // Simulate model predictions (getting better over time)
    let predictions: Vec<Vec<f32>> = unlabeled_data
        .iter()
        .enumerate()
        .map(|(i, _features)| {
            let mut hasher = DefaultHasher::new();
            i.hash(&mut hasher);
            let seed = hasher.finish();

            // Simulate 3-class probabilities
            let base_probs = [
                ((seed % 100) as f32 + confidence_boost * 50.0) / 100.0,
                ((seed.wrapping_shr(8) % 100) as f32) / 100.0,
                ((seed.wrapping_shr(16) % 100) as f32) / 100.0,
            ];

            // Normalize to sum to 1
            let sum: f32 = base_probs.iter().sum();
            base_probs.iter().map(|&p| p / sum).collect()
        })
        .collect();

    // Generate pseudo-labels
    let generator = PseudoLabelGenerator::new(config.clone());
    let pseudo_result = generator.generate_labels(&predictions);

    // Apply threshold selection
    let threshold = 0.8 - confidence_boost; // Decreasing threshold as model improves
    let balancer = ClassBalancer::new();
    let selected = balancer.balance_selection(
        &predictions,
        &pseudo_result.confidences,
        threshold,
        config.max_pseudo_per_class,
    );

    let avg_confidence = if !selected.is_empty() {
        selected
            .iter()
            .map(|&i| pseudo_result.confidences[i])
            .sum::<f32>()
            / selected.len() as f32
    } else {
        0.0
    };

    PseudoLabelingResult {
        supervised_loss: 1.5 - confidence_boost, // Decreasing loss
        pseudo_loss: 0.8 - confidence_boost * 0.5,
        total_loss: 2.3 - confidence_boost * 1.5,
        num_pseudo_labels: selected.len(),
        avg_pseudo_confidence: avg_confidence,
    }
}