zigzag-engine/src/diagram.rs

//! Diagrams in associative n-categories
//!
//! An n-diagram is an object of Zⁿ(ℕ), the n-fold iterated zigzag category
//! over natural numbers. The natural number at each point encodes the dimension
//! of the algebraic generator at that location.
//!
//! The concrete representation uses:
//! - `Diagram`: enum of Diagram0 | DiagramN
//! - `DiagramN`: source diagram + list of cospans (the zigzag structure)
//! - `Cospan`: forward rewrite + backward rewrite
//! - `Rewrite`: zigzag map between diagram slices
//! - `Cone`: atomic rewrite data

use crate::signature::Generator;

/// An n-diagram in the iterated zigzag category.
///
/// This is the main data structure representing diagrams in associative n-categories.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Diagram {
    /// A 0-diagram: a single generator (point in the signature).
    Diagram0(Generator),
    /// An n-diagram for n > 0: source + zigzag of cospans.
    DiagramN(DiagramN),
}

/// An n-dimensional diagram (n > 0).
///
/// Represented as:
/// - A source (n-1)-diagram (the first regular slice)
/// - A sequence of cospans encoding the zigzag structure
///
/// The regular slices r₀, r₁, ..., rₙ are:
/// - r₀ = source
/// - rᵢ₊₁ = backward.target of cospan i (equivalently, forward.target of cospan i)
///
/// The singular slices s₀, s₁, ..., sₙ₋₁ are the apexes of each cospan.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DiagramN {
    /// The source diagram (first regular slice, r₀)
    pub source: Box<Diagram>,
    /// The cospans forming the zigzag structure
    pub cospans: Vec<Cospan>,
}

impl DiagramN {
    /// Create a new n-diagram.
    pub fn new(source: Diagram, cospans: Vec<Cospan>) -> Self {
        Self {
            source: Box::new(source),
            cospans,
        }
    }

    /// Create an identity diagram (zigzag of length 0).
    pub fn identity(source: Diagram) -> Self {
        Self {
            source: Box::new(source),
            cospans: vec![],
        }
    }

    /// The zigzag length (number of cospans / singular heights).
    pub fn length(&self) -> usize {
        self.cospans.len()
    }

    /// Get the source (first regular slice).
    pub fn source(&self) -> &Diagram {
        &self.source
    }

    /// Compute the target (last regular slice).
    ///
    /// For an identity (length 0), this is the same as source.
    /// Otherwise, we traverse the rewrites to find the final regular slice.
    ///
    /// The target is computed by starting with the source and applying
    /// each cospan's rewrites in sequence: for each cospan, apply the
    /// forward rewrite (to reach the apex), then apply the backward
    /// rewrite in reverse (to reach the next regular slice).
    pub fn target(&self) -> Diagram {
        self.regular_slice(self.cospans.len())
            .expect("target should always be computable")
    }

    /// Get the regular slice at height h.
    ///
    /// Regular slices r₀, r₁, ..., rₙ where n = number of cospans:
    /// - r₀ = source
    /// - rᵢ₊₁ is computed by traversing cospan i
    ///
    /// To traverse a cospan (forward: rᵢ → sᵢ, backward: rᵢ₊₁ → sᵢ):
    /// 1. Apply forward rewrite to rᵢ to get sᵢ
    /// 2. Apply backward rewrite in reverse to sᵢ to get rᵢ₊₁
    pub fn regular_slice(&self, h: usize) -> Option<Diagram> {
        if h > self.cospans.len() {
            return None;
        }

        let mut slice = (*self.source).clone();

        // Traverse cospans 0..h to reach regular slice h
        for cospan in &self.cospans[..h] {
            // Apply forward rewrite to get to the apex (singular slice)
            slice = cospan.forward.apply_forward(&slice)?;
            // Apply backward rewrite in reverse to get to the next regular slice
            slice = cospan.backward.apply_backward(&slice)?;
        }

        Some(slice)
    }

    /// Get the singular slice at height h.
    ///
    /// The singular slice sₕ is the apex of cospan h.
    /// It is computed by:
    /// 1. Getting regular slice h (rₕ)
    /// 2. Applying the forward rewrite of cospan h
    pub fn singular_slice(&self, h: usize) -> Option<Diagram> {
        if h >= self.cospans.len() {
            return None;
        }

        // Get the regular slice at height h
        let regular = self.regular_slice(h)?;

        // Apply the forward rewrite to get the apex
        self.cospans[h].forward.apply_forward(&regular)
    }

