//! Gating Mechanism for TGGN
//!
//! Attention-based gating for temporal graph neural networks

use ndarray::{s, Array1, Array2, Axis};
use serde::{Deserialize, Serialize};

use ml_core::MLError;
use rand::prelude::*; // Replace common::rng with standard rand

/// Gradients for attention mechanism components
#[derive(Debug, Clone)]
struct AttentionGradients {
    pub query_weights_grad: Array2<f64>,
    pub key_weights_grad: Array2<f64>,
    pub value_weights_grad: Array2<f64>,
    pub output_weights_grad: Array2<f64>,
    pub bias_grad: Array1<f64>,
}

impl AttentionGradients {
    pub(crate) fn new(hidden_dim: usize) -> Self {
        Self {
            query_weights_grad: Array2::zeros((hidden_dim, hidden_dim)),
            key_weights_grad: Array2::zeros((hidden_dim, hidden_dim)),
            value_weights_grad: Array2::zeros((hidden_dim, hidden_dim)),
            output_weights_grad: Array2::zeros((hidden_dim, hidden_dim)),
            bias_grad: Array1::zeros(hidden_dim),
        }
    }
}

/// Gating mechanism for filtering and weighting messages
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GatingMechanism {
    /// Hidden dimension
    hidden_dim: usize,

    /// Query weights for attention
    query_weights: Array2<f64>,

    /// Key weights for attention
    key_weights: Array2<f64>,

    /// Value weights for attention
    value_weights: Array2<f64>,

    /// Output projection weights
    output_weights: Array2<f64>,

    /// Bias terms
    bias: Array1<f64>,

    /// Temperature for softmax
    temperature: f64,
}

impl GatingMechanism {
    /// Create new gating mechanism
    pub fn new(hidden_dim: usize) -> Result<Self, MLError> {
        // Initialize weights with Xavier initialization
        let scale = (2.0 / hidden_dim as f64).sqrt();

        let query_weights = Array2::from_shape_fn((hidden_dim, hidden_dim), |_| {
            (thread_rng().gen::<f64>() - 0.5) * scale
        });

        let key_weights = Array2::from_shape_fn((hidden_dim, hidden_dim), |_| {
            (thread_rng().gen::<f64>() - 0.5) * scale
        });

        let value_weights = Array2::from_shape_fn((hidden_dim, hidden_dim), |_| {
            (thread_rng().gen::<f64>() - 0.5) * scale
        });

        let output_weights = Array2::from_shape_fn((hidden_dim, hidden_dim), |_| {
            (thread_rng().gen::<f64>() - 0.5) * scale
        });

        let bias = Array1::zeros(hidden_dim);

        Ok(Self {
            hidden_dim,
            query_weights,
            key_weights,
            value_weights,
            output_weights,
            bias,
            temperature: 1.0,
        })
    }

    /// Apply gating mechanism to messages
    pub fn apply(&self, messages: &[Array1<f64>]) -> Result<Vec<Array1<f64>>, MLError> {
        if messages.is_empty() {
            return Ok(vec![]);
        }

        // Ensure all messages have correct dimension
        for msg in messages {
            if msg.len() != self.hidden_dim {
                return Err(MLError::DimensionMismatch {
                    expected: self.hidden_dim,
                    actual: msg.len(),
                });
            }
        }

        let n_messages = messages.len();

        // Stack messages into matrix for batch processing
        let mut message_matrix = Array2::zeros((n_messages, self.hidden_dim));
        for (i, msg) in messages.iter().enumerate() {
            message_matrix.row_mut(i).assign(msg);
        }

        // Compute queries, keys, and values
        let queries = self.compute_linear_transform(&message_matrix, &self.query_weights)?;
        let keys = self.compute_linear_transform(&message_matrix, &self.key_weights)?;
        let values = self.compute_linear_transform(&message_matrix, &self.value_weights)?;

        // Compute attention scores
        let attention_scores = self.compute_attention(&queries, &keys)?;

        // Apply attention to values
        let attended_values = self.apply_attention(&attention_scores, &values)?;

        // Apply output projection
        let output = self.compute_linear_transform(&attended_values, &self.output_weights)?;

        // Add bias and apply activation
        let mut result = Vec::new();
        for i in 0..output.nrows() {
            let mut row = output.row(i).to_owned();
            row += &self.bias;

            // Apply gated linear unit (GLU) activation
            let gated = self.apply_glu(&row)?;
            result.push(gated);
        }

        Ok(result)
    }

