//! Brain-specific GNN layers.
//!
//! These layers extend rtx-geom's GNN layers with brain-specific functionality.

use crate::error::{GnnError, GnnResult};
use crate::graph::BrainGraph;
use ndarray::{Array1, Array2, Axis};
use serde::{Deserialize, Serialize};

/// Configuration for BrainConv layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrainConvConfig {
    /// Input feature dimension
    pub in_features: usize,
    /// Output feature dimension
    pub out_features: usize,
    /// Whether to add self-loops
    pub add_self_loops: bool,
    /// Whether to normalize adjacency
    pub normalize: bool,
    /// Activation function
    pub activation: Activation,
    /// Dropout probability
    pub dropout: f64,
}

impl Default for BrainConvConfig {
    fn default() -> Self {
        Self {
            in_features: 64,
            out_features: 64,
            add_self_loops: true,
            normalize: true,
            activation: Activation::ReLU,
            dropout: 0.0,
        }
    }
}

/// Activation functions
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum Activation {
    /// No activation
    None,
    /// ReLU
    ReLU,
    /// Leaky ReLU
    LeakyReLU,
    /// ELU
    ELU,
    /// Tanh
    Tanh,
    /// Sigmoid
    Sigmoid,
    /// GELU
    GELU,
}

impl Activation {
    /// Apply activation function
    pub fn apply(&self, x: f64) -> f64 {
        match self {
            Activation::None => x,
            Activation::ReLU => x.max(0.0),
            Activation::LeakyReLU => {
                if x > 0.0 {
                    x
                } else {
                    0.01 * x
                }
            }
            Activation::ELU => {
                if x > 0.0 {
                    x
                } else {
                    x.exp() - 1.0
                }
            }
            Activation::Tanh => x.tanh(),
            Activation::Sigmoid => 1.0 / (1.0 + (-x).exp()),
            Activation::GELU => {
                // Approximate GELU
                0.5 * x * (1.0 + (0.7978845608 * (x + 0.044715 * x * x * x)).tanh())
            }
        }
    }

    /// Apply to array
    pub fn apply_array(&self, x: &Array2<f64>) -> Array2<f64> {
        x.mapv(|v| self.apply(v))
    }
}

/// Brain-specific graph convolution layer
///
/// Extends standard GCN with brain topology awareness.
#[derive(Debug, Clone)]
pub struct BrainConv {
    config: BrainConvConfig,
    /// Weight matrix [in_features, out_features]
    weights: Array2<f64>,
    /// Bias vector [out_features]
    bias: Array1<f64>,
}

impl BrainConv {
    /// Create a new BrainConv layer
    pub fn new(config: BrainConvConfig) -> GnnResult<Self> {
        // Xavier initialization
        let scale = (2.0 / (config.in_features + config.out_features) as f64).sqrt();

        let weights = Array2::from_shape_fn((config.in_features, config.out_features), |_| {
            (rand_simple() * 2.0 - 1.0) * scale
        });

        let bias = Array1::zeros(config.out_features);

        Ok(Self {
            config,
            weights,
            bias,
        })
    }

    /// Forward pass
    pub fn forward(&self, x: &Array2<f64>, adj: &Array2<f64>) -> GnnResult<Array2<f64>> {
        let n = x.nrows();
        if adj.nrows() != n || adj.ncols() != n {
            return Err(GnnError::DimensionMismatch(format!(
                "Adjacency {}x{} doesn't match features {}",
                adj.nrows(),
                adj.ncols(),
                n
            )));
        }

        // Normalize adjacency
        let adj_norm = if self.config.normalize {
            normalize_adjacency(adj, self.config.add_self_loops)
        } else if self.config.add_self_loops {
            adj + &Array2::<f64>::eye(n)
        } else {
            adj.clone()
        };

        // Message passing: A * X * W + b
        let h = adj_norm.dot(x).dot(&self.weights);
        let mut output = h + &self.bias;

        // Activation
        output = self.config.activation.apply_array(&output);

        Ok(output)
    }

    /// Forward with brain graph
    pub fn forward_graph(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
        let x = graph.node_feature_matrix()?;
        let adj = graph.adjacency_matrix();
        self.forward(&x, &adj)
    }
}

/// Configuration for BrainAttention layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrainAttentionConfig {
    /// Input feature dimension
    pub in_features: usize,
    /// Output feature dimension
    pub out_features: usize,
    /// Number of attention heads
    pub n_heads: usize,
    /// Dropout probability
    pub dropout: f64,
    /// Whether to concatenate heads (true) or average (false)
    pub concat_heads: bool,
    /// Negative slope for LeakyReLU
    pub negative_slope: f64,
}

