//! Zigzags and zigzag maps
//!
//! A zigzag of length n is a diagram:
//! ```text
//! X(r₀) → X(s₀) ← X(r₁) → X(s₁) ← ... → X(sₙ₋₁) ← X(rₙ)
//! ```
//!
//! - Regular objects X(rⱼ) for j ∈ {0,...,n}
//! - Singular objects X(sᵢ) for i ∈ {0,...,n-1}
//!
//! A zigzag map f: X → Y consists of:
//! - A singular map fˢ: n → m in Δ₊
//! - A derived regular map fʳ = (Rfˢ)ᵒᵖ: m+1 → n+1
//! - Slice maps at each height

use crate::monotone::MonotoneMap;

/// A zigzag in a category C, parameterized by the object type T.
///
/// A zigzag of length n has:
/// - n+1 regular objects
/// - n singular objects
/// - n cospans connecting them
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Zigzag<T> {
    /// Regular objects X(r₀), X(r₁), ..., X(rₙ) — length n+1
    pub regular: Vec<T>,
    /// Singular objects X(s₀), X(s₁), ..., X(sₙ₋₁) — length n
    pub singular: Vec<T>,
    // Note: The cospan structure maps (forward/backward arrows) are implicit
    // in the diagram representation; they're encoded in the Rewrite/Cospan types.
}

impl<T> Zigzag<T> {
    /// Create a new zigzag with the given regular and singular objects.
    ///
    /// # Panics
    /// Panics if regular.len() != singular.len() + 1
    pub fn new(regular: Vec<T>, singular: Vec<T>) -> Self {
        assert_eq!(
            regular.len(),
            singular.len() + 1,
            "Zigzag requires regular.len() = singular.len() + 1, got {} and {}",
            regular.len(),
            singular.len()
        );
        Self { regular, singular }
    }

    /// Create a zigzag of length 0 (single regular object, no singular objects).
    pub fn point(obj: T) -> Self {
        Self {
            regular: vec![obj],
            singular: vec![],
        }
    }

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

    /// Number of regular heights (length + 1).
    pub fn regular_count(&self) -> usize {
        self.regular.len()
    }

    /// Number of singular heights (same as length).
    pub fn singular_count(&self) -> usize {
        self.singular.len()
    }

    /// Get regular object at height h.
    pub fn regular_at(&self, h: usize) -> Option<&T> {
        self.regular.get(h)
    }

    /// Get singular object at height h.
    pub fn singular_at(&self, h: usize) -> Option<&T> {
        self.singular.get(h)
    }
}

impl<T: Clone> Zigzag<T> {
    /// Map a function over all objects in the zigzag.
    pub fn map<U, F: Fn(&T) -> U>(&self, f: F) -> Zigzag<U> {
        Zigzag {
            regular: self.regular.iter().map(&f).collect(),
            singular: self.singular.iter().map(&f).collect(),
        }
    }
}

/// A map between zigzags.
///
/// Given zigzags X (length n) and Y (length m), a zigzag map f: X → Y consists of:
/// - A singular map fˢ: n → m in Δ₊
/// - A derived regular map fʳ = (Rfˢ)ᵒᵖ: m+1 → n+1 in Δ₌
/// - Slice data at each height (stored separately in the category-specific implementation)
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ZigzagMap<S> {
    /// The singular map fˢ: source_length → target_length
    pub singular_map: MonotoneMap,
    /// Slice data at regular heights (one per target regular height)
    pub regular_slices: Vec<S>,
    /// Slice data at singular heights (one per source singular height)
    pub singular_slices: Vec<S>,
}

impl<S> ZigzagMap<S> {
    /// Create a new zigzag map.
    ///
    /// # Arguments
    /// - `singular_map`: The singular map fˢ: n → m
    /// - `regular_slices`: Slice data at each target regular height (length m+1)
    /// - `singular_slices`: Slice data at each source singular height (length n)
    ///
    /// # Panics
    /// Panics if slice counts don't match the singular map dimensions.
    pub fn new(
        singular_map: MonotoneMap,
        regular_slices: Vec<S>,
        singular_slices: Vec<S>,
    ) -> Self {
        let n = singular_map.source_size();
        let m = singular_map.target_size();

        assert_eq!(
            regular_slices.len(),
            m + 1,
            "Expected {} regular slices, got {}",
            m + 1,
            regular_slices.len()
        );
        assert_eq!(
            singular_slices.len(),
            n,
            "Expected {} singular slices, got {}",
            n,
            singular_slices.len()
        );

        Self {
            singular_map,
            regular_slices,
            singular_slices,
        }
    }

    /// The source zigzag length (n).
    pub fn source_length(&self) -> usize {
        self.singular_map.source_size()
    }

    /// The target zigzag length (m).
    pub fn target_length(&self) -> usize {
        self.singular_map.target_size()
    }

