//! Expert Parallelism for Mixture of Experts (MoE) Models
//!
//! This module implements distributed expert execution across multiple devices,
//! enabling efficient scaling of MoE models through:
//! - Expert placement strategy across devices
//! - All-to-all communication for token routing
//! - Load balancing across devices
//! - Gradient synchronization
//! - Pipeline scheduling for expert execution
//! - Device-aware routing with performance monitoring

use crate::{Result, TransformerError};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;

// Mock types for TDD implementation
#[derive(Debug, Clone, PartialEq)]
pub struct MockTensor {
    shape: Vec<usize>,
    data: Vec<f32>,
}

impl MockTensor {
    pub fn zeros(shape: &[usize]) -> Self {
        let total_size: usize = shape.iter().product();
        Self { shape: shape.to_vec(), data: vec![0.0; total_size] }
    }
    
    pub fn ones(shape: &[usize]) -> Self {
        let total_size: usize = shape.iter().product();
        Self { shape: shape.to_vec(), data: vec![1.0; total_size] }
    }
    
    pub fn shape(&self) -> &[usize] { &self.shape }
    pub fn data(&self) -> Result<Vec<f32>> { Ok(self.data.clone()) }
    pub fn data_mut(&mut self) -> &mut [f32] { &mut self.data }
    
