//! Temporal GNN for dynamic connectivity analysis.
//!
//! Provides models for analyzing time-varying brain connectivity graphs.

use crate::error::{GnnError, GnnResult};
use crate::graph::BrainGraph;
use crate::layers::{BrainPool, BrainPoolConfig, PoolMethod};
use crate::models::{BrainGNN, BrainNetCNN, BrainNetCNNConfig};
use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

/// Time window specification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimeWindow {
    /// Start time in seconds
    pub start: f64,
    /// End time in seconds
    pub end: f64,
    /// Window index
    pub index: usize,
}

impl TimeWindow {
    /// Create a new time window
    pub fn new(start: f64, end: f64, index: usize) -> Self {
        Self { start, end, index }
    }

    /// Duration of the window
    pub fn duration(&self) -> f64 {
        self.end - self.start
    }

    /// Center time of the window
    pub fn center(&self) -> f64 {
        f64::midpoint(self.start, self.end)
    }
}

/// Dynamic connectivity sequence
#[derive(Debug, Clone)]
pub struct DynamicConnectivity {
    /// Sequence of brain graphs (one per time window)
    pub graphs: Vec<BrainGraph>,
    /// Time windows
    pub windows: Vec<TimeWindow>,
    /// Sampling frequency
    pub sfreq: f64,
}

impl DynamicConnectivity {
    /// Create from a sequence of connectivity matrices
    pub fn from_matrices(
        matrices: &[Vec<Vec<f64>>],
        channel_names: &[&str],
        window_duration: f64,
        step: f64,
        threshold: f64,
    ) -> GnnResult<Self> {
        let mut graphs = Vec::with_capacity(matrices.len());
        let mut windows = Vec::with_capacity(matrices.len());

        for (i, matrix) in matrices.iter().enumerate() {
            let graph = BrainGraph::from_connectivity_matrix(matrix, channel_names, threshold)?;
            graphs.push(graph);

            let start = i as f64 * step;
            windows.push(TimeWindow::new(start, start + window_duration, i));
        }

        Ok(Self {
            graphs,
            windows,
            sfreq: 1.0 / step,
        })
    }

    /// Number of time points
    pub fn n_times(&self) -> usize {
        self.graphs.len()
    }

    /// Get graph at specific time index
    pub fn get(&self, index: usize) -> Option<&BrainGraph> {
        self.graphs.get(index)
    }

    /// Iterator over (time_window, graph) pairs
    pub fn iter(&self) -> impl Iterator<Item = (&TimeWindow, &BrainGraph)> {
        self.windows.iter().zip(self.graphs.iter())
    }
}

/// Configuration for temporal GNN
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TemporalConfig {
    /// Spatial GNN configuration
    pub spatial: BrainNetCNNConfig,
    /// Number of temporal attention heads
    pub temporal_heads: usize,
    /// Temporal hidden dimension
    pub temporal_hidden: usize,
    /// Whether to use temporal positional encoding
    pub positional_encoding: bool,
    /// Maximum sequence length
    pub max_seq_len: usize,
}

impl Default for TemporalConfig {
    fn default() -> Self {
        Self {
            spatial: BrainNetCNNConfig {
                hidden_dims: vec![32, 64],
                n_classes: 2,
                ..Default::default()
            },
            temporal_heads: 4,
            temporal_hidden: 64,
            positional_encoding: true,
            max_seq_len: 100,
        }
    }
}

/// Temporal attention mechanism
#[derive(Debug, Clone)]
pub struct TemporalAttention {
    /// Number of attention heads
    n_heads: usize,
    /// Dimension per head
    head_dim: usize,
    /// Query weights per head
    wq: Vec<Array2<f64>>,
    /// Key weights per head
    wk: Vec<Array2<f64>>,
    /// Value weights per head
    wv: Vec<Array2<f64>>,
    /// Output projection
    wo: Array2<f64>,
    /// Positional encoding matrix
    positional_encoding: Option<Array2<f64>>,
}

impl TemporalAttention {
    /// Create a new temporal attention layer
    pub fn new(d_model: usize, n_heads: usize, max_seq_len: usize, use_pe: bool) -> Self {
        let head_dim = d_model / n_heads;
        let scale = (1.0 / head_dim as f64).sqrt();

        let mut wq = Vec::with_capacity(n_heads);
        let mut wk = Vec::with_capacity(n_heads);
        let mut wv = Vec::with_capacity(n_heads);

        for _ in 0..n_heads {
            wq.push(Array2::from_shape_fn((d_model, head_dim), |_| {
                rand_simple() * scale
            }));
            wk.push(Array2::from_shape_fn((d_model, head_dim), |_| {
                rand_simple() * scale
            }));
            wv.push(Array2::from_shape_fn((d_model, head_dim), |_| {
                rand_simple() * scale
            }));
        }

        let wo = Array2::from_shape_fn((d_model, d_model), |_| rand_simple() * scale);

        // Sinusoidal positional encoding
        let positional_encoding = if use_pe {
            let mut pe = Array2::zeros((max_seq_len, d_model));
            for pos in 0..max_seq_len {
                for i in 0..(d_model / 2) {
                    let angle = pos as f64 / (10000.0_f64).powf(2.0 * i as f64 / d_model as f64);
                    pe[[pos, 2 * i]] = angle.sin();
                    pe[[pos, 2 * i + 1]] = angle.cos();
                }
            }
            Some(pe)
        } else {
            None
        };

        Self {
            n_heads,
            head_dim,
            wq,
            wk,
            wv,
            wo,
            positional_encoding,
        }
    }

