//! Router Z-loss (Regularization for MoE routers) implementation with strict TDD
//!
//! Router Z-loss is a regularization technique for Mixture of Experts (MoE) routers
//! that prevents router collapse by applying penalties based on the magnitude of
//! router logits. This implementation provides:
//!
//! - Z-loss computation from router logits
//! - Configurable loss weight and scaling strategies
//! - Router entropy regularization
//! - Gradient penalty mechanisms for stability
//! - Integration with existing MoE load balancing losses
//! - Comprehensive statistics tracking

use crate::{Result, TransformerError};
use rtx_tensor::{Tensor, Device, DType};
use rtx_autograd::TensorAutograd;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Configuration for Router Z-loss regularization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterZLossConfig {
    /// Weight for Z-loss term in total loss (typically 1e-4 to 1e-2)
    pub z_loss_weight: f32,
    /// Normalization strategy for logits before Z-loss computation
    pub normalization_strategy: NormalizationStrategy,
    /// Strength of entropy regularization (0.0 to disable)
    pub entropy_regularization_weight: f32,
    /// Gradient penalty coefficient for stability
    pub gradient_penalty_weight: f32,
    /// Temperature parameter for softmax normalization
    pub temperature: f32,
    /// Minimum epsilon for numerical stability
    pub epsilon: f32,
    /// Whether to apply auxiliary losses
    pub apply_auxiliary_losses: bool,
}

/// Normalization strategies for router logits
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum NormalizationStrategy {
    /// No normalization applied
    None,
    /// Layer normalization across experts dimension
    LayerNorm,
    /// L2 normalization across experts dimension
    L2Norm,
    /// Z-score normalization (subtract mean, divide by std)
    ZScore,
}

impl Default for RouterZLossConfig {
    fn default() -> Self {
        Self {
            z_loss_weight: 1e-3,
            normalization_strategy: NormalizationStrategy::L2Norm,
            entropy_regularization_weight: 1e-4,
            gradient_penalty_weight: 1e-5,
            temperature: 1.0,
            epsilon: 1e-8,
            apply_auxiliary_losses: true,
        }
    }
}

/// Statistics for Router Z-loss computation and analysis
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterZLossStats {
    /// Current Z-loss value
    pub z_loss: f32,
    /// Entropy regularization loss
    pub entropy_loss: f32,
    /// Gradient penalty loss
    pub gradient_penalty: f32,
    /// Total regularization loss (sum of all components)
    pub total_loss: f32,
    /// Mean logit magnitude across all tokens and experts
    pub mean_logit_magnitude: f32,
    /// Standard deviation of logit magnitudes
    pub logit_magnitude_std: f32,
    /// Maximum logit magnitude
    pub max_logit_magnitude: f32,
    /// Minimum logit magnitude
    pub min_logit_magnitude: f32,
    /// Router entropy (measure of routing diversity)
    pub router_entropy: f32,
    /// Number of tokens processed
    pub num_tokens: usize,
    /// Number of experts
    pub num_experts: usize,
}

/// Router Z-loss regularization for MoE training stability
pub struct RouterZLoss {
    config: RouterZLossConfig,
    device: Device,
    /// Running statistics for monitoring
    stats_history: Vec<RouterZLossStats>,
}

impl RouterZLoss {
    /// Create a new Router Z-loss regularizer
    pub fn new(config: RouterZLossConfig, device: Device) -> Result<Self> {
        config.validate()?;
        
        Ok(Self {
            config,
            device,
            stats_history: Vec::new(),
        })
    }

