geoit 0.0.2

use crate::algebra::blade_new::{grade, BladeMask};
use crate::algebra::mv::Mv;
use crate::algebra::signature::Signature;
use crate::governance::geom_class::GeomClass;
use crate::governance::groebner;
use crate::scalar::Scalar;
use std::collections::BTreeMap;

// ═══════════════════════════════════════════════════════════
// ExtractionMap
// ═══════════════════════════════════════════════════════════

#[derive(Clone, Debug)]
pub enum ExtractionMap {
    Coefficient(BladeMask),
    Linear(Vec<(BladeMask, Scalar)>),
    Rational {
        numerator: Vec<(BladeMask, Scalar)>,
        denominator: Vec<(BladeMask, Scalar)>,
    },
}

#[derive(Clone, Debug)]
pub enum ExtractionError {
    DivisionByZero,
    IndexOutOfRange { index: usize, len: usize },
    ConstructionEvalFailed { context: String },
    EmptyProbeResult,
}

impl ExtractionMap {
    pub fn apply(&self, mv: &Mv) -> Result<Scalar, ExtractionError> {
        match self {
            ExtractionMap::Coefficient(mask) => Ok(mv.coefficient(*mask)),
            ExtractionMap::Linear(terms) => {
                let mut result = Scalar::from(0i64);
                for (mask, weight) in terms {
                    result += mv.coefficient(*mask) * weight.clone();
                }
                Ok(result)
            }
            ExtractionMap::Rational {
                numerator,
                denominator,
            } => {
                let mut num = Scalar::from(0i64);
                for (mask, weight) in numerator {
                    num += mv.coefficient(*mask) * weight.clone();
                }
                let mut den = Scalar::from(0i64);
                for (mask, weight) in denominator {
                    den += mv.coefficient(*mask) * weight.clone();
                }
                if den.is_zero() {
                    return Err(ExtractionError::DivisionByZero);
                }
                Ok(num / den)
            }
        }
    }
}

impl std::fmt::Display for ExtractionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ExtractionError::DivisionByZero => write!(f, "division by zero in extraction"),
            ExtractionError::IndexOutOfRange { index, len } => {
                write!(f, "extraction index {} out of range (have {})", index, len)
            }
            ExtractionError::ConstructionEvalFailed { context } => {
                write!(f, "construction eval failed: {}", context)
            }
            ExtractionError::EmptyProbeResult => write!(
                f,
                "probe produced empty diff — cannot determine blade mapping"
            ),
        }
    }
}

// ═══════════════════════════════════════════════════════════
// ReadingRules
// ═══════════════════════════════════════════════════════════

#[derive(Clone, Debug)]
pub struct ReadingRules {
    pub extractions: Vec<ExtractionMap>,
}

impl ReadingRules {
    pub fn apply(&self, mv: &Mv) -> Result<Vec<Scalar>, ExtractionError> {
        self.extractions.iter().map(|e| e.apply(mv)).collect()
    }

    pub fn read(&self, mv: &Mv, index: usize) -> Result<Scalar, ExtractionError> {
        if index >= self.extractions.len() {
            return Err(ExtractionError::IndexOutOfRange {
                index,
                len: self.extractions.len(),
            });
        }
        self.extractions[index].apply(mv)
    }
}

// ═══════════════════════════════════════════════════════════
// Variable Map
// ═══════════════════════════════════════════════════════════

#[derive(Clone, Debug)]
pub struct VariableMap {
    pub mask_to_var: BTreeMap<BladeMask, usize>,
    pub var_to_mask: Vec<BladeMask>,
    pub num_vars: usize,
}

impl VariableMap {
    pub fn for_all_blades(sig: &Signature) -> Self {
        let n = sig.n();
        let mut mask_to_var = BTreeMap::new();
        let mut var_to_mask = Vec::new();
        for mask in 0..(1u64 << n) {
            mask_to_var.insert(mask, var_to_mask.len());
            var_to_mask.push(mask);
        }
        let num_vars = var_to_mask.len();
        VariableMap {
            mask_to_var,
            var_to_mask,
            num_vars,
        }
    }

    pub fn for_grade(sig: &Signature, k: u8) -> Self {
        let n = sig.n();
        let mut mask_to_var = BTreeMap::new();
        let mut var_to_mask = Vec::new();
        for mask in 0..(1u64 << n) {
            if grade(mask) == k {
                mask_to_var.insert(mask, var_to_mask.len());
                var_to_mask.push(mask);
            }
        }
        let num_vars = var_to_mask.len();
        VariableMap {
            mask_to_var,
            var_to_mask,
            num_vars,
        }
    }

    /// Build a variable map for blades at any grade permitted by the mask.
    pub fn for_grade_mask(sig: &Signature, grade_mask: u64) -> Self {
        let n = sig.n();
        let mut mask_to_var = BTreeMap::new();
        let mut var_to_mask = Vec::new();
        for mask in 0..(1u64 << n) {
            let g = grade(mask);
            if grade_mask & (1u64 << g) != 0 {
                mask_to_var.insert(mask, var_to_mask.len());
                var_to_mask.push(mask);
            }
        }
        let num_vars = var_to_mask.len();
        VariableMap {
            mask_to_var,
            var_to_mask,
            num_vars,
        }
    }
}

