//! Pre-built brain GNN model architectures.
//!
//! Includes BrainNetCNN, BrainGAT, and BrainTransformer models.

use crate::error::{GnnError, GnnResult};
use crate::graph::BrainGraph;
use crate::layers::{
    Activation, BrainAttention, BrainAttentionConfig, BrainConv, BrainConvConfig, BrainPool,
    BrainPoolConfig, PoolMethod,
};
use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

/// Trait for brain GNN models
pub trait BrainGNN {
    /// Forward pass on a brain graph
    fn forward(&self, graph: &BrainGraph) -> GnnResult<Array1<f64>>;

    /// Forward pass returning node embeddings
    fn forward_node_embeddings(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>>;

    /// Get number of parameters
    fn n_parameters(&self) -> usize;
}

/// BrainNetCNN configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrainNetCNNConfig {
    /// Number of input channels
    pub n_channels: usize,
    /// Number of input features per channel
    pub n_features: usize,
    /// Hidden dimensions for each layer
    pub hidden_dims: Vec<usize>,
    /// Number of output classes
    pub n_classes: usize,
    /// Dropout probability
    pub dropout: f64,
    /// Pooling method
    pub pooling: PoolMethod,
}

impl Default for BrainNetCNNConfig {
    fn default() -> Self {
        Self {
            n_channels: 64,
            n_features: 1,
            hidden_dims: vec![32, 64, 128],
            n_classes: 2,
            dropout: 0.5,
            pooling: PoolMethod::Mean,
        }
    }
}

/// BrainNetCNN model
///
/// A graph convolutional network inspired by BrainNetCNN paper.
/// Uses edge-to-edge and edge-to-node convolutions.
#[derive(Debug, Clone)]
pub struct BrainNetCNN {
    config: BrainNetCNNConfig,
    /// Graph convolution layers
    conv_layers: Vec<BrainConv>,
    /// Pooling layer
    pool: BrainPool,
    /// Final classifier weights
    classifier: Array2<f64>,
    /// Classifier bias
    classifier_bias: Array1<f64>,
}

impl BrainNetCNN {
    /// Create a new BrainNetCNN model
    pub fn new(config: BrainNetCNNConfig) -> GnnResult<Self> {
        let mut conv_layers = Vec::with_capacity(config.hidden_dims.len());

        let mut prev_dim = config.n_features;
        for &hidden_dim in &config.hidden_dims {
            conv_layers.push(BrainConv::new(BrainConvConfig {
                in_features: prev_dim,
                out_features: hidden_dim,
                add_self_loops: true,
                normalize: true,
                activation: Activation::ReLU,
                dropout: config.dropout,
            })?);
            prev_dim = hidden_dim;
        }

        let pool = BrainPool::new(BrainPoolConfig {
            method: config.pooling,
            ratio: 0.5,
        });

        // Classifier
        let final_dim = *config.hidden_dims.last().unwrap_or(&config.n_features);
        let scale = (2.0 / (final_dim + config.n_classes) as f64).sqrt();
        let classifier =
            Array2::from_shape_fn((final_dim, config.n_classes), |_| rand_simple() * scale);
        let classifier_bias = Array1::zeros(config.n_classes);

        Ok(Self {
            config,
            conv_layers,
            pool,
            classifier,
            classifier_bias,
        })
    }
}

impl BrainGNN for BrainNetCNN {
    fn forward(&self, graph: &BrainGraph) -> GnnResult<Array1<f64>> {
        let embeddings = self.forward_node_embeddings(graph)?;

        // Global pooling
        let pooled = self.pool.global_pool(&embeddings);

        // Classification
        let mut logits = Array1::zeros(self.config.n_classes);
        for j in 0..self.config.n_classes {
            logits[j] = self.classifier_bias[j];
            for i in 0..pooled.len() {
                logits[j] += pooled[i] * self.classifier[[i, j]];
            }
        }

        Ok(logits)
    }

    fn forward_node_embeddings(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
        let mut x = graph.node_feature_matrix()?;
        let adj = graph.adjacency_matrix();

        // Apply conv layers
        for layer in &self.conv_layers {
            x = layer.forward(&x, &adj)?;
        }

        Ok(x)
    }

    fn n_parameters(&self) -> usize {
        let mut count = 0;
        for _layer in &self.conv_layers {
            // Approximate parameter count
            count += 100; // Placeholder
        }
        count += self.classifier.len() + self.classifier_bias.len();
        count
    }
}

