geoit 0.0.2

use crate::algebra::blade_new::{blade_to_mask, mask_to_blade, BladeKey};
use crate::scalar::{Rat, Scalar};

// Legacy re-export for backward compatibility
use crate::algebra::blade_new::BladeMask;

/// Sparse multivector: only nonzero terms stored, sorted by BladeKey.
/// Invariant: no entry has a zero coefficient. Sorted by BladeKey ordering.
#[derive(Clone, Debug)]
pub struct Mv {
    pub(crate) terms: Vec<(BladeKey, Scalar)>,
}

impl Mv {
    /// The zero multivector.
    pub fn new() -> Self {
        Mv { terms: Vec::new() }
    }

    // ─── BladeKey-native constructors ───

    /// A single-term multivector from a BladeKey.
    pub fn term_bk(blade: BladeKey, coeff: Scalar) -> Self {
        if coeff.is_zero() {
            return Mv::new();
        }
        Mv {
            terms: vec![(blade, coeff)],
        }
    }

    /// A scalar multivector (grade 0).
    pub fn scalar(value: Scalar) -> Self {
        Self::term_bk(BladeKey::SCALAR, value)
    }

    /// A single basis generator g_k.
    pub fn generator(k: u8) -> Self {
        Self::term_bk(BladeKey::generator(k as u16), Scalar::from(1i64))
    }

    // ─── Legacy mask-based constructors ───

    /// A single-term multivector from a legacy bitmask.
    pub fn term(mask: BladeMask, coeff: Scalar) -> Self {
        Self::term_bk(mask_to_blade(mask), coeff)
    }

    /// Construct from a slice of (mask, Rat) pairs.
    pub fn from_rat_terms(pairs: &[(BladeMask, Rat)]) -> Self {
        let mut mv = Mv::new();
        for &(mask, coeff) in pairs {
            if !coeff.is_zero() {
                mv.add_term(mask, Scalar::Rat(coeff));
            }
        }
        mv
    }

    // ─── BladeKey-native mutation ───

    /// Add a term by BladeKey. Combines with existing if present. Removes if result is zero.
    pub fn add_term_bk(&mut self, blade: BladeKey, coeff: Scalar) {
        if coeff.is_zero() {
            return;
        }
        match self.terms.binary_search_by(|(k, _)| k.cmp(&blade)) {
            Ok(idx) => {
                self.terms[idx].1 = self.terms[idx].1.clone() + coeff;
                if self.terms[idx].1.is_zero() {
                    let _ = self.terms.remove(idx);
                }
            }
            Err(idx) => {
                self.terms.insert(idx, (blade, coeff));
            }
        }
    }

    // ─── Legacy mask-based mutation ───

    /// Add a term by legacy bitmask.
    pub fn add_term(&mut self, mask: BladeMask, coeff: Scalar) {
        self.add_term_bk(mask_to_blade(mask), coeff);
    }

    // ─── BladeKey-native accessors ───

    /// Get the coefficient of a BladeKey. Zero if absent.
    pub fn coefficient_bk(&self, blade: &BladeKey) -> Scalar {
        match self.terms.binary_search_by(|(k, _)| k.cmp(blade)) {
            Ok(idx) => self.terms[idx].1.clone(),
            Err(_) => Scalar::from(0i64),
        }
    }

    /// Iterator over (&BladeKey, &Scalar) pairs — the native interface.
    pub fn blades_bk(&self) -> impl Iterator<Item = (&BladeKey, &Scalar)> {
        self.terms.iter().map(|(k, v)| (k, v))
    }

    // ─── Common accessors ───

    pub fn is_zero(&self) -> bool {
        self.terms.is_empty()
    }

    /// Get the coefficient of a legacy bitmask. Zero if absent.
    pub fn coefficient(&self, mask: BladeMask) -> Scalar {
        self.coefficient_bk(&mask_to_blade(mask))
    }

    /// Keep only terms of the specified grade.
    pub fn grade_project(&self, k: u8) -> Self {
        let terms: Vec<_> = self
            .terms
            .iter()
            .filter(|(blade, _)| blade.grade() == k)
            .cloned()
            .collect();
        Mv { terms }
    }

