geoit 0.0.2

use crate::scalar::Rat;
use std::collections::BTreeMap;

/// Multivariate polynomial over Rat.
///
/// Variables are indexed: variable i corresponds to a blade coefficient.
/// Terms are sparse: only nonzero coefficients stored.
/// Monomial ordering: graded reverse lexicographic (grlex).
#[derive(Clone, Debug)]
pub struct Poly {
    /// Exponent vector → coefficient. Key: exponents[j] = power of variable j.
    pub(crate) terms: BTreeMap<Vec<u8>, Rat>,
    /// Number of variables (fixed per polynomial system).
    pub(crate) num_vars: usize,
}

impl Poly {
    pub fn zero(num_vars: usize) -> Self {
        Poly {
            terms: BTreeMap::new(),
            num_vars,
        }
    }

    /// Number of variables in this polynomial system.
    pub fn num_vars(&self) -> usize {
        self.num_vars
    }

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

    pub fn constant(r: Rat, num_vars: usize) -> Self {
        if r.is_zero() {
            return Self::zero(num_vars);
        }
        let mut terms = BTreeMap::new();
        terms.insert(vec![0u8; num_vars], r);
        Poly { terms, num_vars }
    }

    pub fn variable(i: usize, num_vars: usize) -> Self {
        let mut exp = vec![0u8; num_vars];
        exp[i] = 1;
        let mut terms = BTreeMap::new();
        terms.insert(exp, Rat::ONE);
        Poly { terms, num_vars }
    }

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

    /// Add a term to the polynomial. If the exponent already exists,
    /// combines coefficients. Removes if result is zero.
    pub fn add_term(&mut self, exp: Vec<u8>, coeff: Rat) {
        if coeff.is_zero() {
            return;
        }
        let entry = self.terms.entry(exp).or_insert(Rat::ZERO);
        *entry += coeff;
        self.clean();
    }

    /// Create a polynomial from a single term.
    pub fn from_term(exp: Vec<u8>, coeff: Rat, num_vars: usize) -> Self {
        let mut p = Self::zero(num_vars);
        if !coeff.is_zero() {
            p.terms.insert(exp, coeff);
        }
        p
    }

    /// Get the coefficient of a specific variable (degree 1) in this polynomial.
    /// Returns zero if the variable doesn't appear at degree 1.
    pub fn coefficient_of_var(&self, var: usize) -> Rat {
        for (exp, &coeff) in &self.terms {
            if exp[var] == 1 {
                // Check this is purely x_var (no other variables in this term)
                let others: u8 = exp
                    .iter()
                    .enumerate()
                    .filter(|&(i, _)| i != var)
                    .map(|(_, &e)| e)
                    .sum();
                if others == 0 {
                    return coeff;
                }
            }
        }
        Rat::ZERO
    }

    /// Maximum total degree of any term in this polynomial.
    /// Returns 0 for the zero polynomial.
    pub fn max_degree(&self) -> usize {
        self.terms
            .keys()
            .map(|exp| exp.iter().map(|&e| e as usize).sum::<usize>())
            .max()
            .unwrap_or(0)
    }

    fn clean(&mut self) {
        self.terms.retain(|_, c| !c.is_zero());
    }

    /// Total degree of a single monomial exponent vector.
    /// (The public method `max_degree(&self)` returns the max over all terms.)
    fn total_degree(exp: &[u8]) -> u16 {
        exp.iter().map(|&e| e as u16).sum()
    }

    /// Compare two monomials under grlex ordering.
    /// Returns Ordering: Greater means "larger" monomial (comes first in leading term).
    fn cmp_grlex(a: &[u8], b: &[u8]) -> std::cmp::Ordering {
        let da = Self::total_degree(a);
        let db = Self::total_degree(b);
        match da.cmp(&db) {
            std::cmp::Ordering::Equal => {
                // Reverse lex: compare right-to-left, smaller exponent = larger monomial
                for i in (0..a.len()).rev() {
                    match a[i].cmp(&b[i]) {
                        std::cmp::Ordering::Equal => continue,
                        std::cmp::Ordering::Less => return std::cmp::Ordering::Greater,
                        std::cmp::Ordering::Greater => return std::cmp::Ordering::Less,
                    }
                }
                std::cmp::Ordering::Equal
            }
            other => other,
        }
    }

