subsume/

sheaf.rs

//! # Sheaf Diffusion
//!
//! Algebraic structures for enforcing transitivity and consistency in graphs.
//!
//! A **sheaf** on a graph assigns vector spaces (stalks) to nodes and linear maps
//! (restriction maps) to edges. The key insight: if data is "consistent" across
//! the graph, the Dirichlet energy is zero. Non-zero energy indicates inconsistency.
//!
//! This module provides the mathematical primitives (sheaf Laplacian, Euler-step
//! diffusion) from Hansen & Ghrist (2019) and Bodnar et al. (ICLR 2022).
//! It does **not** implement learnable restriction maps or neural architectures --
//! those would be built on top of these primitives.
//!
//! # Why Sheaves for Coreference?
//!
//! Coreference requires transitivity: if A=B and B=C, then A=C.
//! Traditional approaches enforce this post-hoc (transitive closure).
//! A sheaf Laplacian can enforce it structurally: diffusion drives stalks toward
//! consistency across the graph.
//!
//! ```text
//! Mention A ──[restriction]──> Mention B ──[restriction]──> Mention C
//!     │                            │                            │
//!   stalk_A                     stalk_B                      stalk_C
//!     │                            │                            │
//!     └── If A=B=C, stalks should be "compatible" under restrictions
//! ```
//!
//! # Mathematical Background
//!
//! Given a graph G = (V, E), a **cellular sheaf** F assigns:
//! - To each vertex v: a vector space F(v) (the "stalk")
//! - To each edge e = (u,v): a linear map F_{u←e}: F(u) → F(e) (restriction)
//!
//! The **sheaf Laplacian** L_F is defined as:
//!
//! L_F = δ^T · δ
//!
//! where δ is the coboundary operator. The **Dirichlet energy** is:
//!
//! E(x) = x^T L_F x = Σ_{(u,v) ∈ E} ||F_{u←e}(x_u) - F_{v←e}(x_v)||²
//!
//! Low energy means the signal x is "consistent" across the sheaf.
//!
//! # Relationship to Graph Neural Networks
//!
//! | Model | Message | Aggregation | Transitivity |
//! |-------|---------|-------------|--------------|
//! | GCN | Identity | Sum | Implicit |
//! | GAT | Attention-weighted | Sum | Implicit |
//! | Sheaf | Restriction maps | Laplacian diffusion | **Explicit** |
//!
//! Sheaf neural networks generalize GNNs by learning edge-specific linear maps
//! instead of using identity or scalar weights.
//!
//! # Implementation
//!
//! This module provides framework-agnostic traits. Implementations live in:
//! - `subsume`: CPU-based with ndarray
//! - `subsume-candle`: GPU-accelerated with candle
//!
//! # References
//!
//! - Hansen & Ghrist (2019): "Toward a spectral theory of cellular sheaves"
//! - Bodnar et al. (2022): "Neural Sheaf Diffusion" (ICLR)
//! - Barbero et al. (2022): "Sheaf Neural Networks with Connection Laplacians"
//! - Bodnar (2023): "Topological Deep Learning: Graphs, Complexes, Sheaves"
//!   (Cambridge PhD thesis) -- connects sheaf structure to asymptotic behavior of
//!   message passing, providing theoretical grounding
//! - Zaghen (2024): "Nonlinear Sheaf Diffusion in Graph Neural Networks" -- introduces
//!   nonlinear Laplacians for sheaf diffusion; the current linear restriction maps could
//!   be extended with nonlinear variants for heterophilic graphs
//! - Hu (2026): "Sheaf-Theoretic and Topological Perspective on Complex Network Modeling"
//!   -- comprehensive survey of sheaf neural networks and sheaf attention mechanisms

use std::collections::HashMap;
use std::fmt::Debug;

/// Error type for sheaf operations.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, thiserror::Error)]
pub enum SheafError {
    /// Node not found in the graph.
    #[error("Node {0} not found")]
    NodeNotFound(usize),
    /// Edge not found in the graph.
    #[error("Edge ({0}, {1}) not found")]
    EdgeNotFound(usize, usize),
    /// Dimension mismatch in linear map.
    #[error("Dimension mismatch: expected {expected}, got {actual}")]
    DimensionMismatch {
        /// Expected dimension.
        expected: usize,
        /// Actual dimension.
        actual: usize,
    },
    /// Invalid restriction map.
    #[error("Invalid restriction: {0}")]
    InvalidRestriction(String),
}

/// A restriction map (linear transformation) on an edge.
///
/// For edge (u, v), this maps from the stalk at u to the edge space.
/// The restriction map captures "how information flows" along the edge.
///
/// For coreference: if mentions u and v are coreferent, their stalks
/// should map to the same point in the edge space.
pub trait RestrictionMap: Clone + Debug {
    /// Scalar type for the map.
    type Scalar: Clone + Debug;
    /// Vector type for input/output.
    type Vector: Clone + Debug;

    /// Input dimension (stalk dimension at source node).
    fn in_dim(&self) -> usize;

    /// Output dimension (edge space dimension).
    fn out_dim(&self) -> usize;