    /// Forward pass on sequence of embeddings [n_times, d_model]
    pub fn forward(&self, x: &Array2<f64>) -> Array2<f64> {
        let n_times = x.nrows();
        let d_model = x.ncols();

        // Add positional encoding
        let mut h = x.clone();
        if let Some(ref pe) = self.positional_encoding {
            for t in 0..n_times {
                for d in 0..d_model {
                    h[[t, d]] += pe[[t, d]];
                }
            }
        }

        // Multi-head attention
        let mut head_outputs = Vec::with_capacity(self.n_heads);

        for head in 0..self.n_heads {
            let q = h.dot(&self.wq[head]); // [n_times, head_dim]
            let k = h.dot(&self.wk[head]);
            let v = h.dot(&self.wv[head]);

            // Attention scores
            let scale = 1.0 / (self.head_dim as f64).sqrt();
            let scores = q.dot(&k.t()) * scale; // [n_times, n_times]

            // Softmax
            let mut attn = Array2::zeros((n_times, n_times));
            for i in 0..n_times {
                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_times {
                    attn[[i, j]] = (scores[[i, j]] - max_val).exp() / exp_sum;
                }
            }

            let out = attn.dot(&v);
            head_outputs.push(out);
        }

        // Concatenate heads
        let mut concat = Array2::zeros((n_times, d_model));
        for (h, head_out) in head_outputs.iter().enumerate() {
            for t in 0..n_times {
                for d in 0..self.head_dim {
                    concat[[t, h * self.head_dim + d]] = head_out[[t, d]];
                }
            }
        }

        // Output projection + residual
        let output = concat.dot(&self.wo);
        x + &output
    }

    /// Get attention weights for interpretation
    pub fn get_attention_weights(&self, x: &Array2<f64>) -> Array2<f64> {
        let n_times = x.nrows();

        // Use first head for visualization
        let q = x.dot(&self.wq[0]);
        let k = x.dot(&self.wk[0]);

        let scale = 1.0 / (self.head_dim as f64).sqrt();
        let scores = q.dot(&k.t()) * scale;

        // Softmax
        let mut attn = Array2::zeros((n_times, n_times));
        for i in 0..n_times {
            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_times {
                attn[[i, j]] = (scores[[i, j]] - max_val).exp() / exp_sum;
            }
        }

        attn
    }
}

/// Temporal Brain GNN model
///
/// Combines spatial GNN for each time window with temporal attention
/// for capturing dynamic brain state changes.
#[derive(Debug, Clone)]
pub struct TemporalBrainGNN {
    config: TemporalConfig,
    /// Spatial GNN (applied to each time window)
    spatial_gnn: BrainNetCNN,
    /// Temporal attention
    temporal_attention: TemporalAttention,
    /// Temporal convolution weights
    temporal_conv: Array2<f64>,
    /// Pooling
    pool: BrainPool,
    /// Classifier
    classifier: Array2<f64>,
    classifier_bias: Array1<f64>,
}

impl TemporalBrainGNN {
    /// Create a new temporal brain GNN
    pub fn new(config: TemporalConfig) -> GnnResult<Self> {
        // Create spatial GNN
        let spatial_gnn = BrainNetCNN::new(config.spatial.clone())?;

        // Final dimension from spatial GNN
        let spatial_dim = *config
            .spatial
            .hidden_dims
            .last()
            .unwrap_or(&config.spatial.n_features);

        // Temporal attention
        let temporal_attention = TemporalAttention::new(
            spatial_dim,
            config.temporal_heads,
            config.max_seq_len,
            config.positional_encoding,
        );

        // Temporal 1D convolution (simplified as matrix)
        let scale = (1.0 / config.temporal_hidden as f64).sqrt();
        let temporal_conv = Array2::from_shape_fn((spatial_dim, config.temporal_hidden), |_| {
            rand_simple() * scale
        });

        let pool = BrainPool::new(BrainPoolConfig {
            method: PoolMethod::Attention,
            ratio: 1.0,
        });

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

        Ok(Self {
            config,
            spatial_gnn,
            temporal_attention,
            temporal_conv,
            pool,
            classifier,
            classifier_bias,
        })
    }

    /// Forward pass on dynamic connectivity
    pub fn forward(&self, dyn_conn: &DynamicConnectivity) -> GnnResult<Array1<f64>> {
        let temporal_embeddings = self.forward_temporal_embeddings(dyn_conn)?;

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

        // Apply temporal convolution
        let h = Array1::from_iter((0..self.config.temporal_hidden).map(|j| {
            (0..pooled.len())
                .map(|i| pooled[i] * self.temporal_conv[[i, j]])
                .sum::<f64>()
        }));

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

        Ok(logits)
    }

