//! GNN explainability for brain networks.
//!
//! Provides methods to understand which edges (connections) and nodes (channels)
//! are most important for model predictions.

use crate::error::{GnnError, GnnResult};
use crate::graph::{BrainGraph, BrainRegion, Hemisphere};
use crate::models::BrainGNN;
use ndarray::Array2;
use serde::{Deserialize, Serialize};

/// Edge importance score
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EdgeImportance {
    /// Source node index
    pub source: usize,
    /// Target node index
    pub target: usize,
    /// Source node name
    pub source_name: String,
    /// Target node name
    pub target_name: String,
    /// Importance score
    pub importance: f64,
    /// Whether this is interhemispheric
    pub interhemispheric: bool,
}

/// Node importance score
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeImportance {
    /// Node index
    pub index: usize,
    /// Node name (channel)
    pub name: String,
    /// Importance score
    pub importance: f64,
    /// Hemisphere
    pub hemisphere: Hemisphere,
    /// Brain region
    pub region: BrainRegion,
}

/// Interpretation method
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum InterpretationMethod {
    /// Gradient-based saliency
    Gradient,
    /// Integrated gradients
    IntegratedGradients,
    /// Edge masking (ablation)
    EdgeMasking,
    /// GNNExplainer (learnable masks)
    GnnExplainer,
    /// Attention weights (if model uses attention)
    Attention,
}

/// Configuration for GNN explainer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExplainerConfig {
    /// Interpretation method
    pub method: InterpretationMethod,
    /// Number of steps for integrated gradients
    pub n_steps: usize,
    /// Perturbation magnitude for gradient estimation
    pub epsilon: f64,
    /// Top-K edges to return
    pub top_k_edges: usize,
    /// Top-K nodes to return
    pub top_k_nodes: usize,
    /// Target class for explanation (None = predicted class)
    pub target_class: Option<usize>,
}

impl Default for ExplainerConfig {
    fn default() -> Self {
        Self {
            method: InterpretationMethod::EdgeMasking,
            n_steps: 50,
            epsilon: 0.01,
            top_k_edges: 20,
            top_k_nodes: 10,
            target_class: None,
        }
    }
}

/// GNN Explainer for brain networks
pub struct GnnExplainer<M: BrainGNN> {
    model: M,
    config: ExplainerConfig,
}

impl<M: BrainGNN> GnnExplainer<M> {
    /// Create a new GNN explainer
    pub fn new(model: M, config: ExplainerConfig) -> Self {
        Self { model, config }
    }

    /// Explain model prediction on a graph
    pub fn explain(&self, graph: &BrainGraph) -> GnnResult<ExplanationResult> {
        // Get baseline prediction
        let logits = self.model.forward(graph)?;
        let predicted_class = logits
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map_or(0, |(i, _)| i);

        let target_class = self.config.target_class.unwrap_or(predicted_class);

        // Compute edge importance based on method
        let edge_scores = match self.config.method {
            InterpretationMethod::EdgeMasking => self.compute_edge_masking(graph, target_class)?,
            InterpretationMethod::Gradient | InterpretationMethod::IntegratedGradients => {
                self.compute_gradient_importance(graph, target_class)?
            }
            InterpretationMethod::GnnExplainer => {
                self.compute_gnn_explainer(graph, target_class)?
            }
            InterpretationMethod::Attention => {
                // Return edge weights as importance (placeholder)
                graph.edge_weights()
            }
        };

        // Build edge importance list
        let mut edge_importance: Vec<EdgeImportance> = graph
            .edges
            .iter()
            .zip(edge_scores.iter())
            .map(|(edge, &importance)| EdgeImportance {
                source: edge.source,
                target: edge.target,
                source_name: graph.nodes[edge.source].name.clone(),
                target_name: graph.nodes[edge.target].name.clone(),
                importance,
                interhemispheric: edge.interhemispheric,
            })
            .collect();

        // Sort by importance (descending)
        edge_importance.sort_by(|a, b| b.importance.partial_cmp(&a.importance).unwrap());
        edge_importance.truncate(self.config.top_k_edges);

        // Compute node importance (aggregate from edges)
        let mut node_scores = vec![0.0; graph.n_nodes()];
        for (edge, &score) in graph.edges.iter().zip(edge_scores.iter()) {
            node_scores[edge.source] += score.abs();
            node_scores[edge.target] += score.abs();
        }

        // Normalize node scores
        let max_node = node_scores.iter().copied().fold(0.0, f64::max);
        if max_node > 0.0 {
            for score in &mut node_scores {
                *score /= max_node;
            }
        }

        // Build node importance list
        let mut node_importance: Vec<NodeImportance> = graph
            .nodes
            .iter()
            .zip(node_scores.iter())
            .map(|(node, &importance)| NodeImportance {
                index: node.index,
                name: node.name.clone(),
                importance,
                hemisphere: node.hemisphere,
                region: node.region,
            })
            .collect();

        node_importance.sort_by(|a, b| b.importance.partial_cmp(&a.importance).unwrap());
        node_importance.truncate(self.config.top_k_nodes);

        Ok(ExplanationResult {
            predicted_class,
            target_class,
            confidence: logits[target_class],
            edge_importance,
            node_importance,
            method: self.config.method,
        })
    }