    /// Apply the restriction map to a stalk vector.
    fn apply(&self, x: &Self::Vector) -> Result<Self::Vector, SheafError>;

    /// Apply the transpose (adjoint) of the restriction map.
    /// Used in Laplacian computation.
    fn apply_transpose(&self, x: &Self::Vector) -> Result<Self::Vector, SheafError>;

    /// Get the matrix representation (for debugging/serialization).
    fn as_matrix(&self) -> Vec<Vec<Self::Scalar>>;

    /// Frobenius norm of the map (for regularization).
    fn frobenius_norm(&self) -> Self::Scalar;
}

/// A stalk (vector space) at a node.
///
/// For coreference: the stalk at a mention node contains its embedding.
pub trait Stalk: Clone + Debug {
    /// Scalar type.
    type Scalar: Clone + Debug;
    /// Vector type.
    type Vector: Clone + Debug;

    /// Dimension of the stalk.
    fn dim(&self) -> usize;

    /// Get the current value (signal on the stalk).
    fn value(&self) -> &Self::Vector;

    /// Set the value.
    fn set_value(&mut self, v: Self::Vector) -> Result<(), SheafError>;

    /// Zero vector in this stalk.
    fn zero(&self) -> Self::Vector;
}

/// Edge data in a sheaf graph.
#[derive(Debug, Clone)]
pub struct SheafEdge<R: RestrictionMap> {
    /// Source node ID.
    pub source: usize,
    /// Target node ID.
    pub target: usize,
    /// Restriction map from source stalk to edge space.
    pub restriction_source: R,
    /// Restriction map from target stalk to edge space.
    pub restriction_target: R,
    /// Edge weight (optional, for weighted Laplacian).
    pub weight: f32,
}

/// A sheaf on a graph.
///
/// This is the main data structure for sheaf neural networks.
/// It assigns stalks to nodes and restriction maps to edges.
pub trait SheafGraph: Debug {
    /// Scalar type.
    type Scalar: Clone + Debug + Default;
    /// Vector type.
    type Vector: Clone + Debug;
    /// Restriction map type.
    type Restriction: RestrictionMap<Scalar = Self::Scalar, Vector = Self::Vector>;
    /// Stalk type.
    type Stalk: Stalk<Scalar = Self::Scalar, Vector = Self::Vector>;

    /// Number of nodes.
    fn num_nodes(&self) -> usize;

    /// Number of edges.
    fn num_edges(&self) -> usize;

    /// Get stalk at node.
    fn stalk(&self, node: usize) -> Result<&Self::Stalk, SheafError>;

    /// Get mutable stalk at node.
    fn stalk_mut(&mut self, node: usize) -> Result<&mut Self::Stalk, SheafError>;

    /// Get edge data.
    fn edge(
        &self,
        source: usize,
        target: usize,
    ) -> Result<&SheafEdge<Self::Restriction>, SheafError>;

    /// Iterate over all edges.
    fn edges(&self) -> impl Iterator<Item = &SheafEdge<Self::Restriction>>;

    /// Get neighbors of a node.
    fn neighbors(&self, node: usize) -> Result<Vec<usize>, SheafError>;

    /// Compute Dirichlet energy for the current stalk values.
    ///
    /// E(x) = Σ_{(u,v) ∈ E} w_{uv} ||R_u(x_u) - R_v(x_v)||²
    ///
    /// Low energy means consistent signal across the sheaf.
    fn dirichlet_energy(&self) -> Result<Self::Scalar, SheafError>;

    /// Compute the sheaf Laplacian action on a given node.
    ///
    /// (L_F x)_v = Σ_{u ~ v} R_v^T (R_v x_v - R_u x_u)
    fn laplacian_at(&self, node: usize) -> Result<Self::Vector, SheafError>;

    /// Perform one step of sheaf diffusion.
    ///
    /// x_{t+1} = x_t - α * L_F * x_t
    ///
    /// This smooths the signal according to sheaf structure.
    fn diffusion_step(&mut self, step_size: Self::Scalar) -> Result<(), SheafError>;
}

// =============================================================================
// Sheaf Laplacian Types
// =============================================================================

/// Describes the structure of a sheaf Laplacian.
///
/// The Laplacian can be:
/// - **Connection Laplacian**: Uses orthogonal restriction maps (preserves norms)
/// - **General Laplacian**: Arbitrary linear maps
/// - **Diagonal Laplacian**: Scalar weights only (reduces to graph Laplacian)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LaplacianType {
    /// Orthogonal restriction maps (O(d) valued).
    Connection,
    /// General linear maps (GL(d) valued).
    General,
    /// Scalar weights (diagonal, equivalent to weighted graph Laplacian).
    Diagonal,
}

/// Configuration for sheaf diffusion.
#[derive(Debug, Clone)]
pub struct DiffusionConfig {
    /// Number of diffusion steps.
    pub num_steps: usize,
    /// Step size (learning rate for diffusion).
    pub step_size: f32,
    /// Whether to normalize the Laplacian (D^{-1/2} L D^{-1/2}).
    pub normalize: bool,
    /// Type of Laplacian to use (Connection, General, or Diagonal).
    pub laplacian_type: LaplacianType,
}

