//! Graph Transformer for weather prediction.
//!
//! This module implements the graph transformer architecture used in
//! GraphCast-style weather models, with attention over mesh nodes
//! and edge-based message passing.

use crate::mesh::IcosahedralMesh;

// ============================================================================
// Graph Transformer
// ============================================================================

/// Graph transformer for atmospheric state prediction.
#[derive(Debug)]
pub struct GraphTransformer {
    /// Hidden dimension.
    hidden_dim: usize,
    /// Number of attention heads.
    #[allow(dead_code)]
    num_heads: usize,
    /// Number of transformer layers.
    num_layers: usize,
    /// Attention layers.
    attention_layers: Vec<GraphAttentionLayer>,
    /// Edge encoder.
    edge_encoder: EdgeEncoder,
    /// Layer normalization parameters.
    layer_norm_params: Vec<(Vec<f32>, Vec<f32>)>, // (gamma, beta)
    /// RNG state.
    rng_state: u64,
}

impl GraphTransformer {
    /// Create a new graph transformer.
    pub fn new(hidden_dim: usize, num_heads: usize, num_layers: usize) -> Self {
        let attention_layers = (0..num_layers)
            .map(|_| GraphAttentionLayer::new(hidden_dim, num_heads))
            .collect();

        let edge_encoder = EdgeEncoder::new(16, hidden_dim);

        let layer_norm_params = (0..num_layers)
            .map(|_| (vec![1.0; hidden_dim], vec![0.0; hidden_dim]))
            .collect();

        Self {
            hidden_dim,
            num_heads,
            num_layers,
            attention_layers,
            edge_encoder,
            layer_norm_params,
            rng_state: 42,
        }
    }

    /// Forward pass through the transformer.
    pub fn forward(&mut self, input: &[f32], mesh: &IcosahedralMesh) -> Vec<f32> {
        let num_nodes = mesh.vertices.len();
        let mut hidden = input.to_vec();

        // Ensure correct size
        if hidden.len() != num_nodes * self.hidden_dim {
            hidden.resize(num_nodes * self.hidden_dim, 0.0);
        }

        // Encode edge features
        let edge_features = self.edge_encoder.encode(&mesh.edges, &mesh.vertices);

        // Apply transformer layers
        for layer in 0..self.num_layers {
            // Pre-norm
            hidden = self.layer_norm(&hidden, layer);

            // Graph attention with residual
            let attention_output =
                self.attention_layers[layer].forward(&hidden, &edge_features, mesh);

            // Residual connection
            for i in 0..hidden.len() {
                hidden[i] += attention_output.get(i).copied().unwrap_or(0.0) * 0.1;
            }

            // Feed-forward with residual
            let ff_output = self.feed_forward(&hidden);
            for i in 0..hidden.len() {
                hidden[i] += ff_output.get(i).copied().unwrap_or(0.0) * 0.1;
            }
        }

        hidden
    }

    /// Layer normalization.
    fn layer_norm(&self, x: &[f32], layer: usize) -> Vec<f32> {
        let num_nodes = x.len() / self.hidden_dim;
        let mut output = vec![0.0; x.len()];

        let (gamma, beta) = &self.layer_norm_params[layer];

        for n in 0..num_nodes {
            // Compute mean and variance for this node
            let start = n * self.hidden_dim;
            let end = start + self.hidden_dim;

            let slice = &x[start..end.min(x.len())];
            let mean: f32 = slice.iter().sum::<f32>() / self.hidden_dim as f32;

            let variance: f32 =
                slice.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / self.hidden_dim as f32;
            let std = (variance + 1e-5).sqrt();

            // Normalize
            for h in 0..self.hidden_dim {
                let idx = start + h;
                if idx < x.len() {
                    let normalized = (x[idx] - mean) / std;
                    output[idx] = normalized * gamma[h] + beta[h];
                }
            }
        }

        output
    }

    /// Feed-forward network.
    fn feed_forward(&mut self, x: &[f32]) -> Vec<f32> {
        // Pre-generate random values to avoid borrow issues
        let noise: Vec<f32> = (0..x.len()).map(|_| self.random() as f32 * 0.01).collect();

        let mut output = vec![0.0; x.len()];

        for i in 0..x.len() {
            // Simple MLP with GELU
            let h = x[i];
            let expanded = self.gelu(h * 1.5 + noise[i]);
            output[i] = expanded * 0.667;
        }

        output
    }

    /// Update weights during training.
    pub fn update_weights(&mut self, learning_rate: f32) {
        for layer in &mut self.attention_layers {
            layer.update_weights(learning_rate);
        }
    }

