//! Expert dropout standalone test

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

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Test basic functionality
    let device = Device::cuda(0).unwrap_or(Device::default());
    let moe_config = MoEConfig::new(8, 2, 768, 3072);
    let dropout_config = ExpertDropoutConfig::new(0.2, DropoutStrategy::Random);

    let dropout_layer = ExpertDropoutLayer::new(dropout_config, moe_config, &device)?;
    println!("✓ Expert dropout layer created successfully");
    println!("  Number of experts: {}", dropout_layer.num_experts());
    println!("  Training mode: {}", dropout_layer.is_training());

    // Test different strategies
    let strategies = vec![
        DropoutStrategy::Random,
        DropoutStrategy::Block,
        DropoutStrategy::Progressive, 
        DropoutStrategy::LoadAware,
    ];

    for strategy in strategies {
        let config = ExpertDropoutConfig::new(0.1, strategy.clone());
        assert!(config.validate().is_ok());
        println!("✓ Strategy {:?} validation passed", strategy);
    }

    // Test dropout scheduler
    let scheduler = DropoutScheduler::new(0.5, 0.1, 1000);
    println!("✓ Dropout scheduler created");
    println!("  Initial rate: {}", scheduler.get_dropout_rate(0));
    println!("  Mid rate: {}", scheduler.get_dropout_rate(500));
    println!("  Final rate: {}", scheduler.get_dropout_rate(1000));

    // Test expert importance scorer
    let scorer = ExpertImportanceScorer::new(8);
    println!("✓ Expert importance scorer created");
    println!("  Initial scores: {:?}", scorer.get_importance_scores());

    println!("\n🎉 All expert dropout tests passed!");
    Ok(())
}