impl Default for DiffusionConfig {
    fn default() -> Self {
        Self {
            num_steps: 5,
            step_size: 0.1,
            normalize: true,
            laplacian_type: LaplacianType::General,
        }
    }
}

// =============================================================================
// Simple In-Memory Implementation (f32, Vec<f32>)
// =============================================================================

/// Simple restriction map using a dense matrix.
#[derive(Debug, Clone)]
pub struct DenseRestriction {
    /// Matrix data in row-major order.
    pub data: Vec<f32>,
    /// Number of rows (output dimension).
    pub rows: usize,
    /// Number of columns (input dimension).
    pub cols: usize,
}

impl DenseRestriction {
    /// Create a new restriction map.
    pub fn new(data: Vec<f32>, rows: usize, cols: usize) -> Result<Self, SheafError> {
        if data.len() != rows * cols {
            return Err(SheafError::DimensionMismatch {
                expected: rows * cols,
                actual: data.len(),
            });
        }
        Ok(Self { data, rows, cols })
    }

    /// Create an identity restriction (for same-dimension stalks).
    pub fn identity(dim: usize) -> Self {
        let mut data = vec![0.0; dim * dim];
        for i in 0..dim {
            data[i * dim + i] = 1.0;
        }
        Self {
            data,
            rows: dim,
            cols: dim,
        }
    }

    /// Create a random orthogonal restriction map.
    ///
    /// Uses QR decomposition of random matrix.
    /// Useful for connection Laplacians.
    #[cfg_attr(docsrs, doc(cfg(feature = "rand")))]
    #[cfg(feature = "rand")]
    #[allow(deprecated)]
    pub fn random_orthogonal(dim: usize) -> Self {
        use rand::Rng;
        let mut rng = rand::thread_rng();

        // Generate random matrix
        let mut data: Vec<f32> = (0..dim * dim).map(|_| rng.gen_range(-0.5..0.5)).collect();

        // Simple Gram-Schmidt orthogonalization
        for i in 0..dim {
            // Normalize column i
            let mut norm: f32 = 0.0;
            for j in 0..dim {
                norm += data[j * dim + i] * data[j * dim + i];
            }
            norm = norm.sqrt();
            if norm > 1e-6 {
                for j in 0..dim {
                    data[j * dim + i] /= norm;
                }
            }

            // Subtract projections from remaining columns
            for k in (i + 1)..dim {
                let mut dot = 0.0;
                for j in 0..dim {
                    dot += data[j * dim + i] * data[j * dim + k];
                }
                for j in 0..dim {
                    data[j * dim + k] -= dot * data[j * dim + i];
                }
            }
        }

        Self {
            data,
            rows: dim,
            cols: dim,
        }
    }
}

impl RestrictionMap for DenseRestriction {
    type Scalar = f32;
    type Vector = Vec<f32>;

    fn in_dim(&self) -> usize {
        self.cols
    }

    fn out_dim(&self) -> usize {
        self.rows
    }

    fn apply(&self, x: &Self::Vector) -> Result<Self::Vector, SheafError> {
        if x.len() != self.cols {
            return Err(SheafError::DimensionMismatch {
                expected: self.cols,
                actual: x.len(),
            });
        }

        // Matrix multiplication: result = A * x
        // Using explicit indexing for clarity in matrix math
        let mut result = vec![0.0; self.rows];
        #[allow(clippy::needless_range_loop)]
        for i in 0..self.rows {
            for j in 0..self.cols {
                result[i] += self.data[i * self.cols + j] * x[j];
            }
        }
        Ok(result)
    }

    fn apply_transpose(&self, x: &Self::Vector) -> Result<Self::Vector, SheafError> {
        if x.len() != self.rows {
            return Err(SheafError::DimensionMismatch {
                expected: self.rows,
                actual: x.len(),
            });
        }

        // Matrix transpose multiplication: result = A^T * x
        let mut result = vec![0.0; self.cols];
        #[allow(clippy::needless_range_loop)]
        for j in 0..self.cols {
            for i in 0..self.rows {
                result[j] += self.data[i * self.cols + j] * x[i];
            }
        }
        Ok(result)
    }

    fn as_matrix(&self) -> Vec<Vec<Self::Scalar>> {
        // Convert flat storage to row-major matrix
        let mut matrix = vec![vec![0.0; self.cols]; self.rows];
        #[allow(clippy::needless_range_loop)]
        for i in 0..self.rows {
            for j in 0..self.cols {
                matrix[i][j] = self.data[i * self.cols + j];
            }
        }
        matrix
    }

    fn frobenius_norm(&self) -> Self::Scalar {
        self.data.iter().map(|x| x * x).sum::<f32>().sqrt()
    }
}

/// Simple stalk holding a `Vec<f32>`.
#[derive(Debug, Clone)]
pub struct VecStalk {
    value: Vec<f32>,
}

impl VecStalk {
    /// Create a new stalk with given value.
    pub fn new(value: Vec<f32>) -> Self {
        Self { value }
    }
}

impl Stalk for VecStalk {
    type Scalar = f32;
    type Vector = Vec<f32>;