// ═══════════════════════════════════════════════════════════
// Gröbner-Based Derivation
// ═══════════════════════════════════════════════════════════

impl ReadingRules {
    /// Derive reading rules from Gröbner basis analysis.
    ///
    /// The equations are taken directly from GeomClass.equations.
    /// No conversion needed — constraints ARE polynomials.
    pub fn derive_from_groebner(
        class: &GeomClass,
        construction: &super::construction::Construction,
        sig: &Signature,
        derived_gens: &[Mv],
    ) -> Result<Self, ExtractionError> {
        let arity = construction.arity;
        if arity == 0 {
            return Ok(ReadingRules {
                extractions: vec![],
            });
        }

        let var_map = VariableMap::for_grade_mask(sig, class.grade_mask);

        // Equations are already polynomials — use them directly
        let basis = groebner::groebner_basis(class.equations.clone()).map_err(|e| {
            ExtractionError::ConstructionEvalFailed {
                context: format!("Gröbner basis computation: {}", e),
            }
        })?;
        let free_vars = groebner::free_variables(&basis, var_map.num_vars);

        let param_mapping =
            probe_param_mapping(construction, sig, derived_gens, &var_map, &free_vars)?;

        let extractions: Vec<ExtractionMap> = param_mapping
            .iter()
            .map(|(mask, weight)| {
                if *weight == Scalar::from(1i64) {
                    ExtractionMap::Coefficient(*mask)
                } else if *weight == Scalar::from(-1i64) {
                    ExtractionMap::Linear(vec![(*mask, Scalar::from(-1i64))])
                } else {
                    let inv = Scalar::from(1i64) / weight.clone();
                    ExtractionMap::Linear(vec![(*mask, inv)])
                }
            })
            .collect();

        Ok(ReadingRules { extractions })
    }

    /// Derive reading rules by probing the construction formula.
    pub fn derive_from_probing(
        construction: &super::construction::Construction,
        sig: &Signature,
        derived_gens: &[Mv],
    ) -> Result<Self, ExtractionError> {
        let arity = construction.arity;
        if arity == 0 {
            return Ok(ReadingRules {
                extractions: vec![],
            });
        }

        let zero_params: Vec<Scalar> = vec![Scalar::from(0i64); arity];
        let base = construction
            .construct(&zero_params, sig, derived_gens, &[])
            .map_err(|e| ExtractionError::ConstructionEvalFailed {
                context: format!("at zero: {}", e),
            })?;

        let mut extractions = Vec::with_capacity(arity);

        for i in 0..arity {
            let mut probe_params: Vec<Scalar> = vec![Scalar::from(0i64); arity];
            probe_params[i] = Scalar::from(1i64);
            let probe = construction
                .construct(&probe_params, sig, derived_gens, &[])
                .map_err(|e| ExtractionError::ConstructionEvalFailed {
                    context: format!("at probe {}: {}", i, e),
                })?;
            let diff = probe - base.clone();

            if diff.len() == 1 {
                let (mask, weight) = diff
                    .blades()
                    .next()
                    .ok_or(ExtractionError::EmptyProbeResult)?;
                if *weight == Scalar::from(1i64) {
                    extractions.push(ExtractionMap::Coefficient(mask));
                } else if *weight == Scalar::from(-1i64) {
                    extractions.push(ExtractionMap::Linear(vec![(mask, Scalar::from(-1i64))]));
                } else {
                    let inv_weight = Scalar::from(1i64) / weight.clone();
                    extractions.push(ExtractionMap::Linear(vec![(mask, inv_weight)]));
                }
            } else if diff.len() > 1 {
                let simple_blade = diff
                    .blades()
                    .find(|(_, w)| **w == Scalar::from(1i64) || **w == Scalar::from(-1i64));
                if let Some((mask, weight)) = simple_blade {
                    if *weight == Scalar::from(1i64) {
                        extractions.push(ExtractionMap::Coefficient(mask));
                    } else {
                        extractions.push(ExtractionMap::Linear(vec![(mask, Scalar::from(-1i64))]));
                    }
                } else {
                    let (mask, weight) = diff
                        .blades()
                        .next()
                        .ok_or(ExtractionError::EmptyProbeResult)?;
                    let inv_weight = Scalar::from(1i64) / weight.clone();
                    extractions.push(ExtractionMap::Linear(vec![(mask, inv_weight)]));
                }
            } else {
                extractions.push(ExtractionMap::Coefficient(0));
            }
        }

        Ok(ReadingRules { extractions })
    }

    /// Simple derivation from grade mask (no construction or equations needed).
    pub fn derive_from_grade_mask(
        class: &GeomClass,
        construction_arity: usize,
        sig: &Signature,
    ) -> Self {
        let var_map = VariableMap::for_grade_mask(sig, class.grade_mask);
        let extractions = var_map
            .var_to_mask
            .iter()
            .take(construction_arity)
            .map(|&mask| ExtractionMap::Coefficient(mask))
            .collect();
        ReadingRules { extractions }
    }
}

