//! VICReg (Variance-Invariance-Covariance Regularization) Implementation
//!
//! Self-supervised learning method that explicitly avoids collapse through three loss terms:
//! - Invariance: Similar representations for augmented views (MSE loss)
//! - Variance: Maintains variance ≥ γ in each dimension (hinge loss)  
//! - Covariance: Decorrelates different dimensions (Frobenius norm of off-diagonal covariance)
//!
//! Based on "VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning" 
//! (Bardes et al., 2021).
//!
//! Key features:
//! - No momentum encoders or large batches required
//! - Explicit variance and covariance regularization prevents collapse
//! - Expander network (projector) transforms representations
//! - Support for different backbone architectures
//! - Symmetric loss computation on both augmented views

use crate::prelude::*;
use super::byol::Backbone;
use std::sync::Arc;
use parking_lot::RwLock;

/// VICReg configuration parameters
#[derive(Debug, Clone)]
pub struct VICRegConfig {
    /// Dimension of backbone output features
    pub backbone_dim: usize,
    /// Expander network layer dimensions
    pub expander_dims: Vec<usize>,
    /// λ - Coefficient for invariance loss (similarity between views)
    pub sim_coeff: f32,
    /// μ - Coefficient for variance loss (prevent collapse)
    pub std_coeff: f32,
    /// ν - Coefficient for covariance loss (decorrelate features)
    pub cov_coeff: f32,
    /// γ - Target standard deviation (variance regularization)
    pub variance_target: f32,
    /// ε - Epsilon for numerical stability
    pub epsilon: f32,
}

impl Default for VICRegConfig {
    fn default() -> Self {
        Self {
            backbone_dim: 2048,
            expander_dims: vec![8192, 8192, 8192],
            sim_coeff: 25.0,
            std_coeff: 25.0,
            cov_coeff: 1.0,
            variance_target: 1.0,
            epsilon: 1e-4,
        }
    }
}

impl VICRegConfig {
    /// Create new VICReg configuration
    pub fn new(backbone_dim: usize, expander_dims: Vec<usize>) -> Self {
        Self {
            backbone_dim,
            expander_dims,
            ..Default::default()
        }
    }

    /// Set backbone dimension
    pub fn with_backbone_dim(mut self, backbone_dim: usize) -> Self {
        self.backbone_dim = backbone_dim;
        self
    }

    /// Set invariance loss coefficient (λ)
    pub fn with_sim_coeff(mut self, sim_coeff: f32) -> Self {
        self.sim_coeff = sim_coeff;
        self
    }

    /// Set variance loss coefficient (μ)  
    pub fn with_std_coeff(mut self, std_coeff: f32) -> Self {
        self.std_coeff = std_coeff;
        self
    }

    /// Set covariance loss coefficient (ν)
    pub fn with_cov_coeff(mut self, cov_coeff: f32) -> Self {
        self.cov_coeff = cov_coeff;
        self
    }

    /// Set variance target (γ)
    pub fn with_variance_target(mut self, variance_target: f32) -> Self {
        self.variance_target = variance_target;
        self
    }

    /// Set epsilon for numerical stability
    pub fn with_epsilon(mut self, epsilon: f32) -> Self {
        self.epsilon = epsilon;
        self
    }
}

/// Expander network that maps backbone features to embedding space
/// Similar to projector but terminology from VICReg paper
#[derive(Debug)]
pub struct ExpanderNetwork {
    layers: Vec<LinearLayer>,
    input_dim: usize,
    layer_dims: Vec<usize>,
    device: Device,
}

#[derive(Debug)]
struct LinearLayer {
    weight: Arc<RwLock<Tensor>>,
    bias: Arc<RwLock<Tensor>>,
}

impl ExpanderNetwork {
    /// Create new expander network
    pub fn new(
        input_dim: usize,
        layer_dims: Vec<usize>,
        device: &Device,
    ) -> Result<Self> {
        let mut layers = Vec::new();
        let mut current_dim = input_dim;
        
        for &output_dim in layer_dims.iter() {
            // Initialize weights with Xavier/Glorot normal initialization
            let scale = (2.0 / (current_dim + output_dim) as f32).sqrt();
            let weight = Tensor::randn(vec![current_dim, output_dim], DType::F32, device)?
                .mul(&Tensor::full(vec![current_dim, output_dim], scale, DType::F32, device)?)?;
            let bias = Tensor::zeros(vec![output_dim], device)?;
            
            layers.push(LinearLayer {
                weight: Arc::new(RwLock::new(weight)),
                bias: Arc::new(RwLock::new(bias)),
            });
            
            current_dim = output_dim;
        }
        
        Ok(Self {
            layers,
            input_dim,
            layer_dims,
            device: device.clone(),
        })
    }

    /// Forward pass through expander network
    pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
        let mut x = input.clone();
        