    fn dim(&self) -> usize {
        self.value.len()
    }

    fn value(&self) -> &Self::Vector {
        &self.value
    }

    fn set_value(&mut self, v: Self::Vector) -> Result<(), SheafError> {
        if v.len() != self.value.len() {
            return Err(SheafError::DimensionMismatch {
                expected: self.value.len(),
                actual: v.len(),
            });
        }
        self.value = v;
        Ok(())
    }

    fn zero(&self) -> Self::Vector {
        vec![0.0; self.value.len()]
    }
}

/// Simple in-memory sheaf graph.
#[derive(Debug, Clone)]
pub struct SimpleSheafGraph {
    stalks: Vec<VecStalk>,
    edges: Vec<SheafEdge<DenseRestriction>>,
    adjacency: HashMap<usize, Vec<usize>>,
}

impl SimpleSheafGraph {
    /// Create a new empty sheaf graph.
    pub fn new() -> Self {
        Self {
            stalks: Vec::new(),
            edges: Vec::new(),
            adjacency: HashMap::new(),
        }
    }

    /// Add a node with initial stalk value.
    pub fn add_node(&mut self, value: Vec<f32>) -> usize {
        let id = self.stalks.len();
        self.stalks.push(VecStalk::new(value));
        self.adjacency.insert(id, Vec::new());
        id
    }

    /// Add an edge with restriction maps.
    pub fn add_edge(
        &mut self,
        source: usize,
        target: usize,
        restriction_source: DenseRestriction,
        restriction_target: DenseRestriction,
        weight: f32,
    ) -> Result<(), SheafError> {
        if source >= self.stalks.len() {
            return Err(SheafError::NodeNotFound(source));
        }
        if target >= self.stalks.len() {
            return Err(SheafError::NodeNotFound(target));
        }

        // Verify dimensions
        if restriction_source.in_dim() != self.stalks[source].dim() {
            return Err(SheafError::DimensionMismatch {
                expected: self.stalks[source].dim(),
                actual: restriction_source.in_dim(),
            });
        }
        if restriction_target.in_dim() != self.stalks[target].dim() {
            return Err(SheafError::DimensionMismatch {
                expected: self.stalks[target].dim(),
                actual: restriction_target.in_dim(),
            });
        }
        if restriction_source.out_dim() != restriction_target.out_dim() {
            return Err(SheafError::InvalidRestriction(
                "Source and target restrictions must have same output dimension".into(),
            ));
        }

        self.edges.push(SheafEdge {
            source,
            target,
            restriction_source,
            restriction_target,
            weight,
        });

        self.adjacency.entry(source).or_default().push(target);
        self.adjacency.entry(target).or_default().push(source);

        Ok(())
    }
}

impl Default for SimpleSheafGraph {
    fn default() -> Self {
        Self::new()
    }
}

impl SheafGraph for SimpleSheafGraph {
    type Scalar = f32;
    type Vector = Vec<f32>;
    type Restriction = DenseRestriction;
    type Stalk = VecStalk;

    fn num_nodes(&self) -> usize {
        self.stalks.len()
    }

    fn num_edges(&self) -> usize {
        self.edges.len()
    }

    fn stalk(&self, node: usize) -> Result<&Self::Stalk, SheafError> {
        self.stalks.get(node).ok_or(SheafError::NodeNotFound(node))
    }

    fn stalk_mut(&mut self, node: usize) -> Result<&mut Self::Stalk, SheafError> {
        self.stalks
            .get_mut(node)
            .ok_or(SheafError::NodeNotFound(node))
    }

    fn edge(
        &self,
        source: usize,
        target: usize,
    ) -> Result<&SheafEdge<Self::Restriction>, SheafError> {
        self.edges
            .iter()
            .find(|e| {
                (e.source == source && e.target == target)
                    || (e.source == target && e.target == source)
            })
            .ok_or(SheafError::EdgeNotFound(source, target))
    }

    fn edges(&self) -> impl Iterator<Item = &SheafEdge<Self::Restriction>> {
        self.edges.iter()
    }

    fn neighbors(&self, node: usize) -> Result<Vec<usize>, SheafError> {
        self.adjacency
            .get(&node)
            .cloned()
            .ok_or(SheafError::NodeNotFound(node))
    }

    fn dirichlet_energy(&self) -> Result<Self::Scalar, SheafError> {
        let mut energy = 0.0;

        for edge in &self.edges {
            let x_u = self.stalks[edge.source].value();
            let x_v = self.stalks[edge.target].value();

            let r_u = edge.restriction_source.apply(x_u)?;
            let r_v = edge.restriction_target.apply(x_v)?;

            // ||R_u(x_u) - R_v(x_v)||²
            let diff_sq: f32 = r_u
                .iter()
                .zip(r_v.iter())
                .map(|(a, b)| (a - b) * (a - b))
                .sum();

            energy += edge.weight * diff_sq;
        }

        Ok(energy)
    }

