//! Graph/Network Analysis for financial market structure.
//!
//! Implements:
//! - Correlation graph construction with thresholding
//! - Minimum spanning tree (Kruskal's algorithm with union-find)
//! - Community detection (Louvain greedy modularity optimization)
//! - Centrality metrics (betweenness, closeness, degree)
//! - Contagion risk propagation via shock cascades
//!
//! All algorithms from scratch in Rust -- no external crate dependencies beyond PyO3.

use pyo3::prelude::*;
use std::collections::VecDeque;

// ===========================================================================
// Frozen result types
// ===========================================================================

/// A node in the market graph with centrality metrics.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct GraphNode {
    /// Name of the asset/market.
    #[pyo3(get)]
    pub name: String,
    /// Number of edges connected to this node.
    #[pyo3(get)]
    pub degree: usize,
    /// Betweenness centrality (fraction of shortest paths through this node).
    #[pyo3(get)]
    pub betweenness: f64,
    /// Closeness centrality (inverse mean distance to all other nodes).
    #[pyo3(get)]
    pub closeness: f64,
}

#[pymethods]
impl GraphNode {
    fn __repr__(&self) -> String {
        format!(
            "GraphNode('{}', degree={}, betweenness={:.4}, closeness={:.4})",
            self.name, self.degree, self.betweenness, self.closeness
        )
    }
}

/// An edge in the market graph.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct GraphEdge {
    /// Source node name.
    #[pyo3(get)]
    pub source: String,
    /// Target node name.
    #[pyo3(get)]
    pub target: String,
    /// Edge weight (correlation or distance).
    #[pyo3(get)]
    pub weight: f64,
}

#[pymethods]
impl GraphEdge {
    fn __repr__(&self) -> String {
        format!(
            "GraphEdge('{}' -> '{}', weight={:.4})",
            self.source, self.target, self.weight
        )
    }
}

/// A community detected in the market graph.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct Community {
    /// Community identifier.
    #[pyo3(get)]
    pub id: usize,
    /// Names of assets in this community.
    #[pyo3(get)]
    pub members: Vec<String>,
    /// Modularity contribution of this community.
    #[pyo3(get)]
    pub modularity: f64,
}

#[pymethods]
impl Community {
    fn __repr__(&self) -> String {
        format!(
            "Community(id={}, members={:?}, modularity={:.4})",
            self.id, self.members, self.modularity
        )
    }
}

/// Full market graph with nodes, edges, communities, and MST.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct MarketGraph {
    /// All nodes with centrality metrics.
    #[pyo3(get)]
    pub nodes: Vec<GraphNode>,
    /// All edges above the correlation threshold.
    #[pyo3(get)]
    pub edges: Vec<GraphEdge>,
    /// Detected communities.
    #[pyo3(get)]
    pub communities: Vec<Community>,
    /// Minimum spanning tree edges.
    #[pyo3(get)]
    pub mst_edges: Vec<GraphEdge>,
    /// Graph density (fraction of possible edges present).
    #[pyo3(get)]
    pub density: f64,
}

#[pymethods]
impl MarketGraph {
    fn __repr__(&self) -> String {
        format!(
            "MarketGraph(nodes={}, edges={}, communities={}, density={:.4})",
            self.nodes.len(),
            self.edges.len(),
            self.communities.len(),
            self.density
        )
    }
}

// ===========================================================================
// Private helpers
// ===========================================================================

/// Clamp a value to finite, returning 0.0 for NaN/Inf.
#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

// ---------------------------------------------------------------------------
// Union-Find (Disjoint Set Union)
// ---------------------------------------------------------------------------

struct UnionFind {
    parent: Vec<usize>,
    rank: Vec<usize>,
}

impl UnionFind {
    fn new(n: usize) -> Self {
        Self {
            parent: (0..n).collect(),
            rank: vec![0; n],
        }
    }

    fn find(&mut self, x: usize) -> usize {
        if self.parent[x] != x {
            self.parent[x] = self.find(self.parent[x]); // path compression
        }
        self.parent[x]
    }

    fn union(&mut self, x: usize, y: usize) -> bool {
        let rx = self.find(x);
        let ry = self.find(y);
        if rx == ry {
            return false; // already in same set
        }
        // union by rank
        if self.rank[rx] < self.rank[ry] {
            self.parent[rx] = ry;
        } else if self.rank[rx] > self.rank[ry] {
            self.parent[ry] = rx;
        } else {
            self.parent[ry] = rx;
            self.rank[rx] += 1;
        }
        true
    }
}