    /// Leading monomial (largest under grlex).
    pub fn leading_monomial(&self) -> Option<&Vec<u8>> {
        self.terms.keys().max_by(|a, b| Self::cmp_grlex(a, b))
    }

    /// Leading coefficient.
    pub fn leading_coefficient(&self) -> Rat {
        self.leading_monomial()
            .and_then(|m| self.terms.get(m))
            .copied()
            .unwrap_or(Rat::ZERO)
    }

    /// Leading term as (exponent vector, coefficient).
    pub fn leading_term(&self) -> Option<(Vec<u8>, Rat)> {
        self.leading_monomial().map(|m| (m.clone(), self.terms[m]))
    }

    /// Make monic: divide all coefficients by the leading coefficient.
    pub fn make_monic(&mut self) {
        let lc = self.leading_coefficient();
        if lc.is_zero() || lc == Rat::ONE {
            return;
        }
        let inv = lc.recip();
        for c in self.terms.values_mut() {
            *c *= inv;
        }
    }

    /// Does monomial `a` divide monomial `b`? (a\[i\] <= b\[i\] for all i)
    pub fn monomial_divides(a: &[u8], b: &[u8]) -> bool {
        a.iter().zip(b.iter()).all(|(&ai, &bi)| ai <= bi)
    }

    /// Quotient monomial: b / a (entry-wise subtraction). Assumes a divides b.
    fn monomial_quotient(a: &[u8], b: &[u8]) -> Vec<u8> {
        a.iter().zip(b.iter()).map(|(&ai, &bi)| bi - ai).collect()
    }

    /// LCM of two monomials (entry-wise max).
    fn monomial_lcm(a: &[u8], b: &[u8]) -> Vec<u8> {
        a.iter()
            .zip(b.iter())
            .map(|(&ai, &bi)| ai.max(bi))
            .collect()
    }

    /// Multiply this polynomial by a monomial (exponent vector) with coefficient.
    pub fn mul_monomial(&self, exp: &[u8], coeff: Rat) -> Self {
        let mut result = Poly::zero(self.num_vars);
        for (e, c) in &self.terms {
            let new_exp: Vec<u8> = e.iter().zip(exp.iter()).map(|(&a, &b)| a + b).collect();
            result.terms.insert(new_exp, *c * coeff);
        }
        result.clean();
        result
    }

    /// S-polynomial of self and other.
    pub fn s_polynomial(&self, other: &Self) -> Self {
        let (lt_a, lc_a) = match self.leading_term() {
            Some(t) => t,
            None => return Poly::zero(self.num_vars),
        };
        let (lt_b, lc_b) = match other.leading_term() {
            Some(t) => t,
            None => return Poly::zero(self.num_vars),
        };
        let lcm = Self::monomial_lcm(&lt_a, &lt_b);
        let qa = Self::monomial_quotient(&lt_a, &lcm);
        let qb = Self::monomial_quotient(&lt_b, &lcm);
        let left = self.mul_monomial(&qa, lc_b);
        let right = other.mul_monomial(&qb, lc_a);
        left.sub(&right)
    }

