//! Barlow Twins Implementation
//!
//! Self-supervised learning method using redundancy reduction to learn representations.
//! Based on "Barlow Twins: Self-Supervised Learning via Redundancy Reduction" (Zbontar et al., 2021).
//!
//! Key features:
//! - Cross-correlation matrix computation between twin representations
//! - Redundancy reduction loss (making correlation matrix close to identity)
//! - Projector network for embedding transformation
//! - Support for different backbone architectures
//! - Batch normalization in projector
//! - Symmetric loss computation

use super::byol::Backbone; // Reuse the Backbone trait
use crate::prelude::*;
use parking_lot::RwLock;
use std::sync::Arc;

/// Barlow Twins configuration parameters
#[derive(Debug, Clone)]
pub struct BarlowTwinsConfig {
    /// Dimension of backbone output features
    pub backbone_dim: usize,
    /// Projector network layer dimensions  
    pub projector_dims: Vec<usize>,
    /// Lambda coefficient for redundancy reduction term
    pub lambda_coeff: f32,
    /// Whether to use batch normalization in projector
    pub batch_norm: bool,
    /// Scale factor for numerical stability
    pub scale_loss: f32,
}

impl Default for BarlowTwinsConfig {
    fn default() -> Self {
        Self {
            backbone_dim: 2048, // Will be overridden based on backbone
            projector_dims: vec![8192, 8192, 8192],
            lambda_coeff: 0.005,
            batch_norm: true,
            scale_loss: 1.0 / 32.0, // For numerical stability
        }
    }
}

impl BarlowTwinsConfig {
    /// Create new Barlow Twins configuration
    pub fn new(backbone_dim: usize, projector_dims: Vec<usize>) -> Self {
        Self {
            backbone_dim,
            projector_dims,
            ..Default::default()
        }
    }

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

    /// Set lambda coefficient for redundancy reduction
    pub fn with_lambda_coeff(mut self, lambda_coeff: f32) -> Self {
        self.lambda_coeff = lambda_coeff;
        self
    }

    /// Set whether to use batch normalization
    pub fn with_batch_norm(mut self, batch_norm: bool) -> Self {
        self.batch_norm = batch_norm;
        self
    }

    /// Set loss scale factor
    pub fn with_scale_loss(mut self, scale_loss: f32) -> Self {
        self.scale_loss = scale_loss;
        self
    }
}

/// Projector network that maps backbone features to embedding space
#[derive(Debug)]
pub struct ProjectorNetwork {
    layers: Vec<LinearLayer>,
    batch_norms: Vec<Option<BatchNorm>>,
    input_dim: usize,
    layer_dims: Vec<usize>,
    use_batch_norm: bool,
    device: Device,
}

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

#[derive(Debug)]
struct BatchNorm {
    weight: Arc<RwLock<Tensor>>,
    bias: Arc<RwLock<Tensor>>,
    running_mean: Arc<RwLock<Tensor>>,
    running_var: Arc<RwLock<Tensor>>,
    eps: f32,
    momentum: f32,
}

impl ProjectorNetwork {
    /// Create new projector network
    pub fn new(
        input_dim: usize,
        layer_dims: Vec<usize>,
        use_batch_norm: bool,
        device: &Device,
    ) -> Result<Self> {
        let mut layers = Vec::new();
        let mut batch_norms = Vec::new();

        let mut current_dim = input_dim;

        for (i, &output_dim) in layer_dims.iter().enumerate() {
            // Linear layer
            let weight = Tensor::randn(&[current_dim, output_dim], device)?;
            let bias = Tensor::zeros(&[output_dim], device)?;

            layers.push(LinearLayer {
                weight: Arc::new(RwLock::new(weight)),
                bias: Arc::new(RwLock::new(bias)),
            });

            // Batch normalization (except for the last layer)
            if use_batch_norm && i < layer_dims.len() - 1 {
                let bn_weight = Tensor::ones(&[output_dim], device)?;
                let bn_bias = Tensor::zeros(&[output_dim], device)?;
                let running_mean = Tensor::zeros(&[output_dim], device)?;
                let running_var = Tensor::ones(&[output_dim], device)?;

                batch_norms.push(Some(BatchNorm {
                    weight: Arc::new(RwLock::new(bn_weight)),
                    bias: Arc::new(RwLock::new(bn_bias)),
                    running_mean: Arc::new(RwLock::new(running_mean)),
                    running_var: Arc::new(RwLock::new(running_var)),
                    eps: 1e-5,
                    momentum: 0.1,
                }));
            } else {
                batch_norms.push(None);
            }

            current_dim = output_dim;
        }

        Ok(Self {
            layers,
            batch_norms,
            input_dim,
            layer_dims,
            use_batch_norm,
            device: device.clone(),
        })
    }