/// BrainGAT configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrainGATConfig {
    /// Number of input channels
    pub n_channels: usize,
    /// Number of input features per channel
    pub n_features: usize,
    /// Hidden dimensions
    pub hidden_dims: Vec<usize>,
    /// Number of attention heads per layer
    pub n_heads: Vec<usize>,
    /// Number of output classes
    pub n_classes: usize,
    /// Dropout probability
    pub dropout: f64,
    /// Whether to concatenate attention heads
    pub concat_heads: bool,
}

impl Default for BrainGATConfig {
    fn default() -> Self {
        Self {
            n_channels: 64,
            n_features: 1,
            hidden_dims: vec![64, 64],
            n_heads: vec![4, 4],
            n_classes: 2,
            dropout: 0.5,
            concat_heads: true,
        }
    }
}

/// Brain Graph Attention Network
#[derive(Debug, Clone)]
pub struct BrainGAT {
    config: BrainGATConfig,
    /// Attention layers
    attn_layers: Vec<BrainAttention>,
    /// Pooling layer
    pool: BrainPool,
    /// Final classifier
    classifier: Array2<f64>,
    classifier_bias: Array1<f64>,
}

impl BrainGAT {
    /// Create a new BrainGAT model
    pub fn new(config: BrainGATConfig) -> GnnResult<Self> {
        if config.hidden_dims.len() != config.n_heads.len() {
            return Err(GnnError::InvalidConfig(
                "hidden_dims and n_heads must have same length".into(),
            ));
        }

        let mut attn_layers = Vec::with_capacity(config.hidden_dims.len());

        let mut prev_dim = config.n_features;
        for (&hidden_dim, &n_heads) in config.hidden_dims.iter().zip(&config.n_heads) {
            attn_layers.push(BrainAttention::new(BrainAttentionConfig {
                in_features: prev_dim,
                out_features: hidden_dim,
                n_heads,
                dropout: config.dropout,
                concat_heads: config.concat_heads,
                ..Default::default()
            })?);
            prev_dim = hidden_dim;
        }

        let pool = BrainPool::new(BrainPoolConfig::default());

        // Classifier
        let final_dim = *config.hidden_dims.last().unwrap_or(&config.n_features);
        let scale = (2.0 / (final_dim + config.n_classes) as f64).sqrt();
        let classifier =
            Array2::from_shape_fn((final_dim, config.n_classes), |_| rand_simple() * scale);
        let classifier_bias = Array1::zeros(config.n_classes);

        Ok(Self {
            config,
            attn_layers,
            pool,
            classifier,
            classifier_bias,
        })
    }
}

impl BrainGNN for BrainGAT {
    fn forward(&self, graph: &BrainGraph) -> GnnResult<Array1<f64>> {
        let embeddings = self.forward_node_embeddings(graph)?;
        let pooled = self.pool.global_pool(&embeddings);

        let mut logits = Array1::zeros(self.config.n_classes);
        for j in 0..self.config.n_classes {
            logits[j] = self.classifier_bias[j];
            for i in 0..pooled.len() {
                logits[j] += pooled[i] * self.classifier[[i, j]];
            }
        }

        Ok(logits)
    }

    fn forward_node_embeddings(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
        let mut x = graph.node_feature_matrix()?;
        let adj = graph.adjacency_matrix();

        for layer in &self.attn_layers {
            x = layer.forward(&x, &adj)?;
        }

        Ok(x)
    }

    fn n_parameters(&self) -> usize {
        self.classifier.len() + self.classifier_bias.len() + 1000 // Approximate
    }
}

/// BrainTransformer configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrainTransformerConfig {
    /// Number of input channels
    pub n_channels: usize,
    /// Number of input features per channel
    pub n_features: usize,
    /// Model dimension
    pub d_model: usize,
    /// Number of transformer layers
    pub n_layers: usize,
    /// Number of attention heads
    pub n_heads: usize,
    /// Feed-forward dimension
    pub d_ff: usize,
    /// Number of output classes
    pub n_classes: usize,
    /// Dropout probability
    pub dropout: f64,
    /// Whether to use graph structure in attention
    pub use_graph_mask: bool,
}

impl Default for BrainTransformerConfig {
    fn default() -> Self {
        Self {
            n_channels: 64,
            n_features: 1,
            d_model: 64,
            n_layers: 4,
            n_heads: 4,
            d_ff: 256,
            n_classes: 2,
            dropout: 0.1,
            use_graph_mask: true,
        }
    }
}