    /// Compute Z-loss and auxiliary regularization terms from router logits
    ///
    /// # Arguments
    /// * `router_logits` - Tensor of shape (batch_size * seq_len, num_experts) containing raw router logits
    ///
    /// # Returns
    /// Tuple containing:
    /// - Total regularization loss tensor
    /// - Detailed statistics for monitoring
    pub fn compute_loss(&mut self, router_logits: &Tensor) -> Result<(Tensor, RouterZLossStats)> {
        let shape = router_logits.shape();
        if shape.len() != 2 {
            return Err(TransformerError::config(
                "Router logits must be 2D tensor (batch_size * seq_len, num_experts)".to_string()
            ));
        }
        
        let num_tokens = shape[0];
        let num_experts = shape[1];
        
        // Apply normalization to logits
        let normalized_logits = self.normalize_logits(router_logits)?;
        
        // Compute core Z-loss: log(sum(exp(z_i^2)))
        let z_loss_raw = self.compute_z_loss_core(&normalized_logits)?;
        let z_loss_weighted = z_loss_raw.mul_scalar(self.config.z_loss_weight)?;
        
        // Compute auxiliary losses if enabled
        let (entropy_loss, gradient_penalty) = if self.config.apply_auxiliary_losses {
            let entropy = self.compute_entropy_regularization(&normalized_logits)?;
            let grad_penalty = self.compute_gradient_penalty(&normalized_logits)?;
            (entropy, grad_penalty)
        } else {
            (Tensor::zeros(&[], router_logits.dtype(), &self.device)?, 
             Tensor::zeros(&[], router_logits.dtype(), &self.device)?)
        };
        
        // Combine all loss terms
        let total_loss = z_loss_weighted.add(&entropy_loss)?.add(&gradient_penalty)?;
        
        // Calculate statistics
        let z_loss_val = z_loss_raw.mean()?.item::<f32>()?;
        let entropy_loss_val = entropy_loss.mean()?.item::<f32>()?;
        let gradient_penalty_val = gradient_penalty.mean()?.item::<f32>()?;
        
        let stats = self.calculate_statistics(
            &normalized_logits, 
            z_loss_val, 
            entropy_loss_val, 
            gradient_penalty_val
        )?;
        
        // Store statistics
        self.stats_history.push(stats.clone());
        
        Ok((total_loss, stats))
    }

    /// Apply normalization to router logits according to configuration
    fn normalize_logits(&self, logits: &Tensor) -> Result<Tensor> {
        match self.config.normalization_strategy {
            NormalizationStrategy::None => Ok(logits.clone()),
            NormalizationStrategy::L2Norm => {
                // L2 normalize across experts dimension (dim=1)
                let norm = logits.pow_scalar(2.0)?.sum_keepdim(1)?.sqrt()?.add_scalar(self.config.epsilon)?;
                logits.div(&norm)
            },
            NormalizationStrategy::LayerNorm => {
                // Layer normalization: (x - mean) / std
                let mean = logits.mean_keepdim(1)?;
                let centered = logits.sub(&mean)?;
                let variance = centered.pow_scalar(2.0)?.mean_keepdim(1)?;
                let std = variance.sqrt()?.add_scalar(self.config.epsilon)?;
                centered.div(&std)
            },
            NormalizationStrategy::ZScore => {
                // Z-score normalization across all elements
                let mean = logits.mean()?;
                let centered = logits.sub_scalar(mean.item::<f32>()?)?;
                let std = centered.pow_scalar(2.0)?.mean()?.sqrt()?.add_scalar(self.config.epsilon)?;
                centered.div_scalar(std.item::<f32>()?)
            }
        }
    }

    /// Compute the core Z-loss term: log(sum(exp(z_i^2))) where z_i are logits
    fn compute_z_loss_core(&self, normalized_logits: &Tensor) -> Result<Tensor> {
        // Z-loss: log(sum(exp(z_i^2))) per token
        // This prevents logit magnitude from growing too large
        let squared_logits = normalized_logits.pow_scalar(2.0)?;
        let exp_squared = squared_logits.exp()?;
        let sum_exp = exp_squared.sum(1)?; // Sum across experts dimension
        let z_loss = sum_exp.log()?; // log(sum(exp(z_i^2)))
        z_loss.mean() // Average across tokens
    }

    /// Compute entropy regularization to encourage diverse routing
    fn compute_entropy_regularization(&self, logits: &Tensor) -> Result<Tensor> {
        if self.config.entropy_regularization_weight <= 0.0 {
            return Tensor::zeros(&[], logits.dtype(), &self.device);
        }
        
        // Compute routing probabilities with temperature scaling
        let scaled_logits = logits.div_scalar(self.config.temperature)?;
        let probabilities = scaled_logits.softmax(1)?; // Softmax across experts
        
        // Entropy: -sum(p * log(p + epsilon))
        let log_probs = probabilities.add_scalar(self.config.epsilon)?.log()?;
        let entropy_per_token = probabilities.mul(&log_probs)?.sum(1)?.neg()?;
        let mean_entropy = entropy_per_token.mean()?;
        
        // We want to maximize entropy (encourage diversity), so we minimize -entropy
        let entropy_loss = mean_entropy.neg()?.mul_scalar(self.config.entropy_regularization_weight)?;
        Ok(entropy_loss)
    }

