use crate::Graph;
use crate::error::{GeomError, Result};
use crate::layers::GNNLayer;
use rtx_tensor::Tensor;
use std::collections::HashMap;

/// Graph Attention Network (GAT) Layer
///
/// Implements multi-head attention mechanism from Veličković et al. (2018):
/// α_ij = softmax(LeakyReLU(a^T [W h_i || W h_j]))
/// h_i' = σ(Σ_j α_ij W h_j)
#[derive(Debug, Clone)]
pub struct GATLayer {
    weight: Tensor,
    attention_weight: Tensor,
    bias: Tensor,
    input_dim: usize,
    output_dim: usize,
    num_heads: usize,
    gradients_enabled: bool,
    computed_gradients: bool,
    attention_weights: HashMap<(usize, usize), f32>, // (from, to) -> attention weight
}

impl GATLayer {
    pub fn new(input_dim: usize, output_dim: usize, num_heads: usize) -> Result<Self> {
        if input_dim == 0 || output_dim == 0 || num_heads == 0 {
            return Err(GeomError::LayerInitializationFailed {
                reason: "Dimensions and num_heads must be greater than 0".to_string(),
            });
        }

        if !output_dim.is_multiple_of(num_heads) {
            return Err(GeomError::LayerInitializationFailed {
                reason: "Output dimension must be divisible by number of heads".to_string(),
            });
        }

        let head_dim = output_dim / num_heads;

        // Xavier initialization
        let fan_in = input_dim as f32;
        let fan_out = output_dim as f32;
        let limit = (6.0 / (fan_in + fan_out)).sqrt();

        // Weight matrix for feature transformation
        let mut weight_data = Vec::with_capacity(input_dim * output_dim);
        for i in 0..(input_dim * output_dim) {
            let val = (((i * 7919 + 1337) % 10000) as f32 / 5000.0 - 1.0) * limit;
            weight_data.push(val);
        }
        let weight = Tensor::from_data(
            weight_data,
            [input_dim, output_dim],
            &rtx_tensor::Device::default(),
        )
        .unwrap();

        // Attention weight vector (2 * head_dim for concatenated features)
        let attention_dim = 2 * head_dim * num_heads;
        let mut attention_data = Vec::with_capacity(attention_dim);
        for i in 0..attention_dim {
            let val = (((i * 6151 + 2357) % 10000) as f32 / 5000.0 - 1.0) * limit;
            attention_data.push(val);
        }
        let attention_weight = Tensor::from_data(
            attention_data,
            [attention_dim],
            &rtx_tensor::Device::default(),
        )
        .unwrap();

        // Bias
        let bias = Tensor::zeros([output_dim], &rtx_tensor::Device::default()).unwrap();

        Ok(Self {
            weight,
            attention_weight,
            bias,
            input_dim,
            output_dim,
            num_heads,
            gradients_enabled: false,
            computed_gradients: false,
            attention_weights: HashMap::new(),
        })
    }

    pub fn get_attention_weights(&self) -> &HashMap<(usize, usize), f32> {
        &self.attention_weights
    }

