#!/usr/bin/env rust-script
//! Simple GREEN phase verification for Router Z-loss implementation
//! This tests minimal functionality to verify the GREEN phase is working

// Minimal tensor simulation for GREEN phase testing
#[derive(Debug, Clone)]
pub struct Tensor {
    data: Vec<f32>,
    shape: Vec<usize>,
}

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

#[derive(Debug, Clone)]
pub enum DType {
    F32,
}

impl Device {
    pub const CPU: Device = Device;
}

impl Tensor {
    pub fn randn(shape: &[usize], _dtype: DType, _device: &Device) -> Result<Self, String> {
        let total_elements: usize = shape.iter().product();
        let data = (0..total_elements).map(|i| (i as f32) * 0.1 - 0.5).collect();
        Ok(Tensor { data, shape: shape.to_vec() })
    }
    
    pub fn zeros(shape: &[usize], _dtype: DType, _device: &Device) -> Result<Self, String> {
        let total_elements: usize = shape.iter().product();
        let data = vec![0.0; total_elements];
        Ok(Tensor { data, shape: shape.to_vec() })
    }
    
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    
    pub fn dtype(&self) -> DType {
        DType::F32
    }
    
    pub fn clone(&self) -> Self {
        Self {
            data: self.data.clone(),
            shape: self.shape.clone(),
        }
    }
    
    // Minimal tensor operations for GREEN phase
    pub fn add(&self, other: &Self) -> Result<Self, String> {
        if self.data.len() != other.data.len() {
            return Err("Tensor size mismatch".to_string());
        }
        let data: Vec<f32> = self.data.iter().zip(other.data.iter()).map(|(a, b)| a + b).collect();
        Ok(Tensor { data, shape: self.shape.clone() })
    }
    
    pub fn mul_scalar(&self, scalar: f32) -> Result<Self, String> {
        let data: Vec<f32> = self.data.iter().map(|x| x * scalar).collect();
        Ok(Tensor { data, shape: self.shape.clone() })
    }
    
    pub fn mean(&self) -> Result<Self, String> {
        let mean_val = self.data.iter().sum::<f32>() / self.data.len() as f32;
        Ok(Tensor { data: vec![mean_val], shape: vec![] })
    }
    
    pub fn item<T: From<f32>>(&self) -> Result<T, String> {
        if self.data.len() != 1 {
            return Err("Tensor must be scalar for item()".to_string());
        }
        Ok(T::from(self.data[0]))
    }
    
    // More operations needed for minimal GREEN phase
    pub fn pow_scalar(&self, power: f32) -> Result<Self, String> {
        let data: Vec<f32> = self.data.iter().map(|x| x.powf(power)).collect();
        Ok(Tensor { data, shape: self.shape.clone() })
    }
    
    pub fn sum_keepdim(&self, _dim: usize) -> Result<Self, String> {
        // Simplified: just return sum as scalar
        let sum_val = self.data.iter().sum::<f32>();
        Ok(Tensor { data: vec![sum_val], shape: vec![1, 1] })
    }
    
    pub fn sqrt(&self) -> Result<Self, String> {
        let data: Vec<f32> = self.data.iter().map(|x| x.sqrt()).collect();
        Ok(Tensor { data, shape: self.shape.clone() })
    }
    
    pub fn add_scalar(&self, scalar: f32) -> Result<Self, String> {
        let data: Vec<f32> = self.data.iter().map(|x| x + scalar).collect();
        Ok(Tensor { data, shape: self.shape.clone() })
    }
    
    pub fn div(&self, other: &Self) -> Result<Self, String> {
        if self.data.len() != other.data.len() {
            return Err("Tensor size mismatch".to_string());
        }
        let data: Vec<f32> = self.data.iter().zip(other.data.iter()).map(|(a, b)| a / (b + 1e-8)).collect();
        Ok(Tensor { data, shape: self.shape.clone() })
    }
}

// Minimal error types
#[derive(Debug)]
pub struct TransformerError(String);

impl TransformerError {
    pub fn config(msg: String) -> Self {
        Self(msg)
    }
}

type Result<T> = std::result::Result<T, TransformerError>;

/// Configuration for Router Z-loss regularization - GREEN phase minimal version
#[derive(Debug, Clone)]
pub struct RouterZLossConfig {
    pub z_loss_weight: f32,
    pub normalization_strategy: NormalizationStrategy,
    pub entropy_regularization_weight: f32,
    pub gradient_penalty_weight: f32,
    pub temperature: f32,
    pub epsilon: f32,
    pub apply_auxiliary_losses: bool,
}

/// Normalization strategies for router logits
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NormalizationStrategy {
    None,
    L2Norm,
    LayerNorm,
    ZScore,
}

impl Default for RouterZLossConfig {
    fn default() -> Self {
        Self {
            z_loss_weight: 1e-3,
            normalization_strategy: NormalizationStrategy::None, // Simplified for GREEN phase
            entropy_regularization_weight: 0.0, // Disabled for GREEN phase
            gradient_penalty_weight: 0.0, // Disabled for GREEN phase
            temperature: 1.0,
            epsilon: 1e-8,
            apply_auxiliary_losses: false, // Disabled for GREEN phase
        }
    }
}