    /// Iterator over all slices (interleaved regular and singular).
    ///
    /// Returns slices in order: r₀, s₀, r₁, s₁, ..., sₙ₋₁, rₙ
    /// Total of 2n + 1 slices for a diagram of length n.
    pub fn slices(&self) -> Slices<'_> {
        Slices::new(self)
    }

    /// Iterator over regular slices only.
    pub fn regular_slices(&self) -> impl Iterator<Item = Diagram> + '_ {
        (0..=self.cospans.len()).filter_map(|h| self.regular_slice(h))
    }

    /// Iterator over singular slices only.
    pub fn singular_slices(&self) -> impl Iterator<Item = Diagram> + '_ {
        (0..self.cospans.len()).filter_map(|h| self.singular_slice(h))
    }
}

/// A cospan in a zigzag: rₕ → sₕ ← rₕ₊₁
///
/// The forward rewrite goes from the left regular to the singular.
/// The backward rewrite goes from the right regular to the singular.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Cospan {
    /// Rewrite from left regular slice to singular apex
    pub forward: Rewrite,
    /// Rewrite from right regular slice to singular apex
    pub backward: Rewrite,
}

impl Cospan {
    /// Create a new cospan.
    pub fn new(forward: Rewrite, backward: Rewrite) -> Self {
        Self { forward, backward }
    }

    /// Check if this is an identity cospan (both legs are trivially identity).
    ///
    /// Uses `is_trivial()` which recognizes both `Rewrite::Identity` and
    /// `RewriteN` with empty cones (as serialized by homotopy-rs).
    pub fn is_identity(&self) -> bool {
        self.forward.is_trivial() && self.backward.is_trivial()
    }
}

/// A rewrite (zigzag map) between diagram slices.
///
/// This encodes how one (n-1)-diagram maps to another within the zigzag structure.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Rewrite {
    /// Identity rewrite (does nothing).
    Identity,
    /// A 0-dimensional rewrite between generators.
    Rewrite0 {
        source: Generator,
        target: Generator,
    },
    /// An n-dimensional rewrite (n > 0), encoded as a sequence of cones.
    RewriteN(RewriteN),
}

impl Rewrite {
    /// Check if this is explicitly marked as an identity rewrite.
    ///
    /// Returns true only for `Rewrite::Identity`. This is conservative and
    /// used for degeneracy tracking where we need to distinguish between
    /// a true identity and a parallel degeneracy with non-identity slices.
    pub fn is_identity(&self) -> bool {
        matches!(self, Rewrite::Identity)
    }

    /// Check if this rewrite is trivially identity (makes no changes).
    ///
    /// Returns true for:
    /// - `Rewrite::Identity` (explicit identity marker)
    /// - `RewriteN` with empty cones (no structural changes at this level)
    /// - `Rewrite0` with source == target
    ///
    /// This is used for detecting identity cospans that should be removed
    /// during normalisation.
    pub fn is_trivial(&self) -> bool {
        match self {
            Rewrite::Identity => true,
            Rewrite::RewriteN(r) => r.cones.is_empty(),
            Rewrite::Rewrite0 { source, target } => source == target,
        }
    }

    /// The dimension of this rewrite.
    pub fn dimension(&self) -> usize {
        match self {
            Rewrite::Identity => 0,
            Rewrite::Rewrite0 { .. } => 0,
            Rewrite::RewriteN(r) => r.dimension,
        }
    }

    /// Apply this rewrite in the forward direction.
    ///
    /// Given a rewrite f: A → B and a diagram matching A, returns B.
    /// For 0-dimensional rewrites, this replaces the generator.
    /// For n-dimensional rewrites, this modifies the cospan structure.
    pub fn apply_forward(&self, diagram: &Diagram) -> Option<Diagram> {
        match (self, diagram) {
            // Identity rewrite: return the diagram unchanged
            (Rewrite::Identity, d) => Some(d.clone()),

            // 0-dimensional rewrite: source must match
            (Rewrite::Rewrite0 { source, target }, Diagram::Diagram0(g)) => {
                if g == source {
                    Some(Diagram::Diagram0(target.clone()))
                } else {
                    // If source doesn't match, the rewrite doesn't apply
                    None
                }
            }

            // n-dimensional rewrite on an n-diagram
            (Rewrite::RewriteN(r), Diagram::DiagramN(d)) => {
                r.apply_forward(d).map(Diagram::DiagramN)
            }

            // Dimension mismatch
            _ => None,
        }
    }

