//! Expert Dropout Demo
//!
//! This example demonstrates the usage of expert dropout with MoE layers,
//! showcasing different dropout strategies and their effects on expert utilization.

use anyhow::Result;
use rtx_tensor::{DType, Device, Tensor};
use rtx_transformers::layers::{
    DropoutScheduler, DropoutStatistics, DropoutStrategy, ExpertDropoutConfig, ExpertDropoutLayer,
    ExpertImportanceScorer, ExpertOutputs, MoEConfig, Router,
};

fn main() -> Result<()> {
    println!("🚀 Expert Dropout Demo");
    println!("======================");

    let device = Device::cuda(0).unwrap_or(Device::default());
    let moe_config = MoEConfig::new(8, 2, 768, 3072);
    let batch_size = 4;
    let seq_len = 16;

    // Demo 1: Random Dropout Strategy
    println!("\n📊 Demo 1: Random Dropout Strategy");
    demo_random_dropout(&moe_config, &device, batch_size, seq_len)?;

    // Demo 2: Block Dropout Strategy
    println!("\n🧱 Demo 2: Block Dropout Strategy");
    demo_block_dropout(&moe_config, &device, batch_size, seq_len)?;

    // Demo 3: Load-Aware Dropout Strategy
    println!("\n⚖️  Demo 3: Load-Aware Dropout Strategy");
    demo_load_aware_dropout(&moe_config, &device, batch_size, seq_len)?;

    // Demo 4: Progressive Dropout Strategy
    println!("\n📈 Demo 4: Progressive Dropout Strategy");
    demo_progressive_dropout(&moe_config, &device, batch_size, seq_len)?;

    // Demo 5: Dropout Scheduling
    println!("\n⏰ Demo 5: Dropout Scheduling");
    demo_dropout_scheduling(&moe_config, &device, batch_size, seq_len)?;

    // Demo 6: Expert Importance Scoring
    println!("\n⭐ Demo 6: Expert Importance Scoring");
    demo_importance_scoring()?;

    // Demo 7: MoE Integration
    println!("\n🔗 Demo 7: MoE Integration");
    demo_moe_integration(&moe_config, &device, batch_size, seq_len)?;

    println!("\n✅ Expert Dropout Demo Complete!");
    Ok(())
}

fn demo_random_dropout(
    moe_config: &MoEConfig,
    device: &Device,
    batch_size: usize,
    seq_len: usize,
) -> Result<()> {
    let dropout_config = ExpertDropoutConfig::new(0.25, DropoutStrategy::Random);
    let mut dropout_layer = ExpertDropoutLayer::new(dropout_config, moe_config.clone(), device)?;
    dropout_layer.set_training(true);

    let expert_outputs = create_mock_expert_outputs(
        batch_size,
        seq_len,
        moe_config.num_experts,
        moe_config.hidden_dim,
        device,
    );

    // Run multiple iterations to show randomness
    for i in 1..=3 {
        let result = dropout_layer.forward(&expert_outputs)?;
        let active_count = result.active_experts.iter().filter(|&&x| x).count();
        println!(
            "  Iteration {}: {}/{} experts active",
            i, active_count, moe_config.num_experts
        );
        println!("    Dropout rate: {:.2}", result.dropout_stats.dropout_rate);
    }

    Ok(())
}

fn demo_block_dropout(
    moe_config: &MoEConfig,
    device: &Device,
    batch_size: usize,
    seq_len: usize,
) -> Result<()> {
    let mut dropout_config = ExpertDropoutConfig::new(0.25, DropoutStrategy::Block);
    dropout_config.block_size = Some(2);

    let mut dropout_layer = ExpertDropoutLayer::new(dropout_config, moe_config.clone(), device)?;
    dropout_layer.set_training(true);

    let expert_outputs = create_mock_expert_outputs(
        batch_size,
        seq_len,
        moe_config.num_experts,
        moe_config.hidden_dim,
        device,
    );

    let result = dropout_layer.forward(&expert_outputs)?;
    let active_count = result.active_experts.iter().filter(|&&x| x).count();

    println!(
        "  Block dropout result: {}/{} experts active",
        active_count, moe_config.num_experts
    );
    println!("  Active experts: {:?}", result.active_experts);
    println!("  Dropout rate: {:.2}", result.dropout_stats.dropout_rate);

    Ok(())
}

fn demo_load_aware_dropout(
    moe_config: &MoEConfig,
    device: &Device,
    batch_size: usize,
    seq_len: usize,
) -> Result<()> {
    let dropout_config = ExpertDropoutConfig::new(0.3, DropoutStrategy::LoadAware);
    let mut dropout_layer = ExpertDropoutLayer::new(dropout_config, moe_config.clone(), device)?;
    dropout_layer.set_training(true);

    // Simulate expert loads (some experts are underutilized)
    let expert_loads = vec![0.9, 0.8, 0.1, 0.2, 0.7, 0.05, 0.15, 0.6];
    dropout_layer.update_expert_loads(&expert_loads);

    let expert_outputs = create_mock_expert_outputs(
        batch_size,
        seq_len,
        moe_config.num_experts,
        moe_config.hidden_dim,
        device,
    );

    let result = dropout_layer.forward(&expert_outputs)?;

    println!("  Expert loads: {:?}", expert_loads);
    println!("  Active experts: {:?}", result.active_experts);
    println!("  Experts 2, 5, 6 (low load) should be more likely to be dropped");

    Ok(())
}