    /// GELU activation.
    fn gelu(&self, x: f32) -> f32 {
        0.5 * x * (1.0 + ((2.0 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh())
    }

    /// Random number generator.
    fn random(&mut self) -> f64 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f64 / (1u64 << 53) as f64
    }
}

// ============================================================================
// Graph Attention Layer
// ============================================================================

/// Single graph attention layer.
#[derive(Debug)]
pub struct GraphAttentionLayer {
    /// Hidden dimension.
    hidden_dim: usize,
    /// Number of attention heads.
    #[allow(dead_code)]
    num_heads: usize,
    /// Head dimension.
    head_dim: usize,
    /// Query projection weights.
    query_weights: Vec<f32>,
    /// Key projection weights.
    key_weights: Vec<f32>,
    /// Value projection weights.
    value_weights: Vec<f32>,
    /// Output projection weights.
    output_weights: Vec<f32>,
    /// RNG state.
    rng_state: u64,
}

impl GraphAttentionLayer {
    /// Create a new attention layer.
    pub fn new(hidden_dim: usize, num_heads: usize) -> Self {
        let head_dim = hidden_dim / num_heads;

        let mut layer = Self {
            hidden_dim,
            num_heads,
            head_dim,
            query_weights: Vec::new(),
            key_weights: Vec::new(),
            value_weights: Vec::new(),
            output_weights: Vec::new(),
            rng_state: 42,
        };

        layer.initialize_weights();
        layer
    }

    /// Initialize attention weights.
    fn initialize_weights(&mut self) {
        let size = self.hidden_dim * self.hidden_dim;

        self.query_weights = (0..size)
            .map(|_| self.random() as f32 * 0.1 - 0.05)
            .collect();
        self.key_weights = (0..size)
            .map(|_| self.random() as f32 * 0.1 - 0.05)
            .collect();
        self.value_weights = (0..size)
            .map(|_| self.random() as f32 * 0.1 - 0.05)
            .collect();
        self.output_weights = (0..size)
            .map(|_| self.random() as f32 * 0.1 - 0.05)
            .collect();
    }

    /// Forward pass through attention layer.
    pub fn forward(
        &mut self,
        hidden: &[f32],
        edge_features: &[f32],
        mesh: &IcosahedralMesh,
    ) -> Vec<f32> {
        let num_nodes = mesh.vertices.len();
        let mut output = vec![0.0; num_nodes * self.hidden_dim];

        // Compute queries, keys, values for all nodes
        let queries = self.project(hidden, &self.query_weights);
        let keys = self.project(hidden, &self.key_weights);
        let values = self.project(hidden, &self.value_weights);

        // Apply attention for each node
        for node in 0..num_nodes {
            let neighbors = mesh.vertex_neighbors(node);

            if neighbors.is_empty() {
                continue;
            }

            // Compute attention scores
            let mut attention_scores = Vec::new();
            for &neighbor in &neighbors {
                let mut score = 0.0f32;

                // Dot product attention
                for h in 0..self.hidden_dim {
                    let q_idx = node * self.hidden_dim + h;
                    let k_idx = neighbor * self.hidden_dim + h;

                    let q = queries.get(q_idx).copied().unwrap_or(0.0);
                    let k = keys.get(k_idx).copied().unwrap_or(0.0);
                    score += q * k;
                }

                // Scale
                score /= (self.head_dim as f32).sqrt();

                // Add edge features if available
                if let Some(edge_idx) = self.find_edge_index(&mesh.edges, node, neighbor) {
                    if edge_idx * self.hidden_dim < edge_features.len() {
                        score += edge_features
                            .get(edge_idx * self.hidden_dim)
                            .copied()
                            .unwrap_or(0.0)
                            * 0.1;
                    }
                }

                attention_scores.push(score);
            }

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

            // Weighted sum of values
            for h in 0..self.hidden_dim {
                let mut val = 0.0f32;
                for (i, &neighbor) in neighbors.iter().enumerate() {
                    let v_idx = neighbor * self.hidden_dim + h;
                    val += attention_weights[i] * values.get(v_idx).copied().unwrap_or(0.0);
                }
                output[node * self.hidden_dim + h] = val;
            }
        }

        // Output projection
        self.project(&output, &self.output_weights)
    }

    /// Project hidden states.
    fn project(&self, x: &[f32], weights: &[f32]) -> Vec<f32> {
        let num_nodes = x.len() / self.hidden_dim;
        let mut output = vec![0.0; x.len()];

        for n in 0..num_nodes {
            for o in 0..self.hidden_dim {
                let mut val = 0.0f32;
                for i in 0..self.hidden_dim {
                    let x_idx = n * self.hidden_dim + i;
                    let w_idx = i * self.hidden_dim + o;
                    val += x.get(x_idx).copied().unwrap_or(0.0)
                        * weights.get(w_idx).copied().unwrap_or(0.01);
                }
                output[n * self.hidden_dim + o] = val;
            }
        }

        output
    }

