//! Message Passing Layer for TGGN
//!
//! Graph neural network message passing implementation

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

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

/// Cache for forward pass computations needed for backpropagation
#[derive(Debug, Clone)]
struct MessagePassingCache {
    pub node_features: Array1<f64>,
    pub neighbor_messages: Vec<Array1<f64>>,
    pub computed_messages: Vec<Array1<f64>>,
    pub aggregated_message: Array1<f64>,
    pub updated_features: Array1<f64>,
    pub output: Array1<f64>,
}

impl MessagePassingCache {
    pub(crate) fn new() -> Self {
        Self {
            node_features: Array1::zeros(0),
            neighbor_messages: Vec::new(),
            computed_messages: Vec::new(),
            aggregated_message: Array1::zeros(0),
            updated_features: Array1::zeros(0),
            output: Array1::zeros(0),
        }
    }
}

/// Gradients for message passing layer components
#[derive(Debug, Clone)]
struct MessagePassingGradients {
    pub message_weights_grad: Array2<f64>,
    pub message_bias_grad: Array1<f64>,
    pub update_weights_grad: Array2<f64>,
    pub update_bias_grad: Array1<f64>,
    pub layer_norm_gamma_grad: Array1<f64>,
    pub layer_norm_beta_grad: Array1<f64>,
}

impl MessagePassingGradients {
    pub(crate) fn new(input_dim: usize, output_dim: usize) -> Self {
        Self {
            message_weights_grad: Array2::zeros((output_dim, input_dim)),
            message_bias_grad: Array1::zeros(output_dim),
            update_weights_grad: Array2::zeros((output_dim, input_dim + output_dim)),
            update_bias_grad: Array1::zeros(output_dim),
            layer_norm_gamma_grad: Array1::zeros(output_dim),
            layer_norm_beta_grad: Array1::zeros(output_dim),
        }
    }
}

/// Message passing layer for graph neural networks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessagePassing {
    /// Input feature dimension
    input_dim: usize,

    /// Output feature dimension  
    output_dim: usize,

    /// Message function weights
    message_weights: Array2<f64>,

    /// Message function bias
    message_bias: Array1<f64>,

    /// Update function weights
    update_weights: Array2<f64>,

    /// Update function bias
    update_bias: Array1<f64>,

    /// Aggregation method
    aggregation: AggregationType,

    /// Layer normalization parameters
    layer_norm_gamma: Array1<f64>,
    layer_norm_beta: Array1<f64>,

    /// Dropout probability (for training)
    dropout_rate: f64,
}

/// Types of message aggregation
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum AggregationType {
    /// Sum aggregation
    Sum,
    /// Mean aggregation
    Mean,
    /// Max aggregation
    Max,
    /// Attention-weighted aggregation
    Attention,
}

impl MessagePassing {
    /// Create new message passing layer
    pub fn new(input_dim: usize, output_dim: usize) -> Result<Self, MLError> {
        // Xavier initialization scale
        let message_scale = (2.0 / (input_dim + output_dim) as f64).sqrt();
        let update_scale = (2.0 / (input_dim + output_dim) as f64).sqrt();

        let message_weights = Array2::from_shape_fn((output_dim, input_dim), |_| {
            (thread_rng().gen::<f64>() - 0.5) * message_scale
        });

        let message_bias = Array1::zeros(output_dim);

        let update_weights = Array2::from_shape_fn((output_dim, input_dim + output_dim), |_| {
            (thread_rng().gen::<f64>() - 0.5) * update_scale
        });

        let update_bias = Array1::zeros(output_dim);

        // Layer normalization parameters
        let layer_norm_gamma = Array1::ones(output_dim);
        let layer_norm_beta = Array1::zeros(output_dim);

        Ok(Self {
            input_dim,
            output_dim,
            message_weights,
            message_bias,
            update_weights,
            update_bias,
            aggregation: AggregationType::Mean,
            layer_norm_gamma,
            layer_norm_beta,
            dropout_rate: 0.1,
        })
    }