impl Default for BrainAttentionConfig {
    fn default() -> Self {
        Self {
            in_features: 64,
            out_features: 64,
            n_heads: 4,
            dropout: 0.0,
            concat_heads: true,
            negative_slope: 0.2,
        }
    }
}

/// Brain-specific graph attention layer
#[derive(Debug, Clone)]
pub struct BrainAttention {
    config: BrainAttentionConfig,
    /// Weight matrices per head [n_heads, in_features, out_features/n_heads]
    weights: Vec<Array2<f64>>,
    /// Attention weights [n_heads, 2 * out_features/n_heads]
    attention: Vec<Array1<f64>>,
}

impl BrainAttention {
    /// Create a new BrainAttention layer
    pub fn new(config: BrainAttentionConfig) -> GnnResult<Self> {
        let head_dim = if config.concat_heads {
            config.out_features / config.n_heads
        } else {
            config.out_features
        };

        let scale = (2.0 / (config.in_features + head_dim) as f64).sqrt();

        let mut weights = Vec::with_capacity(config.n_heads);
        let mut attention = Vec::with_capacity(config.n_heads);

        for _ in 0..config.n_heads {
            weights.push(Array2::from_shape_fn(
                (config.in_features, head_dim),
                |_| (rand_simple() * 2.0 - 1.0) * scale,
            ));
            attention.push(Array1::from_shape_fn(2 * head_dim, |_| {
                (rand_simple() * 2.0 - 1.0) * scale
            }));
        }

        Ok(Self {
            config,
            weights,
            attention,
        })
    }

    /// Forward pass
    pub fn forward(&self, x: &Array2<f64>, adj: &Array2<f64>) -> GnnResult<Array2<f64>> {
        let n = x.nrows();
        let head_dim = self.weights[0].ncols();

        // Process each attention head
        let mut head_outputs = Vec::with_capacity(self.config.n_heads);

        for head in 0..self.config.n_heads {
            // Transform features
            let h = x.dot(&self.weights[head]); // [n, head_dim]

            // Compute attention coefficients
            let mut attn_matrix = Array2::zeros((n, n));

            for i in 0..n {
                for j in 0..n {
                    if adj[[i, j]] > 0.0 || i == j {
                        // Concatenate hi || hj
                        let mut concat: Vec<f64> = Vec::with_capacity(2 * head_dim);
                        concat.extend(h.row(i).iter().copied());
                        concat.extend(h.row(j).iter().copied());

                        // Attention score
                        let score: f64 = concat
                            .iter()
                            .zip(self.attention[head].iter())
                            .map(|(&c, &a)| c * a)
                            .sum();

                        // LeakyReLU
                        attn_matrix[[i, j]] = if score > 0.0 {
                            score
                        } else {
                            self.config.negative_slope * score
                        };
                    }
                }
            }

            // Softmax over neighbors
            for i in 0..n {
                let row = attn_matrix.row(i);
                let max_val = row.iter().copied().fold(f64::NEG_INFINITY, f64::max);
                let exp_sum: f64 = row.iter().map(|&v| (v - max_val).exp()).sum();

                for j in 0..n {
                    if adj[[i, j]] > 0.0 || i == j {
                        attn_matrix[[i, j]] = (attn_matrix[[i, j]] - max_val).exp() / exp_sum;
                    } else {
                        attn_matrix[[i, j]] = 0.0;
                    }
                }
            }

            // Aggregate
            let head_out = attn_matrix.dot(&h);
            head_outputs.push(head_out);
        }

        // Combine heads
        if self.config.concat_heads {
            // Concatenate along feature dimension
            let mut output = Array2::zeros((n, self.config.out_features));
            for (i, head_out) in head_outputs.iter().enumerate() {
                for j in 0..n {
                    for k in 0..head_dim {
                        output[[j, i * head_dim + k]] = head_out[[j, k]];
                    }
                }
            }
            Ok(output)
        } else {
            // Average heads
            let mut output = Array2::zeros((n, self.config.out_features));
            for head_out in &head_outputs {
                output += head_out;
            }
            output /= self.config.n_heads as f64;
            Ok(output)
        }
    }
}