/// Brain Transformer model
///
/// A transformer-based model for brain connectivity analysis.
/// Can optionally mask attention based on graph structure.
#[derive(Debug, Clone)]
pub struct BrainTransformer {
    config: BrainTransformerConfig,
    /// Input projection
    input_proj: Array2<f64>,
    /// Transformer layers
    layers: Vec<TransformerLayer>,
    /// Pooling
    pool: BrainPool,
    /// Classifier
    classifier: Array2<f64>,
    classifier_bias: Array1<f64>,
}

/// A single transformer layer
#[derive(Debug, Clone)]
struct TransformerLayer {
    /// Query projection
    wq: Array2<f64>,
    /// Key projection
    wk: Array2<f64>,
    /// Value projection
    wv: Array2<f64>,
    /// Output projection
    wo: Array2<f64>,
    /// FFN layer 1
    ff1: Array2<f64>,
    /// FFN layer 2
    ff2: Array2<f64>,
    /// Layer norms (simulated with scale)
    scale: f64,
}

impl BrainTransformer {
    /// Create a new BrainTransformer model
    pub fn new(config: BrainTransformerConfig) -> GnnResult<Self> {
        let scale = (1.0 / config.d_model as f64).sqrt();

        // Input projection
        let input_proj = Array2::from_shape_fn((config.n_features, config.d_model), |_| {
            rand_simple() * scale
        });

        // Transformer layers
        let mut layers = Vec::with_capacity(config.n_layers);
        for _ in 0..config.n_layers {
            layers.push(TransformerLayer {
                wq: Array2::from_shape_fn((config.d_model, config.d_model), |_| {
                    rand_simple() * scale
                }),
                wk: Array2::from_shape_fn((config.d_model, config.d_model), |_| {
                    rand_simple() * scale
                }),
                wv: Array2::from_shape_fn((config.d_model, config.d_model), |_| {
                    rand_simple() * scale
                }),
                wo: Array2::from_shape_fn((config.d_model, config.d_model), |_| {
                    rand_simple() * scale
                }),
                ff1: Array2::from_shape_fn((config.d_model, config.d_ff), |_| {
                    rand_simple() * scale
                }),
                ff2: Array2::from_shape_fn((config.d_ff, config.d_model), |_| {
                    rand_simple() * scale
                }),
                scale: 1.0,
            });
        }

        let pool = BrainPool::new(BrainPoolConfig::default());

        // Classifier
        let classifier = Array2::from_shape_fn((config.d_model, config.n_classes), |_| {
            rand_simple() * scale
        });
        let classifier_bias = Array1::zeros(config.n_classes);

        Ok(Self {
            config,
            input_proj,
            layers,
            pool,
            classifier,
            classifier_bias,
        })
    }

    /// Self-attention with optional graph masking
    fn self_attention(
        &self,
        layer: &TransformerLayer,
        x: &Array2<f64>,
        mask: Option<&Array2<f64>>,
    ) -> Array2<f64> {
        let n = x.nrows();
        let d = x.ncols();

        // Compute Q, K, V
        let q = x.dot(&layer.wq);
        let k = x.dot(&layer.wk);
        let v = x.dot(&layer.wv);

        // Attention scores
        let scale = 1.0 / (d as f64).sqrt();
        let mut scores = q.dot(&k.t()) * scale;

        // Apply graph mask if provided
        if let Some(m) = mask {
            for i in 0..n {
                for j in 0..n {
                    if m[[i, j]] < 0.5 && i != j {
                        scores[[i, j]] = f64::NEG_INFINITY;
                    }
                }
            }
        }

        // Softmax
        for i in 0..n {
            let max_val = scores
                .row(i)
                .iter()
                .copied()
                .fold(f64::NEG_INFINITY, f64::max);
            let exp_sum: f64 = scores.row(i).iter().map(|&s| (s - max_val).exp()).sum();
            for j in 0..n {
                scores[[i, j]] = (scores[[i, j]] - max_val).exp() / exp_sum;
            }
        }

        // Apply attention
        let attn_out = scores.dot(&v);
        attn_out.dot(&layer.wo)
    }