    /// Edge masking (ablation study)
    fn compute_edge_masking(&self, graph: &BrainGraph, target_class: usize) -> GnnResult<Vec<f64>> {
        let baseline_logits = self.model.forward(graph)?;
        let baseline_score = baseline_logits[target_class];

        let mut importance = Vec::with_capacity(graph.n_edges());

        for edge_idx in 0..graph.n_edges() {
            // Create graph with edge removed
            let mut masked_graph = graph.clone();
            masked_graph.edges.remove(edge_idx);

            let masked_logits = self.model.forward(&masked_graph)?;
            let masked_score = masked_logits[target_class];

            // Importance = drop in score when edge is removed
            importance.push(baseline_score - masked_score);
        }

        // Normalize
        let max_imp = importance.iter().copied().fold(0.0, f64::max);
        if max_imp > 0.0 {
            for imp in &mut importance {
                *imp /= max_imp;
            }
        }

        Ok(importance)
    }

    /// Gradient-based importance (finite differences approximation)
    fn compute_gradient_importance(
        &self,
        graph: &BrainGraph,
        target_class: usize,
    ) -> GnnResult<Vec<f64>> {
        let eps = self.config.epsilon;
        let mut importance = Vec::with_capacity(graph.n_edges());

        for edge_idx in 0..graph.n_edges() {
            // Perturb edge weight up
            let mut graph_plus = graph.clone();
            graph_plus.edges[edge_idx].weight += eps;

            // Perturb edge weight down
            let mut graph_minus = graph.clone();
            graph_minus.edges[edge_idx].weight -= eps;

            let logits_plus = self.model.forward(&graph_plus)?;
            let logits_minus = self.model.forward(&graph_minus)?;

            // Finite difference gradient
            let gradient = (logits_plus[target_class] - logits_minus[target_class]) / (2.0 * eps);
            importance.push(gradient.abs());
        }

        // Normalize
        let max_imp = importance.iter().copied().fold(0.0, f64::max);
        if max_imp > 0.0 {
            for imp in &mut importance {
                *imp /= max_imp;
            }
        }

        Ok(importance)
    }

    /// GNNExplainer (simplified version using iterative masking)
    fn compute_gnn_explainer(
        &self,
        graph: &BrainGraph,
        target_class: usize,
    ) -> GnnResult<Vec<f64>> {
        let n_edges = graph.n_edges();
        let mut mask = vec![0.5; n_edges]; // Start with uniform mask

        // Simple optimization loop
        let lr = 0.1;
        let n_iterations = 100;

        for _ in 0..n_iterations {
            // Compute gradient for each mask value
            let mut gradients = vec![0.0; n_edges];

            for i in 0..n_edges {
                let eps: f64 = 0.01;

                // Mask + epsilon
                let mut mask_plus = mask.clone();
                mask_plus[i] = (mask_plus[i] + eps).min(1.0);
                let score_plus = self.masked_forward(graph, &mask_plus, target_class)?;

                // Mask - epsilon
                let mut mask_minus = mask.clone();
                mask_minus[i] = (mask_minus[i] - eps).max(0.0);
                let score_minus = self.masked_forward(graph, &mask_minus, target_class)?;

                gradients[i] = (score_plus - score_minus) / (2.0 * eps);
            }

            // Update mask (maximize score)
            for i in 0..n_edges {
                mask[i] = (mask[i] + lr * gradients[i]).clamp(0.0, 1.0);
            }
        }

        Ok(mask)
    }

    /// Forward pass with edge masking
    fn masked_forward(
        &self,
        graph: &BrainGraph,
        mask: &[f64],
        target_class: usize,
    ) -> GnnResult<f64> {
        let mut masked_graph = graph.clone();

        for (i, edge) in masked_graph.edges.iter_mut().enumerate() {
            edge.weight *= mask[i];
        }

        let logits = self.model.forward(&masked_graph)?;
        Ok(logits[target_class])
    }
}