    /// Forward pass through projector 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)?;

            // Batch normalization (if enabled and not last layer)
            if let Some(ref batch_norm) = self.batch_norms[i] {
                x = self.apply_batch_norm(&x, batch_norm)?;
            }

            // ReLU activation (except for the last layer)
            if i < self.layers.len() - 1 {
                x = x.relu()?;
            }
        }

        Ok(x)
    }

    fn apply_batch_norm(&self, x: &Tensor, batch_norm: &BatchNorm) -> Result<Tensor> {
        // Simplified batch normalization implementation
        let weight = batch_norm.weight.read();
        let bias = batch_norm.bias.read();
        let running_mean = batch_norm.running_mean.read();
        let running_var = batch_norm.running_var.read();

        // Normalize: (x - mean) / sqrt(var + eps)
        let eps_tensor = Tensor::full(running_var.dims(), batch_norm.eps, &self.device)?;
        let var_eps = running_var.add(&eps_tensor)?;
        let std = var_eps.sqrt()?;

        let normalized = x.sub(&*running_mean)?.div(&std)?;
        let scaled = normalized.mul(&*weight)?;
        let output = scaled.add(&*bias)?;

        Ok(output)
    }

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

    /// Get output dimension
    pub fn output_dim(&self) -> usize {
        self.layer_dims.last().copied().unwrap_or(self.input_dim)
    }

    /// Get layer dimensions
    pub fn layer_dims(&self) -> &[usize] {
        &self.layer_dims
    }

    /// Check if using batch normalization
    pub fn has_batch_norm(&self) -> bool {
        self.use_batch_norm
    }
}

/// Normalize features to have zero mean and unit variance per feature dimension
pub fn normalize_features(features: &Tensor) -> Result<Tensor> {
    let mean = features.mean(&[0i32], false)?; // Mean across batch dimension
    let centered = features.sub(&mean)?;

    let variance = centered.pow_scalar(2.0)?.mean(&[0i32], false)?;
    let eps = Tensor::full(variance.dims(), 1e-8, variance.device())?;
    let std = variance.add(&eps)?.sqrt()?;

    Ok(centered.div(&std)?)
}

/// Compute cross-correlation matrix between two normalized embeddings
pub fn compute_cross_correlation_matrix(y1: &Tensor, y2: &Tensor) -> Result<Tensor> {
    let batch_size = y1.shape()[0] as f32;

    // Normalize features
    let y1_norm = normalize_features(y1)?;
    let y2_norm = normalize_features(y2)?;

    // Compute cross-correlation: C[i,j] = sum(y1_norm[i] * y2_norm[j]) / batch_size
    let y1_t = y1_norm.transpose(0, 1)?; // [feature_dim, batch_size]
    let cross_corr = y1_t.matmul(&y2_norm)?; // [feature_dim, feature_dim]

    let batch_size_tensor = Tensor::full(&[], batch_size, cross_corr.device())?;
    Ok(cross_corr.div(&batch_size_tensor)?)
}

/// Extract diagonal elements from a square matrix
pub fn extract_diagonal(matrix: &Tensor) -> Result<Tensor> {
    let shape = matrix.shape();
    let size = shape[0];

    let mut diag_values = Vec::with_capacity(size);
    let matrix_data = matrix.to_vec()?;

    for i in 0..size {
        let idx = i * size + i; // Diagonal index in flattened matrix
        diag_values.push(matrix_data[idx]);
    }

    Ok(Tensor::from_vec(diag_values, &[size], matrix.device())?)
}

/// Result of Barlow Twins loss computation
#[derive(Debug)]
pub struct BarlowTwinsLossResult {
    /// Total loss
    pub loss: Tensor,
    /// Invariance loss (diagonal term)
    pub invariance_loss: f32,
    /// Redundancy reduction loss (off-diagonal term)
    pub redundancy_loss: f32,
    /// Cross-correlation matrix
    pub cross_correlation: Tensor,
}