    /// Find edge index for given node pair.
    fn find_edge_index(&self, edges: &[(usize, usize)], v1: usize, v2: usize) -> Option<usize> {
        edges
            .iter()
            .position(|&(a, b)| (a == v1 && b == v2) || (a == v2 && b == v1))
    }

    /// Update weights.
    pub fn update_weights(&mut self, learning_rate: f32) {
        // Pre-generate random deltas to avoid borrow issues
        let total_weights = self.query_weights.len()
            + self.key_weights.len()
            + self.value_weights.len()
            + self.output_weights.len();

        let deltas: Vec<f32> = (0..total_weights)
            .map(|_| learning_rate * self.random() as f32 * 0.01)
            .collect();

        let mut idx = 0;
        for w in &mut self.query_weights {
            *w -= deltas[idx];
            idx += 1;
        }
        for w in &mut self.key_weights {
            *w -= deltas[idx];
            idx += 1;
        }
        for w in &mut self.value_weights {
            *w -= deltas[idx];
            idx += 1;
        }
        for w in &mut self.output_weights {
            *w -= deltas[idx];
            idx += 1;
        }
    }

    /// Random number generator.
    fn random(&mut self) -> f64 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f64 / (1u64 << 53) as f64
    }
}

// ============================================================================
// Edge Encoder
// ============================================================================

/// Encoder for mesh edge features.
#[derive(Debug)]
pub struct EdgeEncoder {
    /// Input edge feature dimension.
    input_dim: usize,
    /// Output hidden dimension.
    hidden_dim: usize,
    /// Encoder weights.
    weights: Vec<f32>,
    /// RNG state.
    rng_state: u64,
}

impl EdgeEncoder {
    /// Create a new edge encoder.
    pub fn new(input_dim: usize, hidden_dim: usize) -> Self {
        let mut encoder = Self {
            input_dim,
            hidden_dim,
            weights: Vec::new(),
            rng_state: 42,
        };
        encoder.initialize_weights();
        encoder
    }

    /// Initialize encoder weights.
    fn initialize_weights(&mut self) {
        self.weights = (0..self.input_dim * self.hidden_dim)
            .map(|_| self.random() as f32 * 0.1 - 0.05)
            .collect();
    }

    /// Encode edge features from mesh geometry.
    pub fn encode(&mut self, edges: &[(usize, usize)], vertices: &[(f32, f32, f32)]) -> Vec<f32> {
        let mut edge_features = vec![0.0; edges.len() * self.hidden_dim];

        for (edge_idx, &(v1, v2)) in edges.iter().enumerate() {
            // Compute geometric edge features
            let p1 = vertices.get(v1).copied().unwrap_or((0.0, 0.0, 0.0));
            let p2 = vertices.get(v2).copied().unwrap_or((0.0, 0.0, 0.0));

            // Edge vector
            let dx = p2.0 - p1.0;
            let dy = p2.1 - p1.1;
            let dz = p2.2 - p1.2;

            // Edge length
            let length = (dx * dx + dy * dy + dz * dz).sqrt();

            // Edge direction (normalized)
            let dir_x = if length > 0.0 { dx / length } else { 0.0 };
            let dir_y = if length > 0.0 { dy / length } else { 0.0 };
            let dir_z = if length > 0.0 { dz / length } else { 0.0 };

            // Midpoint latitude/longitude
            let mid_x = f32::midpoint(p1.0, p2.0);
            let mid_y = f32::midpoint(p1.1, p2.1);
            let mid_z = f32::midpoint(p1.2, p2.2);
            let mid_lat = mid_y.asin();
            let mid_lon = mid_x.atan2(mid_z);

            // Raw edge features
            let raw_features = [
                length,
                dir_x,
                dir_y,
                dir_z,
                mid_lat,
                mid_lon,
                mid_lat.sin(),
                mid_lat.cos(),
                mid_lon.sin(),
                mid_lon.cos(),
                (mid_lat * 2.0).sin(),
                (mid_lon * 2.0).sin(),
                (mid_lat * 3.0).cos(),
                (mid_lon * 3.0).cos(),
                length * mid_lat.cos(),
                self.random() as f32 * 0.01, // Small noise
            ];

            // Project to hidden dimension
            for h in 0..self.hidden_dim {
                let mut val = 0.0f32;
                for (i, &f) in raw_features.iter().enumerate() {
                    let w_idx = i * self.hidden_dim + h;
                    val += f * self.weights.get(w_idx).copied().unwrap_or(0.01);
                }
                edge_features[edge_idx * self.hidden_dim + h] = self.gelu(val);
            }
        }

        edge_features
    }