    /// Compute linear transformation: input * weights^T
    fn compute_linear_transform(
        &self,
        input: &Array2<f64>,
        weights: &Array2<f64>,
    ) -> Result<Array2<f64>, MLError> {
        if input.ncols() != weights.nrows() {
            return Err(MLError::DimensionMismatch {
                expected: weights.nrows(),
                actual: input.ncols(),
            });
        }

        Ok(input.dot(weights))
    }

    /// Compute attention scores using scaled dot-product attention
    fn compute_attention(
        &self,
        queries: &Array2<f64>,
        keys: &Array2<f64>,
    ) -> Result<Array2<f64>, MLError> {
        let scale = 1.0 / (self.hidden_dim as f64).sqrt();

        // Compute Q * K^T
        let scores = queries.dot(&keys.t()) * scale / self.temperature;

        // Apply softmax to each row
        let mut attention = Array2::zeros(scores.dim());
        for (mut row, score_row) in attention
            .axis_iter_mut(Axis(0))
            .zip(scores.axis_iter(Axis(0)))
        {
            let softmax = self.softmax(&score_row.to_owned())?;
            row.assign(&softmax);
        }

        Ok(attention)
    }

    /// Apply attention weights to values
    fn apply_attention(
        &self,
        attention: &Array2<f64>,
        values: &Array2<f64>,
    ) -> Result<Array2<f64>, MLError> {
        if attention.ncols() != values.nrows() {
            return Err(MLError::DimensionMismatch {
                expected: values.nrows(),
                actual: attention.ncols(),
            });
        }

        Ok(attention.dot(values))
    }

    /// Softmax activation function
    #[allow(clippy::unnecessary_wraps)]
    fn softmax(&self, x: &Array1<f64>) -> Result<Array1<f64>, MLError> {
        let max_val = x.iter().fold(f64::NEG_INFINITY, |acc, &val| acc.max(val));

        let exp_x: Array1<f64> = x.mapv(|val| (val - max_val).exp());
        let sum_exp = exp_x.sum();

        if sum_exp == 0.0 || !sum_exp.is_finite() {
            // Return uniform distribution as fallback
            Ok(Array1::from_elem(x.len(), 1.0 / x.len() as f64))
        } else {
            Ok(exp_x / sum_exp)
        }
    }

    /// Gated Linear Unit (GLU) activation
    fn apply_glu(&self, x: &Array1<f64>) -> Result<Array1<f64>, MLError> {
        let n = x.len();
        if n % 2 != 0 {
            return Err(MLError::DimensionMismatch {
                expected: n - (n % 2),
                actual: n,
            });
        }

        let half = n / 2;
        let first_half = x.slice(s![..half]);
        let second_half = x.slice(s![half..]);

        // GLU: first_half * sigmoid(second_half)
        let sigmoid_second = second_half.mapv(|val| 1.0 / (1.0 + (-val).exp()));
        let result = &first_half.to_owned() * &sigmoid_second;

        Ok(result)
    }

    /// Backpropagate gradient through GLU activation
    /// Input: gradient w.r.t. GLU output (dimension n/2)
    /// Output: gradient w.r.t. GLU input (dimension n)
    fn backprop_glu(
        &self,
        x: &Array1<f64>,
        grad_output: &Array1<f64>,
    ) -> Result<Array1<f64>, MLError> {
        let n = x.len();
        if n % 2 != 0 {
            return Err(MLError::DimensionMismatch {
                expected: n - (n % 2),
                actual: n,
            });
        }

        let half = n / 2;
        let first_half = x.slice(s![..half]);
        let second_half = x.slice(s![half..]);

        // Forward pass values needed for gradient
        let sigmoid_second = second_half.mapv(|val| 1.0 / (1.0 + (-val).exp()));

        // GLU: output = first_half * sigmoid(second_half)
        // d(output)/d(first_half) = sigmoid(second_half)
        // d(output)/d(second_half) = first_half * sigmoid(second_half) * (1 - sigmoid(second_half))

        let grad_first_half = grad_output * &sigmoid_second;
        let grad_second_half =
            grad_output * &first_half.to_owned() * &sigmoid_second.mapv(|s| s * (1.0 - s));

        // Concatenate gradients
        let mut grad_input = Array1::zeros(n);
        grad_input.slice_mut(s![..half]).assign(&grad_first_half);
        grad_input.slice_mut(s![half..]).assign(&grad_second_half);

        Ok(grad_input)
    }