    fn laplacian_at(&self, node: usize) -> Result<Self::Vector, SheafError> {
        let stalk = self.stalk(node)?;
        let mut result = stalk.zero();

        for edge in &self.edges {
            let (is_source, other) = if edge.source == node {
                (true, edge.target)
            } else if edge.target == node {
                (false, edge.source)
            } else {
                continue;
            };

            let x_node = self.stalks[node].value();
            let x_other = self.stalks[other].value();

            let (r_node, r_other) = if is_source {
                (&edge.restriction_source, &edge.restriction_target)
            } else {
                (&edge.restriction_target, &edge.restriction_source)
            };

            // R_node(x_node) - R_other(x_other)
            let r_x_node = r_node.apply(x_node)?;
            let r_x_other = r_other.apply(x_other)?;

            let diff: Vec<f32> = r_x_node
                .iter()
                .zip(r_x_other.iter())
                .map(|(a, b)| a - b)
                .collect();

            // R_node^T (diff)
            let contrib = r_node.apply_transpose(&diff)?;

            // Accumulate weighted contribution
            for (i, c) in contrib.iter().enumerate() {
                result[i] += edge.weight * c;
            }
        }

        Ok(result)
    }

    fn diffusion_step(&mut self, step_size: Self::Scalar) -> Result<(), SheafError> {
        // Compute Laplacian at all nodes first (to avoid borrowing issues)
        let laplacians: Vec<Vec<f32>> = (0..self.num_nodes())
            .map(|i| self.laplacian_at(i))
            .collect::<Result<_, _>>()?;

        // Update all stalks: x = x - step_size * L_F * x
        for (i, lap) in laplacians.into_iter().enumerate() {
            let stalk = &mut self.stalks[i];
            let new_value: Vec<f32> = stalk
                .value()
                .iter()
                .zip(lap.iter())
                .map(|(x, l)| x - step_size * l)
                .collect();
            stalk.set_value(new_value)?;
        }

        Ok(())
    }
}

// =============================================================================
// Utility Functions
// =============================================================================

/// Compute consistency score for a sheaf graph.
///
/// Returns 1.0 for perfect consistency (zero energy), decreasing toward 0.
pub fn consistency_score(graph: &impl SheafGraph<Scalar = f32>) -> Result<f32, SheafError> {
    let energy = graph.dirichlet_energy()?;
    // Use exponential decay: exp(-energy)
    Ok((-energy).exp())
}