        for (i, layer) in self.layers.iter().enumerate() {
            // Linear transformation
            let weight = layer.weight.read();
            let bias = layer.bias.read();
            x = x.matmul(&*weight)?.add(&*bias)?;
            
            // ReLU activation (except for the last layer)
            if i < self.layers.len() - 1 {
                x = x.relu()?;
            }
        }
        
        Ok(x)
    }

    pub fn input_dim(&self) -> usize {
        self.input_dim
    }

    pub fn output_dim(&self) -> usize {
        *self.layer_dims.last().unwrap_or(&self.input_dim)
    }

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

/// Result of VICReg loss computation
#[derive(Debug, Clone)]
pub struct VICRegLossResult {
    pub total_loss: Tensor,
    pub invariance_loss: Tensor,
    pub variance_loss_y1: Tensor,
    pub variance_loss_y2: Tensor,
    pub covariance_loss_y1: Tensor,
    pub covariance_loss_y2: Tensor,
}

/// Training result from VICReg training step
#[derive(Debug, Clone)]
pub struct VICRegTrainingResult {
    pub total_loss: f32,
    pub invariance_loss: f32,
    pub variance_loss: f32,
    pub covariance_loss: f32,
}

/// VICReg trainer combining backbone and expander networks
pub struct VICRegTrainer<B: Backbone> {
    backbone: B,
    expander: ExpanderNetwork,
    config: VICRegConfig,
    device: Device,
}

impl<B: Backbone> VICRegTrainer<B> {
    /// Create new VICReg trainer
    pub fn new(
        backbone: B,
        config: VICRegConfig,
        device: &Device,
    ) -> Result<Self> {
        let backbone_dim = backbone.output_dim();
        let mut config = config;
        config.backbone_dim = backbone_dim;
        
        let expander = ExpanderNetwork::new(
            backbone_dim,
            config.expander_dims.clone(),
            device,
        )?;
        
        Ok(Self {
            backbone,
            expander,
            config,
            device: device.clone(),
        })
    }

    /// Forward pass through both networks
    pub fn forward(&self, x1: &Tensor, x2: &Tensor) -> Result<VICRegLossResult> {
        // Get representations from backbone
        let z1 = self.backbone.forward(x1)?;
        let z2 = self.backbone.forward(x2)?;
        
        // Transform through expander network
        let y1 = self.expander.forward(&z1)?;
        let y2 = self.expander.forward(&z2)?;
        
        // Compute VICReg loss
        compute_vicreg_loss(&y1, &y2, &self.config)
    }

    /// Training step with augmentation (simplified for test)
    pub fn train_step(&mut self, images: &Tensor) -> Result<VICRegTrainingResult> {
        // In practice, you would apply different augmentations here
        // For testing, we'll use the same image with small noise as "augmentation"
        let noise = Tensor::randn(images.shape().to_vec(), &self.device)?
            .mul(&Tensor::full(images.shape().to_vec(), 0.01, DType::F32, &self.device)?)?;
        let x2 = images.add(&noise)?;
        
        let loss_result = self.forward(images, &x2)?;
        
        // In real implementation, you would:
        // 1. Compute gradients
        // 2. Update parameters
        // 3. Zero gradients
        
        Ok(VICRegTrainingResult {
            total_loss: loss_result.total_loss.to_scalar::<f32>()?,
            invariance_loss: loss_result.invariance_loss.to_scalar::<f32>()?,
            variance_loss: (loss_result.variance_loss_y1.to_scalar::<f32>()? + 
                           loss_result.variance_loss_y2.to_scalar::<f32>()?) * 0.5,
            covariance_loss: (loss_result.covariance_loss_y1.to_scalar::<f32>()? + 
                             loss_result.covariance_loss_y2.to_scalar::<f32>()?) * 0.5,
        })
    }

    pub fn config(&self) -> &VICRegConfig {
        &self.config
    }
}

/// Compute invariance loss: MSE between representations
pub fn compute_invariance_loss(y1: &Tensor, y2: &Tensor) -> Result<Tensor> {
    let diff = y1.sub(y2)?;
    let squared_diff = diff.mul(&diff)?;
    
    // Mean over all dimensions
    let total_elements = squared_diff.shape().iter().product::<usize>() as f32;
    let sum = squared_diff.sum(None)?;
    sum.div(&Tensor::full(vec![], total_elements, DType::F32, squared_diff.device())?)
}