// ---------------------------------------------------------------------------
// BFS shortest paths for betweenness and closeness centrality
// ---------------------------------------------------------------------------

/// Compute unweighted betweenness and closeness centrality using BFS from every node.
/// adj is an adjacency list. Returns (betweenness[n], closeness[n]).
fn compute_centrality(adj: &[Vec<usize>], n: usize) -> (Vec<f64>, Vec<f64>) {
    let mut betweenness = vec![0.0f64; n];
    let mut closeness = vec![0.0f64; n];

    for s in 0..n {
        // BFS from node s (Brandes' algorithm for betweenness)
        let mut dist = vec![-1i32; n];
        let mut sigma = vec![0.0f64; n]; // number of shortest paths
        let mut pred: Vec<Vec<usize>> = vec![vec![]; n];
        let mut stack = Vec::new();
        let mut queue = VecDeque::new();

        dist[s] = 0;
        sigma[s] = 1.0;
        queue.push_back(s);

        while let Some(v) = queue.pop_front() {
            stack.push(v);
            for &w in &adj[v] {
                if dist[w] < 0 {
                    // First visit
                    dist[w] = dist[v] + 1;
                    queue.push_back(w);
                }
                if dist[w] == dist[v] + 1 {
                    sigma[w] += sigma[v];
                    pred[w].push(v);
                }
            }
        }

        // Closeness: inverse of mean distance
        let total_dist: f64 = dist.iter().filter(|&&d| d > 0).map(|&d| d as f64).sum();
        let reachable = dist.iter().filter(|&&d| d > 0).count();
        if reachable > 0 && total_dist > 0.0 {
            closeness[s] = finite_or_zero(reachable as f64 / total_dist);
        }

        // Accumulate betweenness
        let mut delta = vec![0.0f64; n];
        while let Some(w) = stack.pop() {
            if w == s {
                continue;
            }
            for &v in &pred[w] {
                if sigma[w] > 0.0 {
                    delta[v] += (sigma[v] / sigma[w]) * (1.0 + delta[w]);
                }
            }
            betweenness[w] += delta[w];
        }
    }

    // Normalize betweenness: divide by 2 for undirected graphs,
    // then by (n-1)(n-2) for the standard normalization
    let norm = if n > 2 {
        ((n - 1) * (n - 2)) as f64
    } else {
        1.0
    };
    for b in &mut betweenness {
        *b = finite_or_zero(*b / norm);
    }

    // Normalize closeness by (n-1) for standard definition
    if n > 1 {
        let cn = (n - 1) as f64;
        for c in &mut closeness {
            *c = finite_or_zero(*c / cn * cn); // closeness is already scaled by reachable/dist
        }
    }

    (betweenness, closeness)
}

// ---------------------------------------------------------------------------
// Modularity computation for Louvain
// ---------------------------------------------------------------------------

/// Compute the modularity of a partition.
/// Q = (1/2m) * sum_ij [ A_ij - k_i*k_j/(2m) ] * delta(c_i, c_j)
#[allow(dead_code)]
fn modularity(adj_matrix: &[Vec<f64>], communities: &[usize]) -> f64 {
    let n = adj_matrix.len();
    if n == 0 {
        return 0.0;
    }

    // Total weight (2m)
    let mut total_weight = 0.0;
    for i in 0..n {
        for j in 0..n {
            total_weight += adj_matrix[i][j];
        }
    }
    if total_weight < 1e-15 {
        return 0.0;
    }

    // Degree of each node
    let degrees: Vec<f64> = (0..n)
        .map(|i| adj_matrix[i].iter().sum::<f64>())
        .collect();

    let mut q = 0.0;
    for i in 0..n {
        for j in 0..n {
            if communities[i] == communities[j] {
                q += adj_matrix[i][j] - degrees[i] * degrees[j] / total_weight;
            }
        }
    }
    q /= total_weight;
    finite_or_zero(q)
}

