reefer 0.3.0

use crate::clifford::{
    DotProduct, Duality, ExteriorProduct, Graded, InnerProduct, Involution, LeftContraction, Mask,
    RegressiveProduct, RightContraction, ScalarProduct, metric::Metric, multivector::Multivector,
};
use abstalg::{
    AbelianGroup, BoundedOrder, CommuntativeMonoid, Domain, IntegralDomain, Lattice, Monoid,
    SemiRing, Semigroup, UnitaryRing,
};
use itertools::{EitherOrBoth, Itertools};
use num_traits::One;
use std::{collections::BTreeMap, fmt::Debug};

pub trait Frame: Mask + Ord + PartialOrd + Eq + PartialEq + One {}
impl<T: Mask + Ord + One> Frame for T {}

pub trait BasisScalar: UnitaryRing + Debug {}
impl<T: UnitaryRing + Debug> BasisScalar for T {}

pub type BasisElement<B, S> = Multivector<B, S>;

/// Bundle of a metric frame with the scalar algebra used to model its coefficients.
#[derive(Debug, Clone)]
pub struct BasisSpace<B: Frame, S: BasisScalar> {
    pub metric: Metric<B>,
    pub scalar: S,
}

#[allow(dead_code)]
// Many of these helpers are exercised via downstream runtime modules; keep them available even
// when the proc-macro crate does not reference them directly.
impl<B: Frame, S: BasisScalar> BasisSpace<B, S> {
    /// Create a new basis space from a metric frame and scalar algebra implementation.
    pub fn new(frame: Metric<B>, scalar: S) -> Self {
        Self {
            metric: frame,
            scalar,
        }
    }

    /// Borrow the metric describing the frame's signature.
    pub fn frame(&self) -> &Metric<B> {
        &self.metric
    }

    /// Borrow the scalar algebra used for coefficients.
    pub fn scalar(&self) -> &S {
        &self.scalar
    }

    fn prune_terms(&self, vectors: &mut BTreeMap<B, S::Elem>) {
        vectors.retain(|_, coeff| !self.scalar.is_zero(coeff));
    }

    fn accumulate_term(
        &self,
        storage: &mut BTreeMap<B, S::Elem>,
        blade: B,
        coeff: S::Elem,
        is_negative: bool,
    ) {
        if self.scalar.is_zero(&coeff) {
            return;
        }

        use std::collections::btree_map::Entry;

        match storage.entry(blade) {
            Entry::Occupied(mut occ) => {
                if is_negative {
                    self.scalar.sub_assign(occ.get_mut(), &coeff);
                } else {
                    self.scalar.add_assign(occ.get_mut(), &coeff);
                }
                if self.scalar.is_zero(occ.get()) {
                    occ.remove_entry();
                }
            }
            Entry::Vacant(vacant) => {
                let value = if is_negative {
                    self.scalar.neg(&coeff)
                } else {
                    coeff
                };
                if !self.scalar.is_zero(&value) {
                    vacant.insert(value);
                }
            }
        }
    }

    #[inline]
    fn pseudoscalar_with_mask(&self, mask: B) -> Multivector<B, S> {
        let mut vectors = BTreeMap::from([(mask, self.scalar.one())]);
        self.prune_terms(&mut vectors);
        Multivector { vectors }
    }

    #[inline]
    fn pseudoscalar(&self) -> Multivector<B, S> {
        self.pseudoscalar_with_mask(self.metric.supremum)
    }

    #[inline]
    fn complement_mask(&self, mask: B) -> B {
        self.metric.sym_diff(
            {
                let this = &self.metric;
                this.supremum
            },
            mask,
        )
    }
    #[inline]
    fn complement_sign(&self, mask: B) -> bool {
        self.metric.mul_parity(mask, self.complement_mask(mask))
    }
    #[inline]
    fn uncomplement_sign(&self, mask: B) -> bool {
        self.metric.mul_parity(self.complement_mask(mask), mask)
    }
    #[inline]
    fn metric_sign(&self, mask: B) -> bool {
        (mask & {
            let this = &self.metric;
            this.imagimum
        })
        .parity()
    }

