geoit 0.0.2

use crate::algebra::blade_new::BladeMask;
use crate::algebra::mv::Mv;
use crate::algebra::signature::Signature;
use crate::governance::construction::Construction;
use crate::governance::expr::Expr;
use crate::governance::geom_class::GeomClass;
use crate::governance::governance::Governance;

/// An embedding maps a smaller algebra into a larger one
/// via a generator-to-generator bijection that preserves the quadratic form.
#[derive(Clone, Debug)]
pub struct Embedding {
    pub source_sig: Signature,
    pub target_sig: Signature,
    /// source generator i maps to target generator generator_map\[i\].
    pub generator_map: Vec<u8>,
}

impl Embedding {
    /// Create a new embedding, verifying quadratic form preservation.
    pub fn new(source: Signature, target: Signature, map: Vec<u8>) -> Result<Self, EmbeddingError> {
        if map.len() != source.n() as usize {
            return Err(EmbeddingError::MapLengthMismatch {
                expected: source.n() as usize,
                got: map.len(),
            });
        }
        // Verify each mapped generator lands within target
        for (i, &t) in map.iter().enumerate() {
            if t >= target.n() {
                return Err(EmbeddingError::TargetOutOfRange {
                    source_gen: i as u8,
                    target_gen: t,
                    target_n: target.n(),
                });
            }
            // Verify quadratic form preservation: source gen_sq(i) == target gen_sq(map[i])
            let s_sq = source.generator_square(i as u8);
            let t_sq = target.generator_square(t);
            if s_sq != t_sq {
                return Err(EmbeddingError::QuadraticFormViolation {
                    source_gen: i as u8,
                    source_square: s_sq,
                    target_gen: t,
                    target_square: t_sq,
                });
            }
        }
        // Verify injectivity: no two source gens map to the same target gen
        let mut seen = vec![false; target.n() as usize];
        for &t in &map {
            if seen[t as usize] {
                return Err(EmbeddingError::NotInjective { target_gen: t });
            }
            seen[t as usize] = true;
        }
        Ok(Embedding {
            source_sig: source,
            target_sig: target,
            generator_map: map,
        })
    }

    /// Remap an Mv from the source algebra into the target algebra.
    pub fn embed_mv(&self, mv: &Mv) -> Mv {
        let mut result = Mv::new();
        for (mask, coeff) in mv.blades() {
            let new_mask = self.remap_mask(mask);
            result.add_term(new_mask, coeff.clone());
        }
        result
    }

    /// Remap a blade mask through the generator map.
    fn remap_mask(&self, mask: BladeMask) -> BladeMask {
        let mut result: BladeMask = 0;
        for i in 0..self.source_sig.n() {
            if mask & (1u64 << i) != 0 {
                result |= 1u64 << self.generator_map[i as usize];
            }
        }
        result
    }
}

#[derive(Clone, Debug)]
pub enum EmbeddingError {
    MapLengthMismatch {
        expected: usize,
        got: usize,
    },
    TargetOutOfRange {
        source_gen: u8,
        target_gen: u8,
        target_n: u8,
    },
    QuadraticFormViolation {
        source_gen: u8,
        source_square: i8,
        target_gen: u8,
        target_square: i8,
    },
    NotInjective {
        target_gen: u8,
    },
}

impl std::fmt::Display for EmbeddingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            EmbeddingError::MapLengthMismatch { expected, got } => {
                write!(f, "generator map length {}, expected {}", got, expected)
            }
            EmbeddingError::TargetOutOfRange {
                source_gen,
                target_gen,
                target_n,
            } => write!(
                f,
                "source g{} maps to target g{}, but target has {} generators",
                source_gen, target_gen, target_n
            ),
            EmbeddingError::QuadraticFormViolation {
                source_gen,
                source_square,
                target_gen,
                target_square,
            } => write!(
                f,
                "g{}²={} but target g{}²={}",
                source_gen, source_square, target_gen, target_square
            ),
            EmbeddingError::NotInjective { target_gen } => write!(
                f,
                "target g{} mapped to by multiple source generators",
                target_gen
            ),
        }
    }
}

