//! Load balancer with work stealing for heterogeneous compute.
//!
//! Supports:
//! - Cross-processor-type work migration
//! - Energy-aware balancing
//! - Latency-aware scheduling
//! - Real-time utilization metrics

use crate::device::{DeviceInfo, DeviceRegistry};
use crate::processor::{Operation, OperationType, ProcessorId, ProcessorType};
use crate::task::{Task, TaskId, TaskPriority};
use super::TaskAssignment;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Balancing strategy for the load balancer.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BalancingStrategy {
    /// Optimize for speed (minimize execution time).
    Speed,
    /// Optimize for energy efficiency.
    Energy,
    /// Balance speed and energy.
    Balanced,
    /// Optimize for cost (spot pricing).
    Cost,
    /// Optimize for latency (inference workloads).
    Latency,
}

impl Default for BalancingStrategy {
    fn default() -> Self {
        BalancingStrategy::Balanced
    }
}

/// Real-time processor metrics.
#[derive(Clone, Debug, Default)]
pub struct ProcessorMetrics {
    /// Current utilization (0.0 - 1.0).
    pub utilization: f64,
    /// Queue depth (pending tasks).
    pub queue_depth: u64,
    /// Average task completion time (ms).
    pub avg_completion_ms: f64,
    /// Tasks completed in last minute.
    pub throughput_per_min: u64,
    /// Current power draw (watts).
    pub power_watts: f64,
    /// Temperature (celsius).
    pub temperature: f64,
    /// Last updated timestamp.
    pub last_updated: Option<Instant>,
}

/// Load balancer for heterogeneous compute environments.
pub struct LoadBalancer {
    /// Device registry for processor info.
    device_registry: Option<Arc<DeviceRegistry>>,
    /// Current load per processor (task count).
    loads: RwLock<HashMap<ProcessorId, AtomicU64>>,
    /// Real-time metrics per processor.
    metrics: RwLock<HashMap<ProcessorId, ProcessorMetrics>>,
    /// Processor type mapping.
    processor_types: RwLock<HashMap<ProcessorId, ProcessorType>>,
    /// Work stealing threshold (0.0 - 1.0).
    steal_threshold: f64,
    /// Rebalance threshold (0.0 - 1.0).
    rebalance_threshold: f64,
    /// Current balancing strategy.
    strategy: RwLock<BalancingStrategy>,
    /// Migration history (to prevent thrashing).
    migration_history: RwLock<Vec<MigrationRecord>>,
}

/// Record of a task migration.
#[derive(Clone, Debug)]
struct MigrationRecord {
    task_id: TaskId,
    from: ProcessorId,
    to: ProcessorId,
    timestamp: Instant,
}

impl LoadBalancer {
    /// Creates a new load balancer.
    pub fn new() -> Self {
        Self {
            device_registry: None,
            loads: RwLock::new(HashMap::new()),
            metrics: RwLock::new(HashMap::new()),
            processor_types: RwLock::new(HashMap::new()),
            steal_threshold: 0.3,
            rebalance_threshold: 0.2,
            strategy: RwLock::new(BalancingStrategy::default()),
            migration_history: RwLock::new(Vec::new()),
        }
    }

    /// Creates a load balancer with device registry.
    pub fn with_registry(device_registry: Arc<DeviceRegistry>) -> Self {
        Self {
            device_registry: Some(device_registry),
            loads: RwLock::new(HashMap::new()),
            metrics: RwLock::new(HashMap::new()),
            processor_types: RwLock::new(HashMap::new()),
            steal_threshold: 0.3,
            rebalance_threshold: 0.2,
            strategy: RwLock::new(BalancingStrategy::default()),
            migration_history: RwLock::new(Vec::new()),
        }
    }

    /// Sets the balancing strategy.
    pub fn set_strategy(&self, strategy: BalancingStrategy) {
        *self.strategy.write() = strategy;
    }

    /// Gets the current strategy.
    pub fn strategy(&self) -> BalancingStrategy {
        *self.strategy.read()
    }

    /// Register a processor with its type.
    pub fn register_processor(&self, processor_id: ProcessorId, processor_type: ProcessorType) {
        self.loads.write().insert(processor_id, AtomicU64::new(0));
        self.metrics.write().insert(processor_id, ProcessorMetrics::default());
        self.processor_types.write().insert(processor_id, processor_type);
    }

