optirs-learned 0.3.1

use std::fmt::Debug;
// Transformer encoder layers and components
//
// This module implements the encoder components of the transformer optimizer,
// including the transformer layer, feed-forward network, and layer normalization.

#[allow(dead_code)]
use scirs2_core::ndarray::{s, Array1, Array2};
use scirs2_core::numeric::Float;
use scirs2_core::random::{CoreRandom as Random, Rng as SCRRng};

use super::super::TransformerOptimizerConfig;
use super::attention::MultiHeadAttention;
use crate::error::{OptimError, Result};

/// Activation functions for feed-forward networks
#[derive(Debug, Clone, Copy)]
pub enum ActivationFunction {
    /// ReLU activation
    ReLU,
    /// GELU activation  
    GELU,
    /// Swish/SiLU activation
    Swish,
    /// GLU (Gated Linear Unit)
    GLU,
    /// GeGLU (GELU variant of GLU)
    GeGLU,
}

/// Single transformer encoder layer
#[derive(Debug, Clone)]
pub struct TransformerLayer<T: Float + Debug + Send + Sync + 'static> {
    /// Multi-head self-attention
    self_attention: MultiHeadAttention<T>,

    /// Cross-attention (for multi-task learning)
    cross_attention: Option<MultiHeadAttention<T>>,

    /// Feed-forward network
    feed_forward: FeedForwardNetwork<T>,

    /// Layer normalization layers
    ln1: LayerNorm<T>,
    ln2: LayerNorm<T>,
    ln3: Option<LayerNorm<T>>, // For cross-attention

    /// Dropout layers
    dropout1: DropoutLayer,
    dropout2: DropoutLayer,
    dropout3: Option<DropoutLayer>,

    /// Use pre-layer normalization
    pre_layer_norm: bool,
}

/// Feed-forward network
#[derive(Debug, Clone)]
pub struct FeedForwardNetwork<T: Float + Debug + Send + Sync + 'static> {
    /// First linear layer weights
    linear1: Array2<T>,

    /// First linear layer bias
    bias1: Array1<T>,

    /// Second linear layer weights
    linear2: Array2<T>,

    /// Second linear layer bias
    bias2: Array1<T>,

    /// Activation function
    activation: ActivationFunction,

    /// Dropout layer
    dropout: DropoutLayer,
}

/// Layer normalization
#[derive(Debug, Clone)]
pub struct LayerNorm<T: Float + Debug + Send + Sync + 'static> {
    /// Scale parameters (gamma)
    gamma: Array1<T>,

    /// Shift parameters (beta)
    beta: Array1<T>,

    /// Epsilon for numerical stability
    eps: T,

    /// Dimension
    dim: usize,
}

/// Dropout layer
#[derive(Debug, Clone)]
pub struct DropoutLayer {
    /// Dropout probability
    prob: f64,

    /// Training mode
    training: bool,
}

impl<T: Float + Debug + Default + Clone + std::iter::Sum + Send + Sync> TransformerLayer<T> {
    pub fn new(config: &TransformerOptimizerConfig, _rng: &mut Random) -> Result<Self> {
        let self_attention = MultiHeadAttention::new(config)?;
        let cross_attention = if config.cross_attention {
            Some(MultiHeadAttention::new(config)?)
        } else {
            None
        };

        let feed_forward = FeedForwardNetwork::new(config)?;

        let ln1 = LayerNorm::new(config.modeldim);
        let ln2 = LayerNorm::new(config.modeldim);
        let ln3 = if config.cross_attention {
            Some(LayerNorm::new(config.modeldim))
        } else {
            None
        };

        let dropout1 = DropoutLayer::new(config.attention_dropout);
        let dropout2 = DropoutLayer::new(config.ff_dropout);
        let dropout3 = if config.cross_attention {
            Some(DropoutLayer::new(config.attention_dropout))
        } else {
            None
        };

        Ok(Self {
            self_attention,
            cross_attention,
            feed_forward,
            ln1,
            ln2,
            ln3,
            dropout1,
            dropout2,
            dropout3,
            pre_layer_norm: config.pre_layer_norm,
        })
    }