/// Louvain community detection (greedy modularity optimization).
/// Phase 1: greedily move nodes between communities.
/// Returns community assignment for each node.
fn louvain_phase1(adj_matrix: &[Vec<f64>]) -> Vec<usize> {
    let n = adj_matrix.len();
    if n == 0 {
        return vec![];
    }

    // Each node starts in its own community
    let mut comm: Vec<usize> = (0..n).collect();

    // Total weight (2m)
    let mut total_weight = 0.0;
    for i in 0..n {
        for j in 0..n {
            total_weight += adj_matrix[i][j];
        }
    }
    if total_weight < 1e-15 {
        return comm;
    }

    // Degree of each node
    let degrees: Vec<f64> = (0..n)
        .map(|i| adj_matrix[i].iter().sum::<f64>())
        .collect();

    // Sum of weights within each community
    let mut sigma_tot: Vec<f64> = degrees.clone();

    let max_iters = 100;
    for _iter in 0..max_iters {
        let mut improved = false;

        for i in 0..n {
            let current_comm = comm[i];
            let ki = degrees[i];

            // Compute delta-Q for moving node i to each neighbor's community
            // First, remove i from its community
            let mut ki_in_current = 0.0;
            for j in 0..n {
                if comm[j] == current_comm && j != i {
                    ki_in_current += adj_matrix[i][j];
                }
            }

            let mut best_comm = current_comm;
            let mut best_delta = 0.0;

            // Gather candidate communities from neighbors
            let mut candidate_comms: Vec<usize> = Vec::new();
            for j in 0..n {
                if adj_matrix[i][j] > 0.0 && comm[j] != current_comm {
                    if !candidate_comms.contains(&comm[j]) {
                        candidate_comms.push(comm[j]);
                    }
                }
            }

            for &c in &candidate_comms {
                // Sum of edges from i to nodes in community c
                let mut ki_in_c = 0.0;
                for j in 0..n {
                    if comm[j] == c {
                        ki_in_c += adj_matrix[i][j];
                    }
                }

                let sigma_c = sigma_tot[c];

                // Delta Q = [ki_in_c / m - sigma_c * ki / (2m^2)] - [ki_in_current / m - (sigma_current - ki) * ki / (2m^2)]
                let m = total_weight;
                let delta_q = (ki_in_c - ki_in_current) / m
                    - ki * (sigma_c - (sigma_tot[current_comm] - ki)) / (m * m);

                if delta_q > best_delta {
                    best_delta = delta_q;
                    best_comm = c;
                }
            }

            if best_comm != current_comm {
                // Move node i from current_comm to best_comm
                sigma_tot[current_comm] -= ki;
                sigma_tot[best_comm] += ki;
                comm[i] = best_comm;
                improved = true;
            }
        }

        if !improved {
            break;
        }
    }

    // Renumber communities to be contiguous 0, 1, 2, ...
    let mut mapping = std::collections::HashMap::new();
    let mut next_id = 0usize;
    for c in &mut comm {
        let entry = mapping.entry(*c).or_insert_with(|| {
            let id = next_id;
            next_id += 1;
            id
        });
        *c = *entry;
    }

    comm
}

// ===========================================================================
// Public functions
// ===========================================================================