/// Statistics for Router Z-loss computation and analysis
#[derive(Debug, Clone)]
pub struct RouterZLossStats {
    pub z_loss: f32,
    pub entropy_loss: f32,
    pub gradient_penalty: f32,
    pub total_loss: f32,
    pub mean_logit_magnitude: f32,
    pub logit_magnitude_std: f32,
    pub max_logit_magnitude: f32,
    pub min_logit_magnitude: f32,
    pub router_entropy: f32,
    pub num_tokens: usize,
    pub num_experts: usize,
}

/// Router Z-loss regularization for MoE training stability
pub struct RouterZLoss {
    config: RouterZLossConfig,
    device: Device,
    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 - GREEN phase minimal implementation
    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 (minimal - just return input for None strategy)
        let normalized_logits = match self.config.normalization_strategy {
            NormalizationStrategy::None => router_logits.clone(),
            _ => return Err(TransformerError::config("Only None normalization supported in GREEN phase".to_string()))
        };
        
        // Compute minimal Z-loss: just use mean of squared logits
        let squared_logits = normalized_logits.pow_scalar(2.0)?;
        let z_loss_raw = squared_logits.mean()?;
        let z_loss_weighted = z_loss_raw.mul_scalar(self.config.z_loss_weight)?;
        
        // No auxiliary losses in GREEN phase
        let entropy_loss = Tensor::zeros(&[], router_logits.dtype(), &self.device)?;
        let gradient_penalty = Tensor::zeros(&[], router_logits.dtype(), &self.device)?;
        
        let total_loss = z_loss_weighted.clone();
        
        // Calculate minimal statistics
        let z_loss_val = z_loss_raw.item::<f32>()?;
        let stats = RouterZLossStats {
            z_loss: z_loss_val,
            entropy_loss: 0.0,
            gradient_penalty: 0.0,
            total_loss: z_loss_val * self.config.z_loss_weight,
            mean_logit_magnitude: 1.0, // Placeholder
            logit_magnitude_std: 0.5,  // Placeholder
            max_logit_magnitude: 2.0,  // Placeholder
            min_logit_magnitude: 0.0,  // Placeholder
            router_entropy: 1.5,       // Placeholder
            num_tokens,
            num_experts,
        };
        
        self.stats_history.push(stats.clone());
        
        Ok((total_loss, stats))
    }

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

    /// 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();
    }
}

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.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(())
    }
}

fn main() {
    println!("Router Z-loss TDD GREEN Phase Verification");
    println!("==========================================");
    
    // Test 1: Configuration validation still works
    println!("\n1. Testing RouterZLossConfig validation...");
    let config = RouterZLossConfig::default();
    match config.validate() {
        Ok(_) => println!("✓ Default config validation PASSED"),
        Err(_) => {
            println!("✗ Default config validation FAILED");
            return;
        }
    }
    
    // Test 2: RouterZLoss creation still works
    println!("\n2. Testing RouterZLoss creation...");
    let device = Device::CPU;
    let mut zloss = match RouterZLoss::new(config.clone(), device) {
        Ok(zloss) => {
            println!("✓ RouterZLoss creation PASSED");
            zloss
        },
        Err(_) => {
            println!("✗ RouterZLoss creation FAILED");
            return;
        }
    };
    
    // Test 3: compute_loss now works (GREEN phase)
    println!("\n3. Testing that compute_loss now works in GREEN phase...");
    let router_logits = match Tensor::randn(&[4, 8], DType::F32, &Device::CPU) {
        Ok(tensor) => tensor,
        Err(_) => {
            println!("✗ Failed to create test tensor");
            return;
        }
    };
    
    match zloss.compute_loss(&router_logits) {
        Ok((loss_tensor, stats)) => {
            println!("✓ compute_loss SUCCEEDED (GREEN phase working!)");
            println!("  - Z-loss value: {:.6}", stats.z_loss);
            println!("  - Total loss: {:.6}", stats.total_loss);
            println!("  - Number of tokens: {}", stats.num_tokens);
            println!("  - Number of experts: {}", stats.num_experts);
        },
        Err(e) => {
            println!("✗ compute_loss should work in GREEN phase: {:?}", e);
            return;
        }
    }
    
    // Test 4: Stats history works
    println!("\n4. Testing stats history...");
    assert_eq!(zloss.stats_history().len(), 1);
    println!("✓ Stats history has 1 entry");
    
    zloss.reset_stats();
    assert!(zloss.stats_history().is_empty());
    println!("✓ Reset stats works");
    
    // Test 5: Invalid tensor shape handling
    println!("\n5. Testing invalid tensor shape handling...");
    let invalid_tensor = match Tensor::randn(&[4], DType::F32, &Device::CPU) {
        Ok(tensor) => tensor,
        Err(_) => {
            println!("✗ Failed to create invalid test tensor");
            return;
        }
    };
    
    match zloss.compute_loss(&invalid_tensor) {
        Err(_) => println!("✓ Invalid tensor shape properly rejected"),
        Ok(_) => {
            println!("✗ Should reject 1D tensor");
            return;
        }
    }
    
    println!("\n🎉 GREEN PHASE VERIFICATION COMPLETE!");
    println!("✅ Basic functionality is working:");
    println!("  - Configuration validation works");
    println!("  - Router creation works");  
    println!("  - compute_loss now succeeds with basic Z-loss computation");
    println!("  - Statistics tracking works");
    println!("  - Error handling works");
    println!("\n✅ Ready to proceed to REFACTOR phase!");
}