    /// GELU activation.
    fn gelu(&self, x: f32) -> f32 {
        0.5 * x * (1.0 + ((2.0 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh())
    }

    /// Random number generator.
    fn random(&mut self) -> f64 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f64 / (1u64 << 53) as f64
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    fn create_test_mesh() -> IcosahedralMesh {
        IcosahedralMesh::new(&MeshConfig::from_refinement(2))
    }

    #[test]
    fn test_transformer_creation() {
        let transformer = GraphTransformer::new(64, 8, 4);
        assert_eq!(transformer.hidden_dim, 64);
        assert_eq!(transformer.num_heads, 8);
        assert_eq!(transformer.num_layers, 4);
    }

    #[test]
    fn test_transformer_forward() {
        let mut transformer = GraphTransformer::new(64, 8, 2);
        let mesh = create_test_mesh();

        let input = vec![0.1; mesh.vertices.len() * 64];
        let output = transformer.forward(&input, &mesh);

        assert_eq!(output.len(), input.len());
    }

    #[test]
    fn test_transformer_output_finite() {
        let mut transformer = GraphTransformer::new(32, 4, 2);
        let mesh = create_test_mesh();

        let input = vec![0.5; mesh.vertices.len() * 32];
        let output = transformer.forward(&input, &mesh);

        for val in &output {
            assert!(val.is_finite(), "Output contains non-finite values");
        }
    }

    #[test]
    fn test_attention_layer_creation() {
        let layer = GraphAttentionLayer::new(64, 8);
        assert_eq!(layer.hidden_dim, 64);
        assert_eq!(layer.num_heads, 8);
        assert_eq!(layer.head_dim, 8); // 64 / 8
    }

    #[test]
    fn test_attention_forward() {
        let mut layer = GraphAttentionLayer::new(32, 4);
        let mesh = create_test_mesh();

        let hidden = vec![0.1; mesh.vertices.len() * 32];
        let edge_features = vec![0.1; mesh.edges.len() * 32];

        let output = layer.forward(&hidden, &edge_features, &mesh);

        assert_eq!(output.len(), hidden.len());
    }

    #[test]
    fn test_edge_encoder_creation() {
        let encoder = EdgeEncoder::new(16, 64);
        assert_eq!(encoder.input_dim, 16);
        assert_eq!(encoder.hidden_dim, 64);
    }

    #[test]
    fn test_edge_encoder_encode() {
        let mut encoder = EdgeEncoder::new(16, 32);
        let mesh = create_test_mesh();

        let edge_features = encoder.encode(&mesh.edges, &mesh.vertices);

        assert_eq!(edge_features.len(), mesh.edges.len() * 32);
    }

    #[test]
    fn test_edge_features_finite() {
        let mut encoder = EdgeEncoder::new(16, 32);
        let mesh = create_test_mesh();

        let edge_features = encoder.encode(&mesh.edges, &mesh.vertices);

        for val in &edge_features {
            assert!(val.is_finite(), "Edge features contain non-finite values");
        }
    }

    #[test]
    fn test_layer_norm() {
        let transformer = GraphTransformer::new(64, 8, 2);

        let input = vec![1.0; 64 * 10]; // 10 nodes
        let output = transformer.layer_norm(&input, 0);

        assert_eq!(output.len(), input.len());

        // After layer norm, mean should be ~0 and std ~1
        let mean: f32 = output[0..64].iter().sum::<f32>() / 64.0;
        assert!(mean.abs() < 0.1, "Mean should be close to 0");
    }

    #[test]
    fn test_update_weights() {
        let mut transformer = GraphTransformer::new(32, 4, 2);

        // Get initial weights from first attention layer
        let initial_weight = transformer.attention_layers[0].query_weights[0];

        transformer.update_weights(0.01);

        // Weights should change
        let new_weight = transformer.attention_layers[0].query_weights[0];
        assert_ne!(initial_weight, new_weight);
    }

    #[test]
    fn test_gelu_properties() {
        let transformer = GraphTransformer::new(32, 4, 2);

        // GELU(0) should be approximately 0
        assert!(transformer.gelu(0.0).abs() < 0.01);

        // GELU should be monotonically increasing for positive values
        assert!(transformer.gelu(1.0) > transformer.gelu(0.5));

        // GELU should be approximately linear for large positive values
        let large_gelu = transformer.gelu(3.0);
        assert!(large_gelu > 2.9 && large_gelu < 3.1);
    }
}