/// Build a correlation graph from a correlation matrix.
///
/// Creates a graph where nodes are assets/markets and edges represent
/// correlations above the given threshold. Computes centrality metrics
/// (degree, betweenness, closeness), detects communities via Louvain,
/// and finds the minimum spanning tree.
///
/// Args:
///     corr_matrix: NxN correlation matrix (values in [-1, 1]).
///     names: List of N asset/market names.
///     threshold: Minimum absolute correlation to create an edge (0 to 1).
///
/// Returns:
///     MarketGraph with full analysis results.
#[pyfunction]
pub fn build_correlation_graph(
    corr_matrix: Vec<Vec<f64>>,
    names: Vec<String>,
    threshold: f64,
) -> PyResult<MarketGraph> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let n = names.len();
    if n == 0 {
        return Ok(MarketGraph {
            nodes: vec![],
            edges: vec![],
            communities: vec![],
            mst_edges: vec![],
            density: 0.0,
        });
    }
    if corr_matrix.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "correlation matrix has {} rows, expected {}",
            corr_matrix.len(),
            n
        )));
    }
    for (i, row) in corr_matrix.iter().enumerate() {
        if row.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "correlation matrix row {} has {} columns, expected {}",
                i,
                row.len(),
                n
            )));
        }
    }

    // Validate finite values
    for i in 0..n {
        for j in 0..n {
            if !corr_matrix[i][j].is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "corr_matrix[{}][{}] is not finite",
                    i, j
                )));
            }
        }
    }

    let threshold = threshold.abs().clamp(0.0, 1.0);

    // Build adjacency list and edges
    let mut adj: Vec<Vec<usize>> = vec![vec![]; n];
    let mut edges = Vec::new();

    // Build weighted adjacency matrix for community detection (absolute correlation above threshold)
    let mut adj_weighted = vec![vec![0.0f64; n]; n];

    for i in 0..n {
        for j in (i + 1)..n {
            let corr = corr_matrix[i][j];
            if corr.abs() >= threshold {
                adj[i].push(j);
                adj[j].push(i);
                edges.push(GraphEdge {
                    source: names[i].clone(),
                    target: names[j].clone(),
                    weight: finite_or_zero(corr),
                });
                adj_weighted[i][j] = corr.abs();
                adj_weighted[j][i] = corr.abs();
            }
        }
    }

    // Compute centrality metrics
    let (betweenness, closeness) = compute_centrality(&adj, n);

    // Build nodes
    let nodes: Vec<GraphNode> = (0..n)
        .map(|i| GraphNode {
            name: names[i].clone(),
            degree: adj[i].len(),
            betweenness: betweenness[i],
            closeness: closeness[i],
        })
        .collect();

    // Compute density
    let max_edges = if n > 1 { n * (n - 1) / 2 } else { 1 };
    let density = finite_or_zero(edges.len() as f64 / max_edges as f64);

    // Detect communities
    let comm_assignments = louvain_phase1(&adj_weighted);
    let n_communities = comm_assignments.iter().copied().max().map(|m| m + 1).unwrap_or(0);

    let communities: Vec<Community> = (0..n_communities)
        .map(|c| {
            let members: Vec<String> = (0..n)
                .filter(|&i| comm_assignments[i] == c)
                .map(|i| names[i].clone())
                .collect();

            // Compute modularity contribution for this community
            let comm_mod = community_modularity(&adj_weighted, &comm_assignments, c);

            Community {
                id: c,
                members,
                modularity: finite_or_zero(comm_mod),
            }
        })
        .filter(|c| !c.members.is_empty())
        .collect();

    // Compute MST
    let mst_edges = compute_mst(&corr_matrix, &names)?;

    Ok(MarketGraph {
        nodes,
        edges,
        communities,
        mst_edges,
        density,
    })
}

/// Compute the modularity contribution of a single community.
fn community_modularity(adj_matrix: &[Vec<f64>], communities: &[usize], target: usize) -> f64 {
    let n = adj_matrix.len();
    let mut total_weight = 0.0;
    for i in 0..n {
        for j in 0..n {
            total_weight += adj_matrix[i][j];
        }
    }
    if total_weight < 1e-15 {
        return 0.0;
    }

    let degrees: Vec<f64> = (0..n)
        .map(|i| adj_matrix[i].iter().sum::<f64>())
        .collect();

    let mut q = 0.0;
    for i in 0..n {
        if communities[i] != target {
            continue;
        }
        for j in 0..n {
            if communities[j] == target {
                q += adj_matrix[i][j] - degrees[i] * degrees[j] / total_weight;
            }
        }
    }
    finite_or_zero(q / total_weight)
}

/// Compute minimum spanning tree from a correlation matrix using Kruskal's algorithm.
///
/// Transforms correlations to distances: d(i,j) = sqrt(2 * (1 - corr(i,j)))
/// and finds the MST using union-find.
///
/// Args:
///     corr_matrix: NxN correlation matrix.
///     names: List of N asset/market names.
///
/// Returns:
///     List of GraphEdge forming the minimum spanning tree.
#[pyfunction]
pub fn minimum_spanning_tree(
    corr_matrix: Vec<Vec<f64>>,
    names: Vec<String>,
) -> PyResult<Vec<GraphEdge>> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let n = names.len();
    if n == 0 {
        return Ok(vec![]);
    }
    if corr_matrix.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "correlation matrix rows must match names length",
        ));
    }
    for (i, row) in corr_matrix.iter().enumerate() {
        if row.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "correlation matrix row {} has incorrect length",
                i
            )));
        }
    }

    compute_mst(&corr_matrix, &names)
}