    /// The regular map fʳ = (Rfˢ)ᵒᵖ: m+1 → n+1.
    ///
    /// This is derived from the singular map via Wraith's R equivalence.
    pub fn regular_map(&self) -> MonotoneMap {
        self.singular_map.wraith_r()
    }

    /// Get the regular slice at target height j.
    pub fn regular_slice(&self, j: usize) -> Option<&S> {
        self.regular_slices.get(j)
    }

    /// Get the singular slice at source height i.
    pub fn singular_slice(&self, i: usize) -> Option<&S> {
        self.singular_slices.get(i)
    }
}

impl<S: Clone> ZigzagMap<S> {
    /// Compose two zigzag maps: (g ∘ f) where f: X → Y and g: Y → W.
    ///
    /// Composition rules:
    /// - (g ∘ f)ˢ = gˢ ∘ fˢ
    /// - (g ∘ f)(sⱼ) = g(s_{fˢ(j)}) ∘ f(sⱼ)
    /// - (g ∘ f)(rᵢ) = g(rᵢ) ∘ f(r_{gʳ(i)})
    ///
    /// Note: This requires a way to compose slice data. The `compose_slices` function
    /// is provided to combine slice morphisms.
    pub fn compose<F>(&self, other: &ZigzagMap<S>, compose_slices: F) -> ZigzagMap<S>
    where
        F: Fn(&S, &S) -> S,
    {
        // Compose singular maps
        let composed_singular = self.singular_map.compose(&other.singular_map);

        // Get other's regular map for indexing
        let other_regular = other.regular_map();

        // Compose regular slices: (g ∘ f)(rᵢ) = g(rᵢ) ∘ f(r_{gʳ(i)})
        let composed_regular: Vec<S> = (0..other.target_length() + 1)
            .map(|i| {
                let g_r_i = other_regular.apply(i);
                let f_slice = &self.regular_slices[g_r_i];
                let g_slice = &other.regular_slices[i];
                compose_slices(f_slice, g_slice)
            })
            .collect();

        // Compose singular slices: (g ∘ f)(sⱼ) = g(s_{fˢ(j)}) ∘ f(sⱼ)
        let composed_singular_slices: Vec<S> = (0..self.source_length())
            .map(|j| {
                let f_s_j = self.singular_map.apply(j);
                let f_slice = &self.singular_slices[j];
                let g_slice = &other.singular_slices[f_s_j];
                compose_slices(f_slice, g_slice)
            })
            .collect();

        ZigzagMap {
            singular_map: composed_singular,
            regular_slices: composed_regular,
            singular_slices: composed_singular_slices,
        }
    }
}

/// The π functor: Z(C) → Δ₊, sending zigzags to their lengths.
pub fn pi_length<T>(zigzag: &Zigzag<T>) -> usize {
    zigzag.length()
}

/// The π functor on maps: sends a zigzag map to its singular map.
pub fn pi_map<S>(map: &ZigzagMap<S>) -> &MonotoneMap {
    &map.singular_map
}

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

    #[test]
    fn test_zigzag_point() {
        let z: Zigzag<i32> = Zigzag::point(42);
        assert_eq!(z.length(), 0);
        assert_eq!(z.regular_count(), 1);
        assert_eq!(z.singular_count(), 0);
    }

    #[test]
    fn test_zigzag_construction() {
        let z: Zigzag<char> = Zigzag::new(
            vec!['a', 'b', 'c'],
            vec!['x', 'y'],
        );
        assert_eq!(z.length(), 2);
        assert_eq!(z.regular_at(0), Some(&'a'));
        assert_eq!(z.singular_at(1), Some(&'y'));
    }

    #[test]
    fn test_zigzag_map_identity() {
        // Identity map on a length-2 zigzag
        let id_singular = MonotoneMap::identity(2);
        let map: ZigzagMap<()> = ZigzagMap::new(
            id_singular,
            vec![(), (), ()], // 3 regular slices
            vec![(), ()],     // 2 singular slices
        );

        assert_eq!(map.source_length(), 2);
        assert_eq!(map.target_length(), 2);

        let reg_map = map.regular_map();
        assert!(reg_map.is_identity());
    }

    #[test]
    fn test_zigzag_map_regular_derived() {
        // Singular map: 1 → 2 given by [0] (maps 0 to 0)
        let singular = MonotoneMap::new(vec![0], 2);
        let map: ZigzagMap<()> = ZigzagMap::new(
            singular.clone(),
            vec![(), (), ()], // 3 regular slices (target has length 2)
            vec![()],         // 1 singular slice (source has length 1)
        );

        // Regular map should be R([0]): 3 → 2
        let reg = map.regular_map();
        assert_eq!(reg.source_size(), 3);
        assert_eq!(reg.target_size(), 2);
    }
}