//! Mean Teacher Implementation
//!
//! Semi-supervised learning method using exponential moving average teacher models.
//! Based on "Mean teachers are better role models" (Tarvainen & Valpola, 2017).
//!
//! Key features:
//! - Student model (trainable) and teacher model (EMA of student)
//! - Consistency regularization between student and teacher predictions
//! - Support for labeled and unlabeled data
//! - Noise injection for robustness (Gaussian noise, dropout)
//! - Consistency weight ramp-up scheduling
//! - EMA update mechanism for teacher parameters

use crate::prelude::*;
use std::sync::Arc;
use parking_lot::RwLock;

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

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

/// 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,
    /// Device
    device: Device,
}

/// Teacher model wrapper (non-trainable, EMA of student)
#[derive(Debug)]
pub struct TeacherModel {
    /// Linear layers
    layers: Vec<LinearLayer>,
    /// Input dimension
    input_dim: usize,
    /// Output dimension
    output_dim: usize,
    /// Device
    device: Device,
}

/// Linear layer for neural networks
#[derive(Debug)]
struct LinearLayer {
    weight: Arc<RwLock<Tensor>>,
    bias: Arc<RwLock<Tensor>>,
}

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

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

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

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

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

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

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

    /// 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: &Tensor) -> Result<Tensor> {
        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<Tensor> {
        self.layers.iter().flat_map(|layer| {
            vec![layer.get_weight(), layer.get_bias()]
        }).collect()
    }

    /// Set parameters
    pub fn set_parameters(&self, params: Vec<Tensor>) -> 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(())
    }
}

impl TeacherModel {
    /// Create new teacher model
    pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize, device: &Device) -> Result<Self> {
        let layer1 = LinearLayer::new(input_dim, hidden_dim, device)?;
        let layer2 = LinearLayer::new(hidden_dim, output_dim, device)?;
        
        Ok(Self {
            layers: vec![layer1, layer2],
            input_dim,
            output_dim,
            device: device.clone(),
        })
    }

    /// 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 (always false for teacher)
    pub fn is_training(&self) -> bool {
        false
    }

    /// Forward pass through teacher model (always in eval mode)
    pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
        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<Tensor> {
        self.layers.iter().flat_map(|layer| {
            vec![layer.get_weight(), layer.get_bias()]
        }).collect()
    }

    /// Set parameters
    pub fn set_parameters(&self, params: Vec<Tensor>) -> 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(())
    }
}

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

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

    /// 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: &Tensor, student_param: &Tensor) -> Result<Tensor> {
        // 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)
    }

    /// Update all teacher parameters
    pub fn update_all(&mut self, teacher_params: &[Tensor], student_params: &[Tensor]) -> Result<Vec<Tensor>> {
        assert_eq!(teacher_params.len(), student_params.len());
        
        let mut updated_params = Vec::new();
        for (teacher_param, student_param) in teacher_params.iter().zip(student_params.iter()) {
            updated_params.push(self.update(teacher_param, student_param)?);
        }
        
        Ok(updated_params)
    }
}

/// Noise augmenter for input perturbation
#[derive(Debug)]
pub struct NoiseAugmenter {
    strategy: AugmentationStrategy,
    noise_level: f32,
    device: Device,
}

impl NoiseAugmenter {
    /// Create new noise augmenter
    pub fn new(strategy: AugmentationStrategy, noise_level: f32, device: &Device) -> Result<Self> {
        Ok(Self {
            strategy,
            noise_level,
            device: device.clone(),
        })
    }

    /// Apply noise to input tensor
    pub fn apply_noise(&self, input: &Tensor, seed: Option<u64>) -> Result<Tensor> {
        match self.strategy {
            AugmentationStrategy::Gaussian => self.apply_gaussian_noise(input, seed),
            AugmentationStrategy::Dropout => self.apply_dropout_noise(input, seed),
            AugmentationStrategy::Both => {
                let gaussian = self.apply_gaussian_noise(input, seed)?;
                self.apply_dropout_noise(&gaussian, seed.map(|s| s.wrapping_add(1)))
            }
        }
    }

    fn apply_gaussian_noise(&self, input: &Tensor, _seed: Option<u64>) -> Result<Tensor> {
        let noise = Tensor::randn(input.dims(), &self.device)?
            .mul_scalar(self.noise_level)?;
        Ok(input.add(&noise)?)
    }