    /// Feed-forward network
    fn ffn(&self, layer: &TransformerLayer, x: &Array2<f64>) -> Array2<f64> {
        // FFN with GELU activation
        let h = x.dot(&layer.ff1);
        let h_act =
            h.mapv(|v| 0.5 * v * (1.0 + (0.7978845608 * (v + 0.044715 * v * v * v)).tanh()));
        h_act.dot(&layer.ff2)
    }
}

impl BrainGNN for BrainTransformer {
    fn forward(&self, graph: &BrainGraph) -> GnnResult<Array1<f64>> {
        let embeddings = self.forward_node_embeddings(graph)?;
        let pooled = self.pool.global_pool(&embeddings);

        let mut logits = Array1::zeros(self.config.n_classes);
        for j in 0..self.config.n_classes {
            logits[j] = self.classifier_bias[j];
            for i in 0..pooled.len() {
                logits[j] += pooled[i] * self.classifier[[i, j]];
            }
        }

        Ok(logits)
    }

    fn forward_node_embeddings(&self, graph: &BrainGraph) -> GnnResult<Array2<f64>> {
        let mut x = graph.node_feature_matrix()?;

        // Project to model dimension
        x = x.dot(&self.input_proj);

        // Graph mask
        let mask = if self.config.use_graph_mask {
            Some(graph.adjacency_matrix())
        } else {
            None
        };

        // Apply transformer layers
        for layer in &self.layers {
            // Self-attention with residual
            let attn_out = self.self_attention(layer, &x, mask.as_ref());
            x = x + attn_out;

            // FFN with residual
            let ffn_out = self.ffn(layer, &x);
            x = x + ffn_out;
        }

        Ok(x)
    }

    fn n_parameters(&self) -> usize {
        let mut count = self.input_proj.len();
        for layer in &self.layers {
            count += layer.wq.len()
                + layer.wk.len()
                + layer.wv.len()
                + layer.wo.len()
                + layer.ff1.len()
                + layer.ff2.len();
        }
        count += self.classifier.len() + self.classifier_bias.len();
        count
    }
}

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

    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)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::GraphBuilder;
    use crate::graph::{BrainRegion, Hemisphere};

    fn create_test_graph() -> BrainGraph {
        let mut builder = GraphBuilder::new();

        // 4-node test graph
        builder.add_node("F3", Hemisphere::Left, BrainRegion::Frontal);
        builder.add_node("F4", Hemisphere::Right, BrainRegion::Frontal);
        builder.add_node("O1", Hemisphere::Left, BrainRegion::Occipital);
        builder.add_node("O2", Hemisphere::Right, BrainRegion::Occipital);

        builder.add_edge(0, 1, 0.8);
        builder.add_edge(0, 2, 0.5);
        builder.add_edge(1, 3, 0.6);
        builder.add_edge(2, 3, 0.9);

        let mut graph = builder.build().unwrap();

        // Add features
        for node in &mut graph.nodes {
            node.features = vec![0.5, 0.3, 0.2, 0.1];
        }

        graph
    }

    #[test]
    fn test_brainnetcnn() {
        let graph = create_test_graph();

        let model = BrainNetCNN::new(BrainNetCNNConfig {
            n_channels: 4,
            n_features: 4,
            hidden_dims: vec![8, 16],
            n_classes: 3,
            ..Default::default()
        })
        .unwrap();

        let logits = model.forward(&graph).unwrap();
        assert_eq!(logits.len(), 3);
    }

    #[test]
    fn test_braingat() {
        let graph = create_test_graph();

        let model = BrainGAT::new(BrainGATConfig {
            n_channels: 4,
            n_features: 4,
            hidden_dims: vec![8, 8],
            n_heads: vec![2, 2],
            n_classes: 2,
            ..Default::default()
        })
        .unwrap();

        let logits = model.forward(&graph).unwrap();
        assert_eq!(logits.len(), 2);
    }

    #[test]
    fn test_brain_transformer() {
        let graph = create_test_graph();

        let model = BrainTransformer::new(BrainTransformerConfig {
            n_channels: 4,
            n_features: 4,
            d_model: 16,
            n_layers: 2,
            n_heads: 2,
            d_ff: 32,
            n_classes: 2,
            ..Default::default()
        })
        .unwrap();

        let logits = model.forward(&graph).unwrap();
        assert_eq!(logits.len(), 2);

        let embeddings = model.forward_node_embeddings(&graph).unwrap();
        assert_eq!(embeddings.shape(), &[4, 16]);
    }
}