    /// Update weights during training using real backpropagated gradients
    pub fn update_weights(
        &mut self,
        input_messages: &[Array1<f64>],
        target_outputs: &[Array1<f64>],
        learning_rate: f64,
    ) -> Result<(), MLError> {
        if input_messages.is_empty() || target_outputs.is_empty() {
            return Ok(()); // No data to learn from
        }

        if input_messages.len() != target_outputs.len() {
            return Err(MLError::DimensionMismatch {
                expected: input_messages.len(),
                actual: target_outputs.len(),
            });
        }

        // Forward pass to get current outputs
        let current_outputs = self.apply(input_messages)?;

        // Compute gradients using backpropagation through attention mechanism
        let gradients = self.compute_gradients(input_messages, &current_outputs, target_outputs)?;

        // Update weights using computed gradients
        self.apply_gradients(&gradients, learning_rate)?;

        Ok(())
    }

    /// Compute gradients for attention weights using backpropagation
    fn compute_gradients(
        &self,
        input_messages: &[Array1<f64>],
        current_outputs: &[Array1<f64>],
        target_outputs: &[Array1<f64>],
    ) -> Result<AttentionGradients, MLError> {
        let mut gradients = AttentionGradients::new(self.hidden_dim);

        // Compute output error (gradient of loss w.r.t. output)
        let mut output_errors = Vec::new();
        for (current, target) in current_outputs.iter().zip(target_outputs.iter()) {
            let error = current - target; // MSE gradient: 2(y_pred - y_true), factor of 2 absorbed into learning rate
            output_errors.push(error);
        }

        // Stack input messages for batch processing
        let n_messages = input_messages.len();
        let mut message_matrix = Array2::zeros((n_messages, self.hidden_dim));
        for (i, msg) in input_messages.iter().enumerate() {
            message_matrix.row_mut(i).assign(msg);
        }

        // Forward pass components for gradient computation
        let queries = self.compute_linear_transform(&message_matrix, &self.query_weights)?;
        let keys = self.compute_linear_transform(&message_matrix, &self.key_weights)?;
        let values = self.compute_linear_transform(&message_matrix, &self.value_weights)?;
        let attention_scores = self.compute_attention(&queries, &keys)?;
        let attended_values = self.apply_attention(&attention_scores, &values)?;

        // Apply output projection to get pre-GLU activations
        let pre_glu = self.compute_linear_transform(&attended_values, &self.output_weights)?;

        // Backpropagate through GLU activation
        // output_errors has dimension hidden_dim/2 (32) due to GLU activation
        // Need to backprop through GLU to get gradients w.r.t. pre-GLU activations (64)
        let mut pre_glu_grad_matrix = Array2::zeros((n_messages, self.hidden_dim));
        for (i, error) in output_errors.into_iter().enumerate() {
            // Get pre-GLU activation with bias
            let mut pre_glu_with_bias = pre_glu.row(i).to_owned();
            pre_glu_with_bias += &self.bias;

            // Backprop through GLU
            let grad_pre_glu = self.backprop_glu(&pre_glu_with_bias, &error)?;
            pre_glu_grad_matrix.row_mut(i).assign(&grad_pre_glu);
        }

        // Gradient w.r.t. output weights: pre_glu_grad^T * attended_values
        gradients.output_weights_grad = pre_glu_grad_matrix.t().dot(&attended_values);

        // Gradient w.r.t. bias: sum of pre-GLU gradients
        for i in 0..n_messages {
            for j in 0..self.hidden_dim {
                gradients.bias_grad[j] += pre_glu_grad_matrix[[i, j]];
            }
        }

        // Backpropagate through attention mechanism
        let attended_values_grad = pre_glu_grad_matrix.dot(&self.output_weights.t());

        // Gradient w.r.t. values: attention_scores^T * attended_values_grad
        let values_grad = attention_scores.t().dot(&attended_values_grad);

        // Gradient w.r.t. attention scores: attended_values_grad * values^T
        let attention_grad = attended_values_grad.dot(&values.t());

        // Backpropagate through softmax and attention computation
        let (queries_grad, keys_grad) =
            self.backprop_attention(&queries, &keys, &attention_grad)?;

        // Gradient w.r.t. query weights: queries_grad^T * input_messages
        gradients.query_weights_grad = queries_grad.t().dot(&message_matrix);

        // Gradient w.r.t. key weights: keys_grad^T * input_messages
        gradients.key_weights_grad = keys_grad.t().dot(&message_matrix);

        // Gradient w.r.t. value weights: values_grad^T * input_messages
        gradients.value_weights_grad = values_grad.t().dot(&message_matrix);

        Ok(gradients)
    }