/// Internal MST computation.
fn compute_mst(corr_matrix: &[Vec<f64>], names: &[String]) -> PyResult<Vec<GraphEdge>> {
    let n = names.len();
    if n <= 1 {
        return Ok(vec![]);
    }

    // Build edge list with distances
    let mut edge_list: Vec<(f64, usize, usize)> = Vec::with_capacity(n * (n - 1) / 2);
    for i in 0..n {
        for j in (i + 1)..n {
            let corr = corr_matrix[i][j];
            // Distance: d = sqrt(2 * (1 - corr)), clamp corr to [-1, 1]
            let clamped = corr.clamp(-1.0, 1.0);
            let dist = (2.0 * (1.0 - clamped)).max(0.0).sqrt();
            let dist = finite_or_zero(dist);
            edge_list.push((dist, i, j));
        }
    }

    // Sort by distance (ascending)
    edge_list.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

    // Kruskal's algorithm with union-find
    let mut uf = UnionFind::new(n);
    let mut mst = Vec::with_capacity(n - 1);

    for (_dist, i, j) in edge_list {
        if uf.union(i, j) {
            mst.push(GraphEdge {
                source: names[i].clone(),
                target: names[j].clone(),
                weight: finite_or_zero(corr_matrix[i][j]),
            });
            if mst.len() == n - 1 {
                break;
            }
        }
    }

    Ok(mst)
}

/// Detect communities in an adjacency matrix using Louvain modularity optimization.
///
/// Args:
///     adj_matrix: NxN weighted adjacency matrix (non-negative weights).
///     names: List of N node names.
///
/// Returns:
///     List of Community objects with member assignments and modularity.
#[pyfunction]
pub fn detect_communities(
    adj_matrix: Vec<Vec<f64>>,
    names: Vec<String>,
) -> PyResult<Vec<Community>> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let n = names.len();
    if n == 0 {
        return Ok(vec![]);
    }
    if adj_matrix.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "adjacency matrix rows must match names length",
        ));
    }
    for (i, row) in adj_matrix.iter().enumerate() {
        if row.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "adjacency matrix row {} has incorrect length",
                i
            )));
        }
    }

    // Validate finite, non-negative
    for i in 0..n {
        for j in 0..n {
            if !adj_matrix[i][j].is_finite() {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "adj_matrix[{}][{}] is not finite",
                    i, j
                )));
            }
        }
    }

    let comm_assignments = louvain_phase1(&adj_matrix);
    let n_communities = comm_assignments.iter().copied().max().map(|m| m + 1).unwrap_or(0);

    let communities: Vec<Community> = (0..n_communities)
        .map(|c| {
            let members: Vec<String> = (0..n)
                .filter(|&i| comm_assignments[i] == c)
                .map(|i| names[i].clone())
                .collect();

            let comm_mod = community_modularity(&adj_matrix, &comm_assignments, c);

            Community {
                id: c,
                members,
                modularity: finite_or_zero(comm_mod),
            }
        })
        .filter(|c| !c.members.is_empty())
        .collect();

    Ok(communities)
}

