//! Simple Mean Teacher test
//! 
//! Tests the Mean Teacher implementation independently

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

impl SimpleTensor {
    pub fn randn(shape: Vec<usize>) -> Self {
        let size: usize = shape.iter().product();
        // Simple pseudo-random using linear congruential generator
        let mut seed = 42u64;
        let data: Vec<f32> = (0..size).map(|_| {
            seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
            ((seed >> 16) as f32 / 65536.0) - 0.5
        }).collect();
        Self { data, shape }
    }

    pub fn zeros(shape: Vec<usize>) -> Self {
        let size: usize = shape.iter().product();
        let data = vec![0.0; size];
        Self { data, shape }
    }

    pub fn ones(shape: Vec<usize>) -> Self {
        let size: usize = shape.iter().product();
        let data = vec![1.0; size];
        Self { data, shape }
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn add(&self, other: &Self) -> Result<Self> {
        if self.shape != other.shape {
            return Err("Shape mismatch");
        }
        let data: Vec<f32> = self.data.iter().zip(&other.data)
            .map(|(a, b)| a + b)
            .collect();
        Ok(Self { data, shape: self.shape.clone() })
    }

    pub fn sub(&self, other: &Self) -> Result<Self> {
        if self.shape != other.shape {
            return Err("Shape mismatch");
        }
        let data: Vec<f32> = self.data.iter().zip(&other.data)
            .map(|(a, b)| a - b)
            .collect();
        Ok(Self { data, shape: self.shape.clone() })
    }

    pub fn mul(&self, other: &Self) -> Result<Self> {
        if self.shape != other.shape {
            return Err("Shape mismatch");
        }
        let data: Vec<f32> = self.data.iter().zip(&other.data)
            .map(|(a, b)| a * b)
            .collect();
        Ok(Self { data, shape: self.shape.clone() })
    }

    pub fn mul_scalar(&self, scalar: f32) -> Result<Self> {
        let data: Vec<f32> = self.data.iter().map(|x| x * scalar).collect();
        Ok(Self { data, shape: self.shape.clone() })
    }

    pub fn relu(&self) -> Result<Self> {
        let data: Vec<f32> = self.data.iter().map(|x| x.max(0.0)).collect();
        Ok(Self { data, shape: self.shape.clone() })
    }

    pub fn matmul(&self, other: &Self) -> Result<Self> {
        if self.shape.len() != 2 || other.shape.len() != 2 {
            return Err("Only 2D tensors supported for matmul");
        }
        if self.shape[1] != other.shape[0] {
            return Err("Matrix dimension mismatch");
        }
        
        let m = self.shape[0];
        let n = other.shape[1];
        let k = self.shape[1];
        
        let mut result = vec![0.0; m * n];
        for i in 0..m {
            for j in 0..n {
                for l in 0..k {
                    result[i * n + j] += self.data[i * k + l] * other.data[l * n + j];
                }
            }
        }
        
        Ok(Self {
            data: result,
            shape: vec![m, n],
        })
    }

    pub fn mean(&self) -> Result<Self> {
        let sum: f32 = self.data.iter().sum();
        let mean_val = sum / self.data.len() as f32;
        Ok(Self {
            data: vec![mean_val],
            shape: vec![],
        })
    }

    pub fn to_vec(&self) -> Vec<f32> {
        self.data.clone()
    }
}

type Result<T> = std::result::Result<T, &'static str>;

/// Augmentation strategy for noise injection
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AugmentationStrategy {
    /// Gaussian noise only
    Gaussian,
    /// Dropout noise only
    Dropout,
    /// Both Gaussian noise and dropout
    Both,
}

/// Mean Teacher configuration parameters
#[derive(Debug, Clone)]
pub struct MeanTeacherConfig {
    /// Exponential moving average decay rate for teacher updates
    pub ema_decay: f32,
    /// Maximum weight for consistency loss
    pub consistency_weight: f32,
    /// Number of epochs for consistency weight ramp-up
    pub consistency_rampup: usize,
    /// Noise level for input augmentation
    pub noise_level: f32,
    /// Augmentation strategy (Gaussian, Dropout, or Both)
    pub augmentation_strategy: AugmentationStrategy,
}

impl Default for MeanTeacherConfig {
    fn default() -> Self {
        Self {
            ema_decay: 0.999,
            consistency_weight: 100.0,
            consistency_rampup: 5,
            noise_level: 0.15,
            augmentation_strategy: AugmentationStrategy::Gaussian,
        }
    }
}

impl MeanTeacherConfig {
    /// Create new Mean Teacher configuration
    pub fn new() -> Self {
        Self::default()
    }