    /// Compute gradient penalty for training stability
    fn compute_gradient_penalty(&self, logits: &Tensor) -> Result<Tensor> {
        if self.config.gradient_penalty_weight <= 0.0 {
            return Tensor::zeros(&[], logits.dtype(), &self.device);
        }
        
        // Simple gradient penalty: penalize large gradients by using logit variance
        // This is a proxy for gradient magnitude that doesn't require actual gradients
        let mean = logits.mean_keepdim(1)?;
        let variance = logits.sub(&mean)?.pow_scalar(2.0)?.mean(1)?;
        let gradient_penalty = variance.mean()?.mul_scalar(self.config.gradient_penalty_weight)?;
        
        Ok(gradient_penalty)
    }

    /// Calculate comprehensive statistics for monitoring and analysis
    fn calculate_statistics(
        &self,
        logits: &Tensor,
        z_loss: f32,
        entropy_loss: f32,
        gradient_penalty: f32,
    ) -> Result<RouterZLossStats> {
        let shape = logits.shape();
        let num_tokens = shape[0];
        let num_experts = shape[1];
        
        // Calculate logit magnitude statistics
        let abs_logits = logits.abs()?;
        let mean_magnitude = abs_logits.mean()?.item::<f32>()?;
        let max_magnitude = abs_logits.max()?.item::<f32>()?;
        let min_magnitude = abs_logits.min()?.item::<f32>()?;
        
        // Calculate standard deviation of logit magnitudes
        let mean_tensor = abs_logits.mean()?;
        let variance = abs_logits.sub(&mean_tensor)?.pow_scalar(2.0)?.mean()?;
        let std_magnitude = variance.sqrt()?.item::<f32>()?;
        
        // Calculate router entropy
        let probabilities = logits.div_scalar(self.config.temperature)?.softmax(1)?;
        let log_probs = probabilities.add_scalar(self.config.epsilon)?.log()?;
        let entropy = probabilities.mul(&log_probs)?.sum(1)?.neg()?.mean()?.item::<f32>()?;
        
        let total_loss = z_loss * self.config.z_loss_weight + entropy_loss + gradient_penalty;
        
        Ok(RouterZLossStats {
            z_loss,
            entropy_loss,
            gradient_penalty,
            total_loss,
            mean_logit_magnitude: mean_magnitude,
            logit_magnitude_std: std_magnitude,
            max_logit_magnitude: max_magnitude,
            min_logit_magnitude: min_magnitude,
            router_entropy: entropy,
            num_tokens,
            num_experts,
        })
    }

    /// Get current configuration
    pub fn config(&self) -> &RouterZLossConfig {
        &self.config
    }

    /// Get device
    pub fn device(&self) -> &Device {
        &self.device
    }

    /// Get historical statistics
    pub fn stats_history(&self) -> &[RouterZLossStats] {
        &self.stats_history
    }

    /// Reset statistics history
    pub fn reset_stats(&mut self) {
        self.stats_history.clear();
    }

    /// Get the latest statistics if available
    pub fn latest_stats(&self) -> Option<&RouterZLossStats> {
        self.stats_history.last()
    }