/// Simulate contagion/shock propagation through a network.
///
/// Starting from a shocked node, propagates the impact through the network
/// where each step multiplies the shock by the edge weight times a decay factor.
/// Uses BFS to model waves of contagion.
///
/// Args:
///     adj_matrix: NxN weighted adjacency matrix.
///     names: List of N node names.
///     shock_node: Name of the initially shocked node.
///     decay: Decay factor per hop (0 to 1). Lower = faster attenuation.
///
/// Returns:
///     List of (node_name, impact) tuples sorted by impact descending.
#[pyfunction]
pub fn contagion_risk(
    adj_matrix: Vec<Vec<f64>>,
    names: Vec<String>,
    shock_node: String,
    decay: f64,
) -> PyResult<Vec<(String, f64)>> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let n = names.len();
    if n == 0 {
        return Ok(vec![]);
    }
    if adj_matrix.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "adjacency matrix rows must match names length",
        ));
    }
    for (i, row) in adj_matrix.iter().enumerate() {
        if row.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "adjacency matrix row {} has incorrect length",
                i
            )));
        }
    }

    let decay = decay.clamp(0.0, 1.0);

    // Find the shock node index
    let shock_idx = names
        .iter()
        .position(|name| *name == shock_node)
        .ok_or_else(|| {
            pyo3::exceptions::PyValueError::new_err(format!(
                "shock_node '{}' not found in names",
                shock_node
            ))
        })?;

    // BFS propagation
    let mut impact = vec![0.0f64; n];
    impact[shock_idx] = 1.0;
    let mut visited = vec![false; n];
    visited[shock_idx] = true;

    let mut queue = VecDeque::new();
    queue.push_back((shock_idx, 1.0f64)); // (node, current_shock)

    while let Some((node, shock)) = queue.pop_front() {
        for j in 0..n {
            if j == node {
                continue;
            }
            let weight = adj_matrix[node][j];
            if weight <= 0.0 || !weight.is_finite() {
                continue;
            }

            let propagated = shock * weight * decay;
            if propagated < 1e-10 {
                continue;
            }

            // Accumulate impact even if visited (multiple paths)
            impact[j] = finite_or_zero(impact[j] + propagated);

            if !visited[j] {
                visited[j] = true;
                queue.push_back((j, propagated));
            }
        }
    }

    // Build result sorted by impact descending
    let mut result: Vec<(String, f64)> = (0..n)
        .map(|i| (names[i].clone(), finite_or_zero(impact[i])))
        .collect();
    result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    Ok(result)
}