/// Run sheaf diffusion until convergence or max iterations.
pub fn diffuse_until_convergence(
    graph: &mut SimpleSheafGraph,
    config: &DiffusionConfig,
    tolerance: f32,
) -> Result<usize, SheafError> {
    let mut prev_energy = graph.dirichlet_energy()?;

    for step in 0..config.num_steps {
        graph.diffusion_step(config.step_size)?;
        let energy = graph.dirichlet_energy()?;

        if (prev_energy - energy).abs() < tolerance {
            return Ok(step + 1);
        }
        prev_energy = energy;
    }

    Ok(config.num_steps)
}

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

    #[test]
    fn test_identity_restriction() {
        let r = DenseRestriction::identity(3);
        let x = vec![1.0, 2.0, 3.0];
        let y = r.apply(&x).unwrap();
        assert_eq!(y, x);
    }

    #[test]
    fn test_restriction_transpose() {
        // 2x3 matrix
        let r = DenseRestriction::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 2, 3).unwrap();

        let x = vec![1.0, 2.0, 3.0];
        let y = r.apply(&x).unwrap();
        assert_eq!(y.len(), 2);

        let z = vec![1.0, 1.0];
        let w = r.apply_transpose(&z).unwrap();
        assert_eq!(w.len(), 3);
        // Transpose of [[1,2,3],[4,5,6]] is [[1,4],[2,5],[3,6]]
        // [1,4] · [1,1] = 5, [2,5] · [1,1] = 7, [3,6] · [1,1] = 9
        assert_eq!(w, vec![5.0, 7.0, 9.0]);
    }

    #[test]
    fn test_simple_sheaf_graph() {
        let mut graph = SimpleSheafGraph::new();

        // Two nodes with 2D stalks
        let n0 = graph.add_node(vec![1.0, 0.0]);
        let n1 = graph.add_node(vec![0.0, 1.0]);

        // Identity restrictions (simplest case)
        let r = DenseRestriction::identity(2);
        graph.add_edge(n0, n1, r.clone(), r.clone(), 1.0).unwrap();

        assert_eq!(graph.num_nodes(), 2);
        assert_eq!(graph.num_edges(), 1);

        // Dirichlet energy should be ||[1,0] - [0,1]||² = 2
        let energy = graph.dirichlet_energy().unwrap();
        assert!((energy - 2.0).abs() < 1e-6);
    }

    #[test]
    fn test_diffusion_reduces_energy() {
        let mut graph = SimpleSheafGraph::new();

        // Three nodes forming a chain
        let n0 = graph.add_node(vec![1.0, 0.0]);
        let n1 = graph.add_node(vec![0.5, 0.5]);
        let n2 = graph.add_node(vec![0.0, 1.0]);

        let r = DenseRestriction::identity(2);
        graph.add_edge(n0, n1, r.clone(), r.clone(), 1.0).unwrap();
        graph.add_edge(n1, n2, r.clone(), r.clone(), 1.0).unwrap();

        let initial_energy = graph.dirichlet_energy().unwrap();

        // Run diffusion
        for _ in 0..10 {
            graph.diffusion_step(0.1).unwrap();
        }

        let final_energy = graph.dirichlet_energy().unwrap();
        assert!(
            final_energy < initial_energy,
            "Diffusion should reduce energy"
        );
    }

    #[test]
    fn test_consistency_score() {
        let mut graph = SimpleSheafGraph::new();

        // Two nodes with identical stalks
        graph.add_node(vec![1.0, 2.0]);
        graph.add_node(vec![1.0, 2.0]);

        let r = DenseRestriction::identity(2);
        graph.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();

        // Perfect consistency: score should be 1.0
        let score = consistency_score(&graph).unwrap();
        assert!((score - 1.0).abs() < 1e-6);
    }

    // =========================================================================
    // DenseRestriction edge cases
    // =========================================================================

    #[test]
    fn test_dense_restriction_new_dimension_mismatch() {
        let result = DenseRestriction::new(vec![1.0, 2.0, 3.0], 2, 2);
        assert!(matches!(
            result,
            Err(SheafError::DimensionMismatch {
                expected: 4,
                actual: 3
            })
        ));
    }

    #[test]
    fn test_dense_restriction_1x1() {
        let r = DenseRestriction::new(vec![3.0], 1, 1).unwrap();
        let x = vec![2.0];
        let y = r.apply(&x).unwrap();
        assert_eq!(y, vec![6.0]);

        let yt = r.apply_transpose(&vec![2.0]).unwrap();
        assert_eq!(yt, vec![6.0]); // Transpose of 1x1 is itself
    }

    #[test]
    fn test_dense_restriction_apply_wrong_dim() {
        let r = DenseRestriction::identity(3);
        let x = vec![1.0, 2.0]; // Wrong dimension
        let result = r.apply(&x);
        assert!(matches!(
            result,
            Err(SheafError::DimensionMismatch {
                expected: 3,
                actual: 2
            })
        ));
    }

    #[test]
    fn test_dense_restriction_apply_transpose_wrong_dim() {
        let r = DenseRestriction::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 2, 3).unwrap();
        let x = vec![1.0, 2.0, 3.0]; // 3 elements but rows = 2
        let result = r.apply_transpose(&x);
        assert!(matches!(
            result,
            Err(SheafError::DimensionMismatch {
                expected: 2,
                actual: 3
            })
        ));
    }

    #[test]
    fn test_dense_restriction_as_matrix() {
        let r = DenseRestriction::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 2, 3).unwrap();
        let m = r.as_matrix();
        assert_eq!(m.len(), 2);
        assert_eq!(m[0], vec![1.0, 2.0, 3.0]);
        assert_eq!(m[1], vec![4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_dense_restriction_frobenius_norm() {
        let r = DenseRestriction::new(vec![3.0, 4.0], 1, 2).unwrap();
        let norm = r.frobenius_norm();
        assert!((norm - 5.0).abs() < 1e-6); // sqrt(9 + 16) = 5
    }

    #[test]
    fn test_identity_restriction_is_identity() {
        let r = DenseRestriction::identity(4);
        assert_eq!(r.in_dim(), 4);
        assert_eq!(r.out_dim(), 4);
        let x = vec![1.0, 2.0, 3.0, 4.0];
        assert_eq!(r.apply(&x).unwrap(), x);
        assert_eq!(r.apply_transpose(&x).unwrap(), x); // I^T = I
    }

    // =========================================================================
    // VecStalk edge cases
    // =========================================================================

    #[test]
    fn test_vec_stalk_set_value_dimension_mismatch() {
        let mut s = VecStalk::new(vec![1.0, 2.0]);
        let result = s.set_value(vec![1.0]);
        assert!(matches!(
            result,
            Err(SheafError::DimensionMismatch {
                expected: 2,
                actual: 1
            })
        ));
    }

    #[test]
    fn test_vec_stalk_zero() {
        let s = VecStalk::new(vec![5.0, 6.0, 7.0]);
        assert_eq!(s.zero(), vec![0.0, 0.0, 0.0]);
    }

    #[test]
    fn test_vec_stalk_roundtrip() {
        let mut s = VecStalk::new(vec![1.0, 2.0]);
        s.set_value(vec![3.0, 4.0]).unwrap();
        assert_eq!(s.value(), &vec![3.0, 4.0]);
        assert_eq!(s.dim(), 2);
    }

    // =========================================================================
    // SimpleSheafGraph error paths
    // =========================================================================

    #[test]
    fn test_add_edge_source_not_found() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0]);
        let r = DenseRestriction::identity(1);
        let result = graph.add_edge(5, 0, r.clone(), r.clone(), 1.0);
        assert!(matches!(result, Err(SheafError::NodeNotFound(5))));
    }

    #[test]
    fn test_add_edge_target_not_found() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0]);
        let r = DenseRestriction::identity(1);
        let result = graph.add_edge(0, 99, r.clone(), r.clone(), 1.0);
        assert!(matches!(result, Err(SheafError::NodeNotFound(99))));
    }

    #[test]
    fn test_add_edge_restriction_dim_mismatch_source() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 2.0]); // dim 2
        graph.add_node(vec![1.0, 2.0]); // dim 2
        let r_wrong = DenseRestriction::identity(3); // dim 3
        let r_ok = DenseRestriction::identity(2);
        let result = graph.add_edge(0, 1, r_wrong, r_ok, 1.0);
        assert!(matches!(result, Err(SheafError::DimensionMismatch { .. })));
    }

    #[test]
    fn test_add_edge_restriction_output_dim_mismatch() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 2.0]);
        graph.add_node(vec![1.0, 2.0]);
        // Source restriction: 2->3, target restriction: 2->2 (output dims differ)
        let r_src = DenseRestriction::new(vec![1.0; 6], 3, 2).unwrap();
        let r_tgt = DenseRestriction::identity(2);
        let result = graph.add_edge(0, 1, r_src, r_tgt, 1.0);
        assert!(matches!(result, Err(SheafError::InvalidRestriction(_))));
    }

    #[test]
    fn test_stalk_not_found() {
        let graph = SimpleSheafGraph::new();
        assert!(matches!(graph.stalk(0), Err(SheafError::NodeNotFound(0))));
    }

    #[test]
    fn test_edge_not_found() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0]);
        graph.add_node(vec![1.0]);
        // No edge added
        assert!(matches!(
            graph.edge(0, 1),
            Err(SheafError::EdgeNotFound(0, 1))
        ));
    }

    #[test]
    fn test_neighbors_not_found() {
        let graph = SimpleSheafGraph::new();
        assert!(matches!(
            graph.neighbors(0),
            Err(SheafError::NodeNotFound(0))
        ));
    }

    #[test]
    fn test_edge_lookup_bidirectional() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0]);
        graph.add_node(vec![2.0]);
        let r = DenseRestriction::identity(1);
        graph.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();

        // Both directions should find the edge
        assert!(graph.edge(0, 1).is_ok());
        assert!(graph.edge(1, 0).is_ok());
    }

    #[test]
    fn test_neighbors_bidirectional() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0]);
        graph.add_node(vec![2.0]);
        graph.add_node(vec![3.0]);
        let r = DenseRestriction::identity(1);
        graph.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();
        graph.add_edge(1, 2, r.clone(), r.clone(), 1.0).unwrap();

        let n1 = graph.neighbors(1).unwrap();
        assert_eq!(n1.len(), 2); // connected to both 0 and 2
    }

    // =========================================================================
    // Dirichlet energy and Laplacian
    // =========================================================================

    #[test]
    fn test_dirichlet_energy_zero_for_identical_stalks() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 2.0, 3.0]);
        graph.add_node(vec![1.0, 2.0, 3.0]);
        graph.add_node(vec![1.0, 2.0, 3.0]);
        let r = DenseRestriction::identity(3);
        graph.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();
        graph.add_edge(1, 2, r.clone(), r.clone(), 1.0).unwrap();
        let energy = graph.dirichlet_energy().unwrap();
        assert!(
            (energy - 0.0).abs() < 1e-6,
            "identical stalks should have zero energy"
        );
    }

    #[test]
    fn test_dirichlet_energy_weighted() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 0.0]);
        graph.add_node(vec![0.0, 1.0]);
        let r = DenseRestriction::identity(2);
        // Weight 2.0 should double the energy
        graph.add_edge(0, 1, r.clone(), r.clone(), 2.0).unwrap();
        let energy = graph.dirichlet_energy().unwrap();
        // ||[1,0] - [0,1]||^2 = 2, weight 2.0 => 4.0
        assert!((energy - 4.0).abs() < 1e-6);
    }

    #[test]
    fn test_laplacian_at_zero_for_consistent_signal() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 2.0]);
        graph.add_node(vec![1.0, 2.0]);
        let r = DenseRestriction::identity(2);
        graph.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();

        let lap = graph.laplacian_at(0).unwrap();
        assert!(
            lap.iter().all(|&x| x.abs() < 1e-6),
            "Laplacian should be zero for consistent signal"
        );
    }

    #[test]
    fn test_laplacian_symmetry() {
        // For identity restrictions on an edge (0,1), the Laplacian at 0
        // should be the negative of the Laplacian at 1 (conservation).
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 0.0]);
        graph.add_node(vec![0.0, 1.0]);
        let r = DenseRestriction::identity(2);
        graph.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();

        let lap0 = graph.laplacian_at(0).unwrap();
        let lap1 = graph.laplacian_at(1).unwrap();
        // L(0) + L(1) = 0 (conservation for identity restrictions)
        for i in 0..2 {
            assert!(
                (lap0[i] + lap1[i]).abs() < 1e-6,
                "Laplacian should sum to zero"
            );
        }
    }

    // =========================================================================
    // Diffusion convergence
    // =========================================================================

    #[test]
    fn test_diffuse_until_convergence_identical_stalks() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 1.0]);
        graph.add_node(vec![1.0, 1.0]);
        let r = DenseRestriction::identity(2);
        graph.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();

        let config = DiffusionConfig {
            num_steps: 100,
            step_size: 0.1,
            ..Default::default()
        };

        // Already converged: should return 1 (converges on first step)
        let steps = diffuse_until_convergence(&mut graph, &config, 1e-8).unwrap();
        assert!(
            steps <= 2,
            "already-converged graph should converge immediately, took {steps}"
        );
    }

    #[test]
    fn test_diffuse_until_convergence_reaches_max_steps() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![100.0, 0.0]);
        graph.add_node(vec![0.0, 100.0]);
        let r = DenseRestriction::identity(2);
        graph.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();

        let config = DiffusionConfig {
            num_steps: 3,
            step_size: 0.01, // Very small step: won't converge in 3 steps
            ..Default::default()
        };

        let steps = diffuse_until_convergence(&mut graph, &config, 1e-12).unwrap();
        assert_eq!(steps, 3, "should reach max steps");
    }

    #[test]
    fn test_consistency_score_decreases_with_distance() {
        // Larger stalk differences should produce lower consistency
        let mut g1 = SimpleSheafGraph::new();
        g1.add_node(vec![1.0, 0.0]);
        g1.add_node(vec![0.9, 0.1]);
        let r = DenseRestriction::identity(2);
        g1.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();
        let score1 = consistency_score(&g1).unwrap();

        let mut g2 = SimpleSheafGraph::new();
        g2.add_node(vec![1.0, 0.0]);
        g2.add_node(vec![0.0, 1.0]);
        g2.add_edge(0, 1, r.clone(), r.clone(), 1.0).unwrap();
        let score2 = consistency_score(&g2).unwrap();

        assert!(
            score1 > score2,
            "closer stalks should have higher consistency"
        );
    }

    // =========================================================================
    // Default and Display impls
    // =========================================================================

    #[test]
    fn test_diffusion_config_default() {
        let config = DiffusionConfig::default();
        assert_eq!(config.num_steps, 5);
        assert!((config.step_size - 0.1).abs() < 1e-6);
        assert!(config.normalize);
        assert_eq!(config.laplacian_type, LaplacianType::General);
    }

    #[test]
    fn test_sheaf_error_display() {
        assert_eq!(
            format!("{}", SheafError::NodeNotFound(5)),
            "Node 5 not found"
        );
        assert_eq!(
            format!("{}", SheafError::EdgeNotFound(1, 2)),
            "Edge (1, 2) not found"
        );
        assert_eq!(
            format!(
                "{}",
                SheafError::DimensionMismatch {
                    expected: 3,
                    actual: 2
                }
            ),
            "Dimension mismatch: expected 3, got 2"
        );
        assert!(format!("{}", SheafError::InvalidRestriction("bad".into())).contains("bad"));
    }

    #[test]
    fn test_simple_sheaf_graph_default() {
        let graph = SimpleSheafGraph::default();
        assert_eq!(graph.num_nodes(), 0);
        assert_eq!(graph.num_edges(), 0);
    }

    // =========================================================================
    // Non-square restriction maps
    // =========================================================================

    #[test]
    fn test_non_square_restriction_maps() {
        // Project 3D stalks into 2D edge space
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 0.0, 0.0]);
        graph.add_node(vec![0.0, 1.0, 0.0]);

        // Projection: keep first 2 dims (2x3 matrix)
        let proj = DenseRestriction::new(vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 2, 3).unwrap();
        graph
            .add_edge(0, 1, proj.clone(), proj.clone(), 1.0)
            .unwrap();

        let energy = graph.dirichlet_energy().unwrap();
        // R(x0) = [1,0], R(x1) = [0,1], ||diff||^2 = 2
        assert!((energy - 2.0).abs() < 1e-6);
    }

    #[test]
    fn test_diffusion_with_non_square_restrictions() {
        // Verify diffusion works with non-square maps
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 0.0, 0.0]);
        graph.add_node(vec![0.0, 1.0, 0.0]);

        let proj = DenseRestriction::new(vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0], 2, 3).unwrap();
        graph
            .add_edge(0, 1, proj.clone(), proj.clone(), 1.0)
            .unwrap();

        let initial_energy = graph.dirichlet_energy().unwrap();
        graph.diffusion_step(0.1).unwrap();
        let final_energy = graph.dirichlet_energy().unwrap();
        assert!(
            final_energy < initial_energy,
            "diffusion should reduce energy with non-square maps"
        );
    }

    // =========================================================================
    // Empty and single-node graphs
    // =========================================================================

    #[test]
    fn test_empty_graph_energy() {
        let graph = SimpleSheafGraph::new();
        let energy = graph.dirichlet_energy().unwrap();
        assert_eq!(energy, 0.0);
    }

    #[test]
    fn test_single_node_graph() {
        let mut graph = SimpleSheafGraph::new();
        graph.add_node(vec![1.0, 2.0]);
        assert_eq!(graph.num_nodes(), 1);
        assert_eq!(graph.num_edges(), 0);
        let energy = graph.dirichlet_energy().unwrap();
        assert_eq!(energy, 0.0);
    }
}