neberu 0.0.0

// ============================================================
// EXPR — THE FREE ALGEBRA
// ============================================================
//
// An Expr is a sparse linear combination of Words with Rat coefficients.
// Σ cᵢ · wᵢ  where cᵢ ∈ Rat, wᵢ ∈ Word.
//
// Before governance: unconstrained. After governance: canonical.
// Invariant: no zero coefficients stored (maintained by all operations).

use crate::rat::Rat;
use crate::trit::N;
use crate::word::Word;
use std::collections::BTreeMap;

/// A sparse linear combination of Words over Rat.
#[derive(Clone, PartialEq, Eq)]
pub struct Expr {
    terms: BTreeMap<Word, Rat>,
}

impl Expr {
    pub fn zero() -> Expr {
        Expr {
            terms: BTreeMap::new(),
        }
    }

    pub fn scalar(r: Rat) -> Expr {
        if r.is_zero() {
            return Expr::zero();
        }
        let mut terms = BTreeMap::new();
        terms.insert(Word::scalar(), r);
        Expr { terms }
    }

    pub fn int(n: i64) -> Expr {
        Expr::scalar(Rat::integer(n))
    }

    pub fn gen(g: crate::gen::Gen) -> Expr {
        let mut terms = BTreeMap::new();
        terms.insert(Word::gen(g), Rat::one());
        Expr { terms }
    }

    pub fn term(coeff: Rat, word: Word) -> Expr {
        if coeff.is_zero() {
            return Expr::zero();
        }
        let mut terms = BTreeMap::new();
        terms.insert(word, coeff);
        Expr { terms }
    }

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

    /// Get coefficient of a specific word.
    pub fn coeff(&self, word: &Word) -> Rat {
        self.terms.get(word).copied().unwrap_or(Rat::zero())
    }

    /// Is this a pure scalar? If so, return the scalar.
    pub fn as_scalar(&self) -> Option<Rat> {
        if self.is_zero() {
            return None;
        }
        if self.terms.len() == 1 {
            let (word, coeff) = self.terms.iter().next().unwrap();
            if word.is_scalar() {
                return Some(*coeff);
            }
        }
        None
    }

    /// Iterate over (word, coefficient) pairs.
    pub fn terms(&self) -> impl Iterator<Item = (&Word, &Rat)> {
        self.terms.iter()
    }

    /// Addition: merge terms, combine coefficients.
    pub fn add(&self, other: &Expr) -> Expr {
        let mut terms = self.terms.clone();
        for (word, coeff) in &other.terms {
            let entry = terms.entry(word.clone()).or_insert(Rat::zero());
            *entry = entry.add(*coeff);
        }
        terms.retain(|_, c| !c.is_zero());
        Expr { terms }
    }

    pub fn sub(&self, other: &Expr) -> Expr {
        self.add(&other.neg())
    }

    pub fn neg(&self) -> Expr {
        let terms = self
            .terms
            .iter()
            .map(|(w, c)| (w.clone(), c.neg()))
            .collect();
        Expr { terms }
    }

    pub fn scale(&self, s: Rat) -> Expr {
        if s.is_zero() {
            return Expr::zero();
        }
        let terms = self
            .terms
            .iter()
            .map(|(w, c)| (w.clone(), c.mul(s)))
            .filter(|(_, c)| !c.is_zero())
            .collect();
        Expr { terms }
    }

    /// Free product: distribute and concatenate (no governance).
    pub fn mul_free(&self, other: &Expr) -> Expr {
        let mut result = Expr::zero();
        for (w1, c1) in &self.terms {
            for (w2, c2) in &other.terms {
                let coeff = c1.mul(*c2);
                if !coeff.is_zero() {
                    let word = w1.mul(w2);
                    result = result.add(&Expr::term(coeff, word));
                }
            }
        }
        result
    }

    /// Grade projection: keep only terms of a given grade.
    pub fn grade(&self, k: usize) -> Expr {
        let terms = self
            .terms
            .iter()
            .filter(|(w, _)| w.grade() == k)
            .map(|(w, c)| (w.clone(), *c))
            .collect();
        Expr { terms }
    }

    /// Maximum grade present.
    pub fn max_grade(&self) -> usize {
        self.terms.keys().map(|w| w.grade()).max().unwrap_or(0)
    }

    /// The reverse: reverse each word, apply grade sign.
    pub fn reverse(&self) -> Expr {
        let mut result = Expr::zero();
        for (word, coeff) in &self.terms {
            let (rev_word, sign_flip) = word.reverse_sign();
            let signed_coeff = if sign_flip { coeff.neg() } else { *coeff };
            result = result.add(&Expr::term(signed_coeff, rev_word));
        }
        result
    }
}

impl std::fmt::Debug for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self}")
    }
}

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

impl std::hash::Hash for Expr {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        for (w, c) in &self.terms {
            w.hash(state);
            c.hash(state);
        }
    }
}

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

    fn ei(idx: u32) -> Gen {
        Gen::imaginary(idx)
    }

    #[test]
    fn zero() {
        assert!(Expr::zero().is_zero());
    }

    #[test]
    fn scalar_42() {
        let e = Expr::int(42);
        assert_eq!(e.as_scalar(), Some(Rat::integer(42)));
    }

    #[test]
    fn gen_expr() {
        let e = Expr::gen(ei(0));
        assert_eq!(e.num_terms(), 1);
        assert_eq!(e.coeff(&Word::gen(ei(0))), Rat::one());
    }

    #[test]
    fn add_like_terms() {
        let e = Expr::gen(ei(0)).add(&Expr::gen(ei(0)));
        assert_eq!(e.coeff(&Word::gen(ei(0))), Rat::integer(2));
    }

    #[test]
    fn cancel() {
        let e = Expr::gen(ei(0)).add(&Expr::gen(ei(0)).neg());
        assert!(e.is_zero());
    }

    #[test]
    fn free_product() {
        let e1 = Expr::gen(ei(0));
        let e2 = Expr::gen(ei(1));
        let p = e1.mul_free(&e2);
        let w = Word::from_gens(&[ei(0), ei(1)]);
        assert_eq!(p.coeff(&w), Rat::one());
    }

    #[test]
    fn grade_projection() {
        let e = Expr::int(3)
            .add(&Expr::gen(ei(0)))
            .add(&Expr::term(Rat::one(), Word::from_gens(&[ei(0), ei(1)])));
        assert_eq!(e.grade(0), Expr::int(3));
        assert_eq!(e.grade(1), Expr::gen(ei(0)));
        assert_eq!(
            e.grade(2),
            Expr::term(Rat::one(), Word::from_gens(&[ei(0), ei(1)]))
        );
    }
}