    fn apply_dropout_noise(&self, input: &Tensor, _seed: Option<u64>) -> Result<Tensor> {
        // Simulate dropout by randomly scaling elements
        let keep_prob = 1.0 - self.noise_level;
        // Use randn sigmoid to get uniform [0,1]-like mask
        let mask = Tensor::randn(input.dims(), &self.device)?.sigmoid()?;
        let dropout_mask = mask.gt_scalar(self.noise_level)?
            .to_dtype(DType::F32)?
            .div_scalar(keep_prob)?;
        
        Ok(input.mul(&dropout_mask)?)
    }
}

/// Compute consistency loss (MSE) between student and teacher predictions
pub fn compute_consistency_loss(student_pred: &Tensor, teacher_pred: &Tensor) -> Result<Tensor> {
    let diff = student_pred.sub(teacher_pred)?;
    let squared_diff = diff.mul(&diff)?;
    // compute mean over all elements by summing and dividing
    let n = squared_diff.dims().iter().product::<usize>() as f32;
    let total = squared_diff.sum(None)?;
    Ok(total.div_scalar(n)?)
}

/// 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 (simple and predictable)
            let progress = epoch as f32 / self.rampup_epochs as f32;
            progress * self.max_weight
        }
    }
}

/// Training result from Mean Teacher step
#[derive(Debug, Clone)]
pub struct MeanTeacherTrainingResult {
    /// Supervised loss (classification loss on labeled data)
    pub supervised_loss: f32,
    /// Consistency loss (between student and teacher predictions)
    pub consistency_loss: f32,
    /// Total loss (supervised + weighted consistency)
    pub total_loss: f32,
}

/// Mean Teacher trainer
#[derive(Debug)]
pub struct MeanTeacherTrainer {
    config: MeanTeacherConfig,
    student: StudentModel,
    teacher: TeacherModel,
    ema_updater: EMAUpdater,
    noise_augmenter: NoiseAugmenter,
    rampup_scheduler: ConsistencyRampUp,
    current_epoch: usize,
    device: Device,
}

impl MeanTeacherTrainer {
    /// Create new Mean Teacher trainer
    pub fn new(
        config: MeanTeacherConfig,
        input_dim: usize,
        hidden_dim: usize,
        output_dim: usize,
        device: &Device,
    ) -> Result<Self> {
        let student = StudentModel::new(input_dim, hidden_dim, output_dim, device)?;
        let mut teacher = TeacherModel::new(input_dim, hidden_dim, output_dim, device)?;
        
        // Initialize teacher parameters with student parameters
        let student_params = student.parameters();
        teacher.set_parameters(student_params)?;
        
        let ema_updater = EMAUpdater::new(config.ema_decay, device)?;
        let noise_augmenter = NoiseAugmenter::new(
            config.augmentation_strategy.clone(),
            config.noise_level,
            device,
        )?;
        let rampup_scheduler = ConsistencyRampUp::new(
            config.consistency_rampup,
            config.consistency_weight,
        );
        
        Ok(Self {
            config: config.clone(),
            student,
            teacher,
            ema_updater,
            noise_augmenter,
            rampup_scheduler,
            current_epoch: 0,
            device: device.clone(),
        })
    }

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

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

    /// Set current epoch
    pub fn set_epoch(&mut self, epoch: usize) {
        self.current_epoch = epoch;
    }

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

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

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

    /// Get student parameters
    pub fn get_student_parameters(&self) -> Result<Vec<Tensor>> {
        Ok(self.student.parameters())
    }

    /// Get teacher parameters
    pub fn get_teacher_parameters(&self) -> Result<Vec<Tensor>> {
        Ok(self.teacher.parameters())
    }

    /// Supervised training step (labeled data only)
    pub fn supervised_step(
        &mut self,
        images: &Tensor,
        labels: &Tensor,
        epoch: usize,
    ) -> Result<MeanTeacherTrainingResult> {
        self.set_epoch(epoch);
        
        // Forward pass through student
        let student_pred = self.student.forward(images)?;
        
        // Compute supervised loss (simplified cross-entropy)
        let supervised_loss = self.compute_supervised_loss(&student_pred, labels)?;
        
        // Update teacher parameters
        self.update_teacher_parameters()?;
        
        Ok(MeanTeacherTrainingResult {
            supervised_loss: supervised_loss.to_vec()?[0],
            consistency_loss: 0.0,
            total_loss: supervised_loss.to_vec()?[0],
        })
    }