    /// Forward pass through message passing layer
    pub fn forward(
        &self,
        node_features: &Array1<f64>,
        neighbor_messages: &[Array1<f64>],
    ) -> Result<Array1<f64>, MLError> {
        // Validate input dimensions
        if node_features.len() != self.input_dim {
            return Err(MLError::DimensionMismatch {
                expected: self.input_dim,
                actual: node_features.len(),
            });
        }

        // If no neighbors, just apply update function to self
        if neighbor_messages.is_empty() {
            let self_message = self.compute_message(node_features)?;
            return self.apply_update(node_features, &self_message);
        }

        // Compute messages from all neighbors
        let mut messages = Vec::new();
        for neighbor_features in neighbor_messages {
            if neighbor_features.len() != self.input_dim {
                return Err(MLError::DimensionMismatch {
                    expected: self.input_dim,
                    actual: neighbor_features.len(),
                });
            }
            let message = self.compute_message(neighbor_features)?;
            messages.push(message);
        }

        // Add self-message
        let self_message = self.compute_message(node_features)?;
        messages.push(self_message);

        // Aggregate messages
        let aggregated = self.aggregate_messages(&messages)?;

        // Apply update function
        let updated = self.apply_update(node_features, &aggregated)?;

        // Apply layer normalization
        let normalized = self.layer_normalize(&updated)?;

        Ok(normalized)
    }

    /// Compute message from neighbor features
    #[allow(clippy::unnecessary_wraps)]
    fn compute_message(&self, features: &Array1<f64>) -> Result<Array1<f64>, MLError> {
        // Linear transformation: message_weights * features + message_bias
        let message = self.message_weights.dot(features) + &self.message_bias;

        // Apply ReLU activation
        let activated = message.mapv(|x| x.max(0.0));

        Ok(activated)
    }

    /// Aggregate messages from all neighbors
    fn aggregate_messages(&self, messages: &[Array1<f64>]) -> Result<Array1<f64>, MLError> {
        if messages.is_empty() {
            return Ok(Array1::zeros(self.output_dim));
        }

        match self.aggregation {
            AggregationType::Sum => {
                let mut sum = Array1::zeros(self.output_dim);
                for message in messages {
                    sum += message;
                }
                Ok(sum)
            },

            AggregationType::Mean => {
                let mut sum = Array1::zeros(self.output_dim);
                for message in messages {
                    sum += message;
                }
                Ok(sum / messages.len() as f64)
            },

            AggregationType::Max => {
                let mut max_message = messages[0].clone();
                for message in messages.iter().skip(1) {
                    for i in 0..max_message.len().min(message.len()) {
                        if message[i] > max_message[i] {
                            max_message[i] = message[i];
                        }
                    }
                }
                Ok(max_message)
            },

            AggregationType::Attention => self.attention_aggregate(messages),
        }
    }

    /// Attention-based message aggregation
    #[allow(clippy::unnecessary_wraps)]
    fn attention_aggregate(&self, messages: &[Array1<f64>]) -> Result<Array1<f64>, MLError> {
        if messages.is_empty() {
            return Ok(Array1::zeros(self.output_dim));
        }

        if messages.len() == 1 {
            return Ok(messages[0].clone());
        }

        // Compute attention scores (simplified)
        let mut attention_scores = Vec::new();
        for message in messages {
            // Simple attention: sum of absolute values
            let score = message.iter().map(|&x| x.abs()).sum::<f64>();
            attention_scores.push(score);
        }

        // Softmax normalization
        let max_score = attention_scores
            .iter()
            .fold(f64::NEG_INFINITY, |acc, &x| acc.max(x));
        let exp_scores: Vec<f64> = attention_scores
            .iter()
            .map(|&score| (score - max_score).exp())
            .collect();
        let sum_exp: f64 = exp_scores.iter().sum();

        let normalized_scores: Vec<f64> = if sum_exp > 0.0 {
            exp_scores.iter().map(|&score| score / sum_exp).collect()
        } else {
            vec![1.0 / messages.len() as f64; messages.len()]
        };

        // Weighted aggregation
        let mut weighted_sum = Array1::zeros(self.output_dim);
        for (message, weight) in messages.iter().zip(normalized_scores.iter()) {
            weighted_sum += &(message.mapv(|x| x * weight));
        }

        Ok(weighted_sum)
    }

    /// Apply update function to combine node features with aggregated messages
    fn apply_update(
        &self,
        node_features: &Array1<f64>,
        aggregated_message: &Array1<f64>,
    ) -> Result<Array1<f64>, MLError> {
        // Concatenate node features with aggregated message
        let mut combined = Vec::new();
        combined.extend_from_slice(node_features.as_slice().ok_or(MLError::ValidationError {
            message: "Failed to get slice from tensor".to_owned(),
        })?);
        combined.extend_from_slice(aggregated_message.as_slice().ok_or(
            MLError::ValidationError {
                message: "Failed to get slice from tensor".to_owned(),
            },
        )?);
        let combined_array = Array1::from(combined);

        // Apply update transformation
        let updated = self.update_weights.dot(&combined_array) + &self.update_bias;

        // Apply activation function (Swish/SiLU)
        let activated = updated.mapv(|x| x / (1.0 + (-x).exp()));

        // Apply dropout during training (simplified - always apply factor)
        let dropout_factor = 1.0 - self.dropout_rate;
        let mut rng = thread_rng();
        let with_dropout = activated.mapv(|x| {
            if rng.gen::<f64>() < dropout_factor {
                x / dropout_factor
            } else {
                0.0
            }
        });

        Ok(with_dropout)
    }