/// Register graph analysis types and functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<GraphNode>()?;
    m.add_class::<GraphEdge>()?;
    m.add_class::<Community>()?;
    m.add_class::<MarketGraph>()?;
    m.add_function(wrap_pyfunction!(build_correlation_graph, m)?)?;
    m.add_function(wrap_pyfunction!(minimum_spanning_tree, m)?)?;
    m.add_function(wrap_pyfunction!(detect_communities, m)?)?;
    m.add_function(wrap_pyfunction!(contagion_risk, m)?)?;
    Ok(())
}

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

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

    fn sample_corr_3x3() -> Vec<Vec<f64>> {
        vec![
            vec![1.0, 0.8, 0.2],
            vec![0.8, 1.0, 0.3],
            vec![0.2, 0.3, 1.0],
        ]
    }

    fn sample_names_3() -> Vec<String> {
        vec!["A".to_string(), "B".to_string(), "C".to_string()]
    }

    #[test]
    fn test_finite_or_zero() {
        assert_eq!(finite_or_zero(1.5), 1.5);
        assert_eq!(finite_or_zero(f64::NAN), 0.0);
        assert_eq!(finite_or_zero(f64::INFINITY), 0.0);
    }

    #[test]
    fn test_union_find() {
        let mut uf = UnionFind::new(5);
        assert!(uf.union(0, 1));
        assert!(uf.union(2, 3));
        assert!(!uf.union(0, 1)); // already connected
        assert!(uf.union(1, 2));  // merge the two components
        assert_eq!(uf.find(0), uf.find(3)); // now all connected
    }

    #[test]
    fn test_mst_simple() {
        let corr = sample_corr_3x3();
        let names = sample_names_3();
        let mst = minimum_spanning_tree(corr, names).unwrap();
        // MST of 3 nodes has 2 edges
        assert_eq!(mst.len(), 2);
    }

    #[test]
    fn test_mst_empty() {
        let result = minimum_spanning_tree(vec![], vec![]).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_mst_single_node() {
        let result = minimum_spanning_tree(vec![vec![1.0]], vec!["X".to_string()]).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_mst_dimension_mismatch() {
        let result = minimum_spanning_tree(
            vec![vec![1.0, 0.5], vec![0.5, 1.0]],
            vec!["A".to_string()],
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_build_correlation_graph() {
        let corr = sample_corr_3x3();
        let names = sample_names_3();
        let graph = build_correlation_graph(corr, names, 0.5).unwrap();
        assert_eq!(graph.nodes.len(), 3);
        // Only A-B (0.8) should be above threshold 0.5
        // A-C (0.2) and B-C (0.3) are below
        assert_eq!(graph.edges.len(), 1);
        assert!(graph.density > 0.0);
    }

    #[test]
    fn test_build_correlation_graph_low_threshold() {
        let corr = sample_corr_3x3();
        let names = sample_names_3();
        let graph = build_correlation_graph(corr, names, 0.1).unwrap();
        // All correlations above 0.1
        assert_eq!(graph.edges.len(), 3);
    }

    #[test]
    fn test_build_correlation_graph_empty() {
        let graph = build_correlation_graph(vec![], vec![], 0.5).unwrap();
        assert!(graph.nodes.is_empty());
    }

    #[test]
    fn test_detect_communities_basic() {
        // Two clear groups: {A, B} highly connected, {C, D} highly connected
        let adj = vec![
            vec![0.0, 0.9, 0.1, 0.1],
            vec![0.9, 0.0, 0.1, 0.1],
            vec![0.1, 0.1, 0.0, 0.9],
            vec![0.1, 0.1, 0.9, 0.0],
        ];
        let names = vec![
            "A".to_string(),
            "B".to_string(),
            "C".to_string(),
            "D".to_string(),
        ];
        let comms = detect_communities(adj, names).unwrap();
        assert!(comms.len() >= 1); // Should detect at least 1 community
    }

    #[test]
    fn test_detect_communities_empty() {
        let comms = detect_communities(vec![], vec![]).unwrap();
        assert!(comms.is_empty());
    }

    #[test]
    fn test_contagion_risk_basic() {
        let adj = vec![
            vec![0.0, 0.8, 0.2],
            vec![0.8, 0.0, 0.5],
            vec![0.2, 0.5, 0.0],
        ];
        let names = sample_names_3();
        let result = contagion_risk(adj, names, "A".to_string(), 0.9).unwrap();
        assert_eq!(result.len(), 3);
        // A should have impact 1.0 (initial shock)
        let a_impact = result.iter().find(|(n, _)| n == "A").unwrap();
        assert!((a_impact.1 - 1.0).abs() < 1e-10);
        // B should have higher impact than C (directly connected with weight 0.8 vs 0.2)
        let b_impact = result.iter().find(|(n, _)| n == "B").unwrap().1;
        let c_impact = result.iter().find(|(n, _)| n == "C").unwrap().1;
        assert!(b_impact > c_impact);
    }

    #[test]
    fn test_contagion_risk_unknown_node() {
        let adj = vec![vec![0.0, 0.5], vec![0.5, 0.0]];
        let names = vec!["A".to_string(), "B".to_string()];
        let result = contagion_risk(adj, names, "Z".to_string(), 0.5);
        assert!(result.is_err());
    }

    #[test]
    fn test_contagion_zero_decay() {
        let adj = vec![
            vec![0.0, 0.8],
            vec![0.8, 0.0],
        ];
        let names = vec!["A".to_string(), "B".to_string()];
        let result = contagion_risk(adj, names, "A".to_string(), 0.0).unwrap();
        // With decay=0, no propagation beyond the initial node
        let b_impact = result.iter().find(|(n, _)| n == "B").unwrap().1;
        assert!(b_impact < 1e-10);
    }

    #[test]
    fn test_modularity_single_community() {
        // All nodes in one community
        let adj = vec![
            vec![0.0, 1.0, 1.0],
            vec![1.0, 0.0, 1.0],
            vec![1.0, 1.0, 0.0],
        ];
        let comm = vec![0, 0, 0];
        let q = modularity(&adj, &comm);
        // All in one community: Q = 0 (no improvement over random)
        assert!((q).abs() < 1e-10, "q={}", q);
    }

    #[test]
    fn test_centrality_line_graph() {
        // A -- B -- C (line graph)
        let adj: Vec<Vec<usize>> = vec![vec![1], vec![0, 2], vec![1]];
        let (betweenness, closeness) = compute_centrality(&adj, 3);
        // B has highest betweenness (all A-C paths go through B)
        assert!(betweenness[1] > betweenness[0]);
        assert!(betweenness[1] > betweenness[2]);
        // B has highest closeness
        assert!(closeness[1] >= closeness[0]);
    }

    #[test]
    fn test_mst_highest_corr_first() {
        // MST should connect highest correlation edges first (shortest distance)
        let corr = vec![
            vec![1.0, 0.95, 0.1],
            vec![0.95, 1.0, 0.5],
            vec![0.1, 0.5, 1.0],
        ];
        let names = sample_names_3();
        let mst = minimum_spanning_tree(corr, names).unwrap();
        // First MST edge should be A-B (highest corr = 0.95)
        assert_eq!(mst[0].source, "A");
        assert_eq!(mst[0].target, "B");
    }
}