    fn right_dual_with_mask(&self, elem: &Multivector<B, S>, mask: B) -> Multivector<B, S> {
        let pseudoscalar = self.pseudoscalar_with_mask(mask);
        let reversed = self.reverse(elem.clone());
        self.mul(&reversed, &pseudoscalar)
    }

    fn left_dual_with_mask(&self, elem: &Multivector<B, S>, mask: B) -> Multivector<B, S> {
        let pseudoscalar = self.pseudoscalar_with_mask(mask);
        let reversed = self.reverse(elem.clone());
        self.mul(&pseudoscalar, &reversed)
    }

    /// Apply a sandwich product `rotor * target * rotor.reverse()` to transform `target`.
    pub fn sandwich(
        &self,
        rotor: &Multivector<B, S>,
        target: &Multivector<B, S>,
    ) -> Multivector<B, S> {
        let left = self.mul(rotor, target);
        let rotor_reverse = self.reverse(rotor.clone());
        self.mul(&left, &rotor_reverse)
    }

    #[inline]
    /// Compute the right dual of a multivector with respect to the full pseudoscalar.
    pub fn right_dual(&self, elem: &Multivector<B, S>) -> Multivector<B, S> {
        self.right_dual_with_mask(elem, self.metric.max())
    }

    #[inline]
    /// Compute the left dual of a multivector with respect to the full pseudoscalar.
    pub fn left_dual(&self, elem: &Multivector<B, S>) -> Multivector<B, S> {
        self.left_dual_with_mask(elem, self.metric.max())
    }

    /// Compute the complement (dual) basis representation of a multivector.
    pub fn complement(&self, elem: &Multivector<B, S>) -> Multivector<B, S> {
        let mut result = self.zero();
        for (blade, coeff) in elem.vectors.iter() {
            self.accumulate_term(
                &mut result.vectors,
                self.complement_mask(*blade),
                coeff.clone(),
                self.complement_sign(*blade),
            );
        }
        result
    }

    /// Extract the scalar component of a multivector, if present.
    pub fn scalar_part(&self, elem: &Multivector<B, S>) -> Option<S::Elem> {
        elem.vectors.get(&self.metric.one()).cloned()
    }

    /// Compute the norm squared of a multivector.
    pub fn norm_squared(&self, elem: &Multivector<B, S>) -> S::Elem {
        self.scalar_part(&self.inner(elem, elem))
            .unwrap_or_else(|| self.scalar.zero())
    }

    /// Compute the meet (regressive product) of two multivectors.
    pub fn meet(&self, lhs: &Multivector<B, S>, rhs: &Multivector<B, S>) -> Multivector<B, S> {
        let lhs_dual = self.right_dual(lhs);
        let rhs_dual = self.right_dual(rhs);
        let join_dual = self.wedge(&lhs_dual, &rhs_dual);
        self.right_dual(&join_dual)
    }
}

impl<B: Frame, S: BasisScalar> Domain for BasisSpace<B, S> {
    type Elem = BasisElement<B, S>;
    fn equals(&self, elem1: &Self::Elem, elem2: &Self::Elem) -> bool {
        elem1
            .vectors
            .iter()
            .merge_join_by(elem2.vectors.iter(), |(b1, _), (b2, _)| b1.cmp(b2))
            .all(|either| match either {
                EitherOrBoth::Both((_, v1), (_, v2)) => self.scalar.equals(v1, v2),
                EitherOrBoth::Left((_, v)) => self.scalar.is_zero(v),
                EitherOrBoth::Right((_, v)) => self.scalar.is_zero(v),
            })
    }
    fn contains(&self, elem: &Self::Elem) -> bool {
        elem.vectors
            .iter()
            .all(|(b, v)| self.metric.contains(b) & self.scalar.contains(v))
    }
}