/// Compose two Governances by union.
///
/// The composed Governance has:
/// - sig = Signature(a.i + b.i, a.d + b.d, a.h + b.h)
/// - derived_gens = a's gens (as-is) + b's gens (with generator indices offset by a.n())
/// - geom_classes = a's classes + b's classes (b's constraints get generator index offsets)
/// - constructions = a's + b's (b's class_index offset by a.geom_classes.len(),
///   b's generator references offset by a.n()
/// # Errors
/// Returns `SignatureError` if the combined signature exceeds 64 generators.
pub fn compose(a: &Governance, b: &Governance) -> Result<Governance, crate::error::SignatureError> {
    let offset = a.sig.n();
    let new_sig = Signature::new(a.sig.i + b.sig.i, a.sig.d + b.sig.d, a.sig.h + b.sig.h)?;

    // Offset derived generators from B
    let mut derived_gens = a.derived_gens.clone();
    for dg in &b.derived_gens {
        derived_gens.push(offset_mv(dg, offset));
    }

    // Geom classes: A's as-is, B's with offset blade masks in constraints
    let a_dg_count = a.derived_gens.len();
    let mut geom_classes = a.geom_classes.clone();
    for class in &b.geom_classes {
        geom_classes.push(offset_geom_class(class, offset, a_dg_count));
    }

    // Constructions: A's as-is, B's with offset class_index and generator refs
    let a_class_count = a.geom_classes.len();
    let mut constructions = a.constructions.clone();
    for constr in &b.constructions {
        constructions.push(Construction {
            class_index: constr.class_index + a_class_count,
            arity: constr.arity,
            body: offset_expr(&constr.body, offset, a_dg_count),
        });
    }

    Ok(Governance {
        sig: new_sig,
        derived_gens,
        geom_classes,
        constructions,
        probe: None,
        transform_rules: vec![],
    })
}

/// Offset all generator indices in an Mv by `offset`.
fn offset_mv(mv: &Mv, offset: u8) -> Mv {
    let mut result = Mv::new();
    for (mask, coeff) in mv.blades() {
        let new_mask = mask << offset;
        result.add_term(new_mask, coeff.clone());
    }
    result
}

/// Offset a GeomClass for composition.
/// Grade mask is preserved (grades don't change under generator offset).
/// Polynomial equations/inequalities are cloned as-is.
/// NOTE: polynomial variable indices are NOT remapped — they reference
/// the original algebra's blade ordering. For composed algebras with
/// polynomial constraints, the constraints should be rebuilt.
fn offset_geom_class(class: &GeomClass, _gen_offset: u8, _dg_offset: usize) -> GeomClass {
    GeomClass {
        grade_mask: class.grade_mask,
        equations: class.equations.clone(),
        inequalities: class.inequalities.clone(),
        field_op: class.field_op.clone(),
        expected_profile: class.expected_profile,
    }
}