    /// Unregister a processor.
    pub fn unregister_processor(&self, processor_id: ProcessorId) {
        self.loads.write().remove(&processor_id);
        self.metrics.write().remove(&processor_id);
        self.processor_types.write().remove(&processor_id);
    }

    /// Update real-time metrics for a processor.
    pub fn update_metrics(&self, processor_id: ProcessorId, metrics: ProcessorMetrics) {
        if let Some(existing) = self.metrics.write().get_mut(&processor_id) {
            *existing = ProcessorMetrics {
                last_updated: Some(Instant::now()),
                ..metrics
            };
        }
    }

    /// Get current load for a processor.
    pub fn get_load(&self, processor_id: ProcessorId) -> u64 {
        self.loads.read()
            .get(&processor_id)
            .map(|l| l.load(Ordering::Relaxed))
            .unwrap_or(0)
    }

    /// Increment load for a processor.
    pub fn increment_load(&self, processor_id: ProcessorId) {
        if let Some(load) = self.loads.read().get(&processor_id) {
            load.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Decrement load for a processor.
    pub fn decrement_load(&self, processor_id: ProcessorId) {
        if let Some(load) = self.loads.read().get(&processor_id) {
            load.fetch_sub(1, Ordering::Relaxed);
        }
    }

    /// Check if an operation can run on a processor type.
    pub fn can_execute(&self, op: &Operation, processor_type: &ProcessorType) -> bool {
        let op_type = op.op_type();

        match processor_type {
            // CPUs can handle most sequential operations
            ProcessorType::Cpu(_) => matches!(
                op_type,
                OperationType::MatMul
                | OperationType::Conv2d
                | OperationType::Conv3d
                | OperationType::DepthwiseConv
                | OperationType::BatchNorm
                | OperationType::LayerNorm
                | OperationType::Add
                | OperationType::Mul
                | OperationType::ReLU
                | OperationType::GeLU
                | OperationType::SiLU
                | OperationType::Softmax
                | OperationType::Sum
                | OperationType::Mean
                | OperationType::Max
                | OperationType::ArgMax
                | OperationType::Embedding
                | OperationType::TopK
                | OperationType::Sampling
                | OperationType::Tokenization
                | OperationType::Detokenization
                | OperationType::DataLoad
                | OperationType::DataPreprocess
                | OperationType::Transpose
                | OperationType::Reshape
                | OperationType::Concat
                | OperationType::Split
            ),

            // GPUs excel at parallel operations
            ProcessorType::Gpu(_) => matches!(
                op_type,
                OperationType::MatMul
                | OperationType::Conv2d
                | OperationType::Conv3d
                | OperationType::DepthwiseConv
                | OperationType::BatchNorm
                | OperationType::LayerNorm
                | OperationType::SelfAttention
                | OperationType::CrossAttention
                | OperationType::FlashAttention
                | OperationType::Add
                | OperationType::Mul
                | OperationType::ReLU
                | OperationType::GeLU
                | OperationType::SiLU
                | OperationType::Softmax
                | OperationType::Sum
                | OperationType::Mean
                | OperationType::Max
                | OperationType::ArgMax
                | OperationType::Embedding
                | OperationType::RoPE
                | OperationType::KVCache
                | OperationType::TopK
                | OperationType::Sampling
                | OperationType::Transpose
                | OperationType::Reshape
                | OperationType::Concat
                | OperationType::Split
                | OperationType::Gather
                | OperationType::Scatter
                | OperationType::AllReduce
                | OperationType::AllGather
                | OperationType::ReduceScatter
                | OperationType::Backward
                | OperationType::OptimizerStep
                | OperationType::GradientClip
            ),

            // TPUs optimized for ML
            ProcessorType::Tpu(_) => matches!(
                op_type,
                OperationType::MatMul
                | OperationType::Conv2d
                | OperationType::BatchNorm
                | OperationType::LayerNorm
                | OperationType::SelfAttention
                | OperationType::CrossAttention
                | OperationType::FlashAttention
                | OperationType::Add
                | OperationType::Mul
                | OperationType::ReLU
                | OperationType::GeLU
                | OperationType::SiLU
                | OperationType::Softmax
                | OperationType::Sum
                | OperationType::Mean
                | OperationType::Embedding
                | OperationType::RoPE
                | OperationType::KVCache
                | OperationType::AllReduce
                | OperationType::AllGather
                | OperationType::ReduceScatter
                | OperationType::Backward
                | OperationType::OptimizerStep
            ),

            // NPUs for neural network inference
            ProcessorType::Npu(_) => matches!(
                op_type,
                OperationType::MatMul
                | OperationType::Conv2d
                | OperationType::DepthwiseConv
                | OperationType::BatchNorm
                | OperationType::LayerNorm
                | OperationType::SelfAttention
                | OperationType::Add
                | OperationType::Mul
                | OperationType::ReLU
                | OperationType::GeLU
                | OperationType::SiLU
                | OperationType::Softmax
                | OperationType::Sum
                | OperationType::Mean
            ),

            // LPUs for sequential inference (optimized for LLMs)
            ProcessorType::Lpu => matches!(
                op_type,
                OperationType::MatMul
                | OperationType::LayerNorm
                | OperationType::SelfAttention
                | OperationType::FlashAttention
                | OperationType::Add
                | OperationType::Mul
                | OperationType::ReLU
                | OperationType::GeLU
                | OperationType::SiLU
                | OperationType::Softmax
                | OperationType::Embedding
                | OperationType::RoPE
                | OperationType::KVCache
                | OperationType::TopK
                | OperationType::Sampling
            ),

            // FPGAs can be programmed for anything
            ProcessorType::Fpga(_) => true,

            // DSPs for signal processing
            ProcessorType::Dsp(_) => matches!(
                op_type,
                OperationType::Conv2d
                | OperationType::DepthwiseConv
                | OperationType::Add
                | OperationType::Mul
                | OperationType::Sum
                | OperationType::Mean
                | OperationType::Max
            ),

            // WebGPU has limited operations
            ProcessorType::WebGpu => matches!(
                op_type,
                OperationType::MatMul
                | OperationType::Conv2d
                | OperationType::Add
                | OperationType::Mul
                | OperationType::ReLU
                | OperationType::Softmax
                | OperationType::Sum
                | OperationType::Transpose
                | OperationType::Reshape
            ),

            // WASM for portable compute
            ProcessorType::Wasm => matches!(
                op_type,
                OperationType::MatMul
                | OperationType::Add
                | OperationType::Mul
                | OperationType::ReLU
                | OperationType::Softmax
                | OperationType::Sum
                | OperationType::Mean
                | OperationType::Tokenization
                | OperationType::Detokenization
            ),

            // Custom processors - assume they can handle anything
            ProcessorType::Custom { .. } => true,
        }
    }

    /// Calculate a score for assigning a task to a processor.
    fn calculate_score(
        &self,
        task: &Task,
        processor_id: ProcessorId,
        processor_type: &ProcessorType,
    ) -> f64 {
        let strategy = *self.strategy.read();
        let load = self.get_load(processor_id);
        let metrics = self.metrics.read();
        let proc_metrics = metrics.get(&processor_id);

        // Base score from compatibility
        if !self.can_execute(&task.operation, processor_type) {
            return f64::NEG_INFINITY;
        }

        // Get utilization and metrics
        let utilization = proc_metrics.map(|m| m.utilization).unwrap_or(load as f64 / 100.0);
        let power = proc_metrics.map(|m| m.power_watts).unwrap_or(100.0);
        let avg_completion = proc_metrics.map(|m| m.avg_completion_ms).unwrap_or(100.0);

        // Calculate score based on strategy
        match strategy {
            BalancingStrategy::Speed => {
                // Prioritize low utilization and fast completion
                let speed_score = 1.0 / (avg_completion.max(1.0) * (1.0 + utilization));

                // Bonus for powerful processor types
                let type_bonus = match processor_type {
                    ProcessorType::Gpu(_) => 2.0,
                    ProcessorType::Tpu(_) => 2.5,
                    ProcessorType::Lpu => 3.0, // Fastest for inference
                    ProcessorType::Npu(_) => 1.5,
                    _ => 1.0,
                };

                speed_score * type_bonus
            }

            BalancingStrategy::Energy => {
                // Prioritize low power consumption
                let energy_score = 1.0 / power.max(1.0);

                // Bonus for efficient processor types
                let efficiency_bonus = match processor_type {
                    ProcessorType::Npu(_) => 3.0, // Most efficient
                    ProcessorType::Lpu => 2.0,
                    ProcessorType::Cpu(_) => 1.5,
                    ProcessorType::Wasm => 2.0, // Low overhead
                    _ => 1.0,
                };

                energy_score * efficiency_bonus * (1.0 - utilization * 0.5)
            }

            BalancingStrategy::Balanced => {
                // Balance speed and energy
                let speed = 1.0 / avg_completion.max(1.0);
                let efficiency = 1.0 / power.max(1.0);
                let load_factor = 1.0 - utilization;

                (speed * 0.4 + efficiency * 0.3 + load_factor * 0.3)
            }

            BalancingStrategy::Cost => {
                // Prioritize cheaper resources (consumer devices)
                let cost_factor = match processor_type {
                    ProcessorType::Wasm => 0.1,      // Cheapest (browser)
                    ProcessorType::WebGpu => 0.15,
                    ProcessorType::Cpu(_) => 0.2,
                    ProcessorType::Npu(_) => 0.3,    // Mobile NPUs
                    ProcessorType::Gpu(_) => 0.5,
                    ProcessorType::Lpu => 0.8,
                    ProcessorType::Tpu(_) => 1.0,    // Most expensive
                    _ => 0.5,
                };

                (1.0 - cost_factor) * (1.0 - utilization)
            }

            BalancingStrategy::Latency => {
                // Prioritize low latency for inference
                let latency_score = 1.0 / avg_completion.max(0.1);

                // Bonus for low-latency processors
                let latency_bonus = match processor_type {
                    ProcessorType::Lpu => 5.0,      // Designed for low latency
                    ProcessorType::Npu(_) => 3.0,
                    ProcessorType::Gpu(_) => 2.0,
                    ProcessorType::Tpu(_) => 1.5,
                    _ => 1.0,
                };

                // Priority boost for critical tasks
                let priority_boost = match task.priority {
                    TaskPriority::Critical => 2.0,
                    TaskPriority::High => 1.5,
                    TaskPriority::Normal => 1.0,
                    TaskPriority::Background => 0.5,
                };

                latency_score * latency_bonus * priority_boost * (1.0 - utilization)
            }
        }
    }

    /// Maybe rebalance a task to a different processor.
    pub fn maybe_rebalance(
        &self,
        task: &Task,
        suggested_processor: ProcessorId,
        current_assignment: &TaskAssignment,
    ) -> ProcessorId {
        // Get all registered processors
        let processor_types = self.processor_types.read();

        // If we don't have processor info, use suggested
        let suggested_type = match processor_types.get(&suggested_processor) {
            Some(t) => t.clone(),
            None => return suggested_processor,
        };

        // Calculate score for suggested processor
        let suggested_score = self.calculate_score(task, suggested_processor, &suggested_type);

        // Find best alternative
        let mut best_processor = suggested_processor;
        let mut best_score = suggested_score;

        for (proc_id, proc_type) in processor_types.iter() {
            if *proc_id == suggested_processor {
                continue;
            }

            let score = self.calculate_score(task, *proc_id, proc_type);

            // Only switch if significantly better (prevents thrashing)
            if score > best_score * (1.0 + self.rebalance_threshold) {
                best_score = score;
                best_processor = *proc_id;
            }
        }

        // Record migration if different
        if best_processor != suggested_processor {
            self.migration_history.write().push(MigrationRecord {
                task_id: task.id,
                from: suggested_processor,
                to: best_processor,
                timestamp: Instant::now(),
            });
        }

        best_processor
    }

    /// Check if work stealing should happen between two processors.
    pub fn should_steal(&self, from: ProcessorId, to: ProcessorId) -> bool {
        let from_load = self.get_load(from) as f64;
        let to_load = self.get_load(to) as f64;

        if from_load == 0.0 {
            return false;
        }

        // Check if processor types are compatible for the queued work
        let processor_types = self.processor_types.read();
        let from_type = processor_types.get(&from);
        let to_type = processor_types.get(&to);

        // Only steal between same processor types by default
        // (cross-type stealing requires operation compatibility check)
        match (from_type, to_type) {
            (Some(ft), Some(tt)) if std::mem::discriminant(ft) == std::mem::discriminant(tt) => {
                let diff = (from_load - to_load) / from_load;
                diff > self.steal_threshold
            }
            _ => false,
        }
    }

    /// Get rebalancing suggestions based on current load.
    pub fn get_rebalance_suggestions(&self) -> Vec<(ProcessorId, ProcessorId)> {
        let mut suggestions = Vec::new();
        let loads = self.loads.read();

        let load_values: Vec<_> = loads.iter()
            .map(|(id, load)| (*id, load.load(Ordering::Relaxed)))
            .collect();

        if load_values.is_empty() {
            return suggestions;
        }

        let avg_load: f64 = load_values.iter().map(|(_, l)| *l as f64).sum::<f64>()
            / load_values.len() as f64;

        let processor_types = self.processor_types.read();

        let overloaded: Vec<_> = load_values.iter()
            .filter(|(_, l)| *l as f64 > avg_load * (1.0 + self.rebalance_threshold))
            .collect();

        let underloaded: Vec<_> = load_values.iter()
            .filter(|(_, l)| (*l as f64) < avg_load * (1.0 - self.rebalance_threshold))
            .collect();

        // Only suggest migrations between compatible processor types
        for (over_id, _) in overloaded {
            let over_type = processor_types.get(over_id);

            for (under_id, _) in &underloaded {
                let under_type = processor_types.get(under_id);

                // Check type compatibility
                if let (Some(ot), Some(ut)) = (over_type, under_type) {
                    if std::mem::discriminant(ot) == std::mem::discriminant(ut) {
                        suggestions.push((*over_id, *under_id));
                    }
                }
            }
        }

        suggestions
    }

    /// Get load statistics.
    pub fn get_stats(&self) -> LoadBalancerStats {
        let loads = self.loads.read();
        let metrics = self.metrics.read();

        let total_load: u64 = loads.values().map(|l| l.load(Ordering::Relaxed)).sum();
        let processor_count = loads.len();
        let avg_load = if processor_count > 0 {
            total_load as f64 / processor_count as f64
        } else {
            0.0
        };

        let total_utilization: f64 = metrics.values().map(|m| m.utilization).sum();
        let avg_utilization = if processor_count > 0 {
            total_utilization / processor_count as f64
        } else {
            0.0
        };

        let total_power: f64 = metrics.values().map(|m| m.power_watts).sum();
        let migrations = self.migration_history.read().len();

        LoadBalancerStats {
            total_load,
            avg_load,
            processor_count,
            avg_utilization,
            total_power_watts: total_power,
            total_migrations: migrations,
            strategy: *self.strategy.read(),
        }
    }

    /// Clean up old migration history.
    pub fn cleanup_history(&self, max_age: Duration) {
        let cutoff = Instant::now() - max_age;
        self.migration_history.write().retain(|r| r.timestamp > cutoff);
    }
}

impl Default for LoadBalancer {
    fn default() -> Self {
        Self::new()
    }
}

/// Load balancer statistics.
#[derive(Clone, Debug)]
pub struct LoadBalancerStats {
    /// Total tasks across all processors.
    pub total_load: u64,
    /// Average load per processor.
    pub avg_load: f64,
    /// Number of registered processors.
    pub processor_count: usize,
    /// Average utilization (0.0 - 1.0).
    pub avg_utilization: f64,
    /// Total power consumption (watts).
    pub total_power_watts: f64,
    /// Total migrations performed.
    pub total_migrations: usize,
    /// Current balancing strategy.
    pub strategy: BalancingStrategy,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::processor::{CpuVariant, GpuVariant, Operation, Precision};
    use crate::task::TaskStatus;

    fn create_test_task(priority: TaskPriority) -> Task {
        Task {
            id: TaskId::new(),
            operation: Operation::MatMul {
                m: 1024,
                n: 1024,
                k: 1024,
                precision: Precision::Fp32,
            },
            priority,
            dependencies: vec![],
            status: TaskStatus::Pending,
            deadline: None,
        }
    }

    #[test]
    fn test_load_tracking() {
        let balancer = LoadBalancer::new();

        balancer.register_processor(ProcessorId(0), ProcessorType::Cpu(CpuVariant::default()));
        balancer.register_processor(ProcessorId(1), ProcessorType::Cpu(CpuVariant::default()));

        assert_eq!(balancer.get_load(ProcessorId(0)), 0);

        balancer.increment_load(ProcessorId(0));
        balancer.increment_load(ProcessorId(0));
        balancer.increment_load(ProcessorId(1));

        assert_eq!(balancer.get_load(ProcessorId(0)), 2);
        assert_eq!(balancer.get_load(ProcessorId(1)), 1);

        balancer.decrement_load(ProcessorId(0));
        assert_eq!(balancer.get_load(ProcessorId(0)), 1);
    }

    #[test]
    fn test_should_steal_same_type() {
        let balancer = LoadBalancer::new();

        // Register two CPUs
        balancer.register_processor(ProcessorId(0), ProcessorType::Cpu(CpuVariant::default()));
        balancer.register_processor(ProcessorId(1), ProcessorType::Cpu(CpuVariant::default()));

        // Give processor 0 high load
        for _ in 0..10 {
            balancer.increment_load(ProcessorId(0));
        }
        balancer.increment_load(ProcessorId(1));

        // Should steal between same types
        assert!(balancer.should_steal(ProcessorId(0), ProcessorId(1)));
        assert!(!balancer.should_steal(ProcessorId(1), ProcessorId(0)));
    }

    #[test]
    fn test_should_not_steal_different_types() {
        let balancer = LoadBalancer::new();

        // Register CPU and GPU
        balancer.register_processor(ProcessorId(0), ProcessorType::Cpu(CpuVariant::default()));
        balancer.register_processor(
            ProcessorId(1),
            ProcessorType::Gpu(GpuVariant::NvidiaCuda { compute_capability: (8, 9) }),
        );

        // Give CPU high load
        for _ in 0..10 {
            balancer.increment_load(ProcessorId(0));
        }

        // Should NOT steal between different types
        assert!(!balancer.should_steal(ProcessorId(0), ProcessorId(1)));
    }

    #[test]
    fn test_can_execute() {
        let balancer = LoadBalancer::new();

        let matmul = Operation::MatMul {
            m: 1024,
            n: 1024,
            k: 1024,
            precision: Precision::Fp32,
        };

        let flash_attention = Operation::FlashAttention {
            batch: 32,
            seq_len: 2048,
            num_heads: 32,
            head_dim: 128,
            precision: Precision::Fp16,
        };

        let cpu = ProcessorType::Cpu(CpuVariant::default());
        let gpu = ProcessorType::Gpu(GpuVariant::NvidiaCuda { compute_capability: (8, 9) });
        let lpu = ProcessorType::Lpu;

        // MatMul can run on all
        assert!(balancer.can_execute(&matmul, &cpu));
        assert!(balancer.can_execute(&matmul, &gpu));
        assert!(balancer.can_execute(&matmul, &lpu));

        // FlashAttention only on GPU/TPU/LPU
        assert!(!balancer.can_execute(&flash_attention, &cpu));
        assert!(balancer.can_execute(&flash_attention, &gpu));
    }

    #[test]
    fn test_strategy_affects_scoring() {
        let balancer = LoadBalancer::new();

        let cpu_id = ProcessorId(0);
        let npu_id = ProcessorId(1);

        balancer.register_processor(cpu_id, ProcessorType::Cpu(CpuVariant::default()));
        balancer.register_processor(npu_id, ProcessorType::Npu(crate::processor::NpuVariant::AppleNeuralEngine { cores: 16 }));

        let task = create_test_task(TaskPriority::Normal);

        // Energy strategy should prefer NPU
        balancer.set_strategy(BalancingStrategy::Energy);
        let assignment = TaskAssignment::new();
        let result = balancer.maybe_rebalance(&task, cpu_id, &assignment);

        // NPU should be preferred for energy efficiency
        assert_eq!(result, npu_id);
    }

    #[test]
    fn test_stats() {
        let balancer = LoadBalancer::new();

        balancer.register_processor(ProcessorId(0), ProcessorType::Cpu(CpuVariant::default()));
        balancer.register_processor(ProcessorId(1), ProcessorType::Cpu(CpuVariant::default()));

        balancer.increment_load(ProcessorId(0));
        balancer.increment_load(ProcessorId(0));
        balancer.increment_load(ProcessorId(1));

        let stats = balancer.get_stats();
        assert_eq!(stats.total_load, 3);
        assert_eq!(stats.processor_count, 2);
        assert!((stats.avg_load - 1.5).abs() < 0.01);
    }
}