    /// Reduce this polynomial by a single divisor. Returns remainder.
    pub fn reduce_by(&self, divisor: &Self) -> Self {
        let (div_lm, div_lc) = match divisor.leading_term() {
            Some(t) => t,
            None => return self.clone(),
        };

        let mut remainder = Poly::zero(self.num_vars);
        let mut current = self.clone();

        while !current.is_zero() {
            let (cur_lm, cur_lc) = match current.leading_term() {
                Some(t) => t,
                None => break,
            };
            if Self::monomial_divides(&div_lm, &cur_lm) {
                let q_exp = Self::monomial_quotient(&div_lm, &cur_lm);
                let q_coeff = cur_lc / div_lc;
                let subtrahend = divisor.mul_monomial(&q_exp, q_coeff);
                current = current.sub(&subtrahend);
            } else {
                // Move leading term to remainder
                remainder.terms.insert(cur_lm.clone(), cur_lc);
                current.terms.remove(&cur_lm);
            }
        }
        remainder.clean();
        remainder
    }

    /// Evaluate the polynomial at given variable values.
    pub fn eval(&self, values: &[Rat]) -> Rat {
        let mut result = Rat::ZERO;
        for (exp, coeff) in &self.terms {
            let mut term = *coeff;
            for (i, &e) in exp.iter().enumerate() {
                for _ in 0..e {
                    term *= values[i];
                }
            }
            result += term;
        }
        result
    }
}

// Arithmetic

impl Poly {
    pub fn add(&self, other: &Self) -> Self {
        let mut result = self.clone();
        for (exp, coeff) in &other.terms {
            let entry = result.terms.entry(exp.clone()).or_insert(Rat::ZERO);
            *entry += *coeff;
        }
        result.clean();
        result
    }

    pub fn sub(&self, other: &Self) -> Self {
        let mut result = self.clone();
        for (exp, coeff) in &other.terms {
            let entry = result.terms.entry(exp.clone()).or_insert(Rat::ZERO);
            *entry -= *coeff;
        }
        result.clean();
        result
    }

    pub fn mul(&self, other: &Self) -> Self {
        let mut result = Poly::zero(self.num_vars);
        for (ea, ca) in &self.terms {
            for (eb, cb) in &other.terms {
                let exp: Vec<u8> = ea.iter().zip(eb.iter()).map(|(&a, &b)| a + b).collect();
                let entry = result.terms.entry(exp).or_insert(Rat::ZERO);
                *entry += *ca * *cb;
            }
        }
        result.clean();
        result
    }

    pub fn neg(&self) -> Self {
        let mut result = self.clone();
        for c in result.terms.values_mut() {
            *c = -*c;
        }
        result
    }

