//! Symbolic multivectors: coefficients as [`Expr`] rather than rationals.
//!
//! A `SymMultivector` stores coefficients as symbolic expressions, enabling
//! differentiation, LaTeX output, and integration with the rest of SymClaw.

use num_traits::Zero;
use std::collections::BTreeMap;
use std::sync::Arc;

use crate::ast::Expr;
use crate::differentiate::differentiate;
use crate::interner::Symbol;
use crate::latex::to_latex;
use crate::simplify::simplify;

use super::basis::{BasisBlade, Signature};

/// A multivector whose coefficients are symbolic [`Expr`] values.
#[derive(Debug, Clone)]
pub struct SymMultivector {
    /// Sparse map: blade → symbolic coefficient.
    terms: BTreeMap<BasisBlade, Arc<Expr>>,
    /// Algebra signature.
    pub sig: Signature,
}

impl SymMultivector {
    /// Zero symbolic multivector.
    #[must_use]
    pub fn zero(sig: Signature) -> Self {
        Self {
            terms: BTreeMap::new(),
            sig,
        }
    }

    /// Scalar symbolic multivector.
    #[must_use]
    pub fn scalar(expr: Arc<Expr>, sig: Signature) -> Self {
        let mut mv = Self::zero(sig);
        mv.terms.insert(BasisBlade::SCALAR, expr);
        mv
    }

    /// Basis vector `eₖ₊₁` with coefficient `expr`.
    #[must_use]
    pub fn basis_vector(k: u8, coeff: Arc<Expr>, sig: Signature) -> Self {
        Self::from_terms(std::iter::once((BasisBlade::vector(k), coeff)), sig)
    }

    /// Coefficient of a blade (zero expr if absent).
    #[must_use]
    pub fn coeff(&self, blade: BasisBlade) -> Arc<Expr> {
        self.terms
            .get(&blade)
            .cloned()
            .unwrap_or_else(|| Arc::new(Expr::from(0i64)))
    }

    /// Differentiate all coefficients with respect to `var`.
    #[must_use]
    pub fn diff(&self, var: Symbol) -> Self {
        let terms = self.terms.iter().map(|(&b, c)| {
            let dc = differentiate(c, var);
            (b, simplify(&dc))
        });
        Self::from_terms(terms, self.sig)
    }

    /// Simplify all coefficients.
    #[must_use]
    pub fn simplified(&self) -> Self {
        let terms = self.terms.iter().map(|(&b, c)| (b, simplify(c)));
        Self::from_terms(terms, self.sig)
    }

    /// Geometric product (coefficients are symbolically multiplied via `Expr::Mul`).
    #[must_use]
    pub fn geometric_product(&self, other: &Self) -> Self {
        assert_eq!(self.sig, other.sig);
        let mut result: BTreeMap<BasisBlade, Vec<Arc<Expr>>> = BTreeMap::new();
        for (&a_blade, a_coeff) in &self.terms {
            for (&b_blade, b_coeff) in &other.terms {
                let (sign, res_blade) = a_blade.geometric_product(b_blade, self.sig);
                let term = if sign == 1 {
                    Expr::mul(vec![a_coeff.clone(), b_coeff.clone()])
                } else {
                    Expr::neg(Expr::mul(vec![a_coeff.clone(), b_coeff.clone()]))
                };
                result.entry(res_blade).or_default().push(term);
            }
        }
        let terms = result.into_iter().map(|(b, parts)| {
            let sum = if parts.len() == 1 {
                parts.into_iter().next().expect("non-empty")
            } else {
                Expr::add(parts)
            };
            (b, simplify(&(*sum).clone()))
        });
        Self::from_terms(terms, self.sig)
    }

    fn from_terms(iter: impl Iterator<Item = (BasisBlade, Arc<Expr>)>, sig: Signature) -> Self {
        let mut mv = Self::zero(sig);
        for (b, c) in iter {
            let simplified = simplify(&(*c).clone());
            // Drop zero coefficients
            if !is_zero_expr(&simplified) {
                mv.terms.insert(b, simplified);
            }
        }
        mv
    }

    /// Generate a LaTeX string for this symbolic multivector.
    #[must_use]
    pub fn to_latex(&self) -> String {
        if self.terms.is_empty() {
            return "0".to_owned();
        }
        let mut parts: Vec<String> = Vec::new();
        for (blade, coeff) in &self.terms {
            let coeff_latex = to_latex(coeff);
            let blade_label = blade.label();
            if blade.is_scalar() {
                parts.push(coeff_latex);
            } else {
                parts.push(format!("{coeff_latex} \\mathbf{{{blade_label}}}"));
            }
        }
        parts.join(" + ")
    }
}

fn is_zero_expr(expr: &Expr) -> bool {
    matches!(expr, Expr::Num(r) if r.is_zero())
}

impl std::fmt::Display for SymMultivector {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.terms.is_empty() {
            return write!(f, "0");
        }
        let parts: Vec<String> = self
            .terms
            .iter()
            .map(|(b, c)| {
                if b.is_scalar() {
                    format!("{c}")
                } else {
                    format!("{c}*{b}")
                }
            })
            .collect();
        write!(f, "{}", parts.join(" + "))
    }
}

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

    fn eu(n: u8) -> Signature {
        Signature::euclidean(n)
    }

    fn expr(s: &str) -> Arc<Expr> {
        parse(s).expect("valid expr")
    }

    #[test]
    fn symbolic_scalar_display() {
        let sig = eu(3);
        let mv = SymMultivector::scalar(expr("a"), sig);
        let s = format!("{mv}");
        assert!(s.contains('a'));
    }

    #[test]
    fn symbolic_multivector_diff_wrt_scalar() {
        let sig = eu(3);
        // coefficient: x^2  →  diff w.r.t. x → 2*x
        let mv = SymMultivector::basis_vector(0, expr("x^2"), sig);
        let diff_mv = mv.diff(Symbol::new("x"));
        let c = diff_mv.coeff(BasisBlade::vector(0));
        // Evaluate symbolically: 2*x
        let val = format!("{c}");
        assert!(val.contains('x'), "expected 2*x, got {val}");
    }

    #[test]
    fn symbolic_multivector_latex_output() {
        let sig = eu(2);
        let mv = SymMultivector::basis_vector(0, expr("a + b"), sig);
        let latex = mv.to_latex();
        assert!(latex.contains("\\mathbf{e1}"), "latex: {latex}");
    }

    #[test]
    fn symbolic_geometric_product_e1_squared() {
        // e1 * e1 should give +1 (scalar) in Cl(1,0)
        let sig = Signature::euclidean(1);
        let one = expr("1");
        let e1 = SymMultivector::basis_vector(0, one.clone(), sig);
        let prod = e1.geometric_product(&SymMultivector::basis_vector(0, one, sig));
        let scalar = prod.coeff(BasisBlade::SCALAR);
        assert_eq!(format!("{scalar}"), "1");
    }

    #[test]
    fn symbolic_zero_is_zero() {
        let mv = SymMultivector::zero(eu(3));
        assert!(mv.terms.is_empty());
    }

    #[test]
    fn symbolic_simplification_removes_zero_blade() {
        let sig = eu(2);
        // Use a literal zero expression so simplify() reliably produces Expr::Num(0)
        let zero_expr = Arc::new(Expr::from(0i64));
        let mv = SymMultivector::basis_vector(0, zero_expr, sig);
        // from_terms checks is_zero_expr after simplification → blade should be dropped
        assert!(mv.terms.is_empty(), "zero-coeff blade should be dropped");
    }
}