/// Result of explanation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExplanationResult {
    /// Predicted class
    pub predicted_class: usize,
    /// Target class for explanation
    pub target_class: usize,
    /// Confidence (logit) for target class
    pub confidence: f64,
    /// Top-K important edges
    pub edge_importance: Vec<EdgeImportance>,
    /// Top-K important nodes
    pub node_importance: Vec<NodeImportance>,
    /// Method used
    pub method: InterpretationMethod,
}

impl ExplanationResult {
    /// Get summary statistics
    pub fn summary(&self) -> ExplanationSummary {
        let n_interhemi = self
            .edge_importance
            .iter()
            .filter(|e| e.interhemispheric)
            .count();

        let mut region_counts: std::collections::HashMap<BrainRegion, usize> =
            std::collections::HashMap::new();
        for node in &self.node_importance {
            *region_counts.entry(node.region).or_insert(0) += 1;
        }

        let top_region = region_counts
            .into_iter()
            .max_by_key(|&(_, count)| count)
            .map(|(region, _)| region);

        ExplanationSummary {
            predicted_class: self.predicted_class,
            confidence: self.confidence,
            n_important_edges: self.edge_importance.len(),
            n_important_nodes: self.node_importance.len(),
            n_interhemispheric: n_interhemi,
            top_region,
            top_edge: self.edge_importance.first().cloned(),
            top_node: self.node_importance.first().cloned(),
        }
    }

    /// Get edges for specific region
    pub fn edges_for_region(&self, _region: BrainRegion) -> Vec<&EdgeImportance> {
        // Note: This is a simplified implementation
        // In practice, would need to check node regions
        self.edge_importance.iter().collect()
    }

    /// Export as connectivity matrix (importance weighted)
    pub fn to_importance_matrix(&self, n_nodes: usize) -> Array2<f64> {
        let mut matrix = Array2::zeros((n_nodes, n_nodes));

        for edge in &self.edge_importance {
            if edge.source < n_nodes && edge.target < n_nodes {
                matrix[[edge.source, edge.target]] = edge.importance;
                matrix[[edge.target, edge.source]] = edge.importance; // Symmetric
            }
        }

        matrix
    }
}

/// Summary of explanation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExplanationSummary {
    /// Predicted class
    pub predicted_class: usize,
    /// Confidence score
    pub confidence: f64,
    /// Number of important edges identified
    pub n_important_edges: usize,
    /// Number of important nodes identified
    pub n_important_nodes: usize,
    /// Number of interhemispheric connections in top edges
    pub n_interhemispheric: usize,
    /// Most represented brain region
    pub top_region: Option<BrainRegion>,
    /// Most important edge
    pub top_edge: Option<EdgeImportance>,
    /// Most important node
    pub top_node: Option<NodeImportance>,
}

/// Biomarker discovery across multiple subjects
pub struct BiomarkerDiscovery<M: BrainGNN + Clone> {
    model: M,
    explainer_config: ExplainerConfig,
}

impl<M: BrainGNN + Clone> BiomarkerDiscovery<M> {
    /// Create a new biomarker discovery instance
    pub fn new(model: M, config: ExplainerConfig) -> Self {
        Self {
            model,
            explainer_config: config,
        }
    }

