optirs-learned 0.3.1

use std::fmt::Debug;
// Feed-forward network components for transformer layers
//
// This module implements the feed-forward network (FFN) layers used in
// transformer encoder/decoder blocks, including various activation functions
// and output projection layers.

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

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

/// Output transformation types
#[derive(Debug, Clone, Copy)]
pub enum OutputTransformation {
    /// Linear transformation
    Linear,
    /// Tanh activation
    Tanh,
    /// Sigmoid activation
    Sigmoid,
    /// Learned activation
    LearnedActivation,
    /// Parameter-specific scaling
    ParameterScaling,
}

/// Output projection layer for final transformer output
#[derive(Debug, Clone)]
pub struct OutputProjectionLayer<T: Float + Debug + Send + Sync + 'static> {
    /// Projection weights
    weights: Array2<T>,

    /// Projection bias
    bias: Array1<T>,

    /// Output transformation
    transformation: OutputTransformation,
}

/// Input embedding layer for transformer input processing
#[derive(Debug, Clone)]
pub struct InputEmbedding<T: Float + Debug + Send + Sync + 'static> {
    /// Embedding weights
    weights: Array2<T>,

    /// Input dimension
    input_dim: usize,

    /// Model dimension
    modeldim: usize,
}

impl<T: Float + Debug + Default + Clone + Send + Sync + 'static> OutputProjectionLayer<T> {
    /// Create new output projection layer
    pub fn new(input_dim: usize, output_dim: usize) -> Result<Self> {
        let mut rng = scirs2_core::random::thread_rng();
        let mut weights = Array2::zeros((input_dim, output_dim));

        // Xavier initialization
        let bound = (6.0 / (input_dim + output_dim) as f64).sqrt();
        for elem in weights.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound).expect("unwrap failed");
        }

        let bias = Array1::zeros(output_dim);

        Ok(Self {
            weights,
            bias,
            transformation: OutputTransformation::Linear,
        })
    }

    /// Create with specific transformation
    pub fn new_with_transformation(
        input_dim: usize,
        output_dim: usize,
        transformation: OutputTransformation,
    ) -> Result<Self> {
        let mut layer = Self::new(input_dim, output_dim)?;
        layer.transformation = transformation;
        Ok(layer)
    }

    /// Forward pass through output projection
    pub fn forward(&self, input: &Array2<T>) -> Result<Array2<T>> {
        let (seq_len, input_dim) = input.dim();
        let (weight_in, weight_out) = self.weights.dim();

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

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

        // Linear transformation
        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]] * self.weights[[k, j]];
                }
                output[[i, j]] = sum + self.bias[j];
            }
        }

        // Apply output transformation
        match self.transformation {
            OutputTransformation::Linear => {
                // No additional transformation
            }
            OutputTransformation::Tanh => {
                output.mapv_inplace(|x| x.tanh());
            }
            OutputTransformation::Sigmoid => {
                output.mapv_inplace(|x| T::one() / (T::one() + (-x).exp()));
            }
            OutputTransformation::LearnedActivation => {
                // For now, use a simple learned scaling
                output.mapv_inplace(|x| {
                    x * scirs2_core::numeric::NumCast::from(1.1).unwrap_or_else(|| T::zero())
                });
            }
            OutputTransformation::ParameterScaling => {
                // Apply different scaling per parameter dimension
                for j in 0..weight_out {
                    let scale = T::from(1.0 + 0.1 * (j as f64).sin()).expect("unwrap failed");
                    for i in 0..seq_len {
                        output[[i, j]] = output[[i, j]] * scale;
                    }
                }
            }
        }

        Ok(output)
    }

    /// Get output transformation type
    pub fn transformation(&self) -> OutputTransformation {
        self.transformation
    }

    /// Set output transformation type
    pub fn set_transformation(&mut self, transformation: OutputTransformation) {
        self.transformation = transformation;
    }

    /// Get projection weights
    pub fn weights(&self) -> &Array2<T> {
        &self.weights
    }

    /// Get projection bias
    pub fn bias(&self) -> &Array1<T> {
        &self.bias
    }

    /// Update weights and bias
    pub fn update_parameters(&mut self, weights: Array2<T>, bias: Array1<T>) -> Result<()> {
        let (weight_in, weight_out) = weights.dim();
        let bias_dim = bias.len();

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

        // Update internal dimensions if they match
        if (weight_in, weight_out) == self.weights.dim() && bias_dim == self.bias.len() {
            self.weights = weights;
            self.bias = bias;
            Ok(())
        } else {
            Err(OptimError::InvalidConfig(
                "New parameter dimensions don't match current layer dimensions".to_string(),
            ))
        }
    }
}

impl<T: Float + Debug + Default + Clone + Send + Sync + 'static> InputEmbedding<T> {
    /// Create new input embedding layer
    pub fn new(input_dim: usize, model_dim: usize) -> Result<Self> {
        let mut rng = scirs2_core::random::thread_rng();
        let mut weights = Array2::zeros((input_dim, model_dim));

        // Xavier initialization
        let bound = (6.0 / (input_dim + model_dim) as f64).sqrt();
        for elem in weights.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound).expect("unwrap failed");
        }

        Ok(Self {
            weights,
            input_dim,
            modeldim: model_dim,
        })
    }

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

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

        let mut output = Array2::zeros((seq_len, self.modeldim));

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

        Ok(output)
    }

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

    /// Get model dimension
    pub fn model_dim(&self) -> usize {
        self.modeldim
    }

    /// Get embedding weights
    pub fn weights(&self) -> &Array2<T> {
        &self.weights
    }

    /// Update embedding weights
    pub fn update_weights(&mut self, weights: Array2<T>) -> Result<()> {
        let (weight_in, weight_out) = weights.dim();

        if weight_in != self.input_dim || weight_out != self.modeldim {
            return Err(OptimError::InvalidConfig(
                "New weight dimensions don't match embedding dimensions".to_string(),
            ));
        }

        self.weights = weights;
        Ok(())
    }
}