fn demo_progressive_dropout(
    moe_config: &MoEConfig,
    device: &Device,
    batch_size: usize,
    seq_len: usize,
) -> Result<()> {
    let dropout_config = ExpertDropoutConfig::new(0.25, DropoutStrategy::Progressive);
    let mut dropout_layer = ExpertDropoutLayer::new(dropout_config, moe_config.clone(), device)?;
    dropout_layer.set_training(true);

    let expert_outputs = create_mock_expert_outputs(
        batch_size,
        seq_len,
        moe_config.num_experts,
        moe_config.hidden_dim,
        device,
    );

    // Show how dropout pattern changes with training step
    for step in [100, 200, 300] {
        dropout_layer.set_training_step(step);
        let result = dropout_layer.forward(&expert_outputs)?;
        let active_count = result.active_experts.iter().filter(|&&x| x).count();
        println!(
            "  Step {}: {}/{} experts active, pattern: {:?}",
            step, active_count, moe_config.num_experts, result.active_experts
        );
    }

    Ok(())
}

fn demo_dropout_scheduling(
    moe_config: &MoEConfig,
    device: &Device,
    batch_size: usize,
    seq_len: usize,
) -> Result<()> {
    let dropout_config = ExpertDropoutConfig::new(0.5, DropoutStrategy::Random);
    let mut dropout_layer = ExpertDropoutLayer::new(dropout_config, moe_config.clone(), device)?;
    dropout_layer.set_training(true);

    // Create scheduler: start at 50% dropout, end at 10%
    let scheduler = DropoutScheduler::new(0.5, 0.1, 1000);
    dropout_layer.set_scheduler(scheduler.clone());

    println!("  Scheduler: 50% → 10% dropout over 1000 steps");

    for step in [0, 250, 500, 750, 1000] {
        let rate = scheduler.get_dropout_rate(step);
        println!("  Step {}: scheduled dropout rate = {:.2}", step, rate);
    }

    Ok(())
}

fn demo_importance_scoring() -> Result<()> {
    let mut scorer = ExpertImportanceScorer::new(8);

    println!(
        "  Initial importance scores: {:?}",
        scorer.get_importance_scores()
    );

    // Simulate routing weights where expert 4 is most important
    let routing_weights = vec![0.05, 0.1, 0.08, 0.12, 0.4, 0.03, 0.07, 0.15];
    scorer.update_scores(&routing_weights);

    println!(
        "  Updated importance scores: {:?}",
        scorer.get_importance_scores()
    );

    let max_idx = scorer
        .get_importance_scores()
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.total_cmp(b))
        .unwrap()
        .0;

    println!("  Most important expert: #{} (expected: #4)", max_idx);

    Ok(())
}

fn demo_moe_integration(
    moe_config: &MoEConfig,
    device: &Device,
    batch_size: usize,
    seq_len: usize,
) -> Result<()> {
    let router = Router::new(moe_config.clone(), device)?;
    let dropout_config = ExpertDropoutConfig::new(0.2, DropoutStrategy::Random);
    let mut dropout_layer = ExpertDropoutLayer::new(dropout_config, moe_config.clone(), device)?;
    dropout_layer.set_training(true);

    // Create input tensor
    let input = create_mock_input(batch_size, seq_len, moe_config.hidden_dim, device);

    // Get routing information
    let routing_info = router.route(&input)?;
    println!(
        "  Original routing - {} experts available",
        moe_config.num_experts
    );

    // Apply expert dropout to routing
    let modified_routing = dropout_layer.apply_to_routing(&routing_info)?;
    let active_count = modified_routing
        .active_expert_mask
        .iter()
        .filter(|&&x| x)
        .count();

    println!(
        "  After dropout - {}/{} experts active",
        active_count, moe_config.num_experts
    );
    println!(
        "  Active expert mask: {:?}",
        modified_routing.active_expert_mask
    );

    Ok(())
}

// Helper functions
fn create_mock_expert_outputs(
    batch_size: usize,
    seq_len: usize,
    num_experts: usize,
    hidden_dim: usize,
    device: &Device,
) -> ExpertOutputs {
    let mut outputs = Vec::new();
    for _ in 0..num_experts {
        let expert_output = Tensor::randn(&[batch_size, seq_len, hidden_dim], DType::F32, device)
            .expect("Failed to create tensor");
        outputs.push(expert_output);
    }

    ExpertOutputs { outputs }
}

fn create_mock_input(
    batch_size: usize,
    seq_len: usize,
    hidden_dim: usize,
    device: &Device,
) -> Tensor {
    Tensor::randn(&[batch_size, seq_len, hidden_dim], DType::F32, device)
        .expect("Failed to create input tensor")
}