impl<B: Frame, S: BasisScalar> CommuntativeMonoid for BasisSpace<B, S> {
    fn zero(&self) -> Self::Elem {
        Multivector {
            vectors: BTreeMap::new(),
        }
    }
    fn add(&self, elem1: &Self::Elem, elem2: &Self::Elem) -> Self::Elem {
        let mut vectors = elem1
            .vectors
            .iter()
            .merge_join_by(elem2.vectors.iter(), |(b1, _), (b2, _)| b1.cmp(b2))
            .map(|either| match either {
                EitherOrBoth::Both((b, v1), (_, v2)) => (b.clone(), self.scalar.add(v1, v2)),
                EitherOrBoth::Left((b, v)) | EitherOrBoth::Right((b, v)) => (b.clone(), v.clone()),
            })
            .collect();
        self.prune_terms(&mut vectors);
        Multivector { vectors }
    }
    fn is_zero(&self, elem: &Self::Elem) -> bool {
        elem.vectors.values().all(|s| self.scalar.is_zero(s))
    }
    fn add_assign(&self, elem1: &mut Self::Elem, elem2: &Self::Elem) {
        elem2.vectors.iter().for_each(|(b, s2)| {
            elem1
                .vectors
                .entry(b.clone())
                .and_modify(|s1| self.scalar.add_assign(s1, s2))
                .or_insert(s2.clone());
        });
        self.prune_terms(&mut elem1.vectors);
    }
    fn double(&self, elem: &mut Self::Elem) {
        elem.vectors
            .iter_mut()
            .for_each(|(_, coeff)| self.scalar.double(coeff));
        self.prune_terms(&mut elem.vectors);
    }
    fn times(&self, num: usize, elem: &Self::Elem) -> Self::Elem {
        let mut vectors = elem
            .vectors
            .iter()
            .map(|(b, s)| (b.clone(), CommuntativeMonoid::times(&self.scalar, num, s)))
            .collect();
        self.prune_terms(&mut vectors);
        Multivector { vectors }
    }
}

impl<B: Frame, S: BasisScalar> AbelianGroup for BasisSpace<B, S> {
    fn neg(&self, elem: &Self::Elem) -> Self::Elem {
        let mut vectors = elem
            .vectors
            .iter()
            .map(|(b, s)| (b.clone(), self.scalar.neg(s)))
            .collect();
        self.prune_terms(&mut vectors);
        Multivector { vectors }
    }
    fn neg_assign(&self, elem: &mut Self::Elem) {
        elem.vectors
            .iter_mut()
            .for_each(|(_, s)| self.scalar.neg_assign(s));
        self.prune_terms(&mut elem.vectors);
    }
    fn sub(&self, elem1: &Self::Elem, elem2: &Self::Elem) -> Self::Elem {
        let mut vectors = elem1
            .vectors
            .iter()
            .merge_join_by(elem2.vectors.iter(), |(b1, _), (b2, _)| b1.cmp(b2))
            .map(|either| match either {
                EitherOrBoth::Both((b, v1), (_, v2)) => (b.clone(), self.scalar.sub(v1, v2)),
                EitherOrBoth::Left((b, v)) => (b.clone(), v.clone()),
                EitherOrBoth::Right((b, v)) => (b.clone(), self.scalar.neg(v)),
            })
            .collect();
        self.prune_terms(&mut vectors);
        Multivector { vectors }
    }
    fn sub_assign(&self, elem1: &mut Self::Elem, elem2: &Self::Elem) {
        elem2.vectors.iter().for_each(|(b, s2)| {
            elem1
                .vectors
                .entry(b.clone())
                .and_modify(|s1| self.scalar.sub_assign(s1, s2))
                .or_insert(self.scalar.neg(s2));
        });
        self.prune_terms(&mut elem1.vectors);
    }
    fn times(&self, num: isize, elem: &Self::Elem) -> Self::Elem {
        let mut vectors = elem
            .vectors
            .iter()
            .map(|(b, s)| (b.clone(), AbelianGroup::times(&self.scalar, num, s)))
            .collect();
        self.prune_terms(&mut vectors);
        Multivector { vectors }
    }
}

impl<B: Frame, S: BasisScalar> Semigroup for BasisSpace<B, S> {
    fn mul(&self, elem1: &Self::Elem, elem2: &Self::Elem) -> Self::Elem {
        let mut product = self.zero();
        for (b1, s1) in elem1.vectors.iter() {
            for (b2, s2) in elem2.vectors.iter() {
                self.accumulate_term(
                    &mut product.vectors,
                    self.metric.sym_diff(*b1, *b2),
                    self.scalar.mul(s1, s2),
                    self.metric.mul_parity(*b1, *b2),
                );
            }
        }
        product
    }
    fn mul_assign(&self, elem1: &mut Self::Elem, elem2: &Self::Elem) {
        *elem1 = self.mul(elem1, elem2);
    }
    fn square(&self, elem: &mut Self::Elem) {
        *elem = self.mul(elem, elem);
    }
}