    /// Get average statistics over the last N computations
    pub fn average_stats(&self, last_n: usize) -> Option<RouterZLossStats> {
        if self.stats_history.is_empty() {
            return None;
        }

        let start_idx = if last_n >= self.stats_history.len() {
            0
        } else {
            self.stats_history.len() - last_n
        };

        let stats_slice = &self.stats_history[start_idx..];
        let count = stats_slice.len();

        if count == 0 {
            return None;
        }

        let sum_stats = stats_slice.iter().fold(
            RouterZLossStats {
                z_loss: 0.0,
                entropy_loss: 0.0,
                gradient_penalty: 0.0,
                total_loss: 0.0,
                mean_logit_magnitude: 0.0,
                logit_magnitude_std: 0.0,
                max_logit_magnitude: f32::MIN,
                min_logit_magnitude: f32::MAX,
                router_entropy: 0.0,
                num_tokens: 0,
                num_experts: 0,
            },
            |mut acc, stats| {
                acc.z_loss += stats.z_loss;
                acc.entropy_loss += stats.entropy_loss;
                acc.gradient_penalty += stats.gradient_penalty;
                acc.total_loss += stats.total_loss;
                acc.mean_logit_magnitude += stats.mean_logit_magnitude;
                acc.logit_magnitude_std += stats.logit_magnitude_std;
                acc.max_logit_magnitude = acc.max_logit_magnitude.max(stats.max_logit_magnitude);
                acc.min_logit_magnitude = acc.min_logit_magnitude.min(stats.min_logit_magnitude);
                acc.router_entropy += stats.router_entropy;
                acc.num_tokens += stats.num_tokens;
                acc.num_experts = stats.num_experts; // Use last value
                acc
            }
        );

        let count_f32 = count as f32;
        Some(RouterZLossStats {
            z_loss: sum_stats.z_loss / count_f32,
            entropy_loss: sum_stats.entropy_loss / count_f32,
            gradient_penalty: sum_stats.gradient_penalty / count_f32,
            total_loss: sum_stats.total_loss / count_f32,
            mean_logit_magnitude: sum_stats.mean_logit_magnitude / count_f32,
            logit_magnitude_std: sum_stats.logit_magnitude_std / count_f32,
            max_logit_magnitude: sum_stats.max_logit_magnitude,
            min_logit_magnitude: sum_stats.min_logit_magnitude,
            router_entropy: sum_stats.router_entropy / count_f32,
            num_tokens: sum_stats.num_tokens,
            num_experts: sum_stats.num_experts,
        })
    }

    /// Check if router shows signs of collapse based on statistics
    pub fn detect_router_collapse(&self, threshold_entropy: f32, min_history: usize) -> bool {
        if self.stats_history.len() < min_history {
            return false; // Not enough data
        }

        // Check recent entropy values
        let recent_stats = self.average_stats(min_history).unwrap();
        recent_stats.router_entropy < threshold_entropy
    }

    /// Compute router health score (0.0 = unhealthy, 1.0 = healthy)
    pub fn compute_health_score(&self) -> f32 {
        if let Some(stats) = self.latest_stats() {
            let entropy_score = (stats.router_entropy / (stats.num_experts as f32).ln()).min(1.0).max(0.0);
            let magnitude_score = (1.0 / (1.0 + stats.mean_logit_magnitude)).min(1.0).max(0.0);
            let balance_score = if stats.logit_magnitude_std > 0.0 {
                (1.0 / (1.0 + stats.logit_magnitude_std)).min(1.0).max(0.0)
            } else {
                1.0
            };
            
            // Weighted combination of scores
            0.4 * entropy_score + 0.3 * magnitude_score + 0.3 * balance_score
        } else {
            0.5 // Neutral score if no history
        }
    }

    /// Update configuration (useful for adaptive training)
    pub fn update_config(&mut self, new_config: RouterZLossConfig) -> Result<()> {
        new_config.validate()?;
        self.config = new_config;
        Ok(())
    }

    /// Create an adaptive config based on current router health
    pub fn create_adaptive_config(&self, base_config: &RouterZLossConfig) -> RouterZLossConfig {
        let health_score = self.compute_health_score();
        
        // Adjust loss weights based on health
        let adaptive_z_loss_weight = if health_score < 0.3 {
            // Router is unhealthy, increase Z-loss weight
            base_config.z_loss_weight * 2.0
        } else if health_score > 0.8 {
            // Router is healthy, can reduce Z-loss weight
            base_config.z_loss_weight * 0.8
        } else {
            base_config.z_loss_weight
        };

        let adaptive_entropy_weight = if health_score < 0.5 {
            // Increase entropy regularization for unhealthy router
            base_config.entropy_regularization_weight * 1.5
        } else {
            base_config.entropy_regularization_weight
        };

        RouterZLossConfig {
            z_loss_weight: adaptive_z_loss_weight,
            entropy_regularization_weight: adaptive_entropy_weight,
            ..base_config.clone()
        }
    }
}