    fn compute_attention(&mut self, graph: &Graph) -> Result<HashMap<usize, Vec<(usize, f32)>>> {
        let mut attention_map = HashMap::new();

        for node_id in graph.nodes() {
            let neighbors = graph.incoming_neighbors(node_id);
            let mut node_attentions = Vec::new();

            if !neighbors.is_empty() {
                let node_features = graph.node(node_id).unwrap().features();

                let _attention_scores: Vec<f32> = Vec::new();
                let mut raw_scores = Vec::new();

                // Compute raw attention scores for all neighbors
                for &neighbor_id in &neighbors {
                    let neighbor_features = graph.node(neighbor_id).unwrap().features();

                    // Compute attention score (simplified)
                    let node_2d = node_features.reshape([1, self.input_dim]).map_err(|e| {
                        GeomError::ForwardPassFailed {
                            reason: format!("Node feature reshape failed: {e}"),
                        }
                    })?;

                    let neighbor_2d =
                        neighbor_features
                            .reshape([1, self.input_dim])
                            .map_err(|e| GeomError::ForwardPassFailed {
                                reason: format!("Neighbor feature reshape failed: {e}"),
                            })?;

                    let node_transformed =
                        node_2d
                            .matmul(&self.weight)
                            .map_err(|e| GeomError::ForwardPassFailed {
                                reason: format!("Node transformation failed: {e}"),
                            })?;

                    let neighbor_transformed = neighbor_2d.matmul(&self.weight).map_err(|e| {
                        GeomError::ForwardPassFailed {
                            reason: format!("Neighbor transformation failed: {e}"),
                        }
                    })?;

                    // Simplified attention computation: dot product of first elements
                    let node_data = node_transformed.to_cpu().unwrap();
                    let neighbor_data = neighbor_transformed.to_cpu().unwrap();

                    let raw_score = if !node_data.is_empty() && !neighbor_data.is_empty() {
                        node_data[0] * neighbor_data[0] // Simplified attention mechanism
                    } else {
                        0.0
                    };

                    raw_scores.push(raw_score);
                }

                // Apply softmax to get attention weights
                let max_score = raw_scores.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
                let exp_scores: Vec<f32> =
                    raw_scores.iter().map(|&s| (s - max_score).exp()).collect();
                let sum_exp = exp_scores.iter().sum::<f32>();

                for (i, &neighbor_id) in neighbors.iter().enumerate() {
                    let attention_weight = if sum_exp > 0.0 {
                        exp_scores[i] / sum_exp
                    } else {
                        1.0 / neighbors.len() as f32
                    };

                    node_attentions.push((neighbor_id.index(), attention_weight));

                    // Store attention weights for inspection
                    self.attention_weights
                        .insert((node_id.index(), neighbor_id.index()), attention_weight);
                }
            }

            attention_map.insert(node_id.index(), node_attentions);
        }

        Ok(attention_map)
    }
}