    /// Scale every coefficient by a scalar.
    pub fn scale(&self, s: &Scalar) -> Self {
        if s.is_zero() {
            return Mv::new();
        }
        let terms: Vec<_> = self
            .terms
            .iter()
            .map(|(blade, coeff)| (blade.clone(), coeff.clone() * s.clone()))
            .filter(|(_, c)| !c.is_zero())
            .collect();
        Mv { terms }
    }

    /// Legacy iterator: yields (BladeMask, &Scalar). Prefer blades_bk().
    pub fn blades(&self) -> BladesIter<'_> {
        BladesIter {
            inner: self.terms.iter(),
        }
    }

    /// Number of nonzero terms.
    pub fn len(&self) -> usize {
        self.terms.len()
    }

    /// Whether this Mv has no nonzero terms.
    pub fn is_empty(&self) -> bool {
        self.terms.is_empty()
    }

    /// Construct from an iterator of (mask, coeff) pairs.
    pub fn from_terms(pairs: impl IntoIterator<Item = (BladeMask, Scalar)>) -> Self {
        let mut mv = Mv::new();
        for (mask, coeff) in pairs {
            if !coeff.is_zero() {
                mv.add_term(mask, coeff);
            }
        }
        mv
    }

    /// Consume into a Vec of (mask, coeff) pairs — legacy format.
    pub fn into_terms(self) -> Vec<(BladeMask, Scalar)> {
        self.terms
            .into_iter()
            .map(|(blade, coeff)| (blade_to_mask(&blade), coeff))
            .collect()
    }

    /// Iterate over (mask, &coeff) pairs — owned mask, borrowed coeff.
    pub fn iter(&self) -> impl Iterator<Item = (BladeMask, &Scalar)> {
        self.terms
            .iter()
            .map(|(blade, coeff)| (blade_to_mask(blade), coeff))
    }

    /// Display with semantic generator names from a Signature.
    pub fn display_with<'a>(
        &'a self,
        sig: &'a crate::algebra::signature::Signature,
    ) -> MvDisplay<'a> {
        MvDisplay { mv: self, sig }
    }
}

/// Legacy-compatible iterator that yields (BladeMask, &Scalar).
pub struct BladesIter<'a> {
    inner: std::slice::Iter<'a, (BladeKey, Scalar)>,
}

impl<'a> Iterator for BladesIter<'a> {
    type Item = (BladeMask, &'a Scalar);
    fn next(&mut self) -> Option<Self::Item> {
        self.inner
            .next()
            .map(|(blade, coeff)| (blade_to_mask(blade), coeff))
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

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

impl std::ops::Neg for Mv {
    type Output = Self;
    fn neg(self) -> Self {
        let terms = self
            .terms
            .into_iter()
            .map(|(blade, coeff)| (blade, -coeff))
            .collect();
        Mv { terms }
    }
}

impl std::ops::Add for Mv {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        let mut result = self;
        for (blade, coeff) in rhs.terms {
            result.add_term_bk(blade, coeff);
        }
        result
    }
}

impl std::ops::Sub for Mv {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        self + (-rhs)
    }
}

impl std::ops::AddAssign for Mv {
    fn add_assign(&mut self, rhs: Self) {
        for (blade, coeff) in rhs.terms {
            self.add_term_bk(blade, coeff);
        }
    }
}

impl PartialEq for Mv {
    fn eq(&self, other: &Self) -> bool {
        self.terms == other.terms
    }
}

impl Eq for Mv {}

impl std::fmt::Display for Mv {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.is_zero() {
            return write!(f, "0");
        }
        let mut first = true;
        for (blade, coeff) in &self.terms {
            if !first {
                write!(f, " + ")?;
            }
            first = false;
            if blade.is_scalar() {
                write!(f, "{}", coeff)?;
            } else {
                write!(f, "{}·{}", coeff, blade)?;
            }
        }
        Ok(())
    }
}

/// Wrapper for displaying an Mv with semantic generator names from a Signature.
pub struct MvDisplay<'a> {
    mv: &'a Mv,
    sig: &'a crate::algebra::signature::Signature,
}