impl RouterZLossConfig {
    /// Validate the Router Z-loss configuration
    pub fn validate(&self) -> Result<()> {
        if self.z_loss_weight < 0.0 {
            return Err(TransformerError::config(
                "z_loss_weight must be non-negative".to_string()
            ));
        }

        if self.entropy_regularization_weight < 0.0 {
            return Err(TransformerError::config(
                "entropy_regularization_weight must be non-negative".to_string()
            ));
        }

        if self.gradient_penalty_weight < 0.0 {
            return Err(TransformerError::config(
                "gradient_penalty_weight must be non-negative".to_string()
            ));
        }

        if self.temperature <= 0.0 {
            return Err(TransformerError::config(
                "temperature must be positive".to_string()
            ));
        }

        if self.epsilon <= 0.0 {
            return Err(TransformerError::config(
                "epsilon must be positive".to_string()
            ));
        }

        Ok(())
    }
}

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

    #[test]
    fn test_router_zloss_config_default() {
        let config = RouterZLossConfig::default();
        assert_eq!(config.z_loss_weight, 1e-3);
        assert_eq!(config.normalization_strategy, NormalizationStrategy::L2Norm);
        assert_eq!(config.entropy_regularization_weight, 1e-4);
        assert_eq!(config.gradient_penalty_weight, 1e-5);
        assert_eq!(config.temperature, 1.0);
        assert_eq!(config.epsilon, 1e-8);
        assert!(config.apply_auxiliary_losses);
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_router_zloss_config_validation_failures() {
        // Test negative z_loss_weight
        let mut config = RouterZLossConfig::default();
        config.z_loss_weight = -0.1;
        assert!(config.validate().is_err());

        // Test negative entropy_regularization_weight
        let mut config = RouterZLossConfig::default();
        config.entropy_regularization_weight = -0.1;
        assert!(config.validate().is_err());

        // Test negative gradient_penalty_weight
        let mut config = RouterZLossConfig::default();
        config.gradient_penalty_weight = -0.1;
        assert!(config.validate().is_err());

        // Test zero temperature
        let mut config = RouterZLossConfig::default();
        config.temperature = 0.0;
        assert!(config.validate().is_err());

        // Test negative temperature
        let mut config = RouterZLossConfig::default();
        config.temperature = -1.0;
        assert!(config.validate().is_err());

        // Test zero epsilon
        let mut config = RouterZLossConfig::default();
        config.epsilon = 0.0;
        assert!(config.validate().is_err());

        // Test negative epsilon
        let mut config = RouterZLossConfig::default();
        config.epsilon = -1e-8;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_normalization_strategy_variants() {
        // Test all normalization strategy variants
        let strategies = vec![
            NormalizationStrategy::None,
            NormalizationStrategy::LayerNorm,
            NormalizationStrategy::L2Norm,
            NormalizationStrategy::ZScore,
        ];

        for strategy in strategies {
            let mut config = RouterZLossConfig::default();
            config.normalization_strategy = strategy.clone();
            assert!(config.validate().is_ok());
            
            // Test serialization/deserialization
            let serialized = serde_json::to_string(&config).unwrap();
            let deserialized: RouterZLossConfig = serde_json::from_str(&serialized).unwrap();
            assert_eq!(deserialized.normalization_strategy, strategy);
        }
    }

    #[test]
    fn test_router_zloss_creation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        
        let zloss = RouterZLoss::new(config.clone(), device.clone());
        assert!(zloss.is_ok());
        
        let zloss = zloss.unwrap();
        assert_eq!(zloss.config(), &config);
        assert_eq!(zloss.device(), &device);
        assert!(zloss.stats_history().is_empty());
    }

    #[test]
    fn test_router_zloss_creation_with_invalid_config() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut config = RouterZLossConfig::default();
        config.z_loss_weight = -0.1; // Invalid
        
        let zloss = RouterZLoss::new(config, device);
        assert!(zloss.is_err());
    }

    #[test]
    fn test_router_zloss_compute_loss_fails_initially() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let mut zloss = RouterZLoss::new(config, device.clone()).unwrap();
        
        // Create test router logits: (batch_size * seq_len, num_experts)
        let batch_size = 4;
        let seq_len = 8;
        let num_experts = 8;
        let router_logits = Tensor::randn(
            &[batch_size * seq_len, num_experts], 
            DType::F32, 
            &device
        ).unwrap();
        
        // Should fail during RED phase
        let result = zloss.compute_loss(&router_logits);
        assert!(result.is_err());
    }

    #[test]
    fn test_normalize_logits_fails_initially() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let zloss = RouterZLoss::new(config, device.clone()).unwrap();
        
        let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
        let result = zloss.normalize_logits(&logits);
        assert!(result.is_err());
    }

    #[test]
    fn test_compute_z_loss_core_fails_initially() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let zloss = RouterZLoss::new(config, device.clone()).unwrap();
        
        let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
        let result = zloss.compute_z_loss_core(&logits);
        assert!(result.is_err());
    }

    #[test]
    fn test_compute_entropy_regularization_fails_initially() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let zloss = RouterZLoss::new(config, device.clone()).unwrap();
        
        let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
        let result = zloss.compute_entropy_regularization(&logits);
        assert!(result.is_err());
    }

    #[test]
    fn test_compute_gradient_penalty_fails_initially() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let zloss = RouterZLoss::new(config, device.clone()).unwrap();
        
        let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
        let result = zloss.compute_gradient_penalty(&logits);
        assert!(result.is_err());
    }

    #[test]
    fn test_calculate_statistics_fails_initially() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let zloss = RouterZLoss::new(config, device.clone()).unwrap();
        
        let logits = Tensor::randn(&[32, 8], DType::F32, &device).unwrap();
        let result = zloss.calculate_statistics(&logits, 0.1, 0.05, 0.01);
        assert!(result.is_err());
    }

    #[test]
    fn test_router_zloss_stats_structure() {
        // Test that RouterZLossStats has all required fields and can be serialized
        let stats = RouterZLossStats {
            z_loss: 0.1,
            entropy_loss: 0.05,
            gradient_penalty: 0.01,
            total_loss: 0.16,
            mean_logit_magnitude: 2.5,
            logit_magnitude_std: 0.8,
            max_logit_magnitude: 4.2,
            min_logit_magnitude: 0.3,
            router_entropy: 1.8,
            num_tokens: 256,
            num_experts: 8,
        };

        // Test serialization
        let serialized = serde_json::to_string(&stats).unwrap();
        let deserialized: RouterZLossStats = serde_json::from_str(&serialized).unwrap();
        
        assert_abs_diff_eq!(deserialized.z_loss, stats.z_loss);
        assert_abs_diff_eq!(deserialized.entropy_loss, stats.entropy_loss);
        assert_abs_diff_eq!(deserialized.gradient_penalty, stats.gradient_penalty);
        assert_abs_diff_eq!(deserialized.total_loss, stats.total_loss);
        assert_abs_diff_eq!(deserialized.mean_logit_magnitude, stats.mean_logit_magnitude);
        assert_abs_diff_eq!(deserialized.logit_magnitude_std, stats.logit_magnitude_std);
        assert_abs_diff_eq!(deserialized.max_logit_magnitude, stats.max_logit_magnitude);
        assert_abs_diff_eq!(deserialized.min_logit_magnitude, stats.min_logit_magnitude);
        assert_abs_diff_eq!(deserialized.router_entropy, stats.router_entropy);
        assert_eq!(deserialized.num_tokens, stats.num_tokens);
        assert_eq!(deserialized.num_experts, stats.num_experts);
    }

    #[test]
    fn test_router_zloss_reset_stats() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let mut zloss = RouterZLoss::new(config, device).unwrap();
        
        // Manually add some stats to history (simulating previous computations)
        zloss.stats_history.push(RouterZLossStats {
            z_loss: 0.1,
            entropy_loss: 0.05,
            gradient_penalty: 0.01,
            total_loss: 0.16,
            mean_logit_magnitude: 2.5,
            logit_magnitude_std: 0.8,
            max_logit_magnitude: 4.2,
            min_logit_magnitude: 0.3,
            router_entropy: 1.8,
            num_tokens: 256,
            num_experts: 8,
        });
        
        assert_eq!(zloss.stats_history().len(), 1);
        
        zloss.reset_stats();
        assert!(zloss.stats_history().is_empty());
    }

    #[test]
    fn test_router_zloss_latest_stats() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let mut zloss = RouterZLoss::new(config, device).unwrap();
        
        // Initially no stats
        assert!(zloss.latest_stats().is_none());
        
        // Add a stat
        let test_stats = RouterZLossStats {
            z_loss: 0.1,
            entropy_loss: 0.05,
            gradient_penalty: 0.01,
            total_loss: 0.16,
            mean_logit_magnitude: 2.5,
            logit_magnitude_std: 0.8,
            max_logit_magnitude: 4.2,
            min_logit_magnitude: 0.3,
            router_entropy: 1.8,
            num_tokens: 256,
            num_experts: 8,
        };
        zloss.stats_history.push(test_stats.clone());
        
        let latest = zloss.latest_stats().unwrap();
        assert_abs_diff_eq!(latest.z_loss, test_stats.z_loss);
        assert_eq!(latest.num_tokens, test_stats.num_tokens);
    }

    #[test]
    fn test_router_zloss_average_stats() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let mut zloss = RouterZLoss::new(config, device).unwrap();
        
        assert!(zloss.average_stats(5).is_none());
        
        // Add test stats
        for i in 0..3 {
            zloss.stats_history.push(RouterZLossStats {
                z_loss: 0.1 + i as f32 * 0.01, entropy_loss: 0.05, gradient_penalty: 0.01,
                total_loss: 0.16, mean_logit_magnitude: 2.0 + i as f32, logit_magnitude_std: 0.8,
                max_logit_magnitude: 4.0, min_logit_magnitude: 0.0, router_entropy: 1.5,
                num_tokens: 256, num_experts: 8,
            });
        }
        
        let avg_stats = zloss.average_stats(2).unwrap();
        assert_abs_diff_eq!(avg_stats.z_loss, (0.11 + 0.12) / 2.0, epsilon = 1e-6);
    }

    #[test]
    fn test_router_collapse_detection() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let mut zloss = RouterZLoss::new(config, device).unwrap();
        
        assert!(!zloss.detect_router_collapse(0.5, 5)); // Not enough history
        
        // Add stats with low entropy
        for _ in 0..5 {
            zloss.stats_history.push(RouterZLossStats {
                z_loss: 0.1, entropy_loss: 0.05, gradient_penalty: 0.01, total_loss: 0.16,
                mean_logit_magnitude: 2.5, logit_magnitude_std: 0.8, max_logit_magnitude: 4.2,
                min_logit_magnitude: 0.3, router_entropy: 0.2, num_tokens: 256, num_experts: 8,
            });
        }
        
        assert!(zloss.detect_router_collapse(0.5, 5));
        assert!(!zloss.detect_router_collapse(0.1, 5));
    }

    #[test]
    fn test_router_health_score() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let mut zloss = RouterZLoss::new(config, device).unwrap();
        
        // No history - should return neutral score
        assert_abs_diff_eq!(zloss.compute_health_score(), 0.5);
        
        // Add healthy stats and verify score > 0.5
        zloss.stats_history.push(RouterZLossStats {
            z_loss: 0.01, entropy_loss: 0.001, gradient_penalty: 0.0001, total_loss: 0.01,
            mean_logit_magnitude: 0.5, logit_magnitude_std: 0.1, max_logit_magnitude: 1.0,
            min_logit_magnitude: 0.0, router_entropy: 2.0, num_tokens: 256, num_experts: 8,
        });
        
        let health_score = zloss.compute_health_score();
        assert!(health_score > 0.5 && health_score <= 1.0);
    }

    #[test]
    fn test_adaptive_config_creation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let base_config = RouterZLossConfig::default();
        let mut zloss = RouterZLoss::new(base_config.clone(), device).unwrap();
        
        // Add unhealthy stats to trigger adaptations
        zloss.stats_history.push(RouterZLossStats {
            z_loss: 0.5, entropy_loss: 0.1, gradient_penalty: 0.05, total_loss: 0.65,
            mean_logit_magnitude: 10.0, logit_magnitude_std: 5.0, max_logit_magnitude: 20.0,
            min_logit_magnitude: 0.0, router_entropy: 0.1, num_tokens: 256, num_experts: 8,
        });
        
        let adaptive_config = zloss.create_adaptive_config(&base_config);
        assert!(adaptive_config.z_loss_weight >= base_config.z_loss_weight);
        assert!(adaptive_config.entropy_regularization_weight >= base_config.entropy_regularization_weight);
    }

    #[test]
    fn test_config_update() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = RouterZLossConfig::default();
        let mut zloss = RouterZLoss::new(config.clone(), device).unwrap();
        
        let mut new_config = config.clone();
        new_config.z_loss_weight = 0.01;
        new_config.temperature = 2.0;
        
        // Should succeed with valid config
        assert!(zloss.update_config(new_config.clone()).is_ok());
        assert_eq!(zloss.config().z_loss_weight, 0.01);
        assert_eq!(zloss.config().temperature, 2.0);
        
        // Should fail with invalid config
        let mut invalid_config = new_config;
        invalid_config.z_loss_weight = -1.0; // Invalid
        assert!(zloss.update_config(invalid_config).is_err());
        
        // Config should remain unchanged after failed update
        assert_eq!(zloss.config().z_loss_weight, 0.01);
    }
}