/// Offset generator and derived-gen references in an Expr.
fn offset_expr(expr: &Expr, gen_offset: u8, dg_offset: usize) -> Expr {
    match expr {
        Expr::Param(i) => Expr::Param(*i),
        Expr::Generator(k) => Expr::Generator(k + gen_offset),
        Expr::DerivedGen(i) => Expr::DerivedGen(i + dg_offset),
        Expr::Literal(s) => Expr::Literal(s.clone()),
        Expr::Add(a, b) => Expr::Add(
            Box::new(offset_expr(a, gen_offset, dg_offset)),
            Box::new(offset_expr(b, gen_offset, dg_offset)),
        ),
        Expr::Mul(a, b) => Expr::Mul(
            Box::new(offset_expr(a, gen_offset, dg_offset)),
            Box::new(offset_expr(b, gen_offset, dg_offset)),
        ),
        Expr::Neg(a) => Expr::Neg(Box::new(offset_expr(a, gen_offset, dg_offset))),
        Expr::Pow(a, n) => Expr::Pow(Box::new(offset_expr(a, gen_offset, dg_offset)), *n),
        Expr::Construct(idx, args) => Expr::Construct(
            *idx, // construction index NOT offset — stays within the same governance
            args.iter()
                .map(|a| Box::new(offset_expr(a, gen_offset, dg_offset)))
                .collect(),
        ),
        Expr::Outer(a, b) => Expr::Outer(
            Box::new(offset_expr(a, gen_offset, dg_offset)),
            Box::new(offset_expr(b, gen_offset, dg_offset)),
        ),
        Expr::Inner(a, b) => Expr::Inner(
            Box::new(offset_expr(a, gen_offset, dg_offset)),
            Box::new(offset_expr(b, gen_offset, dg_offset)),
        ),
        Expr::Reverse(a) => Expr::Reverse(Box::new(offset_expr(a, gen_offset, dg_offset))),
        Expr::Dual(a) => Expr::Dual(Box::new(offset_expr(a, gen_offset, dg_offset))),
        Expr::Sandwich(a, b) => Expr::Sandwich(
            Box::new(offset_expr(a, gen_offset, dg_offset)),
            Box::new(offset_expr(b, gen_offset, dg_offset)),
        ),
        Expr::ValueRef(idx) => Expr::ValueRef(*idx),
        Expr::GradeProject(a, k) => {
            Expr::GradeProject(Box::new(offset_expr(a, gen_offset, dg_offset)), *k)
        }
        Expr::LeftContract(a, b) => Expr::LeftContract(
            Box::new(offset_expr(a, gen_offset, dg_offset)),
            Box::new(offset_expr(b, gen_offset, dg_offset)),
        ),
        Expr::ScalarProduct(a, b) => Expr::ScalarProduct(
            Box::new(offset_expr(a, gen_offset, dg_offset)),
            Box::new(offset_expr(b, gen_offset, dg_offset)),
        ),
        Expr::Read(inner, idx) => {
            Expr::Read(Box::new(offset_expr(inner, gen_offset, dg_offset)), *idx)
        }
        Expr::WithGov(gov_idx, inner) => Expr::WithGov(
            *gov_idx,
            Box::new(offset_expr(inner, gen_offset, dg_offset)),
        ),
        Expr::Embed(inner, emb_idx) => Expr::Embed(
            Box::new(offset_expr(inner, gen_offset, dg_offset)),
            *emb_idx,
        ),
        Expr::Morph(inner, morph_idx) => Expr::Morph(
            Box::new(offset_expr(inner, gen_offset, dg_offset)),
            *morph_idx,
        ),
        Expr::Probe => Expr::Probe,
        Expr::Object => Expr::Object,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::governance::govern;
    use crate::scalar::{Rat, Scalar};

    /// VGA(2): Cl(0,0,2) with one vector class
    fn vga2() -> Governance {
        Governance {
            sig: Signature::new(0, 0, 2).unwrap(),
            derived_gens: vec![],
            geom_classes: vec![GeomClass::grades_only(&[1])],
            constructions: vec![Construction {
                class_index: 0,
                arity: 2,
                body: Expr::Add(
                    Expr::mul(Expr::param(0), Expr::gen(0)),
                    Expr::mul(Expr::param(1), Expr::gen(1)),
                ),
            }],
            probe: None,
            transform_rules: vec![],
        }
    }

    /// VGA(1): Cl(0,0,1) with one vector class
    fn vga1() -> Governance {
        Governance {
            sig: Signature::new(0, 0, 1).unwrap(),
            derived_gens: vec![],
            geom_classes: vec![GeomClass::grades_only(&[1])],
            constructions: vec![Construction {
                class_index: 0,
                arity: 1,
                body: *Expr::mul(Expr::param(0), Expr::gen(0)),
            }],
            probe: None,
            transform_rules: vec![],
        }
    }

    #[test]
    fn compose_vga2_vga1_is_vga3() {
        let a = vga2();
        let b = vga1();
        let c = compose(&a, &b).unwrap();

        assert_eq!(c.sig, Signature::new(0, 0, 3).unwrap());
        assert_eq!(c.geom_classes.len(), 2); // one from each
        assert_eq!(c.constructions.len(), 2);
    }

    #[test]
    fn compose_preserves_factor_a() {
        let a = vga2();
        let b = vga1();
        let c = compose(&a, &b).unwrap();

        // Factor A's construction (class 0) still works in the composed algebra
        let params = vec![Scalar::from(3i64), Scalar::from(4i64)];
        let mv = c.construct(0, &params).unwrap();
        let geoit = govern(&mv, &c, 0).unwrap();
        let extracted = geoit.read_all().unwrap();
        assert_eq!(extracted, params);
    }

    #[test]
    fn compose_preserves_factor_b_offset() {
        let a = vga2();
        let b = vga1();
        let c = compose(&a, &b).unwrap();

        // Factor B's construction (class 1) uses offset generators
        let params = vec![Scalar::from(7i64)];
        let mv = c.construct(1, &params).unwrap();

        // The Mv should use g₂ (offset by 2 from B's g₀)
        assert_eq!(mv.coefficient(0b100), Scalar::from(7i64)); // g₂
        assert_eq!(mv.len(), 1);

        let geoit = govern(&mv, &c, 1).unwrap();
        let extracted = geoit.read_all().unwrap();
        assert_eq!(extracted, params);
    }

    #[test]
    fn compose_combined_signature() {
        let a = Governance {
            sig: Signature::new(1, 0, 2).unwrap(), // Cl(1,0,2)
            derived_gens: vec![],
            geom_classes: vec![],
            constructions: vec![],
            probe: None,
            transform_rules: vec![],
        };
        let b = Governance {
            sig: Signature::new(0, 1, 3).unwrap(), // Cl(0,1,3)
            derived_gens: vec![],
            geom_classes: vec![],
            constructions: vec![],
            probe: None,
            transform_rules: vec![],
        };
        let c = compose(&a, &b).unwrap();
        assert_eq!(c.sig, Signature::new(1, 1, 5).unwrap()); // combined
        assert_eq!(c.sig.n(), 7);
    }

    // === Embedding tests ===

    #[test]
    fn embedding_vga3_into_cga3() {
        // VGA(3) = Cl(0,0,3), generators g₀, g₁, g₂
        // CGA(3) = Cl(1,0,4), generators g₀(i-type), g₁..g₄(h-type)
        // Embed VGA's g₀→g₂, g₁→g₃, g₂→g₄ (the Euclidean subspace)
        let vga = Signature::new(0, 0, 3).unwrap();
        let cga = Signature::new(1, 0, 4).unwrap();
        let emb = Embedding::new(vga, cga, vec![2, 3, 4]).unwrap();

        // Embed a vector
        let v = Mv::from_rat_terms(&[
            (0b001, Rat::from(3)), // g₀
            (0b010, Rat::from(4)), // g₁
            (0b100, Rat::from(5)), // g₂
        ]);
        let embedded = emb.embed_mv(&v);

        // Should now be at g₂, g₃, g₄ in CGA
        assert_eq!(embedded.coefficient(0b00100), Scalar::from(3i64)); // g₂
        assert_eq!(embedded.coefficient(0b01000), Scalar::from(4i64)); // g₃
        assert_eq!(embedded.coefficient(0b10000), Scalar::from(5i64)); // g₄
        assert_eq!(embedded.len(), 3);
    }

    #[test]
    fn embedding_preserves_quadratic_form() {
        let vga = Signature::new(0, 0, 3).unwrap();
        let cga = Signature::new(1, 0, 4).unwrap();
        let emb = Embedding::new(vga, cga, vec![2, 3, 4]).unwrap();

        // Verify: VGA g₀²=+1, mapped to CGA g₂²=+1 ✓
        assert_eq!(vga.generator_square(0), cga.generator_square(2));
        assert_eq!(vga.generator_square(1), cga.generator_square(3));
        assert_eq!(vga.generator_square(2), cga.generator_square(4));

        // Embed a bivector and verify norm is preserved
        let bv = Mv::from_rat_terms(&[(0b011, Rat::from(1))]); // g₀₁ in VGA
        let embedded = emb.embed_mv(&bv);

        let vga_norm = crate::algebra::ops::norm_squared(&bv, &vga);
        let cga_norm = crate::algebra::ops::norm_squared(&embedded, &cga);
        assert_eq!(vga_norm, cga_norm, "embedding should preserve norm");
    }

    #[test]
    fn embedding_quadratic_form_violation() {
        // Try to embed Cl(0,0,1) g₀²=+1 into Cl(1,0,0) g₀²=-1
        let source = Signature::new(0, 0, 1).unwrap();
        let target = Signature::new(1, 0, 0).unwrap();
        let result = Embedding::new(source, target, vec![0]);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            EmbeddingError::QuadraticFormViolation { .. }
        ));
    }

    #[test]
    fn embedding_not_injective() {
        let source = Signature::new(0, 0, 2).unwrap();
        let target = Signature::new(0, 0, 3).unwrap();
        // Both source gens map to target g₀
        let result = Embedding::new(source, target, vec![0, 0]);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            EmbeddingError::NotInjective { .. }
        ));
    }

    #[test]
    fn embedding_target_out_of_range() {
        let source = Signature::new(0, 0, 1).unwrap();
        let target = Signature::new(0, 0, 2).unwrap();
        let result = Embedding::new(source, target, vec![5]); // 5 > target.n()=2
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            EmbeddingError::TargetOutOfRange { .. }
        ));
    }

    #[test]
    fn embedding_map_length_mismatch() {
        let source = Signature::new(0, 0, 3).unwrap();
        let target = Signature::new(0, 0, 5).unwrap();
        let result = Embedding::new(source, target, vec![0, 1]); // need 3, gave 2
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            EmbeddingError::MapLengthMismatch { .. }
        ));
    }

    #[test]
    fn embedding_identity() {
        let sig = Signature::new(0, 0, 3).unwrap();
        let emb = Embedding::new(sig, sig, vec![0, 1, 2]).unwrap();
        let v = Mv::from_rat_terms(&[
            (0b001, Rat::from(3)),
            (0b010, Rat::from(4)),
            (0b100, Rat::from(5)),
        ]);
        let embedded = emb.embed_mv(&v);
        assert_eq!(embedded, v);
    }
}