    /// Apply layer normalization
    #[allow(clippy::unnecessary_wraps)]
    fn layer_normalize(&self, input: &Array1<f64>) -> Result<Array1<f64>, MLError> {
        let mean = input.mean().unwrap_or(0.0);
        let variance = input.mapv(|x| (x - mean).powi(2)).mean().unwrap_or(1.0);
        let std_dev = (variance + 1e-8).sqrt(); // Add epsilon for numerical stability

        let normalized = input.mapv(|x| (x - mean) / std_dev);
        let output = &normalized * &self.layer_norm_gamma + &self.layer_norm_beta;

        Ok(output)
    }

    /// Set aggregation type
    pub const fn set_aggregation(&mut self, aggregation: AggregationType) {
        self.aggregation = aggregation;
    }

    /// Set dropout rate
    pub const fn set_dropout_rate(&mut self, rate: f64) {
        self.dropout_rate = rate.clamp(0.0, 1.0);
    }

    /// Train weights using proper backpropagation through message passing
    pub fn train_weights(
        &mut self,
        node_features_batch: &[Array1<f64>],
        neighbor_messages_batch: &[Vec<Array1<f64>>],
        targets: &[Array1<f64>],
        learning_rate: f64,
    ) -> Result<(), MLError> {
        if node_features_batch.len() != targets.len()
            || node_features_batch.len() != neighbor_messages_batch.len()
        {
            return Err(MLError::DimensionMismatch {
                expected: node_features_batch.len(),
                actual: targets.len(),
            });
        }

        let batch_size = node_features_batch.len();
        if batch_size == 0 {
            return Ok(());
        }

        // Forward pass for all samples in batch
        let mut predictions = Vec::new();
        let mut forward_cache = Vec::new();

        for (node_features, neighbor_messages) in node_features_batch
            .iter()
            .zip(neighbor_messages_batch.iter())
        {
            let (prediction, cache) = self.forward_with_cache(node_features, neighbor_messages)?;
            predictions.push(prediction);
            forward_cache.push(cache);
        }

        // Compute gradients using backpropagation
        let gradients = self.compute_message_passing_gradients(
            &predictions,
            targets,
            &forward_cache,
            node_features_batch,
            neighbor_messages_batch,
        )?;

        // Apply gradients with proper scaling
        self.apply_gradients(&gradients, learning_rate, batch_size)?;

        Ok(())
    }

    /// Forward pass with caching for gradient computation
    fn forward_with_cache(
        &self,
        node_features: &Array1<f64>,
        neighbor_messages: &[Array1<f64>],
    ) -> Result<(Array1<f64>, MessagePassingCache), MLError> {
        // Validate input dimensions
        if node_features.len() != self.input_dim {
            return Err(MLError::DimensionMismatch {
                expected: self.input_dim,
                actual: node_features.len(),
            });
        }

        let mut cache = MessagePassingCache::new();

        // Cache input
        cache.node_features = node_features.clone();
        cache.neighbor_messages = neighbor_messages.to_vec(); // cpu-side slice clone

        // Compute messages from neighbors
        let mut computed_messages = Vec::new();
        for neighbor_features in neighbor_messages {
            let message = self.compute_message(neighbor_features)?;
            computed_messages.push(message.clone());
            cache.computed_messages.push(message);
        }

        // Add self-message
        let self_message = self.compute_message(node_features)?;
        computed_messages.push(self_message.clone());
        cache.computed_messages.push(self_message);

        // Aggregate messages
        let aggregated = self.aggregate_messages(&computed_messages)?;
        cache.aggregated_message = aggregated.clone();

        // Apply update function
        let updated = self.apply_update(node_features, &aggregated)?;
        cache.updated_features = updated.clone();

        // Apply layer normalization
        let normalized = self.layer_normalize(&updated)?;
        cache.output = normalized.clone();

        Ok((normalized, cache))
    }