    /// Set EMA decay rate
    pub fn with_ema_decay(mut self, ema_decay: f32) -> Self {
        self.ema_decay = ema_decay;
        self
    }

    /// Set consistency weight
    pub fn with_consistency_weight(mut self, consistency_weight: f32) -> Self {
        self.consistency_weight = consistency_weight;
        self
    }

    /// Set consistency ramp-up epochs
    pub fn with_consistency_rampup(mut self, consistency_rampup: usize) -> Self {
        self.consistency_rampup = consistency_rampup;
        self
    }

    /// Set noise level
    pub fn with_noise_level(mut self, noise_level: f32) -> Self {
        self.noise_level = noise_level;
        self
    }

    /// Set augmentation strategy
    pub fn with_augmentation_strategy(mut self, augmentation_strategy: AugmentationStrategy) -> Self {
        self.augmentation_strategy = augmentation_strategy;
        self
    }
}

/// Linear layer for neural networks
#[derive(Debug)]
struct LinearLayer {
    weight: SimpleTensor,
    bias: SimpleTensor,
}

impl LinearLayer {
    fn new(input_dim: usize, output_dim: usize) -> Result<Self> {
        // Xavier initialization
        let scale = (2.0 / (input_dim + output_dim) as f32).sqrt();
        let weight = SimpleTensor::randn(vec![input_dim, output_dim]).mul_scalar(scale)?;
        let bias = SimpleTensor::zeros(vec![output_dim]);
        
        Ok(Self { weight, bias })
    }

    fn forward(&self, input: &SimpleTensor) -> Result<SimpleTensor> {
        let output = input.matmul(&self.weight)?;
        // Bias broadcasting: add bias to each row
        let mut result_data = output.data.clone();
        let batch_size = output.shape[0];
        let output_dim = output.shape[1];
        
        for i in 0..batch_size {
            for j in 0..output_dim {
                result_data[i * output_dim + j] += self.bias.data[j];
            }
        }
        
        Ok(SimpleTensor {
            data: result_data,
            shape: output.shape.clone(),
        })
    }

    fn get_weight(&self) -> &SimpleTensor {
        &self.weight
    }

    fn get_bias(&self) -> &SimpleTensor {
        &self.bias
    }
}

/// Student model wrapper (trainable)
#[derive(Debug)]
pub struct StudentModel {
    /// Linear layers
    layers: Vec<LinearLayer>,
    /// Input dimension
    input_dim: usize,
    /// Output dimension
    output_dim: usize,
    /// Training mode flag
    training: bool,
}

impl StudentModel {
    /// Create new student model
    pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize) -> Result<Self> {
        let layer1 = LinearLayer::new(input_dim, hidden_dim)?;
        let layer2 = LinearLayer::new(hidden_dim, output_dim)?;
        
        Ok(Self {
            layers: vec![layer1, layer2],
            input_dim,
            output_dim,
            training: false,
        })
    }

    /// Get input dimension
    pub fn input_dim(&self) -> usize {
        self.input_dim
    }

    /// Get output dimension
    pub fn output_dim(&self) -> usize {
        self.output_dim
    }

    /// Check if in training mode
    pub fn is_training(&self) -> bool {
        self.training
    }

    /// Forward pass through student model
    pub fn forward(&self, input: &SimpleTensor) -> Result<SimpleTensor> {
        let mut x = input.clone();
        
        // First layer with ReLU activation
        x = self.layers[0].forward(&x)?;
        x = x.relu()?;
        
        // Output layer
        x = self.layers[1].forward(&x)?;
        
        Ok(x)
    }
}

/// EMA updater for teacher parameters
#[derive(Debug)]
pub struct EMAUpdater {
    /// Decay rate for EMA
    decay: f32,
    /// Current step count
    step: usize,
}

impl EMAUpdater {
    /// Create new EMA updater
    pub fn new(decay: f32) -> Result<Self> {
        Ok(Self {
            decay,
            step: 0,
        })
    }

    /// Get decay rate
    pub fn decay(&self) -> f32 {
        self.decay
    }

    /// Get current step
    pub fn step(&self) -> usize {
        self.step
    }

    /// Update teacher parameter with student parameter
    pub fn update(&mut self, teacher_param: &SimpleTensor, student_param: &SimpleTensor) -> Result<SimpleTensor> {
        if teacher_param.shape() != student_param.shape() {
            return Err("Parameter shape mismatch");
        }
        
        // EMA update: teacher = decay * teacher + (1 - decay) * student
        let teacher_scaled = teacher_param.mul_scalar(self.decay)?;
        let student_scaled = student_param.mul_scalar(1.0 - self.decay)?;
        let updated = teacher_scaled.add(&student_scaled)?;
        
        self.step += 1;
        Ok(updated)
    }
}

