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

use std::sync::Arc;
use parking_lot::RwLock;

// 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();
        let data: Vec<f32> = (0..size).map(|_| rand::random::<f32>() - 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, Box<dyn std::error::Error>> {
        if self.shape != other.shape {
            return Err("Shape mismatch".into());
        }
        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, Box<dyn std::error::Error>> {
        if self.shape != other.shape {
            return Err("Shape mismatch".into());
        }
        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, Box<dyn std::error::Error>> {
        if self.shape != other.shape {
            return Err("Shape mismatch".into());
        }
        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, Box<dyn std::error::Error>> {
        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, Box<dyn std::error::Error>> {
        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, Box<dyn std::error::Error>> {
        if self.shape.len() != 2 || other.shape.len() != 2 {
            return Err("Only 2D tensors supported for matmul".into());
        }
        if self.shape[1] != other.shape[0] {
            return Err("Matrix dimension mismatch".into());
        }
        
        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, Box<dyn std::error::Error>> {
        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, Box<dyn std::error::Error>>;

/// 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: Arc<RwLock<SimpleTensor>>,
    bias: Arc<RwLock<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: Arc::new(RwLock::new(weight)),
            bias: Arc::new(RwLock::new(bias)),
        })
    }

    fn forward(&self, input: &SimpleTensor) -> Result<SimpleTensor> {
        let weight = self.weight.read();
        let bias = self.bias.read();
        
        let output = input.matmul(&*weight)?;
        output.add(&*bias)
    }

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

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

    fn set_weight(&self, new_weight: SimpleTensor) -> Result<()> {
        *self.weight.write() = new_weight;
        Ok(())
    }

    fn set_bias(&self, new_bias: SimpleTensor) -> Result<()> {
        *self.bias.write() = new_bias;
        Ok(())
    }
}

/// 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
    }

    /// Set training mode
    pub fn train(&mut self) {
        self.training = true;
    }

    /// Set evaluation mode
    pub fn eval(&mut self) {
        self.training = false;
    }

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

    /// Get all parameters
    pub fn parameters(&self) -> Vec<SimpleTensor> {
        self.layers.iter().flat_map(|layer| {
            vec![layer.get_weight(), layer.get_bias()]
        }).collect()
    }

    /// Set parameters
    pub fn set_parameters(&self, params: Vec<SimpleTensor>) -> Result<()> {
        assert_eq!(params.len(), self.layers.len() * 2);
        
        for (i, layer) in self.layers.iter().enumerate() {
            layer.set_weight(params[i * 2].clone())?;
            layer.set_bias(params[i * 2 + 1].clone())?;
        }
        
        Ok(())
    }
}

/// 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> {
        // EMA update: teacher = decay * teacher + (1 - decay) * student
        let updated = teacher_param.mul_scalar(self.decay)?
            .add(&student_param.mul_scalar(1.0 - self.decay)?)?;
        
        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> {
    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 {
            // Sigmoid ramp-up function
            let progress = epoch as f32 / self.rampup_epochs as f32;
            let sigmoid = 1.0 / (1.0 + (-5.0 * (progress - 0.5)).exp());
            sigmoid * 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]);
    let output = student.forward(&input)?;
    assert_eq!(output.shape(), &[2, 3]);
    println!("✅ Forward pass test passed");
    
    // 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]);
    let updated = ema.update(&teacher_param, &student_param)?;
    assert_eq!(updated.shape(), &[5, 5]);
    println!("✅ EMA updater test passed");
    
    // Test 5: Consistency loss
    let pred1 = SimpleTensor::randn(vec![4, 3]);
    let pred2 = SimpleTensor::randn(vec![4, 3]);
    let loss = compute_consistency_loss(&pred1, &pred2)?;
    assert_eq!(loss.shape(), &[]);
    println!("✅ Consistency loss test passed");
    
    // Test 6: Ramp-up scheduler
    let scheduler = ConsistencyRampUp::new(5, 100.0);
    assert_eq!(scheduler.get_weight(0), 0.0);
    assert!((scheduler.get_weight(5) - 100.0).abs() < 1e-6);
    println!("✅ Ramp-up scheduler test passed");
    
    println!("🎉 All tests passed! Mean Teacher implementation is working correctly.");
    
    Ok(())
}