    /// Compute gradients for message passing using backpropagation
    fn compute_message_passing_gradients(
        &self,
        predictions: &[Array1<f64>],
        targets: &[Array1<f64>],
        forward_caches: &[MessagePassingCache],
        node_features_batch: &[Array1<f64>],
        neighbor_messages_batch: &[Vec<Array1<f64>>],
    ) -> Result<MessagePassingGradients, MLError> {
        let mut gradients = MessagePassingGradients::new(self.input_dim, self.output_dim);

        for (sample_idx, ((prediction, target), cache)) in predictions
            .iter()
            .zip(targets.iter())
            .zip(forward_caches.iter())
            .enumerate()
        {
            // Compute output error (gradient of loss w.r.t. output)
            let output_error = prediction - target;

            // Backpropagate through layer normalization
            let ln_input_grad =
                self.backprop_layer_norm(&output_error, &cache.updated_features, &mut gradients)?;

            // Backpropagate through update function
            let (_node_grad, aggregated_grad) = self.backprop_update(
                &ln_input_grad,
                &cache.node_features,
                &cache.aggregated_message,
                &mut gradients,
            )?;

            // Backpropagate through message aggregation
            let message_grads =
                self.backprop_aggregation(&aggregated_grad, &cache.computed_messages)?;

            // Backpropagate through message computation
            for (message_grad, neighbor_features) in message_grads
                .iter()
                .zip(
                    neighbor_messages_batch[sample_idx]
                        .iter()
                        .chain(std::iter::once(&node_features_batch[sample_idx])),
                )
            {
                self.backprop_message_computation(message_grad, neighbor_features, &mut gradients)?;
            }
        }

        Ok(gradients)
    }

    /// Backpropagate through layer normalization
    #[allow(clippy::unnecessary_wraps)]
    fn backprop_layer_norm(
        &self,
        output_grad: &Array1<f64>,
        input: &Array1<f64>,
        gradients: &mut MessagePassingGradients,
    ) -> Result<Array1<f64>, MLError> {
        let mean = input.mean().unwrap_or(0.0);
        let variance = input.mapv(|x| (x - mean).powi(2)).mean().unwrap_or(1.0);
        let std_dev = (variance + 1e-8).sqrt();

        let normalized = input.mapv(|x| (x - mean) / std_dev);

        // Gradient w.r.t. gamma (scale parameter)
        for (i, (out_grad, norm_val)) in output_grad.iter().zip(normalized.iter()).enumerate() {
            gradients.layer_norm_gamma_grad[i] += out_grad * norm_val;
        }

        // Gradient w.r.t. beta (shift parameter)
        for (i, &out_grad) in output_grad.into_iter().enumerate() {
            gradients.layer_norm_beta_grad[i] += out_grad;
        }

        // Gradient w.r.t. input (chain rule through normalization)
        let n = input.len() as f64;
        let mut input_grad = Array1::zeros(input.len());

        for i in 0..input.len() {
            let x_centered = input[i] - mean;
            let gamma_out_grad = self.layer_norm_gamma[i] * output_grad[i];

            // Gradient through normalization computation
            let grad_normalized = gamma_out_grad;
            let grad_variance = -0.5 * x_centered * grad_normalized / (variance + 1e-8).powf(1.5);
            let grad_mean = -grad_normalized / std_dev - 2.0 * x_centered * grad_variance / n;

            input_grad[i] =
                grad_normalized / std_dev + grad_variance * 2.0 * x_centered / n + grad_mean / n;
        }

        Ok(input_grad)
    }