    /// Apply this rewrite in the backward direction.
    ///
    /// Given a rewrite f: A → B and a diagram matching B, returns A.
    /// This is the inverse direction of apply_forward.
    pub fn apply_backward(&self, diagram: &Diagram) -> Option<Diagram> {
        match (self, diagram) {
            // Identity rewrite: return the diagram unchanged
            (Rewrite::Identity, d) => Some(d.clone()),

            // 0-dimensional rewrite: target must match
            (Rewrite::Rewrite0 { source, target }, Diagram::Diagram0(g)) => {
                if g == target {
                    Some(Diagram::Diagram0(source.clone()))
                } else {
                    None
                }
            }

            // n-dimensional rewrite on an n-diagram
            (Rewrite::RewriteN(r), Diagram::DiagramN(d)) => {
                r.apply_backward(d).map(Diagram::DiagramN)
            }

            // Dimension mismatch
            _ => None,
        }
    }
}

/// An n-dimensional rewrite (n > 0).
///
/// Encoded as a sequence of cones, where each cone describes an atomic
/// transformation at a specific position.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RewriteN {
    /// Dimension of this rewrite (> 0)
    pub dimension: usize,
    /// The cones encoding the rewrite structure
    pub cones: Vec<Cone>,
}

impl RewriteN {
    /// Create a new n-dimensional rewrite.
    pub fn new(dimension: usize, cones: Vec<Cone>) -> Self {
        assert!(dimension > 0, "RewriteN dimension must be > 0");
        Self { dimension, cones }
    }

    /// Create an identity rewrite at dimension n.
    pub fn identity(dimension: usize) -> Self {
        Self {
            dimension,
            cones: vec![],
        }
    }

    /// Apply this rewrite in the forward direction.
    ///
    /// A forward rewrite transforms the cospan structure by:
    /// - For each cone, replacing source[cone.index..cone.index+cone.source.len()]
    ///   with the single target cospan
    ///
    /// The source of the diagram is unchanged; only cospans are modified.
    pub fn apply_forward(&self, diagram: &DiagramN) -> Option<DiagramN> {
        let mut cospans = diagram.cospans.clone();
        let mut offset: isize = 0;

        for cone in &self.cones {
            let start = (cone.index as isize + offset) as usize;
            let end = start + cone.source.len();

            // Verify the source cospans match
            if cospans.get(start..end) != Some(&cone.source[..]) {
                return None;
            }

            // Replace source cospans with target cospan
            cospans.splice(start..end, std::iter::once(cone.target.clone()));

            // Update offset: we removed cone.source.len() cospans and added 1
            offset -= cone.source.len() as isize - 1;
        }

        Some(DiagramN::new((*diagram.source).clone(), cospans))
    }

    /// Apply this rewrite in the backward direction.
    ///
    /// A backward rewrite is the inverse: for each cone, we replace
    /// the single target cospan with the source cospans.
    pub fn apply_backward(&self, diagram: &DiagramN) -> Option<DiagramN> {
        let mut cospans = diagram.cospans.clone();

        for cone in &self.cones {
            let start = cone.index;
            let end = start + 1;

            // Verify the target cospan matches
            if cospans.get(start) != Some(&cone.target) {
                return None;
            }

            // Replace target cospan with source cospans
            cospans.splice(start..end, cone.source.iter().cloned());
        }

        Some(DiagramN::new((*diagram.source).clone(), cospans))
    }
}

/// A cone: atomic rewrite data.
///
/// A cone describes a single "bubble" in a string diagram — a region where
/// a source configuration of cospans is replaced by a target cospan.
///
/// The singular map structure of the parent rewrite is determined by the
/// indices of the cones.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Cone {
    /// Index in the target zigzag where this cone maps to
    pub index: usize,
    /// Source configuration (sequence of cospans being contracted)
    pub source: Vec<Cospan>,
    /// Target cospan (what the source contracts to)
    pub target: Cospan,
    /// Slice rewrites for each source singular height.
    /// Length equals source.len() (one per source cospan apex).
    /// Matches homotopy-rs singular_slices convention.
    pub slices: Vec<Rewrite>,
}

impl Cone {
    /// Create a new cone.
    pub fn new(index: usize, source: Vec<Cospan>, target: Cospan, slices: Vec<Rewrite>) -> Self {
        Self {
            index,
            source,
            target,
            slices,
        }
    }

    /// The number of source cospans this cone contracts.
    pub fn source_size(&self) -> usize {
        self.source.len()
    }

    /// The number of singular slices in this cone.
    /// Same as source_size() - one singular slice per source cospan.
    pub fn len(&self) -> usize {
        self.source.len()
    }