/// Compute Barlow Twins loss
pub fn compute_barlow_twins_loss(
    y1: &Tensor,
    y2: &Tensor,
    lambda_coeff: f32,
    scale_loss: f32,
) -> Result<BarlowTwinsLossResult> {
    // Compute cross-correlation matrix
    let cross_corr = compute_cross_correlation_matrix(y1, y2)?;
    let feature_dim = cross_corr.shape()[0];

    // Create identity matrix
    let identity = Tensor::eye(feature_dim, cross_corr.device())?;

    // Invariance loss: sum((1 - C[i,i])^2) - diagonal should be 1
    let diag_diff = identity.sub(&cross_corr)?;
    let diag_squared = diag_diff.pow_scalar(2.0)?;

    // Extract diagonal elements for invariance loss
    let diagonal = extract_diagonal(&cross_corr)?;
    let ones = Tensor::ones(diagonal.dims(), diagonal.device())?;
    let diag_loss_vec = ones.sub(&diagonal)?.pow_scalar(2.0)?;
    let invariance_loss_tensor = diag_loss_vec.sum(None)?;
    let invariance_loss = invariance_loss_tensor.to_vec()?[0];

    // Redundancy reduction loss: sum(C[i,j]^2 for i≠j) - off-diagonal should be 0
    let cross_corr_squared = cross_corr.pow_scalar(2.0)?;
    let total_squared = cross_corr_squared.sum(None)?;
    let diag_squared_sum = diag_squared.sum(None)?;
    let redundancy_loss_tensor = total_squared.sub(&diag_squared_sum)?;
    let redundancy_loss = redundancy_loss_tensor.to_vec()?[0];

    // Total loss = invariance_loss + lambda * redundancy_loss
    let lambda_tensor = Tensor::full(&[], lambda_coeff, cross_corr.device())?;
    let weighted_redundancy = redundancy_loss_tensor.mul(&lambda_tensor)?;
    let total_loss = invariance_loss_tensor.add(&weighted_redundancy)?;

    // Apply scaling
    let scale_tensor = Tensor::full(&[], scale_loss, cross_corr.device())?;
    let scaled_loss = total_loss.mul(&scale_tensor)?;

    Ok(BarlowTwinsLossResult {
        loss: scaled_loss,
        invariance_loss,
        redundancy_loss,
        cross_correlation: cross_corr,
    })
}

/// Barlow Twins trainer
pub struct BarlowTwinsTrainer {
    backbone: Box<dyn Backbone>,
    projector: ProjectorNetwork,
    config: BarlowTwinsConfig,
    device: Device,
    training: bool,
}

/// Result of a Barlow Twins training step
#[derive(Debug)]
pub struct BarlowTwinsTrainingResult {
    /// Total loss
    pub loss: Tensor,
    /// Invariance loss component
    pub invariance_loss: f32,
    /// Redundancy reduction loss component
    pub redundancy_loss: f32,
    /// Cross-correlation matrix
    pub cross_correlation: Tensor,
}

impl BarlowTwinsTrainer {
    /// Create new Barlow Twins trainer
    pub fn new(
        backbone: Box<dyn Backbone>,
        config: BarlowTwinsConfig,
        device: &Device,
    ) -> Result<Self> {
        let backbone_dim = backbone.output_dim();

        // Update config with actual backbone dimension
        let mut config = config;
        config.backbone_dim = backbone_dim;

        let projector = ProjectorNetwork::new(
            backbone_dim,
            config.projector_dims.clone(),
            config.batch_norm,
            device,
        )?;

        Ok(Self {
            backbone,
            projector,
            config,
            device: device.clone(),
            training: true,
        })
    }

    /// Perform one training step with two augmented views
    pub fn train_step(
        &mut self,
        images: &Tensor,
        _seed: Option<u64>,
    ) -> Result<BarlowTwinsTrainingResult> {
        if !self.training {
            return Err(TransformerError::InvalidInput(
                "Trainer must be in training mode".to_string(),
            ));
        }

        // For now, create two simple "augmented" views by adding noise
        // In practice, this would use proper augmentation pipeline
        let noise1 = Tensor::randn(images.dims(), &self.device)?.mul_scalar(0.01)?;
        let noise2 = Tensor::randn(images.dims(), &self.device)?.mul_scalar(0.01)?;
        let view1 = images.add(&noise1)?;
        let view2 = images.add(&noise2)?;

        // Forward pass through backbone and projector
        let z1 = self.backbone.forward(&view1)?;
        let z2 = self.backbone.forward(&view2)?;

        let y1 = self.projector.forward(&z1)?;
        let y2 = self.projector.forward(&z2)?;

        // Compute Barlow Twins loss
        let loss_result =
            compute_barlow_twins_loss(&y1, &y2, self.config.lambda_coeff, self.config.scale_loss)?;

        Ok(BarlowTwinsTrainingResult {
            loss: loss_result.loss,
            invariance_loss: loss_result.invariance_loss,
            redundancy_loss: loss_result.redundancy_loss,
            cross_correlation: loss_result.cross_correlation,
        })
    }

    /// Extract features for evaluation (using backbone only)
    pub fn extract_features(&self, images: &Tensor) -> Result<Tensor> {
        self.backbone.forward(images)
    }

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

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

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