    /// Backpropagate through update function
    #[allow(clippy::unnecessary_wraps)]
    fn backprop_update(
        &self,
        output_grad: &Array1<f64>,
        node_features: &Array1<f64>,
        aggregated_message: &Array1<f64>,
        gradients: &mut MessagePassingGradients,
    ) -> Result<(Array1<f64>, Array1<f64>), MLError> {
        // The update function concatenates node features with aggregated message
        let combined_dim = node_features.len() + aggregated_message.len();

        // Gradient w.r.t. update weights: output_grad^T * combined_input
        for i in 0..self.output_dim {
            for j in 0..combined_dim {
                let combined_input_j = if j < node_features.len() {
                    node_features[j]
                } else {
                    aggregated_message[j - node_features.len()]
                };
                gradients.update_weights_grad[[i, j]] += output_grad[i] * combined_input_j;
            }
        }

        // Gradient w.r.t. update bias
        for i in 0..self.output_dim {
            gradients.update_bias_grad[i] += output_grad[i];
        }

        // Gradient w.r.t. combined input: weights^T * output_grad
        let mut combined_input_grad = Array1::zeros(combined_dim);
        for j in 0..combined_dim {
            for i in 0..self.output_dim {
                combined_input_grad[j] += self.update_weights[[i, j]] * output_grad[i];
            }
        }

        // Apply Swish/SiLU derivative (d/dx[x * sigmoid(x)] = sigmoid(x) + x * sigmoid(x) * (1 - sigmoid(x)))
        for j in 0..combined_dim {
            let combined_input_j = if j < node_features.len() {
                node_features[j]
            } else {
                aggregated_message[j - node_features.len()]
            };

            let sigmoid_val = 1.0 / (1.0 + (-combined_input_j).exp());
            let swish_derivative =
                sigmoid_val + combined_input_j * sigmoid_val * (1.0 - sigmoid_val);
            combined_input_grad[j] *= swish_derivative;
        }

        // Split gradient back to node features and aggregated message
        let node_grad = combined_input_grad
            .slice(s![..node_features.len()])
            .to_owned();
        let aggregated_grad = combined_input_grad
            .slice(s![node_features.len()..])
            .to_owned();

        Ok((node_grad, aggregated_grad))
    }

    /// Backpropagate through message aggregation
    fn backprop_aggregation(
        &self,
        aggregated_grad: &Array1<f64>,
        messages: &[Array1<f64>],
    ) -> Result<Vec<Array1<f64>>, MLError> {
        match self.aggregation {
            AggregationType::Sum => {
                // For sum aggregation, gradient is just passed through to all messages
                Ok(vec![aggregated_grad.clone(); messages.len()])
            },

            AggregationType::Mean => {
                // For mean aggregation, gradient is divided by number of messages
                let mean_grad = aggregated_grad / messages.len() as f64;
                Ok(vec![mean_grad; messages.len()])
            },

            AggregationType::Max => {
                // For max aggregation, gradient goes only to the message that was maximum
                let mut message_grads = vec![Array1::zeros(self.output_dim); messages.len()];

                for i in 0..self.output_dim {
                    let mut max_val = f64::NEG_INFINITY;
                    let mut max_idx = 0;

                    for (j, message) in messages.iter().enumerate() {
                        if message[i] > max_val {
                            max_val = message[i];
                            max_idx = j;
                        }
                    }

                    message_grads[max_idx][i] = aggregated_grad[i];
                }

                Ok(message_grads)
            },

            AggregationType::Attention => {
                // For attention aggregation, need to backprop through attention weights
                self.backprop_attention_aggregation(aggregated_grad, messages)
            },
        }
    }

    /// Backpropagate through attention-based aggregation
    #[allow(clippy::unnecessary_wraps)]
    fn backprop_attention_aggregation(
        &self,
        aggregated_grad: &Array1<f64>,
        messages: &[Array1<f64>],
    ) -> Result<Vec<Array1<f64>>, MLError> {
        // Recompute attention scores for backward pass
        let mut attention_scores = Vec::new();
        for message in messages {
            let score = message.iter().map(|&x| x.abs()).sum::<f64>();
            attention_scores.push(score);
        }

        // Softmax normalization (recompute for gradient)
        let max_score = attention_scores
            .iter()
            .fold(f64::NEG_INFINITY, |acc, &x| acc.max(x));
        let exp_scores: Vec<f64> = attention_scores
            .iter()
            .map(|&score| (score - max_score).exp())
            .collect();
        let sum_exp: f64 = exp_scores.iter().sum();

        let normalized_scores: Vec<f64> = if sum_exp > 0.0 {
            exp_scores.iter().map(|&score| score / sum_exp).collect()
        } else {
            vec![1.0 / messages.len() as f64; messages.len()]
        };

        // Gradient w.r.t. messages through attention weights
        let mut message_grads = Vec::new();
        for weight in normalized_scores.iter() {
            let message_grad = aggregated_grad.mapv(|x| x * weight);
            message_grads.push(message_grad);
        }

        Ok(message_grads)
    }

    /// Backpropagate through message computation
    #[allow(clippy::unnecessary_wraps)]
    fn backprop_message_computation(
        &self,
        message_grad: &Array1<f64>,
        input_features: &Array1<f64>,
        gradients: &mut MessagePassingGradients,
    ) -> Result<(), MLError> {
        // Apply ReLU derivative (1 if input > 0, 0 otherwise)
        let mut adjusted_grad = message_grad.clone();
        for i in 0..adjusted_grad.len() {
            let linear_output =
                self.message_weights.row(i).dot(input_features) + self.message_bias[i];
            if linear_output <= 0.0 {
                adjusted_grad[i] = 0.0; // ReLU derivative
            }
        }

        // Gradient w.r.t. message weights: adjusted_grad^T * input_features
        for i in 0..self.output_dim {
            for j in 0..self.input_dim {
                gradients.message_weights_grad[[i, j]] += adjusted_grad[i] * input_features[j];
            }
        }

        // Gradient w.r.t. message bias
        for i in 0..self.output_dim {
            gradients.message_bias_grad[i] += adjusted_grad[i];
        }

        Ok(())
    }