/// Configuration for BrainPool layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrainPoolConfig {
    /// Pooling method
    pub method: PoolMethod,
    /// Pooling ratio (for TopK, SAGPool)
    pub ratio: f64,
}

impl Default for BrainPoolConfig {
    fn default() -> Self {
        Self {
            method: PoolMethod::Mean,
            ratio: 0.5,
        }
    }
}

/// Pooling methods
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum PoolMethod {
    /// Mean pooling (global)
    Mean,
    /// Max pooling (global)
    Max,
    /// Sum pooling (global)
    Sum,
    /// Attention-weighted pooling
    Attention,
    /// Region-based pooling (pools by brain region)
    Region,
}

/// Brain-specific graph pooling layer
#[derive(Debug, Clone)]
pub struct BrainPool {
    config: BrainPoolConfig,
    /// Attention weights for attention pooling
    attention_weights: Option<Array1<f64>>,
}

impl BrainPool {
    /// Create a new BrainPool layer
    pub fn new(config: BrainPoolConfig) -> Self {
        Self {
            config,
            attention_weights: None,
        }
    }

    /// Initialize attention weights
    pub fn init_attention(&mut self, n_features: usize) {
        self.attention_weights = Some(Array1::from_shape_fn(n_features, |_| rand_simple() * 0.1));
    }

    /// Global pooling (for graph-level predictions)
    pub fn global_pool(&self, x: &Array2<f64>) -> Array1<f64> {
        match self.config.method {
            PoolMethod::Mean => x.mean_axis(Axis(0)).unwrap(),
            PoolMethod::Max => {
                let mut result = Array1::zeros(x.ncols());
                for j in 0..x.ncols() {
                    result[j] = x
                        .column(j)
                        .iter()
                        .copied()
                        .fold(f64::NEG_INFINITY, f64::max);
                }
                result
            }
            PoolMethod::Sum => x.sum_axis(Axis(0)),
            PoolMethod::Attention => {
                if let Some(ref attn) = self.attention_weights {
                    // Compute attention scores
                    let scores: Vec<f64> = x
                        .rows()
                        .into_iter()
                        .map(|row| {
                            row.iter()
                                .zip(attn.iter())
                                .map(|(&r, &a)| r * a)
                                .sum::<f64>()
                        })
                        .collect();

                    // Softmax
                    let max_score = scores.iter().copied().fold(f64::NEG_INFINITY, f64::max);
                    let exp_scores: Vec<f64> =
                        scores.iter().map(|&s| (s - max_score).exp()).collect();
                    let sum_exp: f64 = exp_scores.iter().sum();
                    let attn_probs: Vec<f64> = exp_scores.iter().map(|&e| e / sum_exp).collect();

                    // Weighted sum
                    let mut result = Array1::zeros(x.ncols());
                    for (i, &prob) in attn_probs.iter().enumerate() {
                        for j in 0..x.ncols() {
                            result[j] += prob * x[[i, j]];
                        }
                    }
                    result
                } else {
                    x.mean_axis(Axis(0)).unwrap()
                }
            }
            PoolMethod::Region => {
                // Falls back to mean for now
                x.mean_axis(Axis(0)).unwrap()
            }
        }
    }

    /// Region-based pooling using brain graph
    pub fn pool_by_region(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
        use crate::graph::BrainRegion;

        let x = graph.node_feature_matrix()?;
        let n_features = x.ncols();

        // Pool by region
        let regions = [
            BrainRegion::Frontal,
            BrainRegion::Central,
            BrainRegion::Temporal,
            BrainRegion::Parietal,
            BrainRegion::Occipital,
        ];

        let mut pooled = Array2::zeros((regions.len(), n_features));

        for (r_idx, &region) in regions.iter().enumerate() {
            let indices: Vec<usize> = graph
                .nodes
                .iter()
                .filter(|n| n.region == region)
                .map(|n| n.index)
                .collect();

            if indices.is_empty() {
                continue;
            }

            // Mean pooling over region nodes
            for &i in &indices {
                for j in 0..n_features {
                    pooled[[r_idx, j]] += x[[i, j]];
                }
            }
            for j in 0..n_features {
                pooled[[r_idx, j]] /= indices.len() as f64;
            }
        }

        Ok(pooled)
    }
}

/// Edge convolution for learning edge representations
#[derive(Debug, Clone)]
pub struct EdgeConv {
    /// Input edge feature dimension
    in_features: usize,
    /// Output edge feature dimension
    out_features: usize,
    /// Weight matrix
    weights: Array2<f64>,
    /// Bias
    bias: Array1<f64>,
    /// Activation
    activation: Activation,
}