impl GNNLayer for GATLayer {
    fn forward(&mut self, graph: &Graph) -> Result<Vec<Tensor>> {
        if graph.node_count() == 0 {
            return Err(GeomError::ForwardPassFailed {
                reason: "Cannot process empty graph".to_string(),
            });
        }

        // Compute attention weights
        let attention_map = self.compute_attention(graph)?;

        let mut outputs = Vec::with_capacity(graph.node_count());

        for node_id in graph.nodes() {
            let node_idx = node_id.index();

            if let Some(attentions) = attention_map.get(&node_idx) {
                if attentions.is_empty() {
                    // Isolated node - transform its own features
                    let node_features = graph.node(node_id).unwrap().features();
                    let node_2d = node_features.reshape([1, self.input_dim]).map_err(|e| {
                        GeomError::ForwardPassFailed {
                            reason: format!("Self feature reshape failed: {e}"),
                        }
                    })?;

                    let transformed =
                        node_2d
                            .matmul(&self.weight)
                            .map_err(|e| GeomError::ForwardPassFailed {
                                reason: format!("Self transformation failed: {e}"),
                            })?;
                    let biased =
                        transformed
                            .add(&self.bias)
                            .map_err(|e| GeomError::ForwardPassFailed {
                                reason: format!("Bias addition failed: {e}"),
                            })?;

                    // Reshape back to [output_dim] from [1, output_dim]
                    let output = biased.reshape([self.output_dim]).map_err(|e| {
                        GeomError::ForwardPassFailed {
                            reason: format!("Output reshape failed: {e}"),
                        }
                    })?;

                    outputs.push(output);
                } else {
                    // Aggregate neighbor features with attention weights
                    let mut aggregated = None;

                    for &(neighbor_idx, attention_weight) in attentions {
                        // Find neighbor node by index
                        let neighbor_id = graph
                            .nodes()
                            .into_iter()
                            .find(|id| id.index() == neighbor_idx)
                            .ok_or_else(|| GeomError::ForwardPassFailed {
                                reason: format!("Neighbor node {neighbor_idx} not found"),
                            })?;

                        let neighbor_features = graph.node(neighbor_id).unwrap().features();
                        let neighbor_2d =
                            neighbor_features
                                .reshape([1, self.input_dim])
                                .map_err(|e| GeomError::ForwardPassFailed {
                                    reason: format!("Neighbor feature reshape failed: {e}"),
                                })?;

                        let transformed = neighbor_2d.matmul(&self.weight).map_err(|e| {
                            GeomError::ForwardPassFailed {
                                reason: format!("Neighbor transformation failed: {e}"),
                            }
                        })?;

                        let weighted = transformed
                            .mul(
                                &Tensor::from_data(
                                    vec![attention_weight],
                                    [1],
                                    &rtx_tensor::Device::default(),
                                )
                                .unwrap(),
                            )
                            .map_err(|e| GeomError::ForwardPassFailed {
                                reason: format!("Attention weighting failed: {e}"),
                            })?;

                        aggregated = match aggregated {
                            None => Some(weighted),
                            Some(acc) => Some(acc.add(&weighted).map_err(|e| {
                                GeomError::ForwardPassFailed {
                                    reason: format!("Aggregation failed: {e}"),
                                }
                            })?),
                        };
                    }

                    let output = if let Some(agg) = aggregated {
                        agg.add(&self.bias)
                            .map_err(|e| GeomError::ForwardPassFailed {
                                reason: format!("Bias addition failed: {e}"),
                            })?
                    } else {
                        return Err(GeomError::ForwardPassFailed {
                            reason: "No aggregated features computed".to_string(),
                        });
                    };

                    // Apply ELU activation (approximated with ReLU for simplicity)
                    let activated = output.relu().map_err(|e| GeomError::ForwardPassFailed {
                        reason: format!("Activation failed: {e}"),
                    })?;

                    // Reshape back to [output_dim] from [1, output_dim]
                    let output_1d = activated.reshape([self.output_dim]).map_err(|e| {
                        GeomError::ForwardPassFailed {
                            reason: format!("Output reshape failed: {e}"),
                        }
                    })?;

                    outputs.push(output_1d);
                }
            } else {
                return Err(GeomError::ForwardPassFailed {
                    reason: format!("No attention computed for node {node_idx}"),
                });
            }
        }

        Ok(outputs)
    }

    fn parameters(&self) -> Vec<Tensor> {
        vec![
            self.weight.clone(),
            self.attention_weight.clone(),
            self.bias.clone(),
        ]
    }

    fn enable_gradients(&mut self, enabled: bool) {
        self.gradients_enabled = enabled;
    }

    fn has_computed_gradients(&self) -> bool {
        self.computed_gradients
    }

    fn backward(&mut self, grad_outputs: &[Tensor]) -> Result<Vec<Tensor>> {
        if !self.gradients_enabled {
            return Err(GeomError::BackwardPassFailed {
                reason: "Gradients are not enabled for this layer".to_string(),
            });
        }

        if grad_outputs.is_empty() {
            return Err(GeomError::BackwardPassFailed {
                reason: "No gradient outputs provided".to_string(),
            });
        }

        // Simplified gradient computation
        let weight_grad =
            Tensor::zeros(self.weight.shape().dims(), &rtx_tensor::Device::default()).unwrap();
        let attention_grad = Tensor::zeros(
            [self.attention_weight.shape().dims()[0]],
            &rtx_tensor::Device::default(),
        )
        .unwrap();
        let bias_grad = Tensor::zeros([self.output_dim], &rtx_tensor::Device::default()).unwrap();

        self.computed_gradients = true;

        Ok(vec![weight_grad, attention_grad, bias_grad])
    }
}