    /// Apply computed gradients to weights
    #[allow(clippy::unnecessary_wraps)]
    fn apply_gradients(
        &mut self,
        gradients: &MessagePassingGradients,
        learning_rate: f64,
        batch_size: usize,
    ) -> Result<(), MLError> {
        let scale = learning_rate / batch_size as f64;

        // Update message weights
        for ((i, j), &grad) in gradients.message_weights_grad.indexed_iter() {
            self.message_weights[[i, j]] -= scale * grad;
        }

        // Update message bias
        for (i, &grad) in gradients.message_bias_grad.indexed_iter() {
            self.message_bias[i] -= scale * grad;
        }

        // Update update weights
        for ((i, j), &grad) in gradients.update_weights_grad.indexed_iter() {
            self.update_weights[[i, j]] -= scale * grad;
        }

        // Update update bias
        for (i, &grad) in gradients.update_bias_grad.indexed_iter() {
            self.update_bias[i] -= scale * grad;
        }

        // Update layer norm parameters
        for (i, &grad) in gradients.layer_norm_gamma_grad.indexed_iter() {
            self.layer_norm_gamma[i] -= scale * grad;
        }

        for (i, &grad) in gradients.layer_norm_beta_grad.indexed_iter() {
            self.layer_norm_beta[i] -= scale * 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 parameter norm
        total_norm_sq += self.message_weights.mapv(|x| x * x).sum();
        total_norm_sq += self.message_bias.mapv(|x| x * x).sum();
        total_norm_sq += self.update_weights.mapv(|x| x * x).sum();
        total_norm_sq += self.update_bias.mapv(|x| x * x).sum();
        total_norm_sq += self.layer_norm_gamma.mapv(|x| x * x).sum();
        total_norm_sq += self.layer_norm_beta.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.message_weights.mapv_inplace(|x| x * clip_factor);
            self.message_bias.mapv_inplace(|x| x * clip_factor);
            self.update_weights.mapv_inplace(|x| x * clip_factor);
            self.update_bias.mapv_inplace(|x| x * clip_factor);
            self.layer_norm_gamma.mapv_inplace(|x| x * clip_factor);
            self.layer_norm_beta.mapv_inplace(|x| x * clip_factor);
        }
    }

    /// Get layer parameters for serialization
    pub fn get_parameters(&self) -> MessagePassingParameters {
        MessagePassingParameters {
            input_dim: self.input_dim,
            output_dim: self.output_dim,
            message_weights: self.message_weights.clone(),
            message_bias: self.message_bias.clone(),
            update_weights: self.update_weights.clone(),
            update_bias: self.update_bias.clone(),
            layer_norm_gamma: self.layer_norm_gamma.clone(),
            layer_norm_beta: self.layer_norm_beta.clone(),
            aggregation: self.aggregation,
            dropout_rate: self.dropout_rate,
        }
    }

    /// Load parameters from serialized form
    pub fn load_parameters(&mut self, params: MessagePassingParameters) -> Result<(), MLError> {
        if params.input_dim != self.input_dim || params.output_dim != self.output_dim {
            return Err(MLError::DimensionMismatch {
                expected: self.input_dim * self.output_dim,
                actual: params.input_dim * params.output_dim,
            });
        }

        self.message_weights = params.message_weights;
        self.message_bias = params.message_bias;
        self.update_weights = params.update_weights;
        self.update_bias = params.update_bias;
        self.layer_norm_gamma = params.layer_norm_gamma;
        self.layer_norm_beta = params.layer_norm_beta;
        self.aggregation = params.aggregation;
        self.dropout_rate = params.dropout_rate;

        Ok(())
    }
}

/// Serializable parameters for `MessagePassing` layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessagePassingParameters {
    pub input_dim: usize,
    pub output_dim: usize,
    pub message_weights: Array2<f64>,
    pub message_bias: Array1<f64>,
    pub update_weights: Array2<f64>,
    pub update_bias: Array1<f64>,
    pub layer_norm_gamma: Array1<f64>,
    pub layer_norm_beta: Array1<f64>,
    pub aggregation: AggregationType,
    pub dropout_rate: f64,
}