    /// Check if this cone is empty (no source cospans).
    pub fn is_empty(&self) -> bool {
        self.source.is_empty()
    }
}

impl RewriteN {
    /// Compute where a regular height in the source maps to in the target.
    ///
    /// For a rewrite f: A → B, given a regular height h in A,
    /// returns the corresponding regular height in B.
    ///
    /// Based on homotopy-rs implementation.
    pub fn regular_image(&self, h: usize) -> usize {
        let mut height = h;
        for cone in &self.cones {
            // Only affect heights that are completely AFTER this cone's source range
            if height >= cone.index + cone.source_size() {
                // Shift down by the contraction: source_size cospans become 1
                height -= cone.source_size().saturating_sub(1);
            }
        }
        height
    }

    /// Compute the preimage of a regular height from target back to source.
    ///
    /// For a rewrite f: A → B, given a regular height h in B,
    /// returns the corresponding regular height in A.
    ///
    /// This is the inverse of regular_image.
    pub fn regular_preimage(&self, target_height: usize) -> usize {
        let mut source_height = target_height;
        for cone in &self.cones {
            // For each cone that starts before or at this target height,
            // we need to account for the expansion
            if target_height > cone.index {
                source_height += cone.source_size().saturating_sub(1);
            }
        }
        source_height
    }

    /// Compute the preimage of a singular height h in the target.
    ///
    /// For a rewrite f: A → B, given a singular height h in B,
    /// returns all singular heights in A that map to h.
    ///
    /// This handles three cases:
    /// 1. **Contraction**: A cone targets h and has len > 0. Returns all source
    ///    heights consumed by that cone.
    /// 2. **Insertion**: A cone targets h but has len == 0. Returns empty (no
    ///    source heights map to this insertion point).
    /// 3. **Passthrough**: No cone targets h. Returns the single source height
    ///    that passes through to h (computed from the cone structure).
    pub fn singular_preimage(&self, target_h: usize) -> Vec<usize> {
        let mut current_source = 0;
        let mut current_target = 0;

        for cone in &self.cones {
            // Handle passthroughs before this cone
            while current_target < cone.index {
                if current_target == target_h {
                    // Found passthrough at target_h
                    return vec![current_source];
                }
                current_source += 1;
                current_target += 1;
            }

            // Handle the cone itself
            if current_target == target_h {
                // This cone targets our height
                // Return all source heights in the cone's range
                // (empty if len == 0, i.e., insertion)
                return (current_source..current_source + cone.len()).collect();
            }

            current_source += cone.len();
            current_target += 1; // Cone produces 1 target cospan
        }

        // Handle passthroughs after all cones
        // target_h is beyond all cone indices
        let offset = target_h - current_target;
        vec![current_source + offset]
    }

    /// Get the target heights (where cones map to).
    pub fn targets(&self) -> impl Iterator<Item = usize> + '_ {
        self.cones.iter().map(|c| c.index)
    }

    /// Get the slice rewrite at a source singular height.
    ///
    /// For a rewrite f: A → B, given a source singular height h,
    /// returns the (n-1)-dimensional rewrite between source's singular
    /// slice at h and the corresponding target singular slice.
    ///
    /// Based on homotopy-rs: finds the cone containing this source height,
    /// then indexes into its singular_slices.
    pub fn slice(&self, source_height: usize) -> Rewrite {
        // Find which cone contains this source height
        let mut source_offset = 0;
        for cone in &self.cones {
            let source_end = source_offset + cone.len();
            if source_height >= source_offset && source_height < source_end {
                // Found the cone - index into its slices
                let local_idx = source_height - source_offset;
                if local_idx < cone.slices.len() {
                    return cone.slices[local_idx].clone();
                }
            }
            source_offset = source_end;
        }
        // Height is outside all cones (passthrough), return identity
        Rewrite::Identity
    }

    /// Get the cone that targets a specific height, if any.
    pub fn cone_over_target(&self, target_height: usize) -> Option<&Cone> {
        self.cones.iter().find(|c| c.index == target_height)
    }
}

// === Diagram methods ===

impl Diagram {
    /// The dimension of this diagram.
    pub fn dimension(&self) -> usize {
        match self {
            Diagram::Diagram0(_) => 0,
            Diagram::DiagramN(d) => 1 + d.source.dimension(),
        }
    }

    /// The zigzag length at the top level.
    ///
    /// For a 0-diagram, this is 0.
    /// For an n-diagram, this is the number of cospans.
    pub fn length(&self) -> usize {
        match self {
            Diagram::Diagram0(_) => 0,
            Diagram::DiagramN(d) => d.length(),
        }
    }