    /// Backpropagate through attention computation
    fn backprop_attention(
        &self,
        queries: &Array2<f64>,
        keys: &Array2<f64>,
        attention_grad: &Array2<f64>,
    ) -> Result<(Array2<f64>, Array2<f64>), MLError> {
        let scale = 1.0 / (self.hidden_dim as f64).sqrt() / self.temperature;

        // Recompute attention scores for gradient computation
        let scores = queries.dot(&keys.t()) * scale;
        let attention_weights = self.compute_softmax_matrix(&scores)?;

        // Gradient through softmax: softmax_grad = attention_weights * (grad - (grad * attention_weights).sum())
        let mut softmax_grad = Array2::zeros(scores.dim());
        for i in 0..attention_grad.nrows() {
            let grad_row = attention_grad.row(i);
            let attn_row = attention_weights.row(i);

            // Compute gradient of softmax
            let grad_sum = grad_row.dot(&attn_row);
            for j in 0..softmax_grad.ncols() {
                softmax_grad[[i, j]] = attn_row[j] * (grad_row[j] - grad_sum);
            }
        }

        // Scale the gradient
        let scores_grad = &softmax_grad * scale;

        // Gradient w.r.t. queries: scores_grad * keys
        let queries_grad = scores_grad.dot(keys);

        // Gradient w.r.t. keys: scores_grad^T * queries
        let keys_grad = scores_grad.t().dot(queries);

        Ok((queries_grad, keys_grad))
    }

    /// Apply computed gradients to weights
    #[allow(clippy::unnecessary_wraps)]
    fn apply_gradients(
        &mut self,
        gradients: &AttentionGradients,
        learning_rate: f64,
    ) -> Result<(), MLError> {
        // Update query weights
        for ((i, j), &grad) in gradients.query_weights_grad.indexed_iter() {
            self.query_weights[[i, j]] -= learning_rate * grad;
        }

        // Update key weights
        for ((i, j), &grad) in gradients.key_weights_grad.indexed_iter() {
            self.key_weights[[i, j]] -= learning_rate * grad;
        }

        // Update value weights
        for ((i, j), &grad) in gradients.value_weights_grad.indexed_iter() {
            self.value_weights[[i, j]] -= learning_rate * grad;
        }

        // Update output weights
        for ((i, j), &grad) in gradients.output_weights_grad.indexed_iter() {
            self.output_weights[[i, j]] -= learning_rate * grad;
        }

        // Update bias
        for (i, &grad) in gradients.bias_grad.indexed_iter() {
            self.bias[i] -= learning_rate * grad;
        }

        // Apply gradient clipping to prevent exploding gradients
        self.clip_gradients(1.0);

        Ok(())
    }