impl EdgeConv {
    /// Create a new EdgeConv layer
    pub fn new(in_features: usize, out_features: usize) -> Self {
        let scale = (2.0 / (in_features + out_features) as f64).sqrt();

        Self {
            in_features,
            out_features,
            weights: Array2::from_shape_fn((in_features, out_features), |_| {
                (rand_simple() * 2.0 - 1.0) * scale
            }),
            bias: Array1::zeros(out_features),
            activation: Activation::ReLU,
        }
    }

    /// Forward pass on edge features
    pub fn forward(&self, edge_features: &Array2<f64>) -> Array2<f64> {
        let h = edge_features.dot(&self.weights) + &self.bias;
        self.activation.apply_array(&h)
    }
}

// Simple pseudo-random for initialization (deterministic for reproducibility)
fn rand_simple() -> f64 {
    use std::sync::atomic::{AtomicU64, Ordering};
    static SEED: AtomicU64 = AtomicU64::new(12345);

    let mut s = SEED.fetch_add(1, Ordering::Relaxed);
    s ^= s >> 12;
    s ^= s << 25;
    s ^= s >> 27;
    s = s.wrapping_mul(0x2545F4914F6CDD1D);
    (s as f64) / (u64::MAX as f64)
}

/// Normalize adjacency matrix with symmetric normalization
fn normalize_adjacency(adj: &Array2<f64>, add_self_loops: bool) -> Array2<f64> {
    let n = adj.nrows();
    let mut a = adj.clone();

    if add_self_loops {
        for i in 0..n {
            a[[i, i]] = 1.0;
        }
    }

    // Compute degree
    let degrees: Vec<f64> = a.sum_axis(Axis(1)).to_vec();

    // D^{-1/2} A D^{-1/2}
    let mut norm = Array2::zeros((n, n));
    for i in 0..n {
        for j in 0..n {
            if a[[i, j]] > 0.0 {
                let d_i = degrees[i].max(1e-10);
                let d_j = degrees[j].max(1e-10);
                norm[[i, j]] = a[[i, j]] / (d_i * d_j).sqrt();
            }
        }
    }

    norm
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_brain_conv() {
        let config = BrainConvConfig {
            in_features: 4,
            out_features: 8,
            ..Default::default()
        };

        let layer = BrainConv::new(config).unwrap();

        let x = Array2::from_shape_fn((5, 4), |_| rand_simple());
        let adj = Array2::from_shape_fn((5, 5), |_| if rand_simple() > 0.5 { 1.0 } else { 0.0 });

        let output = layer.forward(&x, &adj).unwrap();
        assert_eq!(output.shape(), &[5, 8]);
    }

    #[test]
    fn test_brain_attention() {
        let config = BrainAttentionConfig {
            in_features: 4,
            out_features: 8,
            n_heads: 2,
            concat_heads: true,
            ..Default::default()
        };

        let layer = BrainAttention::new(config).unwrap();

        let x = Array2::from_shape_fn((5, 4), |_| rand_simple());
        let mut adj = Array2::zeros((5, 5));
        for i in 0..5 {
            adj[[i, i]] = 1.0;
            if i > 0 {
                adj[[i, i - 1]] = 1.0;
            }
        }

        let output = layer.forward(&x, &adj).unwrap();
        assert_eq!(output.shape(), &[5, 8]);
    }

    #[test]
    fn test_brain_pool() {
        let x = Array2::from_shape_fn((10, 4), |_| rand_simple());

        // Mean pooling
        let pool = BrainPool::new(BrainPoolConfig {
            method: PoolMethod::Mean,
            ..Default::default()
        });
        let pooled = pool.global_pool(&x);
        assert_eq!(pooled.len(), 4);

        // Max pooling
        let pool_max = BrainPool::new(BrainPoolConfig {
            method: PoolMethod::Max,
            ..Default::default()
        });
        let pooled_max = pool_max.global_pool(&x);
        assert_eq!(pooled_max.len(), 4);
    }

    #[test]
    fn test_activation() {
        assert_eq!(Activation::ReLU.apply(-1.0), 0.0);
        assert_eq!(Activation::ReLU.apply(1.0), 1.0);
        assert!((Activation::Sigmoid.apply(0.0) - 0.5).abs() < 1e-10);
        assert!(Activation::Tanh.apply(0.0).abs() < 1e-10);
    }
}