/// Compute variance loss: Hinge loss to maintain std ≥ γ  
pub fn compute_variance_loss(y: &Tensor, gamma: f32, epsilon: f32) -> Result<Tensor> {
    let batch_size = y.shape()[0] as f32;
    let feature_dim = y.shape()[1];
    
    // Compute mean along batch dimension
    let mean = y.sum_axis(0)?.div(&Tensor::full(vec![feature_dim], batch_size, DType::F32, y.device())?)?;
    
    // Center the features: y - mean
    let centered = y.sub(&mean.unsqueeze(0)?)?;
    
    // Compute variance: E[(y - mean)²]
    let squared_diff = centered.mul(&centered)?;
    let variance = squared_diff.sum_axis(0)?.div(&Tensor::full(vec![feature_dim], batch_size, DType::F32, y.device())?)?;
    
    // Compute standard deviation with epsilon for stability
    let eps_tensor = Tensor::full(vec![feature_dim], epsilon, DType::F32, y.device())?;
    let variance_eps = variance.add(&eps_tensor)?;
    let std = variance_eps.sqrt()?;
    
    // Hinge loss: ReLU(γ - std)
    let gamma_tensor = Tensor::full(vec![feature_dim], gamma, DType::F32, y.device())?;
    let hinge = gamma_tensor.sub(&std)?;
    let hinge_loss = hinge.relu()?;
    
    // Mean over all features
    hinge_loss.mean_all()
}

/// Compute covariance loss: Frobenius norm of off-diagonal covariance matrix
pub fn compute_covariance_loss(y: &Tensor, epsilon: f32) -> Result<Tensor> {
    let batch_size = y.shape()[0] as f32;
    let feature_dim = y.shape()[1];
    
    // Center the features
    let mean = y.sum_axis(0)?.div(&Tensor::full(vec![feature_dim], batch_size, DType::F32, y.device())?)?;
    let centered = y.sub(&mean.unsqueeze(0)?)?;
    
    // Compute covariance matrix: (1/N) * X^T * X
    let covariance = centered.transpose(0, 1)?.matmul(&centered)?
        .div(&Tensor::full(vec![feature_dim, feature_dim], batch_size, DType::F32, y.device())?)?;
    
    // Add epsilon to diagonal for numerical stability
    let eps_tensor = Tensor::full(vec![feature_dim], epsilon, DType::F32, y.device())?;
    let eye = create_identity_matrix(feature_dim, y.device())?;
    let eps_eye = eye.mul(&eps_tensor.unsqueeze(1)?)?;
    let stable_cov = covariance.add(&eps_eye)?;
    
    // Zero out diagonal elements (we only want off-diagonal)
    let off_diagonal_mask = create_off_diagonal_mask(feature_dim, y.device())?;
    let off_diagonal_cov = stable_cov.mul(&off_diagonal_mask)?;
    
    // Compute Frobenius norm of off-diagonal elements
    let squared = off_diagonal_cov.mul(&off_diagonal_cov)?;
    squared.sum(None)
}

/// Compute complete VICReg loss with all three terms
pub fn compute_vicreg_loss(
    y1: &Tensor, 
    y2: &Tensor, 
    config: &VICRegConfig
) -> Result<VICRegLossResult> {
    // Invariance loss: similarity between representations
    let invariance_loss = compute_invariance_loss(y1, y2)?;
    
    // Variance losses: prevent dimensional collapse
    let variance_loss_y1 = compute_variance_loss(y1, config.variance_target, config.epsilon)?;
    let variance_loss_y2 = compute_variance_loss(y2, config.variance_target, config.epsilon)?;
    
    // Covariance losses: decorrelate features
    let covariance_loss_y1 = compute_covariance_loss(y1, config.epsilon)?;
    let covariance_loss_y2 = compute_covariance_loss(y2, config.epsilon)?;
    
    // Weighted total loss
    let sim_term = invariance_loss.mul(&Tensor::full(vec![], config.sim_coeff, DType::F32, y1.device())?)?;
    let std_term1 = variance_loss_y1.mul(&Tensor::full(vec![], config.std_coeff, DType::F32, y1.device())?)?;
    let std_term2 = variance_loss_y2.mul(&Tensor::full(vec![], config.std_coeff, DType::F32, y1.device())?)?;
    let cov_term1 = covariance_loss_y1.mul(&Tensor::full(vec![], config.cov_coeff, DType::F32, y1.device())?)?;
    let cov_term2 = covariance_loss_y2.mul(&Tensor::full(vec![], config.cov_coeff, DType::F32, y1.device())?)?;
    
    let total_loss = sim_term
        .add(&std_term1)?
        .add(&std_term2)?
        .add(&cov_term1)?
        .add(&cov_term2)?;
    
    Ok(VICRegLossResult {
        total_loss,
        invariance_loss,
        variance_loss_y1,
        variance_loss_y2,
        covariance_loss_y1,
        covariance_loss_y2,
    })
}

/// Helper function to create identity matrix
fn create_identity_matrix(size: usize, device: &Device) -> Result<Tensor> {
    let mut data = vec![0.0f32; size * size];
    for i in 0..size {
        data[i * size + i] = 1.0;
    }
    Tensor::from_slice(&data, vec![size, size], DType::F32, device)
}

/// Helper function to create off-diagonal mask (1s off-diagonal, 0s on diagonal)
fn create_off_diagonal_mask(size: usize, device: &Device) -> Result<Tensor> {
    let mut data = vec![1.0f32; size * size];
    for i in 0..size {
        data[i * size + i] = 0.0;
    }
    Tensor::from_slice(&data, vec![size, size], DType::F32, device)
}