/// Probe the construction at unit parameters to determine which
/// blade mask each parameter maps to, and with what weight.
fn probe_param_mapping(
    construction: &super::construction::Construction,
    sig: &Signature,
    derived_gens: &[Mv],
    var_map: &VariableMap,
    free_vars: &[usize],
) -> Result<Vec<(BladeMask, Scalar)>, ExtractionError> {
    let arity = construction.arity;
    let zero_params: Vec<Scalar> = vec![Scalar::from(0i64); arity];
    let base = construction
        .construct(&zero_params, sig, derived_gens, &[])
        .map_err(|e| ExtractionError::ConstructionEvalFailed {
            context: format!("probe at zero: {}", e),
        })?;

    let free_masks: Vec<BladeMask> = free_vars
        .iter()
        .map(|&vi| var_map.var_to_mask[vi])
        .collect();

    let mut result = Vec::with_capacity(arity);

    for i in 0..arity {
        let mut probe_params: Vec<Scalar> = vec![Scalar::from(0i64); arity];
        probe_params[i] = Scalar::from(1i64);
        let probe = construction
            .construct(&probe_params, sig, derived_gens, &[])
            .map_err(|e| ExtractionError::ConstructionEvalFailed {
                context: format!("probe at unit {}: {}", i, e),
            })?;
        let diff = probe - base.clone();

        let mut best_mask: BladeMask = 0;
        let mut best_weight = Scalar::from(0i64);
        for &mask in &free_masks {
            let coeff = diff.coefficient(mask);
            if !coeff.is_zero() {
                best_mask = mask;
                best_weight = coeff;
                break;
            }
        }

        if !best_weight.is_zero() {
            result.push((best_mask, best_weight));
        } else if !diff.is_zero() {
            let (mask, weight) = diff
                .blades()
                .next()
                .ok_or(ExtractionError::EmptyProbeResult)?;
            result.push((mask, weight.clone()));
        } else {
            result.push((0, Scalar::from(1i64)));
        }
    }

    Ok(result)
}

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

    #[test]
    fn coefficient_extraction() {
        let mv = Mv::from_rat_terms(&[
            (0b001, Rat::from(3)),
            (0b010, Rat::from(4)),
            (0b100, Rat::from(5)),
        ]);
        assert_eq!(
            ExtractionMap::Coefficient(0b010).apply(&mv).unwrap(),
            Scalar::from(4i64)
        );
    }

    #[test]
    fn coefficient_absent() {
        let mv = Mv::from_rat_terms(&[(0b001, Rat::from(3))]);
        assert_eq!(
            ExtractionMap::Coefficient(0b010).apply(&mv).unwrap(),
            Scalar::from(0i64)
        );
    }

    #[test]
    fn linear_extraction() {
        let mv = Mv::from_rat_terms(&[(0b01, Rat::from(10)), (0b10, Rat::from(20))]);
        let e = ExtractionMap::Linear(vec![(0b01, Scalar::from(2i64)), (0b10, Scalar::from(3i64))]);
        assert_eq!(e.apply(&mv).unwrap(), Scalar::from(80i64));
    }

    #[test]
    fn rational_division_by_zero() {
        let mv = Mv::from_rat_terms(&[(0b01, Rat::from(5))]);
        let e = ExtractionMap::Rational {
            numerator: vec![(0b01, Scalar::from(1i64))],
            denominator: vec![(0b10, Scalar::from(1i64))],
        };
        assert!(matches!(e.apply(&mv), Err(ExtractionError::DivisionByZero)));
    }

    #[test]
    fn reading_rules_apply() {
        let mv = Mv::from_rat_terms(&[
            (0b001, Rat::from(3)),
            (0b010, Rat::from(4)),
            (0b100, Rat::from(5)),
        ]);
        let rules = ReadingRules {
            extractions: vec![
                ExtractionMap::Coefficient(0b001),
                ExtractionMap::Coefficient(0b010),
                ExtractionMap::Coefficient(0b100),
            ],
        };
        assert_eq!(
            rules.apply(&mv).unwrap(),
            vec![Scalar::from(3i64), Scalar::from(4i64), Scalar::from(5i64)]
        );
    }

    #[test]
    fn read_out_of_range() {
        let rules = ReadingRules {
            extractions: vec![],
        };
        assert!(matches!(
            rules.read(&Mv::new(), 0),
            Err(ExtractionError::IndexOutOfRange { .. })
        ));
    }

    #[test]
    fn variable_map_grade_mask() {
        let sig = Signature::new(0, 0, 3).unwrap();
        // Grade 1 only
        let vm = VariableMap::for_grade_mask(&sig, 0b10);
        assert_eq!(vm.num_vars, 3);
        assert_eq!(vm.var_to_mask, vec![0b001, 0b010, 0b100]);
        // Grades 0 and 2
        let vm2 = VariableMap::for_grade_mask(&sig, 0b101);
        assert_eq!(vm2.num_vars, 4); // 1 scalar + 3 bivectors
    }
}