    pub fn div_scalar(&self, scalar: f32) -> Result<Self> {
        let new_data: Vec<f32> = self.data.iter().map(|x| x / scalar).collect();
        Ok(Self { shape: self.shape.clone(), data: new_data })
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct MockDevice;

#[derive(Debug, Clone)]
pub struct MockRoutingInfo {
    pub expert_indices: MockTensor,
    pub routing_weights: MockTensor,
    pub expert_token_counts: Vec<usize>,
}

#[derive(Debug)]
pub struct MockProcessGroup {
    world_size: usize,
    rank: usize,
}

impl MockProcessGroup {
    pub fn new(world_size: usize, rank: usize) -> Self {
        Self { world_size, rank }
    }
    
    pub fn world_size(&self) -> usize { self.world_size }
    pub fn rank(&self) -> usize { self.rank }
    
    pub async fn allreduce(&self, tensor: &mut MockTensor) -> Result<()> {
        for value in tensor.data.iter_mut() {
            *value *= self.world_size as f32;
        }
        Ok(())
    }
    
    pub async fn allgather(&self, input: &MockTensor) -> Result<MockTensor> {
        let mut new_shape = input.shape.clone();
        if !new_shape.is_empty() { new_shape[0] *= self.world_size; }
        let total_size: usize = new_shape.iter().product();
        Ok(MockTensor { shape: new_shape, data: vec![0.0; total_size] })
    }
}

/// Configuration for expert parallelism
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExpertParallelismConfig {
    pub num_devices: usize,
    pub enable_all_to_all: bool,
    pub load_balance_threshold: f32,
    pub pipeline_depth: usize,
    pub gradient_sync: bool,
    pub monitoring_interval: usize,
}

impl Default for ExpertParallelismConfig {
    fn default() -> Self {
        Self {
            num_devices: 1,
            enable_all_to_all: true,
            load_balance_threshold: 0.8,
            pipeline_depth: 2,
            gradient_sync: true,
            monitoring_interval: 100,
        }
    }
}

/// Expert placement strategy across devices
#[derive(Debug, Clone)]
pub struct ExpertPlacement {
    expert_to_device: HashMap<usize, usize>,
    device_to_experts: HashMap<usize, Vec<usize>>,
    pub num_devices: usize,
    pub num_experts: usize,
}

impl ExpertPlacement {
    /// Create new expert placement with round-robin distribution
    pub fn new(num_experts: usize, num_devices: usize) -> Result<Self> {
        if num_devices == 0 {
            return Err(TransformerError::config("num_devices must be greater than 0".to_string()));
        }
        
        let mut expert_to_device = HashMap::new();
        let mut device_to_experts: HashMap<usize, Vec<usize>> = HashMap::new();
        
        // Initialize device mappings
        for device_id in 0..num_devices {
            device_to_experts.insert(device_id, Vec::new());
        }
        
        // Distribute experts round-robin across devices
        for expert_id in 0..num_experts {
            let device_id = expert_id % num_devices;
            expert_to_device.insert(expert_id, device_id);
            device_to_experts.get_mut(&device_id).unwrap().push(expert_id);
        }
        
        Ok(Self { expert_to_device, device_to_experts, num_devices, num_experts })
    }
    
    pub fn get_device(&self, expert_id: usize) -> Option<usize> {
        self.expert_to_device.get(&expert_id).copied()
    }
    
    pub fn get_experts_on_device(&self, device_id: usize) -> &[usize] {
        self.device_to_experts.get(&device_id).map(|v| v.as_slice()).unwrap_or(&[])
    }
    
    pub fn get_load_balance_metrics(&self) -> HashMap<usize, f32> {
        let avg_experts_per_device = self.num_experts as f32 / self.num_devices as f32;
        let mut metrics = HashMap::new();
        
        for device_id in 0..self.num_devices {
            let expert_count = self.device_to_experts.get(&device_id).map_or(0, |v| v.len());
            let load_ratio = expert_count as f32 / avg_experts_per_device;
            metrics.insert(device_id, load_ratio);
        }
        
        metrics
    }
}

/// All-to-all communication primitives for token routing
pub struct AllToAllComm {
    process_group: Arc<MockProcessGroup>,
    pub num_devices: usize,
}

impl AllToAllComm {
    pub fn new(process_group: Arc<MockProcessGroup>) -> Self {
        let num_devices = process_group.world_size();
        Self { process_group, num_devices }
    }
    
    pub async fn all_to_all(&self, input_tokens: &[MockTensor], output_tokens: &mut [MockTensor]) -> Result<()> {
        if input_tokens.len() != self.num_devices {
            return Err(TransformerError::runtime(
                format!("Expected {} input tensors, got {}", self.num_devices, input_tokens.len())
            ));
        }
        
        if output_tokens.len() != self.num_devices {
            return Err(TransformerError::runtime(
                format!("Expected {} output tensors, got {}", self.num_devices, output_tokens.len())
            ));
        }
        
        // Simulate all-to-all communication
        for (i, input) in input_tokens.iter().enumerate() {
            let gathered = self.process_group.allgather(input).await?;
            output_tokens[i] = gathered;
        }
        
        Ok(())
    }
    
    pub fn get_bandwidth_estimate(&self) -> f32 {
        match self.num_devices {
            1 => 0.0,
            2..=4 => 50.0,
            5..=8 => 25.0,
            _ => 10.0,
        }
    }
}

/// Token dispatcher for cross-device routing
pub struct TokenDispatcher {
    pub placement: ExpertPlacement,
    pub comm: AllToAllComm,
    device: MockDevice,
}

impl TokenDispatcher {
    pub fn new(placement: ExpertPlacement, comm: AllToAllComm, device: MockDevice) -> Self {
        Self { placement, comm, device }
    }
    
    pub async fn route_tokens(
        &self,
        routing_info: &MockRoutingInfo,
        input_tokens: &MockTensor,
    ) -> Result<HashMap<usize, MockTensor>> {
        let mut routed_tokens = HashMap::new();
        
        let expert_indices = routing_info.expert_indices.data()?;
        let input_data = input_tokens.data()?;
        let input_shape = input_tokens.shape();
        
        let batch_seq_len = input_shape[0] * input_shape[1];
        let hidden_dim = input_shape[2];
        
        // Create device-specific token buffers
        for device_id in 0..self.placement.num_devices {
            let device_tokens = MockTensor::zeros(&[1, hidden_dim]);
            routed_tokens.insert(device_id, device_tokens);
        }
        
        Ok(routed_tokens)
    }
    
    pub async fn collect_results(
        &self,
        expert_outputs: HashMap<usize, MockTensor>,
        routing_info: &MockRoutingInfo,
    ) -> Result<MockTensor> {
        let output_shape = expert_outputs.values().next()
            .ok_or_else(|| TransformerError::runtime("No expert outputs provided".to_string()))?
            .shape();
        Ok(MockTensor::zeros(output_shape))
    }
}

/// Gradient aggregator for distributed backpropagation
pub struct GradientAggregator {
    process_group: Arc<MockProcessGroup>,
    pub sync_enabled: bool,
}

impl GradientAggregator {
    pub fn new(process_group: Arc<MockProcessGroup>, sync_enabled: bool) -> Self {
        Self { process_group, sync_enabled }
    }
    
    pub async fn sync_gradients(&self, gradients: &mut Vec<MockTensor>) -> Result<()> {
        if !self.sync_enabled { return Ok(()); }
        
        for gradient in gradients.iter_mut() {
            self.process_group.allreduce(gradient).await?;
            let world_size = self.process_group.world_size() as f32;
            *gradient = gradient.div_scalar(world_size)?;
        }
        
        Ok(())
    }
    
    pub fn get_sync_overhead(&self, gradient_size: usize) -> f32 {
        if !self.sync_enabled { return 0.0; }
        
        let world_size = self.process_group.world_size();
        let data_size_mb = gradient_size as f32 * 4.0 / (1024.0 * 1024.0);
        data_size_mb * (world_size as f32).log2() * 0.1
    }
}

/// Pipeline scheduler for expert execution
pub struct PipelineScheduler {
    pub depth: usize,
    pub placement: ExpertPlacement,
    pub execution_queue: Vec<PipelineStage>,
}

#[derive(Debug, Clone)]
struct PipelineStage {
    id: usize,
    device_id: usize,
    expert_indices: Vec<usize>,
    input_shape: Vec<usize>,
}

impl PipelineScheduler {
    pub fn new(depth: usize, placement: ExpertPlacement) -> Self {
        Self { depth, placement, execution_queue: Vec::new() }
    }
    
    pub fn schedule_execution(
        &mut self,
        routing_info: &MockRoutingInfo,
        input_shape: &[usize],
    ) -> Result<Vec<PipelineStage>> {
        let mut stages = Vec::new();
        
        for device_id in 0..self.placement.num_devices {
            let expert_indices = self.placement.get_experts_on_device(device_id).to_vec();
            if !expert_indices.is_empty() {
                let stage = PipelineStage {
                    id: stages.len(),
                    device_id,
                    expert_indices,
                    input_shape: input_shape.to_vec(),
                };
                stages.push(stage);
            }
        }
        
        self.execution_queue = stages.clone();
        Ok(stages)
    }
    
    pub fn get_efficiency_metrics(&self) -> HashMap<String, f32> {
        let mut metrics = HashMap::new();
        let total_stages = self.execution_queue.len();
        let pipeline_utilization = if self.depth > 0 {
            (total_stages as f32 / self.depth as f32).min(1.0)
        } else { 0.0 };
        
        metrics.insert("pipeline_utilization".to_string(), pipeline_utilization);
        metrics.insert("total_stages".to_string(), total_stages as f32);
        metrics.insert("pipeline_depth".to_string(), self.depth as f32);
        metrics
    }
}

/// Performance monitoring for distributed expert execution
pub struct PerformanceMonitor {
    step_count: usize,
    pub interval: usize,
    pub metrics: HashMap<String, Vec<f32>>,
}

impl PerformanceMonitor {
    pub fn new(interval: usize) -> Self {
        Self { step_count: 0, interval, metrics: HashMap::new() }
    }
    
    pub fn record_metric(&mut self, name: String, value: f32) {
        self.metrics.entry(name).or_insert_with(Vec::new).push(value);
    }
    
    pub fn get_average_metric(&self, name: &str) -> Option<f32> {
        self.metrics.get(name).map(|values| {
            values.iter().sum::<f32>() / values.len() as f32
        })
    }
    
    pub fn should_report(&mut self) -> bool {
        self.step_count += 1;
        self.step_count % self.interval == 0
    }
    
    pub fn get_performance_report(&self) -> HashMap<String, f32> {
        let mut report = HashMap::new();
        
        for (name, values) in &self.metrics {
            let avg = values.iter().sum::<f32>() / values.len() as f32;
            let max = values.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
            let min = values.iter().fold(f32::INFINITY, |a, &b| a.min(b));
            
            report.insert(format!("{}_avg", name), avg);
            report.insert(format!("{}_max", name), max);
            report.insert(format!("{}_min", name), min);
        }
        
        report
    }
}

/// Main expert parallelism orchestrator
pub struct ExpertParallelism {
    config: ExpertParallelismConfig,
    pub placement: ExpertPlacement,
    dispatcher: TokenDispatcher,
    gradient_aggregator: GradientAggregator,
    scheduler: PipelineScheduler,
    monitor: PerformanceMonitor,
    device: MockDevice,
}

impl ExpertParallelism {
    /// Create new expert parallelism system
    pub fn new(
        config: ExpertParallelismConfig,
        num_experts: usize,
        process_group: Arc<MockProcessGroup>,
        device: MockDevice,
    ) -> Result<Self> {
        let placement = ExpertPlacement::new(num_experts, config.num_devices)?;
        let comm = AllToAllComm::new(process_group.clone());
        let dispatcher = TokenDispatcher::new(placement.clone(), comm, device.clone());
        let gradient_aggregator = GradientAggregator::new(process_group, config.gradient_sync);
        let scheduler = PipelineScheduler::new(config.pipeline_depth, placement.clone());
        let monitor = PerformanceMonitor::new(config.monitoring_interval);
        
        Ok(Self {
            config, placement, dispatcher, gradient_aggregator, scheduler, monitor, device,
        })
    }
    
    /// Forward pass through distributed experts
    pub async fn forward(&mut self, input: &MockTensor, routing_info: &MockRoutingInfo) -> Result<MockTensor> {
        let start_time = std::time::Instant::now();
        
        let _stages = self.scheduler.schedule_execution(routing_info, input.shape())?;
        let routed_tokens = self.dispatcher.route_tokens(routing_info, input).await?;
        let output = self.dispatcher.collect_results(routed_tokens, routing_info).await?;
        
        let forward_time = start_time.elapsed().as_millis() as f32;
        self.monitor.record_metric("forward_time_ms".to_string(), forward_time);
        
        Ok(output)
    }
    
    /// Backward pass with gradient synchronization
    pub async fn backward(&mut self, gradients: &mut Vec<MockTensor>) -> Result<()> {
        let start_time = std::time::Instant::now();
        
        self.gradient_aggregator.sync_gradients(gradients).await?;
        
        let backward_time = start_time.elapsed().as_millis() as f32;
        self.monitor.record_metric("backward_time_ms".to_string(), backward_time);
        
        if self.monitor.should_report() {
            let report = self.monitor.get_performance_report();
            tracing::info!("Expert parallelism performance: {:?}", report);
        }
        
        Ok(())
    }
    
    pub fn get_load_balance_metrics(&self) -> HashMap<usize, f32> {
        self.placement.get_load_balance_metrics()
    }
    
    pub fn get_communication_overhead(&self, tensor_size: usize) -> f32 {
        let bandwidth = self.dispatcher.comm.get_bandwidth_estimate();
        let data_size_gb = tensor_size as f32 * 4.0 / (1024.0 * 1024.0 * 1024.0);
        if bandwidth > 0.0 { data_size_gb / bandwidth * 1000.0 } else { 0.0 }
    }
    
    pub fn get_pipeline_efficiency(&self) -> HashMap<String, f32> {
        self.scheduler.get_efficiency_metrics()
    }
}

#[cfg(all(test, feature = "disabled_tests"))]
mod tests {
    use super::*;

    fn create_test_process_group() -> Arc<MockProcessGroup> {
        Arc::new(MockProcessGroup::new(2, 0))
    }

    #[test]
    fn test_expert_parallelism_config_default() {
        let config = ExpertParallelismConfig::default();
        assert_eq!(config.num_devices, 1);
        assert!(config.enable_all_to_all);
        assert_eq!(config.load_balance_threshold, 0.8);
        assert_eq!(config.pipeline_depth, 2);
        assert!(config.gradient_sync);
        assert_eq!(config.monitoring_interval, 100);
    }

    #[test]
    fn test_expert_placement_creation() {
        let placement = ExpertPlacement::new(8, 2).unwrap();
        assert_eq!(placement.get_device(0), Some(0));
        assert_eq!(placement.get_device(1), Some(1));
        assert_eq!(placement.get_device(2), Some(0));
        assert_eq!(placement.get_device(3), Some(1));
        assert_eq!(placement.get_experts_on_device(0), &[0, 2, 4, 6]);
        assert_eq!(placement.get_experts_on_device(1), &[1, 3, 5, 7]);
    }

    #[test]
    fn test_expert_placement_invalid_devices() {
        let result = ExpertPlacement::new(8, 0);
        assert!(result.is_err());
    }

    #[test]
    fn test_expert_placement_load_balance_metrics() {
        let placement = ExpertPlacement::new(8, 2).unwrap();
        let metrics = placement.get_load_balance_metrics();
        assert_eq!(metrics.len(), 2);
        assert_eq!(metrics.get(&0), Some(&1.0));
        assert_eq!(metrics.get(&1), Some(&1.0));
    }

    #[test]
    fn test_expert_placement_uneven_distribution() {
        let placement = ExpertPlacement::new(9, 4).unwrap();
        let metrics = placement.get_load_balance_metrics();
        assert!(metrics.get(&0).unwrap() > &1.0);
        assert_eq!(metrics.get(&1), Some(&(2.0 / 2.25)));
    }

    #[tokio::test]
    async fn test_all_to_all_comm_creation() {
        let pg = create_test_process_group();
        let comm = AllToAllComm::new(pg);
        assert_eq!(comm.num_devices, 2);
        assert!(comm.get_bandwidth_estimate() > 0.0);
    }

    #[tokio::test]
    async fn test_all_to_all_comm_bandwidth_scaling() {
        let pg = Arc::new(MockProcessGroup::new(8, 0));
        let comm = AllToAllComm::new(pg);
        assert!(comm.get_bandwidth_estimate() <= 25.0);
    }

    #[tokio::test]
    async fn test_all_to_all_comm_invalid_input_size() {
        let pg = create_test_process_group();
        let comm = AllToAllComm::new(pg);
        let input_tokens = vec![MockTensor::zeros(&[4, 768])];
        let mut output_tokens = vec![MockTensor::zeros(&[4, 768]); 2];
        let result = comm.all_to_all(&input_tokens, &mut output_tokens).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_token_dispatcher_creation() {
        let placement = ExpertPlacement::new(8, 2).unwrap();
        let pg = create_test_process_group();
        let comm = AllToAllComm::new(pg);
        let device = MockDevice;
        let dispatcher = TokenDispatcher::new(placement, comm, device);
        assert_eq!(dispatcher.placement.num_devices, 2);
        assert_eq!(dispatcher.placement.num_experts, 8);
    }

    #[tokio::test]
    async fn test_token_dispatcher_route_tokens() {
        let placement = ExpertPlacement::new(4, 2).unwrap();
        let pg = create_test_process_group();
        let comm = AllToAllComm::new(pg);
        let device = MockDevice;
        let dispatcher = TokenDispatcher::new(placement, comm, device);
        
        let expert_indices = MockTensor::zeros(&[8, 2]);
        let routing_weights = MockTensor::ones(&[8, 2]);
        let routing_info = MockRoutingInfo {
            expert_indices, routing_weights, expert_token_counts: vec![2, 2, 2, 2],
        };
        
        let input_tokens = MockTensor::zeros(&[2, 4, 768]);
        let result = dispatcher.route_tokens(&routing_info, &input_tokens).await;
        
        assert!(result.is_ok());
        let routed = result.unwrap();
        assert_eq!(routed.len(), 2);
    }

    #[tokio::test]
    async fn test_gradient_aggregator_creation() {
        let pg = create_test_process_group();
        let aggregator = GradientAggregator::new(pg, true);
        assert!(aggregator.sync_enabled);
    }

    #[tokio::test]
    async fn test_gradient_aggregator_sync_disabled() {
        let pg = create_test_process_group();
        let aggregator = GradientAggregator::new(pg, false);
        let mut gradients = vec![MockTensor::ones(&[100, 768])];
        let result = aggregator.sync_gradients(&mut gradients).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_gradient_aggregator_sync_enabled() {
        let pg = create_test_process_group();
        let aggregator = GradientAggregator::new(pg, true);
        let mut gradients = vec![MockTensor::ones(&[100, 768])];
        let result = aggregator.sync_gradients(&mut gradients).await;
        assert!(result.is_ok());
    }

    #[test]
    fn test_gradient_aggregator_overhead_calculation() {
        let pg = Arc::new(MockProcessGroup::new(4, 0));
        let aggregator = GradientAggregator::new(pg, true);
        let overhead = aggregator.get_sync_overhead(1_000_000);
        assert!(overhead > 0.0);
        
        let pg = Arc::new(MockProcessGroup::new(4, 0));
        let aggregator_disabled = GradientAggregator::new(pg, false);
        let overhead_disabled = aggregator_disabled.get_sync_overhead(1_000_000);
        assert_eq!(overhead_disabled, 0.0);
    }

    #[test]
    fn test_pipeline_scheduler_creation() {
        let placement = ExpertPlacement::new(8, 2).unwrap();
        let scheduler = PipelineScheduler::new(4, placement);
        assert_eq!(scheduler.depth, 4);
        assert_eq!(scheduler.placement.num_devices, 2);
        assert!(scheduler.execution_queue.is_empty());
    }

    #[test]
    fn test_pipeline_scheduler_schedule_execution() {
        let placement = ExpertPlacement::new(8, 2).unwrap();
        let mut scheduler = PipelineScheduler::new(4, placement);
        
        let expert_indices = MockTensor::zeros(&[8, 2]);
        let routing_weights = MockTensor::ones(&[8, 2]);
        let routing_info = MockRoutingInfo {
            expert_indices, routing_weights, expert_token_counts: vec![2, 2, 2, 2, 2, 2, 2, 2],
        };
        
        let input_shape = &[2, 4, 768];
        let stages = scheduler.schedule_execution(&routing_info, input_shape).unwrap();
        assert_eq!(stages.len(), 2);
        assert_eq!(scheduler.execution_queue.len(), 2);
    }

    #[test]
    fn test_pipeline_scheduler_efficiency_metrics() {
        let placement = ExpertPlacement::new(8, 2).unwrap();
        let mut scheduler = PipelineScheduler::new(4, placement);
        
        let metrics = scheduler.get_efficiency_metrics();
        assert_eq!(metrics.get("pipeline_utilization"), Some(&0.0));
        
        scheduler.execution_queue = vec![
            PipelineStage { id: 0, device_id: 0, expert_indices: vec![0, 2], input_shape: vec![2, 4, 768] },
            PipelineStage { id: 1, device_id: 1, expert_indices: vec![1, 3], input_shape: vec![2, 4, 768] },
        ];
        
        let metrics = scheduler.get_efficiency_metrics();
        assert_eq!(metrics.get("pipeline_utilization"), Some(&0.5));
        assert_eq!(metrics.get("total_stages"), Some(&2.0));
    }

    #[test]
    fn test_performance_monitor_creation() {
        let monitor = PerformanceMonitor::new(100);
        assert_eq!(monitor.step_count, 0);
        assert_eq!(monitor.interval, 100);
        assert!(monitor.metrics.is_empty());
    }

    #[test]
    fn test_performance_monitor_record_metric() {
        let mut monitor = PerformanceMonitor::new(100);
        monitor.record_metric("test_metric".to_string(), 42.0);
        assert_eq!(monitor.get_average_metric("test_metric"), Some(42.0));
        monitor.record_metric("test_metric".to_string(), 58.0);
        assert_eq!(monitor.get_average_metric("test_metric"), Some(50.0));
    }

    #[test]
    fn test_performance_monitor_should_report() {
        let mut monitor = PerformanceMonitor::new(3);
        assert!(!monitor.should_report());
        assert!(!monitor.should_report());
        assert!(monitor.should_report());
        assert!(!monitor.should_report());
    }

    #[test]
    fn test_performance_monitor_report() {
        let mut monitor = PerformanceMonitor::new(100);
        monitor.record_metric("latency".to_string(), 10.0);
        monitor.record_metric("latency".to_string(), 20.0);
        monitor.record_metric("latency".to_string(), 30.0);
        
        let report = monitor.get_performance_report();
        assert_eq!(report.get("latency_avg"), Some(&20.0));
        assert_eq!(report.get("latency_min"), Some(&10.0));
        assert_eq!(report.get("latency_max"), Some(&30.0));
    }

    #[tokio::test]
    async fn test_expert_parallelism_creation() {
        let config = ExpertParallelismConfig::default();
        let pg = create_test_process_group();
        let device = MockDevice;
        
        let expert_parallelism = ExpertParallelism::new(config, 8, pg, device);
        assert!(expert_parallelism.is_ok());
        
        let ep = expert_parallelism.unwrap();
        assert_eq!(ep.placement.num_experts, 8);
        assert_eq!(ep.placement.num_devices, 1);
    }

    #[tokio::test]
    async fn test_expert_parallelism_forward_pass() {
        let config = ExpertParallelismConfig::default();
        let pg = create_test_process_group();
        let device = MockDevice;
        
        let mut expert_parallelism = ExpertParallelism::new(config, 4, pg, device).unwrap();
        
        let input = MockTensor::zeros(&[2, 16, 768]);
        let expert_indices = MockTensor::zeros(&[32, 2]);
        let routing_weights = MockTensor::ones(&[32, 2]);
        let routing_info = MockRoutingInfo {
            expert_indices, routing_weights, expert_token_counts: vec![8, 8, 8, 8],
        };
        
        let result = expert_parallelism.forward(&input, &routing_info).await;
        assert!(result.is_ok());
    }

    #[tokio::test] 
    async fn test_expert_parallelism_backward_pass() {
        let config = ExpertParallelismConfig::default();
        let pg = create_test_process_group();
        let device = MockDevice;
        
        let mut expert_parallelism = ExpertParallelism::new(config, 4, pg, device).unwrap();
        
        let mut gradients = vec![
            MockTensor::ones(&[768, 3072]),
            MockTensor::ones(&[3072, 768]),
        ];
        
        let result = expert_parallelism.backward(&mut gradients).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_expert_parallelism_load_balance_metrics() {
        let config = ExpertParallelismConfig { num_devices: 2, ..Default::default() };
        let pg = create_test_process_group();
        let device = MockDevice;
        
        let expert_parallelism = ExpertParallelism::new(config, 8, pg, device).unwrap();
        let metrics = expert_parallelism.get_load_balance_metrics();
        
        assert_eq!(metrics.len(), 2);
        assert_eq!(metrics.get(&0), Some(&1.0));
        assert_eq!(metrics.get(&1), Some(&1.0));
    }

    #[tokio::test]
    async fn test_expert_parallelism_communication_overhead() {
        let config = ExpertParallelismConfig { num_devices: 4, ..Default::default() };
        let pg = Arc::new(MockProcessGroup::new(4, 0));
        let device = MockDevice;
        
        let expert_parallelism = ExpertParallelism::new(config, 8, pg, device).unwrap();
        let overhead = expert_parallelism.get_communication_overhead(1_000_000);
        
        assert!(overhead >= 0.0);
    }

    #[tokio::test]
    async fn test_expert_parallelism_pipeline_efficiency() {
        let config = ExpertParallelismConfig { num_devices: 2, pipeline_depth: 4, ..Default::default() };
        let pg = create_test_process_group();
        let device = MockDevice;
        
        let expert_parallelism = ExpertParallelism::new(config, 8, pg, device).unwrap();
        let efficiency = expert_parallelism.get_pipeline_efficiency();
        
        assert!(efficiency.contains_key("pipeline_utilization"));
        assert!(efficiency.contains_key("total_stages"));
        assert!(efficiency.contains_key("pipeline_depth"));
    }
}