    /// Unsupervised training step (unlabeled data only)
    pub fn unsupervised_step(
        &mut self,
        unlabeled_images: &Tensor,
        epoch: usize,
    ) -> Result<MeanTeacherTrainingResult> {
        self.set_epoch(epoch);
        
        // Apply different noise to inputs for student and teacher
        let student_input = self.noise_augmenter.apply_noise(unlabeled_images, Some(42))?;
        let teacher_input = self.noise_augmenter.apply_noise(unlabeled_images, Some(24))?;
        
        // Forward passes
        let student_pred = self.student.forward(&student_input)?;
        let teacher_pred = self.teacher.forward(&teacher_input)?;
        
        // Compute consistency loss
        let consistency_loss = compute_consistency_loss(&student_pred, &teacher_pred)?;
        
        // Apply consistency weight ramp-up
        let consistency_weight = self.rampup_scheduler.get_weight(epoch);
        let weighted_consistency_loss = consistency_loss.mul_scalar(consistency_weight)?;
        
        // Update teacher parameters
        self.update_teacher_parameters()?;
        
        Ok(MeanTeacherTrainingResult {
            supervised_loss: 0.0,
            consistency_loss: weighted_consistency_loss.to_vec()?[0],
            total_loss: weighted_consistency_loss.to_vec()?[0],
        })
    }

    /// Mixed training step (labeled + unlabeled data)
    pub fn mixed_step(
        &mut self,
        labeled_images: &Tensor,
        labels: &Tensor,
        unlabeled_images: &Tensor,
        epoch: usize,
    ) -> Result<MeanTeacherTrainingResult> {
        self.set_epoch(epoch);
        
        // Supervised loss on labeled data
        let student_pred_labeled = self.student.forward(labeled_images)?;
        let supervised_loss = self.compute_supervised_loss(&student_pred_labeled, labels)?;
        
        // Consistency loss on unlabeled data
        let student_input = self.noise_augmenter.apply_noise(unlabeled_images, Some(42))?;
        let teacher_input = self.noise_augmenter.apply_noise(unlabeled_images, Some(24))?;
        
        let student_pred_unlabeled = self.student.forward(&student_input)?;
        let teacher_pred_unlabeled = self.teacher.forward(&teacher_input)?;
        
        let consistency_loss = compute_consistency_loss(&student_pred_unlabeled, &teacher_pred_unlabeled)?;
        let consistency_weight = self.rampup_scheduler.get_weight(epoch);
        let weighted_consistency_loss = consistency_loss.mul_scalar(consistency_weight)?;
        
        // Total loss
        let total_loss = supervised_loss.add(&weighted_consistency_loss)?;
        
        // Update teacher parameters
        self.update_teacher_parameters()?;
        
        Ok(MeanTeacherTrainingResult {
            supervised_loss: supervised_loss.to_vec()?[0],
            consistency_loss: weighted_consistency_loss.to_vec()?[0],
            total_loss: total_loss.to_vec()?[0],
        })
    }

    /// Extract features using student model (evaluation mode)
    pub fn extract_features(&self, input: &Tensor) -> Result<Tensor> {
        self.student.forward(input)
    }

    fn compute_supervised_loss(&self, predictions: &Tensor, labels: &Tensor) -> Result<Tensor> {
        // Simplified cross-entropy loss
        let log_probs = predictions.log_softmax(-1)?;
        let labels_one_hot = self.to_one_hot(labels, predictions.shape()[1])?;
        let loss = log_probs.mul(&labels_one_hot)?.sum(None)?.neg()?;
        Ok(loss.div_scalar(predictions.shape()[0] as f32)?)
    }

    fn to_one_hot(&self, labels: &Tensor, num_classes: usize) -> Result<Tensor> {
        let batch_size = labels.shape()[0];
        // Simplified one-hot: build the data manually then create tensor
        let labels_data = labels.to_vec()?;
        let mut one_hot_data = vec![0.0f32; batch_size * num_classes];
        for (i, &label_val) in labels_data.iter().enumerate().take(batch_size) {
            let idx = label_val as usize;
            if idx < num_classes {
                one_hot_data[i * num_classes + idx] = 1.0;
            }
        }
        Ok(Tensor::from_vec(one_hot_data, &[batch_size, num_classes], &self.device)?)
    }

    fn update_teacher_parameters(&mut self) -> Result<()> {
        let student_params = self.student.parameters();
        let teacher_params = self.teacher.parameters();
        
        let updated_params = self.ema_updater.update_all(&teacher_params, &student_params)?;
        self.teacher.set_parameters(updated_params)?;
        
        Ok(())
    }
}