    pub fn forward(&mut self, input: &Array2<T>) -> Result<Array2<T>> {
        let mut x = input.clone();

        // Self-attention with residual connection
        let residual = x.clone();
        if self.pre_layer_norm {
            x = self.ln1.forward(&x)?;
        }

        x = self.self_attention.forward(&x, &x, &x)?;
        x = self.dropout1.forward(&x)?;
        x = x + &residual;

        if !self.pre_layer_norm {
            x = self.ln1.forward(&x)?;
        }

        // Cross-attention (if enabled)
        if let Some(ref mut cross_attn) = self.cross_attention {
            let residual = x.clone();
            if self.pre_layer_norm {
                if let Some(ref ln3) = self.ln3 {
                    x = ln3.forward(&x)?;
                }
            }

            // For now, use same input as key/value for cross-attention
            x = cross_attn.forward(&x, &x, &x)?;
            if let Some(ref dropout3) = self.dropout3 {
                x = dropout3.forward(&x)?;
            }
            x = x + &residual;

            if !self.pre_layer_norm {
                if let Some(ref ln3) = self.ln3 {
                    x = ln3.forward(&x)?;
                }
            }
        }

        // Feed-forward with residual connection
        let residual = x.clone();
        if self.pre_layer_norm {
            x = self.ln2.forward(&x)?;
        }

        x = self.feed_forward.forward(&x)?;
        x = self.dropout2.forward(&x)?;
        x = x + &residual;

        if !self.pre_layer_norm {
            x = self.ln2.forward(&x)?;
        }

        Ok(x)
    }

    /// Get attention patterns for analysis
    pub fn get_attention_patterns(&self) -> Option<&scirs2_core::ndarray::Array3<T>> {
        self.self_attention.get_attention_patterns()
    }
}

impl<T: Float + Debug + Default + Clone + Send + Sync + 'static> FeedForwardNetwork<T> {
    pub fn new(config: &TransformerOptimizerConfig) -> Result<Self> {
        let modeldim = config.modeldim;
        let ff_dim = config.ff_dim;
        let mut rng = scirs2_core::random::thread_rng();

        // Initialize weights with Xavier initialization
        let bound1 = (6.0 / (modeldim + ff_dim) as f64).sqrt();
        let bound2 = (6.0 / (ff_dim + modeldim) as f64).sqrt();

        let mut linear1 = Array2::zeros((modeldim, ff_dim));
        let mut linear2 = Array2::zeros((ff_dim, modeldim));

        for elem in linear1.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound1).expect("unwrap failed");
        }
        for elem in linear2.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound2).expect("unwrap failed");
        }

        let bias1 = Array1::zeros(ff_dim);
        let bias2 = Array1::zeros(modeldim);

        Ok(Self {
            linear1,
            bias1,
            linear2,
            bias2,
            activation: ActivationFunction::GELU,
            dropout: DropoutLayer::new(config.ff_dropout),
        })
    }

    pub fn forward(&self, input: &Array2<T>) -> Result<Array2<T>> {
        // First linear layer
        let x1 = self.linear_transform(input, &self.linear1, &self.bias1)?;

        // Activation
        let x2 = self.apply_activation(&x1)?;

        // Dropout
        let x3 = self.dropout.forward(&x2)?;

        // Second linear layer
        let output = self.linear_transform(&x3, &self.linear2, &self.bias2)?;

        Ok(output)
    }

    fn linear_transform(
        &self,
        input: &Array2<T>,
        weights: &Array2<T>,
        bias: &Array1<T>,
    ) -> Result<Array2<T>> {
        let (seq_len, input_dim) = input.dim();
        let (weight_in, weight_out) = weights.dim();

        if input_dim != weight_in {
            return Err(OptimError::InvalidConfig(
                "Input dimension doesn't match weight matrix".to_string(),
            ));
        }

        if bias.len() != weight_out {
            return Err(OptimError::InvalidConfig(
                "Bias dimension doesn't match output dimension".to_string(),
            ));
        }

        let mut output = Array2::zeros((seq_len, weight_out));

        for i in 0..seq_len {
            for j in 0..weight_out {
                let mut sum = T::zero();
                for k in 0..input_dim {
                    sum = sum + input[[i, k]] * weights[[k, j]];
                }
                output[[i, j]] = sum + bias[j];
            }
        }

        Ok(output)
    }

    fn apply_activation(&self, input: &Array2<T>) -> Result<Array2<T>> {
        let mut output = input.clone();

        match self.activation {
            ActivationFunction::ReLU => {
                output.mapv_inplace(|x| if x > T::zero() { x } else { T::zero() });
            }
            ActivationFunction::GELU => {
                // Approximate GELU: 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
                output.mapv_inplace(|x| {
                    let sqrt_2_pi = scirs2_core::numeric::NumCast::from(0.7978845608)
                        .unwrap_or_else(|| T::zero()); // sqrt(2/π)
                    let coeff =
                        scirs2_core::numeric::NumCast::from(0.044715).unwrap_or_else(|| T::zero());
                    let x_cubed = x * x * x;
                    let inner = sqrt_2_pi * (x + coeff * x_cubed);
                    let tanh_val = inner.tanh();
                    scirs2_core::numeric::NumCast::from(0.5).unwrap_or_else(|| T::zero())
                        * x
                        * (T::one() + tanh_val)
                });
            }
            ActivationFunction::Swish => {
                output.mapv_inplace(|x| x * x.exp() / (T::one() + x.exp()));
            }
            ActivationFunction::GLU => {
                // For simplicity, treating as GELU for now
                output.mapv_inplace(|x| {
                    let sqrt_2_pi = scirs2_core::numeric::NumCast::from(0.7978845608)
                        .unwrap_or_else(|| T::zero());
                    let coeff =
                        scirs2_core::numeric::NumCast::from(0.044715).unwrap_or_else(|| T::zero());
                    let x_cubed = x * x * x;
                    let inner = sqrt_2_pi * (x + coeff * x_cubed);
                    let tanh_val = inner.tanh();
                    scirs2_core::numeric::NumCast::from(0.5).unwrap_or_else(|| T::zero())
                        * x
                        * (T::one() + tanh_val)
                });
            }
            ActivationFunction::GeGLU => {
                // For simplicity, treating as GELU for now
                output.mapv_inplace(|x| {
                    let sqrt_2_pi = scirs2_core::numeric::NumCast::from(0.7978845608)
                        .unwrap_or_else(|| T::zero());
                    let coeff =
                        scirs2_core::numeric::NumCast::from(0.044715).unwrap_or_else(|| T::zero());
                    let x_cubed = x * x * x;
                    let inner = sqrt_2_pi * (x + coeff * x_cubed);
                    let tanh_val = inner.tanh();
                    scirs2_core::numeric::NumCast::from(0.5).unwrap_or_else(|| T::zero())
                        * x
                        * (T::one() + tanh_val)
                });
            }
        }

        Ok(output)
    }

    /// Set activation function
    pub fn set_activation(&mut self, activation: ActivationFunction) {
        self.activation = activation;
    }
}

