//! Pipeline Parallelism Demo
//!
//! This example demonstrates how to use the RTX Transformers Pipeline Parallelism
//! implementation for distributed training of large transformer models.

use rtx_tensor::{DType, Device, Tensor};
use rtx_transformers::distributed::{
    ModelPartitionable, PipelineConfig, PipelineParallelism, PipelineSchedule,
};
use rtx_transformers::prelude::*;
use std::error::Error;

/// Mock transformer model for demonstration
#[derive(Debug, Clone)]
struct DemoTransformerModel {
    layers: Vec<DemoLayer>,
    vocab_size: usize,
    hidden_dim: usize,
}

#[derive(Debug, Clone)]
struct DemoLayer {
    id: usize,
    hidden_dim: usize,
}

impl DemoTransformerModel {
    fn new(num_layers: usize, hidden_dim: usize, vocab_size: usize) -> Self {
        let layers = (0..num_layers)
            .map(|i| DemoLayer { id: i, hidden_dim })
            .collect();

        Self {
            layers,
            vocab_size,
            hidden_dim,
        }
    }

    fn num_layers(&self) -> usize {
        self.layers.len()
    }
}

impl ModelPartitionable for DemoTransformerModel {
    fn num_layers(&self) -> usize {
        self.num_layers()
    }