    /// Clip gradients to prevent exploding gradients
    fn clip_gradients(&mut self, max_norm: f64) {
        let mut total_norm_sq = 0.0;

        // Calculate total gradient norm (we'll use the weights as proxy)
        total_norm_sq += self.query_weights.mapv(|x| x * x).sum();
        total_norm_sq += self.key_weights.mapv(|x| x * x).sum();
        total_norm_sq += self.value_weights.mapv(|x| x * x).sum();
        total_norm_sq += self.output_weights.mapv(|x| x * x).sum();
        total_norm_sq += self.bias.mapv(|x| x * x).sum();

        let total_norm = total_norm_sq.sqrt();

        if total_norm > max_norm {
            let clip_factor = max_norm / total_norm;
            self.query_weights.mapv_inplace(|x| x * clip_factor);
            self.key_weights.mapv_inplace(|x| x * clip_factor);
            self.value_weights.mapv_inplace(|x| x * clip_factor);
            self.output_weights.mapv_inplace(|x| x * clip_factor);
            self.bias.mapv_inplace(|x| x * clip_factor);
        }
    }

    /// Compute softmax for matrix (row-wise)
    fn compute_softmax_matrix(&self, scores: &Array2<f64>) -> Result<Array2<f64>, MLError> {
        let mut softmax_matrix = Array2::zeros(scores.dim());

        for (i, score_row) in scores.axis_iter(Axis(0)).enumerate() {
            let softmax_row = self.softmax(&score_row.to_owned())?;
            softmax_matrix.row_mut(i).assign(&softmax_row);
        }

        Ok(softmax_matrix)
    }

    /// Set temperature for attention softmax
    pub const fn set_temperature(&mut self, temperature: f64) {
        self.temperature = temperature.max(0.01); // Prevent division by zero
    }

    /// Get current temperature
    pub const fn temperature(&self) -> f64 {
        self.temperature
    }

    /// Reset weights to random initialization
    pub fn reset_weights(&mut self) -> Result<(), MLError> {
        *self = Self::new(self.hidden_dim)?;
        Ok(())
    }
}

/// Multi-head attention variant of gating mechanism
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultiHeadGating {
    /// Number of attention heads
    num_heads: usize,

    /// Individual gating mechanisms for each head
    heads: Vec<GatingMechanism>,

    /// Output projection layer
    output_projection: Array2<f64>,

    /// Bias for output projection
    output_bias: Array1<f64>,
}

impl MultiHeadGating {
    /// Create new multi-head gating mechanism
    pub fn new(hidden_dim: usize, num_heads: usize) -> Result<Self, MLError> {
        if hidden_dim % num_heads != 0 {
            return Err(MLError::ConfigError(format!(
                    "Hidden dimension {} must be divisible by number of heads {}",
                    hidden_dim, num_heads
                )));
        }

        let head_dim = hidden_dim / num_heads;
        let mut heads = Vec::new();

        for _ in 0..num_heads {
            heads.push(GatingMechanism::new(head_dim)?);
        }

        // After GLU, each head outputs head_dim/2, so concatenated dimension is num_heads * (head_dim/2)
        let concat_dim = num_heads * (head_dim / 2);
        let scale = (2.0 / hidden_dim as f64).sqrt();
        let output_projection = Array2::from_shape_fn((concat_dim, hidden_dim), |_| {
            (thread_rng().gen::<f64>() - 0.5) * scale
        });

        let output_bias = Array1::zeros(hidden_dim);

        Ok(Self {
            num_heads,
            heads,
            output_projection,
            output_bias,
        })
    }