impl<T: Float + Debug + Default + Clone + std::iter::Sum + Send + Sync> LayerNorm<T> {
    pub fn new(dim: usize) -> Self {
        Self {
            gamma: Array1::ones(dim),
            beta: Array1::zeros(dim),
            eps: scirs2_core::numeric::NumCast::from(1e-6).unwrap_or_else(|| T::zero()),
            dim,
        }
    }

    pub fn forward(&self, input: &Array2<T>) -> Result<Array2<T>> {
        let (seq_len, input_dim) = input.dim();

        if input_dim != self.dim {
            return Err(OptimError::InvalidConfig(format!(
                "Input dimension {} doesn't match layer norm dimension {}",
                input_dim, self.dim
            )));
        }

        let mut output = Array2::zeros((seq_len, input_dim));

        for i in 0..seq_len {
            let row = input.slice(s![i, ..]);

            // Compute mean
            let mean = row.iter().cloned().sum::<T>()
                / scirs2_core::numeric::NumCast::from(input_dim).unwrap_or_else(|| T::zero());

            // Compute variance
            let variance = row
                .iter()
                .map(|&x| {
                    let diff = x - mean;
                    diff * diff
                })
                .sum::<T>()
                / scirs2_core::numeric::NumCast::from(input_dim).unwrap_or_else(|| T::zero());

            let std = (variance + self.eps).sqrt();

            // Normalize and scale/shift
            for j in 0..input_dim {
                let normalized = (input[[i, j]] - mean) / std;
                output[[i, j]] = self.gamma[j] * normalized + self.beta[j];
            }
        }

        Ok(output)
    }

    /// Get layer normalization parameters
    pub fn parameters(&self) -> (&Array1<T>, &Array1<T>) {
        (&self.gamma, &self.beta)
    }

    /// Set layer normalization parameters
    pub fn set_parameters(&mut self, gamma: Array1<T>, beta: Array1<T>) -> Result<()> {
        if gamma.len() != self.dim || beta.len() != self.dim {
            return Err(OptimError::InvalidConfig(
                "Parameter dimensions don't match layer norm dimension".to_string(),
            ));
        }
        self.gamma = gamma;
        self.beta = beta;
        Ok(())
    }
}

impl DropoutLayer {
    pub fn new(prob: f64) -> Self {
        Self {
            prob,
            training: true,
        }
    }

    pub fn forward<T: Float + Clone>(&self, input: &Array2<T>) -> Result<Array2<T>> {
        if !self.training || self.prob == 0.0 {
            return Ok(input.clone());
        }

        // For simplicity, just return input during inference/testing
        // In a full implementation, this would apply dropout during training
        Ok(input.clone())
    }

    /// Set training mode
    pub fn set_training(&mut self, training: bool) {
        self.training = training;
    }

    /// Get dropout probability
    pub fn prob(&self) -> f64 {
        self.prob
    }
}