impl<'a> std::fmt::Display for MvDisplay<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.mv.is_zero() {
            return write!(f, "0");
        }
        let mut first = true;
        for (blade, coeff) in &self.mv.terms {
            if !first {
                write!(f, " + ")?;
            }
            first = false;
            if blade.is_scalar() {
                write!(f, "{}", coeff)?;
            } else {
                write!(f, "{}·", coeff)?;
                let mut first_gen = true;
                for &k in blade.indices() {
                    if !first_gen {
                        write!(f, "∧")?;
                    }
                    first_gen = false;
                    write!(f, "{}", self.sig.generator_name(k as u8))?;
                }
            }
        }
        Ok(())
    }
}

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

    #[test]
    fn zero_mv() {
        let mv = Mv::new();
        assert!(mv.is_zero());
        assert_eq!(mv.len(), 0);
    }

    #[test]
    fn single_term() {
        let mv = Mv::term(0b001, Scalar::from(3));
        assert!(!mv.is_zero());
        assert_eq!(mv.len(), 1);
        assert_eq!(mv.coefficient(0b001), Scalar::from(3));
        assert_eq!(mv.coefficient(0b010), Scalar::from(0));
    }

    #[test]
    fn zero_coeff_not_stored() {
        let mv = Mv::term(0b001, Scalar::from(0));
        assert!(mv.is_zero());
    }

    #[test]
    fn addition() {
        let a = Mv::term(0b001, Scalar::from(3));
        let b = Mv::term(0b010, Scalar::from(4));
        let c = a + b;
        assert_eq!(c.len(), 2);
        assert_eq!(c.coefficient(0b001), Scalar::from(3));
        assert_eq!(c.coefficient(0b010), Scalar::from(4));
    }

    #[test]
    fn additive_cancellation() {
        let a = Mv::term(0b001, Scalar::from(3));
        let b = Mv::term(0b001, Scalar::from(-3));
        let c = a + b;
        assert!(c.is_zero());
    }

    #[test]
    fn additive_identity() {
        let a = Mv::term(0b101, Scalar::from(7));
        let z = Mv::new();
        assert_eq!(a.clone() + z, a);
    }

    #[test]
    fn negation() {
        let a = Mv::term(0b001, Scalar::from(5));
        let b = -a;
        assert_eq!(b.coefficient(0b001), Scalar::from(-5));
    }

    #[test]
    fn subtraction() {
        let a = Mv::term(0b001, Scalar::from(5));
        let b = Mv::term(0b001, Scalar::from(2));
        let c = a - b;
        assert_eq!(c.coefficient(0b001), Scalar::from(3));
    }

    #[test]
    fn scale() {
        let a = Mv::term(0b001, Scalar::from(3));
        let b = a.scale(&Scalar::from(4));
        assert_eq!(b.coefficient(0b001), Scalar::from(12));
    }

    #[test]
    fn scale_by_zero() {
        let a = Mv::term(0b001, Scalar::from(3));
        let b = a.scale(&Scalar::from(0));
        assert!(b.is_zero());
    }

    #[test]
    fn grade_project() {
        let mut mv = Mv::new();
        mv.add_term(0b000, Scalar::from(1)); // scalar
        mv.add_term(0b001, Scalar::from(2)); // grade 1
        mv.add_term(0b010, Scalar::from(3)); // grade 1
        mv.add_term(0b011, Scalar::from(4)); // grade 2

        let g1 = mv.grade_project(1);
        assert_eq!(g1.len(), 2);
        assert_eq!(g1.coefficient(0b001), Scalar::from(2));
        assert_eq!(g1.coefficient(0b010), Scalar::from(3));

        let g0 = mv.grade_project(0);
        assert_eq!(g0.len(), 1);
        assert_eq!(g0.coefficient(0b000), Scalar::from(1));
    }

    #[test]
    fn generator() {
        let g2 = Mv::generator(2);
        assert_eq!(g2.coefficient(0b100), Scalar::from(1));
        assert_eq!(g2.len(), 1);
    }

    #[test]
    fn from_rat_terms() {
        let mv = Mv::from_rat_terms(&[(0b01, Rat::new(1, 2)), (0b10, Rat::from(3))]);
        assert_eq!(mv.len(), 2);
    }
}