    /// Apply multi-head gating to messages
    pub fn apply(&self, messages: &[Array1<f64>]) -> Result<Vec<Array1<f64>>, MLError> {
        if messages.is_empty() {
            return Ok(vec![]);
        }

        let n_messages = messages.len();
        let total_dim = messages[0].len();
        let head_dim = total_dim / self.num_heads;

        // Split messages across heads
        let mut head_outputs = Vec::new();

        for head_idx in 0..self.num_heads {
            let start_idx = head_idx * head_dim;
            let end_idx = (head_idx + 1) * head_dim;

            // Extract head-specific features from each message
            let head_messages: Vec<Array1<f64>> = messages
                .iter()
                .map(|msg| msg.slice(s![start_idx..end_idx]).to_owned())
                .collect();

            // Apply head-specific gating
            let head_output = self.heads[head_idx].apply(&head_messages)?;
            head_outputs.push(head_output);
        }

        // Concatenate head outputs
        let mut result = Vec::new();
        for msg_idx in 0..n_messages {
            let mut concatenated = Vec::new();
            for head_output in head_outputs.iter().take(self.num_heads) {
                concatenated.extend_from_slice(head_output[msg_idx].as_slice().ok_or(
                    MLError::ValidationError {
                        message: "Failed to get slice from tensor".to_owned(),
                    },
                )?);
            }

            // Apply output projection
            let concat_array = Array1::from(concatenated);
            let projected = concat_array.dot(&self.output_projection) + &self.output_bias;

            result.push(projected);
        }

        Ok(result)
    }

    /// Update weights for all heads using real gradients
    pub fn update_weights(
        &mut self,
        input_messages: &[Array1<f64>],
        target_outputs: &[Array1<f64>],
        learning_rate: f64,
    ) -> Result<(), MLError> {
        if input_messages.is_empty() || target_outputs.is_empty() {
            return Ok(());
        }

        // Update each head with its portion of the data
        let total_dim = input_messages[0].len();
        let head_dim = total_dim / self.num_heads;

        for head_idx in 0..self.num_heads {
            let start_idx = head_idx * head_dim;
            let end_idx = (head_idx + 1) * head_dim;

            // Extract head-specific data
            let head_inputs: Vec<Array1<f64>> = input_messages
                .iter()
                .map(|msg| msg.slice(s![start_idx..end_idx]).to_owned())
                .collect();

            let head_targets: Vec<Array1<f64>> = target_outputs
                .iter()
                .map(|msg| msg.slice(s![start_idx..end_idx]).to_owned())
                .collect();

            // Update head with real gradients
            self.heads[head_idx].update_weights(&head_inputs, &head_targets, learning_rate)?;
        }

        // Update output projection with real gradients
        self.update_output_projection(input_messages, target_outputs, learning_rate)?;

        Ok(())
    }

    /// Update output projection weights using gradients
    fn update_output_projection(
        &mut self,
        input_messages: &[Array1<f64>],
        target_outputs: &[Array1<f64>],
        learning_rate: f64,
    ) -> Result<(), MLError> {
        // Forward pass through heads to get concatenated features
        let head_outputs = self.compute_head_outputs(input_messages)?;

        // Compute output projection gradients
        let mut projection_grad: Array2<f32> = Array2::zeros(self.output_projection.dim());
        let mut bias_grad: Array1<f32> = Array1::zeros(self.output_bias.len());

        for (head_output, target) in head_outputs
            .into_iter()
            .zip(target_outputs.iter())
        {
            // Error in output
            let output_error = &head_output - target;

            // Gradient w.r.t. projection weights: error * input^T
            for i in 0..projection_grad.nrows() {
                for j in 0..projection_grad.ncols() {
                    projection_grad[[i, j]] += (output_error[i] * head_output[j]) as f32;
                }
            }

            // Gradient w.r.t. bias: sum of errors
            for i in 0..bias_grad.len() {
                bias_grad[i] += output_error[i] as f32;
            }
        }

        // Apply gradients
        let n_samples = input_messages.len() as f64;
        for ((i, j), &grad) in projection_grad.indexed_iter() {
            self.output_projection[[i, j]] -= learning_rate * grad as f64 / n_samples;
        }

        for (i, grad) in bias_grad.iter().enumerate() {
            self.output_bias[i] -= learning_rate * *grad as f64 / n_samples;
        }

        Ok(())
    }