/// Graph Attention Network (GAT) variant of message passing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GATMessagePassing {
    /// Base message passing layer
    base_layer: MessagePassing,

    /// Attention mechanism weights
    attention_weights: Array2<f64>,

    /// Number of attention heads
    num_heads: usize,

    /// Attention dropout rate
    attention_dropout: f64,
}

impl GATMessagePassing {
    /// Create new GAT message passing layer
    pub fn new(input_dim: usize, output_dim: usize, num_heads: usize) -> Result<Self, MLError> {
        let base_layer = MessagePassing::new(input_dim, output_dim)?;

        // Attention weights for computing attention coefficients
        let attention_scale = (2.0 / (2 * output_dim) as f64).sqrt();
        let attention_weights = Array2::from_shape_fn((num_heads, 2 * output_dim), |_| {
            (thread_rng().gen::<f64>() - 0.5) * attention_scale
        });

        Ok(Self {
            base_layer,
            attention_weights,
            num_heads,
            attention_dropout: 0.1,
        })
    }

    /// Forward pass with graph attention
    pub fn forward(
        &self,
        node_features: &Array1<f64>,
        neighbor_features: &[Array1<f64>],
    ) -> Result<Array1<f64>, MLError> {
        if neighbor_features.is_empty() {
            return self.base_layer.forward(node_features, neighbor_features);
        }

        // Apply multi-head attention
        let mut head_outputs = Vec::new();

        for head in 0..self.num_heads {
            let head_output = self.apply_attention_head(node_features, neighbor_features, head)?;
            head_outputs.push(head_output);
        }

        // Average attention heads (could also concatenate)
        let mut averaged = Array1::zeros(self.base_layer.output_dim);
        for head_output in &head_outputs {
            averaged += head_output;
        }
        averaged /= self.num_heads as f64;

        Ok(averaged)
    }

    /// Apply single attention head
    fn apply_attention_head(
        &self,
        node_features: &Array1<f64>,
        neighbor_features: &[Array1<f64>],
        head_idx: usize,
    ) -> Result<Array1<f64>, MLError> {
        // Transform node and neighbor features
        let node_transformed = self.base_layer.compute_message(node_features)?;

        let mut neighbor_transformed = Vec::new();
        for neighbor in neighbor_features {
            let transformed = self.base_layer.compute_message(neighbor)?;
            neighbor_transformed.push(transformed);
        }

        // Compute attention coefficients
        let mut attention_coeffs = Vec::new();
        let attention_head_weights = self.attention_weights.row(head_idx);

        for neighbor_trans in &neighbor_transformed {
            // Concatenate node and neighbor features for attention computation
            let mut combined = Vec::new();
            combined.extend_from_slice(node_transformed.as_slice().ok_or(
                MLError::ValidationError {
                    message: "Failed to get slice from tensor".to_owned(),
                },
            )?);
            combined.extend_from_slice(neighbor_trans.as_slice().ok_or(
                MLError::ValidationError {
                    message: "Failed to get slice from tensor".to_owned(),
                },
            )?);
            let combined_array = Array1::from(combined);

            // Compute attention coefficient
            let attention_logit = attention_head_weights.dot(&combined_array);
            let attention_coeff = attention_logit.tanh(); // Use tanh as activation
            attention_coeffs.push(attention_coeff);
        }

        // Apply softmax to attention coefficients
        let max_coeff = attention_coeffs
            .iter()
            .fold(f64::NEG_INFINITY, |acc, &x| acc.max(x));
        let exp_coeffs: Vec<f64> = attention_coeffs
            .iter()
            .map(|&coeff| (coeff - max_coeff).exp())
            .collect();
        let sum_exp: f64 = exp_coeffs.iter().sum();

        let normalized_coeffs: Vec<f64> = if sum_exp > 0.0 {
            exp_coeffs.iter().map(|&coeff| coeff / sum_exp).collect()
        } else {
            vec![1.0 / neighbor_transformed.len() as f64; neighbor_transformed.len()]
        };

        // Apply attention dropout
        let dropout_coeffs: Vec<f64> = normalized_coeffs
            .iter()
            .map(|&coeff| {
                if thread_rng().gen::<f64>() < (1.0 - self.attention_dropout) {
                    coeff / (1.0 - self.attention_dropout)
                } else {
                    0.0
                }
            })
            .collect();

        // Weighted aggregation
        let mut output = Array1::zeros(self.base_layer.output_dim);
        for (neighbor_trans, weight) in neighbor_transformed.iter().zip(dropout_coeffs.iter()) {
            output += &neighbor_trans.mapv(|x| x * weight);
        }

        Ok(output)
    }
}