    /// Check if this is a 0-diagram.
    pub fn is_zero(&self) -> bool {
        matches!(self, Diagram::Diagram0(_))
    }

    /// Create an identity diagram over this one (adds one dimension).
    pub fn identity(self) -> DiagramN {
        DiagramN::identity(self)
    }

    /// Get the source of this diagram.
    ///
    /// For a 0-diagram, returns None.
    /// For an n-diagram, returns the source (n-1)-diagram.
    pub fn source(&self) -> Option<&Diagram> {
        match self {
            Diagram::Diagram0(_) => None,
            Diagram::DiagramN(d) => Some(&d.source),
        }
    }

    /// Get the target of this diagram.
    pub fn target(&self) -> Option<Diagram> {
        match self {
            Diagram::Diagram0(_) => None,
            Diagram::DiagramN(d) => Some(d.target()),
        }
    }

    /// Check if this diagram is globular.
    ///
    /// A diagram is globular if:
    /// - It's a 0-diagram, OR
    /// - Its source and target are equal globular diagrams
    pub fn is_globular(&self) -> bool {
        match self {
            Diagram::Diagram0(_) => true,
            Diagram::DiagramN(d) => {
                let src = d.source();
                let tgt = d.target();
                src.is_globular() && tgt.is_globular() && src == &tgt
            }
        }
    }
}

impl From<Generator> for Diagram {
    fn from(g: Generator) -> Self {
        Diagram::Diagram0(g)
    }
}

impl From<DiagramN> for Diagram {
    fn from(d: DiagramN) -> Self {
        Diagram::DiagramN(d)
    }
}

/// A map between diagrams (zigzag map in the iterated category).
///
/// This is the full morphism structure, containing the singular map
/// and all slice data recursively.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DiagramMap {
    /// The rewrite encoding this map
    pub rewrite: Rewrite,
}

impl DiagramMap {
    /// Create a diagram map from a rewrite.
    pub fn new(rewrite: Rewrite) -> Self {
        Self { rewrite }
    }

    /// Create an identity map.
    pub fn identity(_diagram: &Diagram) -> Self {
        Self {
            rewrite: Rewrite::Identity,
        }
    }

    /// Check if this is an identity map.
    pub fn is_identity(&self) -> bool {
        self.rewrite.is_identity()
    }

    /// Compose two diagram maps: (g ∘ f) where self = f and other = g.
    ///
    /// For identity maps, composition is trivial.
    /// For rewrites, we need to compose the underlying structure.
    pub fn compose(&self, other: &DiagramMap) -> DiagramMap {
        if self.is_identity() {
            other.clone()
        } else if other.is_identity() {
            self.clone()
        } else {
            // TODO: Implement full rewrite composition
            // For now, return other (this is a simplification)
            other.clone()
        }
    }

    /// Check if a singular height h is in the image of this map's singular component.
    ///
    /// For degeneracy maps, this checks if height h would be preserved
    /// (i.e., is not an inserted identity cospan position).
    pub fn has_singular_height_in_image(&self, h: usize) -> bool {
        match &self.rewrite {
            Rewrite::Identity => true, // Identity maps preserve all heights
            Rewrite::Rewrite0 { .. } => true, // 0-dim has no singular structure
            Rewrite::RewriteN(r) => {
                // Check if h is NOT an insertion point (not in any cone's empty-source positions)
                let insertion_points: std::collections::HashSet<usize> = r.cones
                    .iter()
                    .filter(|c| c.source.is_empty())
                    .map(|c| c.index)
                    .collect();
                !insertion_points.contains(&h)
            }
        }
    }
}

/// Direction for slice iteration.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SliceDirection {
    Forward,
    Backward,
}

/// Iterator over all slices of a diagram (interleaved regular and singular).
///
/// Returns slices in order: r₀, s₀, r₁, s₁, ..., sₙ₋₁, rₙ
pub struct Slices<'a> {
    diagram: &'a DiagramN,
    current: Option<Diagram>,
    direction: SliceDirection,
    cospan_index: usize,
}

impl<'a> Slices<'a> {
    fn new(diagram: &'a DiagramN) -> Self {
        Self {
            diagram,
            current: Some((*diagram.source).clone()),
            direction: SliceDirection::Forward,
            cospan_index: 0,
        }
    }
}

impl<'a> Iterator for Slices<'a> {
    type Item = Diagram;