#[allow(unused_variables)]
impl<B: Frame, S: BasisScalar> Monoid for BasisSpace<B, S> {
    fn one(&self) -> Self::Elem {
        let mut vectors = BTreeMap::from([(self.metric.one(), self.scalar.one())]);
        self.prune_terms(&mut vectors);
        Multivector { vectors }
    }
    fn is_one(&self, elem: &Self::Elem) -> bool {
        self.equals(elem, &self.one())
    }
    fn try_inv(&self, elem: &Self::Elem) -> Option<Self::Elem> {
        todo!("Implement multivector inversion")
    }
    fn invertible(&self, elem: &Self::Elem) -> bool {
        self.try_inv(elem).is_some()
    }
}

impl<B: Frame, S: BasisScalar> SemiRing for BasisSpace<B, S> {}

impl<B: Frame, S: BasisScalar> UnitaryRing for BasisSpace<B, S> {}

#[allow(unused_variables)]
impl<B: Frame, S: BasisScalar> IntegralDomain for BasisSpace<B, S> {
    fn try_div(&self, elem1: &Self::Elem, elem2: &Self::Elem) -> Option<Self::Elem> {
        todo!("Implement multivector division")
    }
    fn associate_repr(&self, elem: &Self::Elem) -> Self::Elem {
        todo!("Implement associate representation")
    }
    fn associate_coef(&self, elem: &Self::Elem) -> Self::Elem {
        todo!("Implement associate coefficient")
    }
    fn divisible(&self, elem1: &Self::Elem, elem2: &Self::Elem) -> bool {
        self.try_div(elem1, elem2).is_some()
    }
    fn associates(&self, elem1: &Self::Elem, elem2: &Self::Elem) -> bool {
        todo!("Implement associate checking")
    }
}

// --

impl<B: Frame, S: BasisScalar> Graded for BasisSpace<B, S> {
    type Grade = Vec<u32>;
    type Output = Self::Elem;
    fn grade_of(&self, elem: &Self::Elem) -> Self::Grade {
        elem.vectors
            .keys()
            .map(|basis| self.metric.grade_of(basis))
            .sorted() // todo could get rid of sort with right Basis Ord
            .dedup()
            .collect()
    }
    fn grade_by(&self, elem: &Self::Elem, grades: Self::Grade) -> Self::Output {
        let vectors = elem
            .vectors
            .iter()
            .filter(|(b, _)| grades.iter().any(|g| self.metric.grade_by(b, *g).is_some()))
            .map(|(b, s)| (b.clone(), s.clone()))
            .collect();
        Multivector { vectors }
    }
}

impl<B: Frame, S: BasisScalar> Duality for BasisSpace<B, S> {
    type Psuedo = Option<B>;
    type Output = Self::Elem;
    fn dual(&self, elem: &Self::Elem, psuedoscalar: &Self::Psuedo) -> Self::Output {
        let w = match psuedoscalar {
            Some(b) => *b,
            None => self.metric.max(),
        };
        let vectors = elem
            .vectors
            .iter()
            .map(|(b, s)| {
                if self.metric.mul_parity(*b, w) {
                    (*b, self.scalar.neg(s))
                } else {
                    (*b, s.clone())
                }
            })
            .collect();
        Multivector { vectors }
    }
    fn undual(&self, elem: &Self::Elem, pseudoscalar: &Self::Psuedo) -> Self::Output {
        let w = match pseudoscalar {
            Some(b) => *b,
            None => self.metric.max(),
        };
        let vectors = elem
            .vectors
            .iter()
            .map(|(b, s)| {
                if self.metric.mul_parity(w, *b) {
                    (*b, self.scalar.neg(s))
                } else {
                    (*b, s.clone())
                }
            })
            .collect();
        Multivector { vectors }
    }
}