    fn layer_range(&self, start: usize, end: usize) -> Result<Box<dyn std::any::Any>> {
        let layers = self.layers[start..end].to_vec();
        Ok(Box::new(layers))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Initialize RTX Transformers
    rtx_transformers::init()?;

    println!("🚀 Pipeline Parallelism Demo");
    println!("===========================\n");

    // 1. Create a large transformer model
    let model = DemoTransformerModel::new(24, 1024, 50000); // 24-layer model
    println!(
        "📊 Created demo model with {} layers, hidden_dim={}",
        model.num_layers(),
        model.hidden_dim
    );

    // 2. Configure pipeline parallelism across 4 devices
    let devices = vec![Device::cuda(0).unwrap_or(Device::default()); 4]; // In practice, use GPUs
    let config = PipelineConfig::new()
        .with_devices(devices)
        .with_micro_batch_size(8)
        .with_schedule(PipelineSchedule::PipeDream1F1B)
        .with_gradient_accumulation_steps(4)
        .with_activation_checkpointing(true)
        .with_checkpointing_segments(2)
        .with_load_balancing(true)
        .with_fault_tolerance(true)
        .with_memory_limit_mb(2048); // 2GB per stage

    println!("⚙️ Pipeline Configuration:");
    println!(
        "   - {} stages across {} devices",
        config.num_stages(),
        config.num_stages()
    );
    println!("   - Micro-batch size: {}", config.micro_batch_size());
    println!("   - Schedule: {:?}", config.schedule());
    println!(
        "   - Gradient accumulation: {} steps",
        config.gradient_accumulation_steps()
    );
    println!(
        "   - Activation checkpointing: {}",
        config.activation_checkpointing_enabled()
    );

    // 3. Initialize pipeline parallelism
    println!("\n🔧 Initializing Pipeline Parallelism...");
    let mut pipeline = PipelineParallelism::new(config, model).await?;

    // 4. Integrate with existing training infrastructure
    let gradient_accumulator = GradientAccumulator::new(4, true);
    let mixed_precision_trainer =
        MixedPrecisionTrainer::new(DType::F16, LossScalingStrategy::default(), true)?;

    let pipeline = pipeline
        .with_gradient_accumulator(gradient_accumulator)?
        .with_mixed_precision(mixed_precision_trainer)?;

    println!(
        "✅ Pipeline initialized with {} stages",
        pipeline.num_stages()
    );
    println!(
        "   - Gradient accumulation: {}",
        pipeline.has_gradient_accumulator()
    );
    println!("   - Mixed precision: {}", pipeline.has_mixed_precision());

    // 5. Demonstrate training workflow
    println!("\n🎯 Starting Training Simulation...");

    // Set training mode
    let mut training_pipeline = pipeline;
    training_pipeline.set_training_mode(true);

    // Create sample training data
    let batch_size = 32;
    let sequence_length = 512;
    let input = Tensor::randn(&[batch_size, sequence_length], DType::F32)?;
    let targets = Tensor::randn(&[batch_size, sequence_length], DType::F32)?;

    println!(
        "📦 Training data: batch_size={}, sequence_length={}",
        batch_size, sequence_length
    );

    // Run several training steps
    for step in 1..=5 {
        println!("\n📈 Training Step {}/5", step);

        let start_time = std::time::Instant::now();
        let result = training_pipeline.training_step(&input, &targets).await?;
        let elapsed = start_time.elapsed();

        println!("   ✅ Loss: {:.6}", result.loss());
        println!(
            "   ⚡ Processed {} micro-batches",
            result.num_micro_batches()
        );
        println!("   ⏱️ Step time: {:.2}ms", elapsed.as_millis());
        println!(
            "   🧠 Used mixed precision: {}",
            result.used_mixed_precision()
        );
        println!(
            "   🔄 Used gradient accumulation: {}",
            result.used_gradient_accumulation()
        );
    }

    // 6. Show comprehensive metrics
    println!("\n📊 Final Pipeline Metrics:");
    let metrics = training_pipeline.get_metrics();

    println!("   Throughput:");
    println!(
        "     - Average: {:.2} samples/sec",
        metrics.samples_per_second()
    );
    println!("     - Peak throughput: {:.2}", metrics.peak_throughput());

    println!("   Efficiency:");
    println!(
        "     - Pipeline efficiency: {:.1}%",
        metrics.pipeline_efficiency() * 100.0
    );
    println!(
        "     - Bubble time ratio: {:.1}%",
        metrics.bubble_time_ratio() * 100.0
    );

    println!("   Memory:");
    println!(
        "     - Peak memory usage: {:.1} MB",
        metrics.peak_memory_usage() as f64 / (1024.0 * 1024.0)
    );
    println!(
        "     - Activation checkpointing: {}",
        metrics.has_activation_checkpoints()
    );

    println!("   Training:");
    println!("     - Total steps: {}", metrics.total_training_steps());
    println!(
        "     - Total micro-batches: {}",
        metrics.total_micro_batches()
    );

    // 7. Demonstrate fault tolerance
    println!("\n🛡️ Testing Fault Tolerance...");
    training_pipeline.inject_stage_failure(1).await?;
    let recovery_stats = training_pipeline.get_recovery_stats();
    println!(
        "   ✅ Recovered from {} failures in {:.2}ms",
        recovery_stats.total_failures(),
        recovery_stats.recovery_time().as_millis()
    );

    // 8. Show load balancing info
    println!("\n⚖️ Load Balancing Statistics:");
    let load_balancer = training_pipeline.load_balancer();
    let balance_stats = load_balancer.get_balance_statistics();
    println!("   - Balanced: {}", balance_stats.is_balanced());
    println!(
        "   - Max stage time: {:.2}ms",
        balance_stats.max_stage_time().as_millis()
    );
    println!(
        "   - Min stage time: {:.2}ms",
        balance_stats.min_stage_time().as_millis()
    );

    // 9. Clean shutdown
    println!("\n🔄 Shutting down pipeline...");
    training_pipeline.shutdown().await?;

    println!("\n🎉 Pipeline Parallelism Demo Completed Successfully!");
    println!("\n📈 Summary of Capabilities Demonstrated:");
    println!("   ✅ Multi-stage pipeline across devices");
    println!("   ✅ PipeDream 1F1B scheduling for efficiency");
    println!("   ✅ Micro-batch processing for memory efficiency");
    println!("   ✅ Gradient accumulation and synchronization");
    println!("   ✅ Activation checkpointing for large models");
    println!("   ✅ Mixed precision training integration");
    println!("   ✅ Load balancing across stages");
    println!("   ✅ Fault tolerance and automatic recovery");
    println!("   ✅ Comprehensive metrics and monitoring");
    println!("   ✅ Clean resource management");

    Ok(())
}

/// Additional helper to show configuration options
fn show_advanced_configurations() {
    println!("\n🔧 Advanced Configuration Options:");

    // GPipe configuration for maximum throughput
    let _gpipe_config = PipelineConfig::new()
        .with_schedule(PipelineSchedule::GPipe)
        .with_micro_batch_size(16)
        .with_max_outstanding_batches(8);
    println!("   📋 GPipe: Fill-drain pattern, high throughput");

    // Memory-constrained configuration
    let _memory_config = PipelineConfig::new()
        .with_activation_checkpointing(true)
        .with_gradient_compression(true)
        .with_memory_limit_mb(1024)
        .with_micro_batch_size(2);
    println!("   💾 Memory-efficient: Checkpointing + compression");

    // High-performance configuration
    let _perf_config = PipelineConfig::new()
        .with_schedule(PipelineSchedule::Interleaved1F1B)
        .with_dynamic_batching(true)
        .with_load_balancing(true)
        .with_communication_backend(rtx_transformers::distributed::CommunicationBackend::Hybrid);
    println!("   🚀 High-performance: Interleaved scheduling + dynamic batching");
}