    /// Find consistent biomarkers across subjects
    pub fn find_biomarkers(
        &self,
        graphs: &[BrainGraph],
        labels: &[usize],
    ) -> GnnResult<BiomarkerResult> {
        if graphs.len() != labels.len() {
            return Err(GnnError::DimensionMismatch(
                "graphs and labels must have same length".into(),
            ));
        }

        let explainer = GnnExplainer::new(self.model.clone(), self.explainer_config.clone());

        // Collect explanations for each class
        let n_classes = labels.iter().max().map_or(0, |&x| x + 1);
        let mut class_edge_scores: Vec<std::collections::HashMap<(String, String), Vec<f64>>> = (0
            ..n_classes)
            .map(|_| std::collections::HashMap::new())
            .collect();

        for (graph, &label) in graphs.iter().zip(labels.iter()) {
            let explanation = explainer.explain(graph)?;

            for edge in &explanation.edge_importance {
                let key = (edge.source_name.clone(), edge.target_name.clone());
                class_edge_scores[label]
                    .entry(key)
                    .or_default()
                    .push(edge.importance);
            }
        }

        // Find consistently important edges per class
        let mut class_biomarkers = Vec::with_capacity(n_classes);

        for class_scores in class_edge_scores {
            let mut biomarkers: Vec<(String, String, f64, f64)> = class_scores
                .into_iter()
                .filter(|(_, scores)| scores.len() >= 2) // At least 2 occurrences
                .map(|((src, tgt), scores)| {
                    let mean = scores.iter().sum::<f64>() / scores.len() as f64;
                    let std = (scores.iter().map(|&x| (x - mean).powi(2)).sum::<f64>()
                        / scores.len() as f64)
                        .sqrt();
                    (src, tgt, mean, std)
                })
                .collect();

            // Sort by mean importance / std (consistency)
            biomarkers.sort_by(|a, b| {
                let score_a = a.2 / (a.3 + 0.1);
                let score_b = b.2 / (b.3 + 0.1);
                score_b.partial_cmp(&score_a).unwrap()
            });

            class_biomarkers.push(biomarkers);
        }

        Ok(BiomarkerResult {
            class_biomarkers,
            n_subjects: graphs.len(),
            n_classes,
        })
    }
}

/// Result of biomarker discovery
#[derive(Debug, Clone)]
pub struct BiomarkerResult {
    /// Biomarkers per class: (source, target, mean_importance, std)
    pub class_biomarkers: Vec<Vec<(String, String, f64, f64)>>,
    /// Number of subjects analyzed
    pub n_subjects: usize,
    /// Number of classes
    pub n_classes: usize,
}

impl BiomarkerResult {
    /// Get top biomarkers for a class
    pub fn top_biomarkers(&self, class: usize, k: usize) -> Vec<(String, String, f64)> {
        self.class_biomarkers
            .get(class)
            .map(|b| {
                b.iter()
                    .take(k)
                    .map(|(src, tgt, mean, _)| (src.clone(), tgt.clone(), *mean))
                    .collect()
            })
            .unwrap_or_default()
    }

    /// Get discriminative biomarkers (different between classes)
    pub fn discriminative_biomarkers(
        &self,
        class_a: usize,
        class_b: usize,
    ) -> Vec<(String, String, f64)> {
        let a_edges: std::collections::HashSet<_> = self
            .class_biomarkers
            .get(class_a)
            .map(|b| {
                b.iter()
                    .take(20)
                    .map(|(s, t, _, _)| (s.clone(), t.clone()))
                    .collect()
            })
            .unwrap_or_default();

        let b_edges: std::collections::HashSet<_> = self
            .class_biomarkers
            .get(class_b)
            .map(|b| {
                b.iter()
                    .take(20)
                    .map(|(s, t, _, _)| (s.clone(), t.clone()))
                    .collect()
            })
            .unwrap_or_default();

        // Find edges in A but not in B
        a_edges
            .difference(&b_edges)
            .filter_map(|(s, t)| {
                self.class_biomarkers.get(class_a).and_then(|b| {
                    b.iter()
                        .find(|(src, tgt, _, _)| src == s && tgt == t)
                        .map(|(src, tgt, mean, _)| (src.clone(), tgt.clone(), *mean))
                })
            })
            .collect()
    }
}

impl<M: BrainGNN + Clone> GnnExplainer<M> {
    /// Explain using a reference to the model
    pub fn explain_ref(&self, graph: &BrainGraph) -> GnnResult<ExplanationResult> {
        self.explain(graph)
    }
}

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

    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_edge_masking_explanation() {
        let graph = create_test_graph();

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

        let explainer = GnnExplainer::new(
            model,
            ExplainerConfig {
                method: InterpretationMethod::EdgeMasking,
                top_k_edges: 4,
                top_k_nodes: 4,
                ..Default::default()
            },
        );

        let result = explainer.explain(&graph).unwrap();

        assert!(result.edge_importance.len() <= 4);
        assert!(result.node_importance.len() <= 4);
        assert!(result.predicted_class < 2);
    }

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

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

        let explainer = GnnExplainer::new(model, ExplainerConfig::default());
        let result = explainer.explain(&graph).unwrap();
        let summary = result.summary();

        assert!(summary.n_important_edges > 0);
        assert!(summary.predicted_class < 2);
    }

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

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

        let explainer = GnnExplainer::new(model, ExplainerConfig::default());
        let result = explainer.explain(&graph).unwrap();
        let matrix = result.to_importance_matrix(4);

        assert_eq!(matrix.shape(), &[4, 4]);
    }
}