    fn next(&mut self) -> Option<Self::Item> {
        // If we've exhausted all cospans, return the final slice
        if self.cospan_index >= self.diagram.cospans.len() {
            return self.current.take();
        }

        let current = self.current.as_ref()?;
        let cospan = &self.diagram.cospans[self.cospan_index];

        let next = match self.direction {
            SliceDirection::Forward => {
                // Apply forward rewrite to get singular slice
                self.direction = SliceDirection::Backward;
                cospan.forward.apply_forward(current)?
            }
            SliceDirection::Backward => {
                // Apply backward rewrite in reverse to get next regular slice
                self.direction = SliceDirection::Forward;
                self.cospan_index += 1;
                cospan.backward.apply_backward(current)?
            }
        };

        std::mem::replace(&mut self.current, Some(next))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = if self.current.is_none() {
            0
        } else {
            let cospans_left = self.diagram.cospans.len() - self.cospan_index;
            let slices_from_cospans = cospans_left * 2;
            let extra = match self.direction {
                SliceDirection::Forward => 1, // Still need to emit current regular + traverse remaining
                SliceDirection::Backward => 0, // Already emitted current, just need backward + remaining
            };
            slices_from_cospans + extra
        };
        (remaining, Some(remaining))
    }
}

impl<'a> ExactSizeIterator for Slices<'a> {
    fn len(&self) -> usize {
        self.size_hint().0
    }
}