    /// Compute outputs from all heads for gradient computation
    fn compute_head_outputs(
        &self,
        input_messages: &[Array1<f64>],
    ) -> Result<Vec<Array1<f64>>, MLError> {
        if input_messages.is_empty() {
            return Ok(vec![]);
        }

        let n_messages = input_messages.len();
        let total_dim = input_messages[0].len();
        let head_dim = total_dim / self.num_heads;

        // Apply each head to its portion of the input
        let mut head_outputs = Vec::new();
        for head_idx in 0..self.num_heads {
            let start_idx = head_idx * head_dim;
            let end_idx = (head_idx + 1) * head_dim;

            let head_messages: Vec<Array1<f64>> = input_messages
                .iter()
                .map(|msg| msg.slice(s![start_idx..end_idx]).to_owned())
                .collect();

            let head_output = self.heads[head_idx].apply(&head_messages)?;
            head_outputs.push(head_output);
        }

        // Concatenate head outputs and apply projection
        let mut result = Vec::new();
        for msg_idx in 0..n_messages {
            let mut concatenated = Vec::new();
            for head_output in head_outputs.iter().take(self.num_heads) {
                concatenated.extend_from_slice(head_output[msg_idx].as_slice().ok_or(
                    MLError::ValidationError {
                        message: "Failed to get slice from tensor".to_owned(),
                    },
                )?);
            }

            // Don't apply projection here - return concatenated features for gradient computation
            result.push(Array1::from(concatenated));
        }

        Ok(result)
    }
}

#[cfg(test)]
#[allow(clippy::assertions_on_result_states, clippy::manual_range_contains)]
mod tests {
    use super::*;
    use ndarray::array;

    #[test]
    fn test_gating_mechanism() -> Result<(), Box<dyn std::error::Error>> {
        let gating = GatingMechanism::new(4)?;

        let messages = vec![array![1.0, 2.0, 3.0, 4.0], array![0.5, 1.5, 2.5, 3.5]];

        let gated = gating.apply(&messages)?;
        assert_eq!(gated.len(), 2);
        assert_eq!(gated[0].len(), 2); // GLU halves dimension: 4 → 2
        Ok(())
    }

    #[test]
    fn test_empty_messages() -> Result<(), Box<dyn std::error::Error>> {
        let gating = GatingMechanism::new(4)?;
        let empty_messages = vec![];

        let result = gating.apply(&empty_messages)?;
        assert!(result.is_empty());
        Ok(())
    }

    #[test]
    fn test_dimension_mismatch() -> Result<(), Box<dyn std::error::Error>> {
        let gating = GatingMechanism::new(4)?;

        let invalid_messages = vec![
            array![1.0, 2.0, 3.0], // Wrong dimension
        ];

        assert!(gating.apply(&invalid_messages).is_err());
        Ok(())
    }

    #[test]
    fn test_softmax() -> Result<(), Box<dyn std::error::Error>> {
        let gating = GatingMechanism::new(4)?;
        let input = array![1.0, 2.0, 3.0, 4.0];

        let softmax_output = gating.softmax(&input)?;
        let sum: f64 = softmax_output.sum();

        assert!((sum - 1.0).abs() < 1e-6);
        assert!(softmax_output.iter().all(|&x| x >= 0.0 && x <= 1.0));
        Ok(())
    }

    #[test]
    fn test_glu_activation() -> Result<(), Box<dyn std::error::Error>> {
        let gating = GatingMechanism::new(4)?;
        let input = array![1.0, 2.0, -1.0, 0.5]; // Even length for GLU

        let glu_output = gating.apply_glu(&input)?;
        assert_eq!(glu_output.len(), 2); // Half the input length
        Ok(())
    }

    #[test]
    fn test_multi_head_gating() -> Result<(), Box<dyn std::error::Error>> {
        let multi_head = MultiHeadGating::new(8, 2)?;

        let messages = vec![
            Array1::from(vec![1.0, 2.0, 3.0, 4.0, 0.5, 1.5, 2.5, 3.5]),
            Array1::from(vec![0.1, 0.2, 0.3, 0.4, 0.05, 0.15, 0.25, 0.35]),
        ];

        let result = multi_head.apply(&messages)?;
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].len(), 8);
        Ok(())
    }

    #[test]
    fn test_temperature_setting() -> Result<(), MLError> {
        let mut gating = GatingMechanism::new(4)?;

        gating.set_temperature(2.0);
        assert_eq!(gating.temperature(), 2.0);

        // Test minimum temperature enforcement
        gating.set_temperature(0.0);
        assert_eq!(gating.temperature(), 0.01);
        Ok(())
    }
}