/// Compute consistency loss (MSE) between student and teacher predictions
pub fn compute_consistency_loss(student_pred: &SimpleTensor, teacher_pred: &SimpleTensor) -> Result<SimpleTensor> {
    if student_pred.shape() != teacher_pred.shape() {
        return Err("Prediction shape mismatch");
    }
    let diff = student_pred.sub(teacher_pred)?;
    let squared_diff = diff.mul(&diff)?;
    squared_diff.mean()
}

/// Consistency weight ramp-up scheduler
#[derive(Debug)]
pub struct ConsistencyRampUp {
    rampup_epochs: usize,
    max_weight: f32,
}

impl ConsistencyRampUp {
    /// Create new consistency ramp-up scheduler
    pub fn new(rampup_epochs: usize, max_weight: f32) -> Self {
        Self {
            rampup_epochs,
            max_weight,
        }
    }

    /// Get rampup epochs
    pub fn rampup_epochs(&self) -> usize {
        self.rampup_epochs
    }

    /// Get maximum weight
    pub fn max_weight(&self) -> f32 {
        self.max_weight
    }

    /// Get current consistency weight based on epoch
    pub fn get_weight(&self, epoch: usize) -> f32 {
        if epoch >= self.rampup_epochs {
            self.max_weight
        } else if self.rampup_epochs == 0 {
            self.max_weight
        } else {
            // Linear ramp-up function (simpler and more predictable)
            let progress = epoch as f32 / self.rampup_epochs as f32;
            progress * self.max_weight
        }
    }
}

fn main() -> Result<()> {
    println!("🧪 Testing Mean Teacher Implementation");
    
    // Test 1: Configuration
    let config = MeanTeacherConfig::default();
    assert_eq!(config.ema_decay, 0.999);
    assert_eq!(config.consistency_weight, 100.0);
    println!("✅ Configuration test passed");
    
    // Test 2: Student model
    let student = StudentModel::new(10, 5, 3)?;
    assert_eq!(student.input_dim(), 10);
    assert_eq!(student.output_dim(), 3);
    println!("✅ Student model test passed");
    
    // Test 3: Forward pass
    let input = SimpleTensor::randn(vec![2, 10]);
    match student.forward(&input) {
        Ok(output) => {
            assert_eq!(output.shape(), &[2, 3]);
            println!("✅ Forward pass test passed");
        },
        Err(e) => {
            println!("❌ Forward pass test failed: {}", e);
            println!("input shape: {:?}", input.shape());
            return Err(e);
        }
    }
    
    // Test 4: EMA updater
    let mut ema = EMAUpdater::new(0.9)?;
    let teacher_param = SimpleTensor::zeros(vec![5, 5]);
    let student_param = SimpleTensor::ones(vec![5, 5]);
    match ema.update(&teacher_param, &student_param) {
        Ok(updated) => {
            assert_eq!(updated.shape(), &[5, 5]);
            println!("✅ EMA updater test passed");
        },
        Err(e) => {
            println!("❌ EMA updater test failed: {}", e);
            println!("teacher shape: {:?}, student shape: {:?}", teacher_param.shape(), student_param.shape());
            return Err(e);
        }
    }
    
    // Test 5: Consistency loss
    let pred1 = SimpleTensor::randn(vec![4, 3]);
    let pred2 = SimpleTensor::randn(vec![4, 3]);
    match compute_consistency_loss(&pred1, &pred2) {
        Ok(loss) => {
            assert_eq!(loss.shape(), &[]);
            println!("✅ Consistency loss test passed");
        },
        Err(e) => {
            println!("❌ Consistency loss test failed: {}", e);
            println!("pred1 shape: {:?}, pred2 shape: {:?}", pred1.shape(), pred2.shape());
            return Err(e);
        }
    }
    
    // Test 6: Ramp-up scheduler
    let scheduler = ConsistencyRampUp::new(5, 100.0);
    let weight_0 = scheduler.get_weight(0);
    let weight_5 = scheduler.get_weight(5);
    assert_eq!(weight_0, 0.0); // Should be exactly 0 with linear ramp-up
    assert!((weight_5 - 100.0).abs() < 1e-6);
    println!("✅ Ramp-up scheduler test passed");
    
    println!("🎉 All tests passed! Mean Teacher implementation is working correctly.");
    
    Ok(())
}