impl<'a> std::iter::FusedIterator for Slices<'a> {}

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

    fn test_generator() -> Generator {
        Generator::new(0, 0, false)
    }

    fn gen(id: usize) -> Generator {
        Generator::new(id, 0, false)
    }

    fn diagram0(id: usize) -> Diagram {
        Diagram::Diagram0(gen(id))
    }

    #[test]
    fn test_diagram_0() {
        let g = test_generator();
        let d = Diagram::Diagram0(g.clone());
        assert_eq!(d.dimension(), 0);
        assert!(d.is_zero());
        assert!(d.is_globular());
    }

    #[test]
    fn test_diagram_n_identity() {
        let g = test_generator();
        let d0 = Diagram::Diagram0(g);
        let d1 = DiagramN::identity(d0);

        assert_eq!(d1.length(), 0);
        assert_eq!(Diagram::DiagramN(d1).dimension(), 1);
    }

    #[test]
    fn test_cospan_identity() {
        let c = Cospan::new(Rewrite::Identity, Rewrite::Identity);
        assert!(c.is_identity());
    }

    // ========== Slice computation tests ==========

    #[test]
    fn test_identity_diagram_slices() {
        // An identity diagram (length 0) has source = target
        let g = test_generator();
        let d0 = Diagram::Diagram0(g.clone());
        let id = DiagramN::identity(d0.clone());

        // Regular slice 0 is the source
        assert_eq!(id.regular_slice(0), Some(d0.clone()));

        // Target should equal source for identity
        assert_eq!(id.target(), d0);

        // No singular slices for identity diagram
        assert_eq!(id.singular_slice(0), None);

        // Out of bounds
        assert_eq!(id.regular_slice(1), None);
    }

    #[test]
    fn test_identity_rewrite_application() {
        // Identity rewrite should not change the diagram
        let d = diagram0(0);
        assert_eq!(Rewrite::Identity.apply_forward(&d), Some(d.clone()));
        assert_eq!(Rewrite::Identity.apply_backward(&d), Some(d.clone()));
    }

    #[test]
    fn test_rewrite0_application() {
        let src = gen(0);
        let tgt = gen(1);
        let rewrite = Rewrite::Rewrite0 {
            source: src.clone(),
            target: tgt.clone(),
        };

        // Forward: source -> target
        let d_src = Diagram::Diagram0(src.clone());
        let d_tgt = Diagram::Diagram0(tgt.clone());
        assert_eq!(rewrite.apply_forward(&d_src), Some(d_tgt.clone()));

        // Backward: target -> source
        assert_eq!(rewrite.apply_backward(&d_tgt), Some(d_src.clone()));

        // Mismatched source should fail
        let d_other = diagram0(2);
        assert_eq!(rewrite.apply_forward(&d_other), None);
        assert_eq!(rewrite.apply_backward(&d_other), None);
    }

    #[test]
    fn test_simple_cospan_slices() {
        // A diagram with one cospan: A -> X <- B
        // Where forward: A -> X and backward: B -> X
        let a = gen(0);
        let b = gen(1);
        let x = gen(2);

        let forward = Rewrite::Rewrite0 {
            source: a.clone(),
            target: x.clone(),
        };
        let backward = Rewrite::Rewrite0 {
            source: b.clone(),
            target: x.clone(),
        };
        let cospan = Cospan::new(forward, backward);

        // Create the 1-diagram with source A
        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source.clone(), vec![cospan]);

        // Regular slices
        assert_eq!(diag.regular_slice(0), Some(Diagram::Diagram0(a.clone())));
        assert_eq!(diag.regular_slice(1), Some(Diagram::Diagram0(b.clone())));
        assert_eq!(diag.regular_slice(2), None);

        // Singular slices
        assert_eq!(diag.singular_slice(0), Some(Diagram::Diagram0(x.clone())));
        assert_eq!(diag.singular_slice(1), None);

        // Target should be the last regular slice
        assert_eq!(diag.target(), Diagram::Diagram0(b.clone()));
    }

    #[test]
    fn test_identity_cospan_slices() {
        // A diagram with identity cospans: A -> A <- A (weak identity)
        let a = gen(0);

        let cospan = Cospan::new(Rewrite::Identity, Rewrite::Identity);

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source.clone(), vec![cospan]);

        // All slices should be A
        assert_eq!(diag.regular_slice(0), Some(Diagram::Diagram0(a.clone())));
        assert_eq!(diag.regular_slice(1), Some(Diagram::Diagram0(a.clone())));
        assert_eq!(diag.singular_slice(0), Some(Diagram::Diagram0(a.clone())));
        assert_eq!(diag.target(), Diagram::Diagram0(a.clone()));
    }

    #[test]
    fn test_multiple_cospans() {
        // A diagram with two cospans: A -> X <- B -> Y <- C
        let a = gen(0);
        let b = gen(1);
        let c = gen(2);
        let x = gen(3);
        let y = gen(4);

        let cospan1 = Cospan::new(
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
            Rewrite::Rewrite0 { source: b.clone(), target: x.clone() },
        );
        let cospan2 = Cospan::new(
            Rewrite::Rewrite0 { source: b.clone(), target: y.clone() },
            Rewrite::Rewrite0 { source: c.clone(), target: y.clone() },
        );

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source, vec![cospan1, cospan2]);

        // Regular slices: A, B, C
        assert_eq!(diag.regular_slice(0), Some(Diagram::Diagram0(a.clone())));
        assert_eq!(diag.regular_slice(1), Some(Diagram::Diagram0(b.clone())));
        assert_eq!(diag.regular_slice(2), Some(Diagram::Diagram0(c.clone())));
        assert_eq!(diag.regular_slice(3), None);

        // Singular slices: X, Y
        assert_eq!(diag.singular_slice(0), Some(Diagram::Diagram0(x.clone())));
        assert_eq!(diag.singular_slice(1), Some(Diagram::Diagram0(y.clone())));
        assert_eq!(diag.singular_slice(2), None);

        // Target is C
        assert_eq!(diag.target(), Diagram::Diagram0(c.clone()));
    }

    #[test]
    fn test_slices_iterator() {
        // A diagram with one cospan: A -> X <- B
        let a = gen(0);
        let b = gen(1);
        let x = gen(2);

        let cospan = Cospan::new(
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
            Rewrite::Rewrite0 { source: b.clone(), target: x.clone() },
        );

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source, vec![cospan]);

        // slices() should yield: r0, s0, r1 = A, X, B
        let slices: Vec<_> = diag.slices().collect();
        assert_eq!(slices.len(), 3);
        assert_eq!(slices[0], Diagram::Diagram0(a.clone()));
        assert_eq!(slices[1], Diagram::Diagram0(x.clone()));
        assert_eq!(slices[2], Diagram::Diagram0(b.clone()));
    }

    #[test]
    fn test_slices_iterator_two_cospans() {
        // A diagram with two cospans: A -> X <- B -> Y <- C
        let a = gen(0);
        let b = gen(1);
        let c = gen(2);
        let x = gen(3);
        let y = gen(4);

        let cospan1 = Cospan::new(
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
            Rewrite::Rewrite0 { source: b.clone(), target: x.clone() },
        );
        let cospan2 = Cospan::new(
            Rewrite::Rewrite0 { source: b.clone(), target: y.clone() },
            Rewrite::Rewrite0 { source: c.clone(), target: y.clone() },
        );

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source, vec![cospan1, cospan2]);

        // slices() should yield: r0, s0, r1, s1, r2 = A, X, B, Y, C
        let slices: Vec<_> = diag.slices().collect();
        assert_eq!(slices.len(), 5);
        assert_eq!(slices[0], Diagram::Diagram0(a.clone()));
        assert_eq!(slices[1], Diagram::Diagram0(x.clone()));
        assert_eq!(slices[2], Diagram::Diagram0(b.clone()));
        assert_eq!(slices[3], Diagram::Diagram0(y.clone()));
        assert_eq!(slices[4], Diagram::Diagram0(c.clone()));
    }

    #[test]
    fn test_slices_iterator_identity() {
        // Identity diagram has just one slice
        let a = gen(0);
        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::identity(source);

        let slices: Vec<_> = diag.slices().collect();
        assert_eq!(slices.len(), 1);
        assert_eq!(slices[0], Diagram::Diagram0(a.clone()));
    }

    #[test]
    fn test_regular_slices_iterator() {
        // A diagram with two cospans: A -> X <- B -> Y <- C
        let a = gen(0);
        let b = gen(1);
        let c = gen(2);
        let x = gen(3);
        let y = gen(4);

        let cospan1 = Cospan::new(
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
            Rewrite::Rewrite0 { source: b.clone(), target: x.clone() },
        );
        let cospan2 = Cospan::new(
            Rewrite::Rewrite0 { source: b.clone(), target: y.clone() },
            Rewrite::Rewrite0 { source: c.clone(), target: y.clone() },
        );

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source, vec![cospan1, cospan2]);

        // regular_slices() should yield: A, B, C
        let regular: Vec<_> = diag.regular_slices().collect();
        assert_eq!(regular.len(), 3);
        assert_eq!(regular[0], Diagram::Diagram0(a.clone()));
        assert_eq!(regular[1], Diagram::Diagram0(b.clone()));
        assert_eq!(regular[2], Diagram::Diagram0(c.clone()));
    }

    #[test]
    fn test_singular_slices_iterator() {
        // A diagram with two cospans: A -> X <- B -> Y <- C
        let a = gen(0);
        let b = gen(1);
        let c = gen(2);
        let x = gen(3);
        let y = gen(4);

        let cospan1 = Cospan::new(
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
            Rewrite::Rewrite0 { source: b.clone(), target: x.clone() },
        );
        let cospan2 = Cospan::new(
            Rewrite::Rewrite0 { source: b.clone(), target: y.clone() },
            Rewrite::Rewrite0 { source: c.clone(), target: y.clone() },
        );

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source, vec![cospan1, cospan2]);

        // singular_slices() should yield: X, Y
        let singular: Vec<_> = diag.singular_slices().collect();
        assert_eq!(singular.len(), 2);
        assert_eq!(singular[0], Diagram::Diagram0(x.clone()));
        assert_eq!(singular[1], Diagram::Diagram0(y.clone()));
    }

    #[test]
    fn test_slices_iterator_len() {
        let a = gen(0);
        let b = gen(1);
        let x = gen(2);

        let cospan = Cospan::new(
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
            Rewrite::Rewrite0 { source: b.clone(), target: x.clone() },
        );

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source, vec![cospan]);

        let mut iter = diag.slices();
        assert_eq!(iter.len(), 3);
        iter.next();
        assert_eq!(iter.len(), 2);
        iter.next();
        assert_eq!(iter.len(), 1);
        iter.next();
        assert_eq!(iter.len(), 0);
    }

    #[test]
    fn test_globular_identity() {
        // An identity diagram over a point is globular
        let g = test_generator();
        let d0 = Diagram::Diagram0(g);
        let d1 = DiagramN::identity(d0.clone());

        assert!(Diagram::DiagramN(d1).is_globular());
    }

    #[test]
    fn test_globular_non_identity() {
        // A non-identity diagram with source != target is not globular
        let a = gen(0);
        let b = gen(1);
        let x = gen(2);

        let cospan = Cospan::new(
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
            Rewrite::Rewrite0 { source: b.clone(), target: x.clone() },
        );

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source, vec![cospan]);

        // Source is A, target is B, so not globular
        assert!(!Diagram::DiagramN(diag).is_globular());
    }

    #[test]
    fn test_globular_loop() {
        // A diagram with source = target is globular (a loop)
        let a = gen(0);
        let x = gen(1);

        // Cospan: A -> X <- A
        let cospan = Cospan::new(
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
            Rewrite::Rewrite0 { source: a.clone(), target: x.clone() },
        );

        let source = Diagram::Diagram0(a.clone());
        let diag = DiagramN::new(source, vec![cospan]);

        // Source is A, target is A, so globular
        assert!(Diagram::DiagramN(diag).is_globular());
    }
}