impl<B: Frame, S: BasisScalar> Involution for BasisSpace<B, S> {
    type Output = Self::Elem;
    fn automorphism(&self, mut elem: Self::Elem) -> Self::Output {
        elem.vectors.iter_mut().for_each(|(b, s)| {
            if self.metric.automorphism(*b) {
                self.scalar.neg_assign(s);
            }
        });
        elem
    }
    fn reverse(&self, mut elem: Self::Elem) -> Self::Output {
        elem.vectors.iter_mut().for_each(|(b, s)| {
            if self.metric.reverse(*b) {
                self.scalar.neg_assign(s);
            }
        });
        elem
    }
    fn conjugate(&self, mut elem: Self::Elem) -> Self::Output {
        elem.vectors.iter_mut().for_each(|(b, s)| {
            if self.metric.conjugate(*b) {
                self.scalar.neg_assign(s);
            }
        });
        elem
    }
}

impl<B: Frame, S: BasisScalar> ExteriorProduct for BasisSpace<B, S> {
    type Output = Self::Elem;
    fn wedge(&self, lhs: &Self::Elem, rhs: &Self::Elem) -> Self::Output {
        let mut result = self.zero();
        for (b1, s1) in lhs.vectors.iter() {
            for (b2, s2) in rhs.vectors.iter() {
                if self.metric.meet(b1, b2).is_zero() {
                    let blade = self.metric.sym_diff(*b1, *b2);
                    let coeff = self.scalar.mul(s1, s2);
                    let is_neg = self.metric.mul_parity(*b1, *b2);
                    self.accumulate_term(&mut result.vectors, blade, coeff, is_neg);
                }
            }
        }
        result
    }
}

impl<B: Frame, S: BasisScalar> InnerProduct for BasisSpace<B, S> {
    type Output = Self::Elem;
    fn inner(&self, lhs: &Self::Elem, rhs: &Self::Elem) -> Self::Output {
        let mut value = self.scalar.zero();
        let mut has_scalar = false;
        for (blade, coeff_l) in lhs.vectors.iter() {
            if let Some(coeff_r) = rhs.vectors.get(blade) {
                let mut coeff = self.scalar.mul(coeff_l, coeff_r);
                match (has_scalar, self.metric_sign(*blade)) {
                    (true, true) => self.scalar.sub_assign(&mut value, &coeff),
                    (true, false) => self.scalar.add_assign(&mut value, &coeff),
                    (false, true) => {
                        self.scalar.neg_assign(&mut coeff);
                        value = coeff;
                        has_scalar = true;
                    }
                    (false, false) => {
                        value = coeff;
                        has_scalar = true;
                    }
                }
            }
        }

        let mut vectors = BTreeMap::new();
        if has_scalar && !self.scalar.is_zero(&value) {
            vectors.insert(self.metric.one(), value);
        }
        Multivector { vectors }
    }
}

impl<B: Frame, S: BasisScalar> LeftContraction for BasisSpace<B, S> {
    type Output = Self::Elem;
    fn contract_onto(&self, lhs: &Self::Elem, rhs: &Self::Elem) -> Self::Output {
        self.antiwedge(&self.dual(lhs, &None), rhs)
    }
}

impl<B: Frame, S: BasisScalar> RightContraction for BasisSpace<B, S> {
    type Output = Self::Elem;
    fn contract_by(&self, lhs: &Self::Elem, rhs: &Self::Elem) -> Self::Output {
        self.antiwedge(lhs, &self.dual(&rhs, &None))
    }
}

impl<B: Frame, S: BasisScalar> ScalarProduct for BasisSpace<B, S> {
    type Output = Self::Elem;
    fn scalar_product(&self, lhs: &Self::Elem, rhs: &Self::Elem) -> Self::Output {
        self.inner(lhs, rhs)
    }
}

impl<B: Frame, S: BasisScalar> DotProduct for BasisSpace<B, S> {
    type Output = Self::Elem;
    fn dot(&self, lhs: &Self::Elem, rhs: &Self::Elem) -> Self::Output {
        self.inner(lhs, rhs)
    }
}