// DISABLED: Tests
// // #[cfg(test)]
// mod tests {
//     use super::*;
//     use ndarray::array;
//
//     #[test]
//     fn test_message_passing_layer() -> Result<(), Box<dyn std::error::Error>> {
//         let layer = MessagePassing::new(4, 8)?;
//
//         let node_features = array![1.0, 2.0, 3.0, 4.0];
//         let messages = vec![array![0.1, 0.2, 0.3, 0.4], array![0.5, 0.6, 0.7, 0.8]];
//
//         let output = layer.forward(&node_features, &messages)?;
//         assert_eq!(output.len(), 8);
//         Ok(())
//     }
//
//     #[test]
//     fn test_empty_messages() -> Result<(), Box<dyn std::error::Error>> {
//         let layer = MessagePassing::new(3, 5)?;
//         let node_features = array![1.0, 2.0, 3.0];
//         let empty_messages = vec![];
//
//         let output = layer.forward(&node_features, &empty_messages)?;
//         assert_eq!(output.len(), 5);
//         Ok(())
//     }
//
//     #[test]
//     fn test_aggregation_types() {
//         let mut layer = MessagePassing::new(2, 2)?;
//
//         let messages = vec![array![1.0, 2.0], array![3.0, 1.0], array![2.0, 4.0]];
//
//         // Test different aggregation methods
//         layer.set_aggregation(AggregationType::Sum);
//         let sum_result = layer.aggregate_messages(&messages)?;
//
//         layer.set_aggregation(AggregationType::Mean);
//         let mean_result = layer.aggregate_messages(&messages)?;
//
//         layer.set_aggregation(AggregationType::Max);
//         let max_result = layer.aggregate_messages(&messages)?;
//
//         // Verify basic properties
//         assert_eq!(sum_result.len(), 2);
//         assert_eq!(mean_result.len(), 2);
//         assert_eq!(max_result.len(), 2);
//
//         // Mean should be sum divided by count
//         assert!((mean_result[0] - sum_result[0] / 3.0).abs() < 1e-6);
//         assert!((mean_result[1] - sum_result[1] / 3.0).abs() < 1e-6);
//     }
//
//     #[test]
//     fn test_layer_normalization() {
//         let layer = MessagePassing::new(3, 3)?;
//         let input = array![1.0, 4.0, 7.0];
//
//         let normalized = layer.layer_normalize(&input)?;
//
//         // Check that output has approximately zero mean and unit variance
//         let mean = normalized.mean()?;
//         let variance = normalized.mapv(|x| (x - mean).powi(2)).mean()?;
//
//         assert!(mean.abs() < 1e-6);
//         assert!((variance - 1.0).abs() < 1e-6);
//     }
//
//     #[test]
//     fn test_gat_message_passing() {
//         let gat_layer = GATMessagePassing::new(4, 6, 2)?;
//
//         let node_features = array![1.0, 2.0, 3.0, 4.0];
//         let neighbor_features = vec![array![0.5, 1.0, 1.5, 2.0], array![2.0, 1.5, 1.0, 0.5]];
//
//         let output = gat_layer.forward(&node_features, &neighbor_features)?;
//         assert_eq!(output.len(), 6);
//     }
//
//     #[test]
//     fn test_attention_aggregation() {
//         let layer = MessagePassing::new(3, 3)?;
//
//         let messages = vec![
//             array![1.0, 0.0, 0.0], // High attention (large magnitude)
//             array![0.1, 0.1, 0.1], // Low attention (small magnitude)
//         ];
//
//         let aggregated = layer.attention_aggregate(&messages)?;
//         assert_eq!(aggregated.len(), 3);
//
//         // First message should have higher weight due to larger magnitude
//         assert!(aggregated[0] > aggregated[1]);
//     }
//
//     #[test]
//     fn test_dropout_setting() {
//         let mut layer = MessagePassing::new(4, 4)?;
//
//         layer.set_dropout_rate(0.5);
//         assert_eq!(layer.dropout_rate, 0.5);
//
//         // Test clamping
//         layer.set_dropout_rate(-0.1);
//         assert_eq!(layer.dropout_rate, 0.0);
//
//         layer.set_dropout_rate(1.5);
//         assert_eq!(layer.dropout_rate, 1.0);
//     }
// }