    /// Get temporal embeddings [n_times, spatial_dim]
    pub fn forward_temporal_embeddings(
        &self,
        dyn_conn: &DynamicConnectivity,
    ) -> GnnResult<Array2<f64>> {
        let n_times = dyn_conn.n_times();
        if n_times == 0 {
            return Err(GnnError::InvalidGraph("No time windows".into()));
        }

        let spatial_dim = *self
            .config
            .spatial
            .hidden_dims
            .last()
            .unwrap_or(&self.config.spatial.n_features);

        // Get spatial embeddings for each time window
        let mut spatial_embeddings = Array2::zeros((n_times, spatial_dim));

        for (t, graph) in dyn_conn.graphs.iter().enumerate() {
            let node_emb = self.spatial_gnn.forward_node_embeddings(graph)?;
            let pooled = self.pool.global_pool(&node_emb);

            for d in 0..spatial_dim.min(pooled.len()) {
                spatial_embeddings[[t, d]] = pooled[d];
            }
        }

        // Apply temporal attention
        let temporal_out = self.temporal_attention.forward(&spatial_embeddings);

        Ok(temporal_out)
    }

    /// Get temporal attention weights for interpretation
    pub fn get_temporal_attention_weights(
        &self,
        dyn_conn: &DynamicConnectivity,
    ) -> GnnResult<Array2<f64>> {
        let temporal_embeddings = self.forward_temporal_embeddings(dyn_conn)?;
        Ok(self
            .temporal_attention
            .get_attention_weights(&temporal_embeddings))
    }

    /// Predict brain state at each time window
    pub fn predict_states(&self, dyn_conn: &DynamicConnectivity) -> GnnResult<Vec<usize>> {
        let temporal_embeddings = self.forward_temporal_embeddings(dyn_conn)?;

        let mut states = Vec::with_capacity(dyn_conn.n_times());

        for t in 0..dyn_conn.n_times() {
            // Get embedding for this time
            let h: Vec<f64> = temporal_embeddings.row(t).to_vec();

            // Apply temporal conv
            let conv_out: Vec<f64> = (0..self.config.temporal_hidden)
                .map(|j| {
                    (0..h.len())
                        .map(|i| h[i] * self.temporal_conv[[i, j]])
                        .sum()
                })
                .collect();

            // Classify
            let mut logits = vec![0.0; self.config.spatial.n_classes];
            for j in 0..self.config.spatial.n_classes {
                logits[j] = self.classifier_bias[j];
                for i in 0..conv_out.len() {
                    logits[j] += conv_out[i] * self.classifier[[i, j]];
                }
            }

            // Argmax
            let state = logits
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                .map_or(0, |(i, _)| i);

            states.push(state);
        }

        Ok(states)
    }
}

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

    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::{BrainRegion, GraphBuilder, Hemisphere};

    fn create_test_graph() -> BrainGraph {
        let mut builder = GraphBuilder::new();
        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();
        for node in &mut graph.nodes {
            node.features = vec![0.5, 0.3, 0.2, 0.1];
        }
        graph
    }

    #[test]
    fn test_dynamic_connectivity() {
        let channels = ["F3", "F4", "O1", "O2"];
        let matrices: Vec<Vec<Vec<f64>>> = (0..5)
            .map(|t| {
                (0..4)
                    .map(|i| {
                        (0..4)
                            .map(|j| if i == j { 1.0 } else { 0.5 + t as f64 * 0.05 })
                            .collect()
                    })
                    .collect()
            })
            .collect();

        let dyn_conn =
            DynamicConnectivity::from_matrices(&matrices, &channels, 1.0, 0.5, 0.4).unwrap();

        assert_eq!(dyn_conn.n_times(), 5);
        assert!(dyn_conn.get(2).is_some());
    }

    #[test]
    fn test_temporal_attention() {
        let attn = TemporalAttention::new(16, 4, 50, true);

        let x = Array2::from_shape_fn((10, 16), |_| rand_simple());
        let out = attn.forward(&x);

        assert_eq!(out.shape(), &[10, 16]);
    }

    #[test]
    fn test_temporal_brain_gnn() {
        // Create dynamic connectivity
        let channels = ["F3", "F4", "O1", "O2"];
        let matrices: Vec<Vec<Vec<f64>>> = (0..5)
            .map(|_| {
                (0..4)
                    .map(|i| (0..4).map(|j| if i == j { 1.0 } else { 0.6 }).collect())
                    .collect()
            })
            .collect();

        let dyn_conn =
            DynamicConnectivity::from_matrices(&matrices, &channels, 1.0, 0.5, 0.4).unwrap();

        // Create model
        let model = TemporalBrainGNN::new(TemporalConfig {
            spatial: BrainNetCNNConfig {
                n_channels: 4,
                n_features: 1, // Will use degree
                hidden_dims: vec![8, 16],
                n_classes: 2,
                ..Default::default()
            },
            temporal_heads: 2,
            temporal_hidden: 8,
            max_seq_len: 10,
            ..Default::default()
        })
        .unwrap();

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

        let states = model.predict_states(&dyn_conn).unwrap();
        assert_eq!(states.len(), 5);
    }
}