    pub fn scale(&self, r: Rat) -> Self {
        if r.is_zero() {
            return Poly::zero(self.num_vars);
        }
        let mut result = self.clone();
        for c in result.terms.values_mut() {
            *c *= r;
        }
        result.clean();
        result
    }
}

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

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

    #[test]
    fn zero_poly() {
        let p = Poly::zero(3);
        assert!(p.is_zero());
    }

    #[test]
    fn constant() {
        let p = Poly::constant(Rat::from(5), 3);
        assert!(!p.is_zero());
        assert_eq!(p.eval(&[Rat::ZERO, Rat::ZERO, Rat::ZERO]), Rat::from(5));
    }

    #[test]
    fn variable() {
        let x = Poly::variable(0, 3);
        assert_eq!(x.eval(&[Rat::from(7), Rat::ZERO, Rat::ZERO]), Rat::from(7));
    }

    #[test]
    fn add_polys() {
        let x = Poly::variable(0, 2);
        let y = Poly::variable(1, 2);
        let sum = x.add(&y);
        assert_eq!(sum.eval(&[Rat::from(3), Rat::from(4)]), Rat::from(7));
    }

    #[test]
    fn mul_polys() {
        // (x + 1)(x - 1) = x² - 1
        let n = 1;
        let x = Poly::variable(0, n);
        let one = Poly::constant(Rat::ONE, n);
        let xp1 = x.add(&one);
        let xm1 = x.sub(&one);
        let product = xp1.mul(&xm1);
        // Eval at x=3: 9 - 1 = 8
        assert_eq!(product.eval(&[Rat::from(3)]), Rat::from(8));
        // Eval at x=1: 0
        assert_eq!(product.eval(&[Rat::from(1)]), Rat::ZERO);
    }

    #[test]
    fn leading_term_grlex() {
        // 3x²y + 2xy² — both total degree 3
        // grlex: x²y > xy² (compare right-to-left: y-exponents 1 vs 2, smaller is larger)
        let mut p = Poly::zero(2);
        p.terms.insert(vec![2, 1], Rat::from(3)); // x²y
        p.terms.insert(vec![1, 2], Rat::from(2)); // xy²
        let (lm, lc) = p.leading_term().unwrap();
        assert_eq!(lm, vec![2, 1]);
        assert_eq!(lc, Rat::from(3));
    }

    #[test]
    fn s_polynomial() {
        // f = x² - 1, g = x*y - 1
        let n = 2;
        let f = {
            let mut p = Poly::zero(n);
            p.terms.insert(vec![2, 0], Rat::ONE); // x²
            p.terms.insert(vec![0, 0], -Rat::ONE); // -1
            p
        };
        let g = {
            let mut p = Poly::zero(n);
            p.terms.insert(vec![1, 1], Rat::ONE); // xy
            p.terms.insert(vec![0, 0], -Rat::ONE); // -1
            p
        };
        let s = f.s_polynomial(&g);
        // lcm(x², xy) = x²y
        // S = (y * f) - (x * g) = x²y - y - x²y + x -= y
        assert_eq!(s.eval(&[Rat::from(5), Rat::from(3)]), Rat::from(2)); // 5 - 3 = 2
    }

    #[test]
    fn reduce_by_simple() {
        // Reduce x² by (x - 1): x² = x*(x-1) + x, so remainder is x
        // Then reduce x by (x - 1): x = 1*(x-1) + 1, so remainder is 1
        let n = 1;
        let x_sq = {
            let mut p = Poly::zero(n);
            p.terms.insert(vec![2], Rat::ONE);
            p
        };
        let x_minus_1 = {
            let mut p = Poly::zero(n);
            p.terms.insert(vec![1], Rat::ONE);
            p.terms.insert(vec![0], -Rat::ONE);
            p
        };
        let rem = x_sq.reduce_by(&x_minus_1);
        // x² mod (x-1) = 1 (since x=1 is the root)
        assert_eq!(rem.eval(&[Rat::from(99)]), Rat::ONE);
    }

    #[test]
    fn monomial_divides() {
        assert!(Poly::monomial_divides(&vec![1, 0], &vec![2, 1])); // x | x²y
        assert!(!Poly::monomial_divides(&vec![0, 2], &vec![1, 1])); // y² does not divide xy
    }

    #[test]
    fn eval_multivariate() {
        // p = x² + 2xy + y²  = (x+y)²
        let mut p = Poly::zero(2);
        p.terms.insert(vec![2, 0], Rat::ONE);
        p.terms.insert(vec![1, 1], Rat::from(2));
        p.terms.insert(vec![0, 2], Rat::ONE);
        assert_eq!(p.eval(&[Rat::from(3), Rat::from(4)]), Rat::from(49)); // (3+4)² = 49
    }

    #[test]
    fn neg_poly() {
        let x = Poly::variable(0, 1);
        let nx = x.neg();
        assert_eq!(nx.eval(&[Rat::from(5)]), Rat::from(-5));
    }

    #[test]
    fn make_monic() {
        let mut p = Poly::zero(1);
        p.terms.insert(vec![2], Rat::from(3)); // 3x²
        p.terms.insert(vec![0], Rat::from(6)); // + 6
        p.make_monic();
        // Leading coeff should now be 1
        assert_eq!(p.leading_coefficient(), Rat::ONE);
        // Constant term should be 2 (6/3)
        assert_eq!(p.eval(&[Rat::ZERO]), Rat::from(2));
    }
}