impl<B: Frame, S: BasisScalar> RegressiveProduct for BasisSpace<B, S> {
    type Output = Self::Elem;
    fn antiwedge(&self, lhs: &Self::Elem, rhs: &Self::Elem) -> Self::Output {
        let mut result = self.zero();
        for (b1, s1) in lhs.vectors.iter() {
            let comp_a = self.complement_mask(*b1);
            for (b2, s2) in rhs.vectors.iter() {
                let comp_b = self.complement_mask(*b2);
                if !self.metric.meet(&comp_a, &comp_b).is_zero() {
                    continue;
                }
                let c_comp = self.metric.sym_diff(comp_a, comp_b);
                let c = self.complement_mask(c_comp);
                let is_neg = self.complement_sign(*b1)
                    ^ self.complement_sign(*b2)
                    ^ self.metric.mul_parity(comp_a, comp_b)
                    ^ self.uncomplement_sign(c_comp);
                let coeff = self.scalar.mul(s1, s2);
                self.accumulate_term(&mut result.vectors, c, coeff, is_neg);
            }
        }
        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use abstalg::{ApproxFloats, F32};

    type Blade = u8;
    type Ring = ApproxFloats<f32>;
    #[allow(dead_code)]
    type ScalarElem = <Ring as Domain>::Elem;

    fn euclidean_space() -> BasisSpace<Blade, Ring> {
        BasisSpace::new(Metric::new(3, 0), F32)
    }

    fn minkowski_space() -> BasisSpace<Blade, Ring> {
        BasisSpace::new(Metric::new(1, 1), F32)
    }

    fn multivector(
        space: &BasisSpace<Blade, Ring>,
        terms: &[(Blade, i64)],
    ) -> Multivector<Blade, Ring> {
        Multivector::from_terms(
            terms
                .iter()
                .map(|(blade, coeff)| (*blade, space.scalar.int(*coeff as isize))),
        )
    }

    #[test]
    fn sandwich_with_identity_returns_target() {
        let space = euclidean_space();
        let rotor = space.one();
        let vector = multivector(&space, &[(0b001, 1)]);

        let transformed = space.sandwich(&rotor, &vector);
        assert!(space.equals(&transformed, &vector));
    }
    fn basis_vector(space: &BasisSpace<Blade, Ring>, mask: Blade) -> Multivector<Blade, Ring> {
        multivector(space, &[(mask, 1)])
    }

    #[test]
    fn complement_wedge_pseudoscalar() {
        let space = euclidean_space();
        let e1 = basis_vector(&space, 0b001);
        let complement = space.complement(&e1);
        let pseudoscalar = space.pseudoscalar();
        let wedge = space.wedge(&e1, &complement);

        assert!(space.equals(&wedge, &pseudoscalar));
    }

    #[test]
    fn dual_relates_scalar_and_pseudoscalar() {
        let space = euclidean_space();
        let scalar = space.one();
        let pseudoscalar = space.pseudoscalar();

        let dual_scalar = space.right_dual(&scalar);
        let dual_ps = space.right_dual(&pseudoscalar);

        assert!(space.equals(&dual_scalar, &pseudoscalar));
        assert!(space.equals(&dual_ps, &scalar));
    }

    #[test]
    fn norm_squared_respects_metric_signature() {
        let space = minkowski_space();
        let e_plus = basis_vector(&space, 0b01); // positive square
        let e_minus = basis_vector(&space, 0b10); // negative square

        let norm_plus = space.norm_squared(&e_plus);
        let norm_minus = space.norm_squared(&e_minus);

        let one = space.scalar.one();
        let neg_one = space.scalar.neg(&one);

        assert!(space.scalar.equals(&norm_plus, &one));
        assert!(space.scalar.equals(&norm_minus, &neg_one));
    }

    #[test]
    fn meet_of_planes_is_common_line() {
        let space = euclidean_space();
        let e1 = basis_vector(&space, 0b001);
        let e2 = basis_vector(&space, 0b010);
        let e3 = basis_vector(&space, 0b100);

        let plane_12 = space.wedge(&e1, &e2);
        let plane_13 = space.wedge(&e1, &e3);

        let meet = space.meet(&plane_12, &plane_13);
        let expected = e1;

        assert!(space.equals(&meet, &expected));
    }
}