//! Semi-naive bottom-up grounder with delta iteration.
//!
//! Instead of full Herbrand enumeration, this grounder works bottom-up from
//! premises, only firing rules when at least one antecedent matches a
//! *newly derived* fact. This avoids producing ground instances that can
//! never participate in any argument, significantly reducing the grounded
//! system size for theories with many constants.

use std::collections::{HashMap, HashSet};

use crate::error::AspicError;
use crate::language::formula::{Formula, Literal, Term};
use crate::language::grounding::{
    apply_substitution, GroundedSystem, GroundingContext, Substitution,
};
use crate::pipeline::traits::Grounder;
use crate::system::knowledge::{KnowledgeBase, Premise};
use crate::system::rules::{DefeasibleRule, RuleId, StrictRule};
use crate::system::AspicSystem;

/// Semi-naive grounder that avoids generating irrelevant ground instances.
///
/// Instead of enumerating all possible variable-to-constant substitutions
/// (Herbrand grounding), this grounder works bottom-up:
/// 1. Start with premises as initial facts.
/// 2. For each rule, only attempt grounding when at least one body literal
///    matches a *newly derived* fact (semi-naive optimisation).
/// 3. Newly derived conclusions become available for the next iteration.
/// 4. Terminate when no new facts are produced (fixpoint).
#[derive(Default)]
pub struct SemiNaiveGrounder;

/// Key for fact indexing: (`predicate_name`, arity).
type FactKey = (String, usize);

/// Index of known ground facts for fast body matching.
struct FactIndex {
    by_key: HashMap<FactKey, Vec<Formula>>,
    all: HashSet<Formula>,
}

impl FactIndex {
    fn new() -> Self {
        FactIndex {
            by_key: HashMap::new(),
            all: HashSet::new(),
        }
    }

    fn insert(&mut self, formula: &Formula) -> bool {
        if !self.all.insert(formula.clone()) {
            return false;
        }
        if let Some(key) = fact_key(formula) {
            self.by_key.entry(key).or_default().push(formula.clone());
        }
        true
    }

    fn contains(&self, formula: &Formula) -> bool {
        self.all.contains(formula)
    }

    fn lookup(&self, key: &FactKey) -> &[Formula] {
        self.by_key.get(key).map_or(&[], std::vec::Vec::as_slice)
    }
}

fn fact_key(formula: &Formula) -> Option<FactKey> {
    match formula {
        Formula::Lit(lit) => {
            let atom = lit.atom();
            Some((atom.predicate.clone(), atom.terms.len()))
        }
        _ => None,
    }
}

impl Grounder for SemiNaiveGrounder {
    /// Ground the system using a bottom-up fixpoint with the semi-naive
    /// optimisation.
    ///
    /// Facts are seeded from ground premises, then rules are fired
    /// iteratively. On each pass only rules where at least one body
    /// literal matches a *newly derived* fact (from the previous
    /// iteration) are considered -- the `prev_new` / `mem::replace`
    /// swap partitions facts into "already known" vs "delta" without
    /// copying.
    ///
    /// Convergence is guaranteed by monotonicity of the immediate-
    /// consequence operator `T_P`: the fact set grows monotonically and
    /// the Herbrand base is finite, so a fixpoint is reached in
    /// finitely many passes.
    fn ground(&self, system: &AspicSystem) -> Result<GroundedSystem, AspicError> {
        let mut facts = FactIndex::new();
        let mut new_facts = FactIndex::new();
        let mut grounded_premises = Vec::new();
        let mut grounded_strict = Vec::new();
        let mut grounded_defeasible = Vec::new();
        let mut seen_rules: HashSet<(Vec<Formula>, Formula)> = HashSet::new();

        let all_constants: Vec<String> = system.collect_constants().into_iter().collect();

        // Phase 1: seed with ground premises.
        for premise in &system.knowledge_base.premises {
            if premise.formula.is_ground() {
                if facts.insert(&premise.formula) {
                    new_facts.insert(&premise.formula);
                    grounded_premises.push(premise.clone());
                }
            } else {
                for gf in ground_formula_all(&premise.formula, &all_constants) {
                    if facts.insert(&gf) {
                        new_facts.insert(&gf);
                        grounded_premises.push(Premise {
                            formula: gf,
                            kind: premise.kind.clone(),
                        });
                    }
                }
            }
        }

        // Phase 2: iterate until fixpoint. Each iteration only considers
        // rules where at least one body literal matches a newly derived fact.
        let mut rule_counter: usize = 0;

        loop {
            if new_facts.all.is_empty() {
                break;
            }

            let prev_new = std::mem::replace(&mut new_facts, FactIndex::new());
            let mut derived_any = false;

            for rule in &system.strict_rules {
                let groundings = ground_rule_body(
                    &rule.antecedents,
                    &rule.consequent,
                    &facts,
                    &prev_new,
                    &all_constants,
                );
                for (ground_ants, ground_cons) in groundings {
                    let key = (ground_ants.clone(), ground_cons.clone());
                    if !seen_rules.insert(key) {
                        continue;
                    }
                    rule_counter += 1;
                    if facts.insert(&ground_cons) {
                        new_facts.insert(&ground_cons);
                        derived_any = true;
                    }
                    grounded_strict.push(StrictRule {
                        id: RuleId(format!("{}_{}", rule.id.0, rule_counter)),
                        antecedents: ground_ants,
                        consequent: ground_cons,
                    });
                }
            }

            for rule in &system.defeasible_rules {
                let groundings = ground_rule_body(
                    &rule.antecedents,
                    &rule.consequent,
                    &facts,
                    &prev_new,
                    &all_constants,
                );
                for (ground_ants, ground_cons) in groundings {
                    let key = (ground_ants.clone(), ground_cons.clone());
                    if !seen_rules.insert(key) {
                        continue;
                    }
                    rule_counter += 1;
                    let ground_name = rule.name.as_ref().map(|n| {
                        let subst = extract_substitution(&rule.antecedents, &ground_ants);
                        apply_substitution(n, &subst)
                    });
                    if facts.insert(&ground_cons) {
                        new_facts.insert(&ground_cons);
                        derived_any = true;
                    }
                    grounded_defeasible.push(DefeasibleRule {
                        id: RuleId(format!("{}_{}", rule.id.0, rule_counter)),
                        antecedents: ground_ants,
                        consequent: ground_cons,
                        name: ground_name,
                    });
                }
            }

            if !derived_any {
                break;
            }
        }

        let ctx = GroundingContext::from_system(system);
        let mut grounded_contrariness = system.contrariness.clone();
        grounded_contrariness.ground(&ctx);

        Ok(GroundedSystem {
            knowledge_base: KnowledgeBase {
                premises: grounded_premises,
            },
            strict_rules: grounded_strict,
            defeasible_rules: grounded_defeasible,
            contrariness: grounded_contrariness,
        })
    }
}

/// Ground a rule body, returning ALL valid (`ground_antecedents`, `ground_consequent`)
/// combinations where every antecedent matches a known fact.
///
/// Semi-naive: only produces groundings where at least one antecedent matched
/// against a fact in `new_facts` (facts derived in the previous iteration).
fn ground_rule_body(
    antecedents: &[Formula],
    consequent: &Formula,
    all_facts: &FactIndex,
    new_facts: &FactIndex,
    all_constants: &[String],
) -> Vec<(Vec<Formula>, Formula)> {
    if antecedents.is_empty() {
        return ground_formula_all(consequent, all_constants)
            .into_iter()
            .map(|c| (vec![], c))
            .collect();
    }

    let vars = collect_vars(antecedents, consequent);

    if vars.is_empty() {
        if antecedents.iter().all(|a| all_facts.contains(a))
            && antecedents.iter().any(|a| new_facts.contains(a))
        {
            return vec![(antecedents.to_vec(), consequent.clone())];
        }
        return vec![];
    }

    let all_substs = find_all_substitutions(antecedents, all_facts, all_constants);

    let mut results = Vec::new();
    for (subst, used_new_in_any) in &all_substs {
        let ground_ants: Vec<Formula> = antecedents
            .iter()
            .map(|a| apply_substitution(a, subst))
            .collect();
        let ground_cons = apply_substitution(consequent, subst);

        let all_ground = ground_ants
            .iter()
            .all(crate::language::formula::Formula::is_ground)
            && ground_cons.is_ground();
        let all_known = ground_ants.iter().all(|a| all_facts.contains(a));
        let touches_new = *used_new_in_any || ground_ants.iter().any(|a| new_facts.contains(a));

        if all_ground && all_known && touches_new {
            results.push((ground_ants, ground_cons));
        }
    }

    results
}

/// Find all valid substitutions by matching antecedents left-to-right against
/// known facts, returning each complete substitution along with whether any
/// matched fact was newly derived.
fn find_all_substitutions(
    antecedents: &[Formula],
    facts: &FactIndex,
    all_constants: &[String],
) -> Vec<(Substitution, bool)> {
    let mut substitutions: Vec<(Substitution, bool)> = vec![(HashMap::new(), false)];

    for ant in antecedents {
        let mut next = Vec::new();
        for (subst, seen_new) in &substitutions {
            let partially_ground = apply_substitution(ant, subst);

            if partially_ground.is_ground() {
                if facts.contains(&partially_ground) {
                    next.push((subst.clone(), *seen_new));
                }
                continue;
            }

            let matches = match_pattern_against_facts(&partially_ground, facts, all_constants);
            for m in matches {
                if let Some(merged) = merge_substitutions(subst, &m) {
                    next.push((merged, *seen_new));
                }
            }
        }
        substitutions = next;
    }

    substitutions
}

/// Match a (possibly non-ground) formula pattern against indexed facts.
fn match_pattern_against_facts(
    pattern: &Formula,
    facts: &FactIndex,
    all_constants: &[String],
) -> Vec<Substitution> {
    let Some(key) = fact_key(pattern) else {
        return enumerate_substitutions_for(pattern, all_constants);
    };

    let mut results = Vec::new();
    for fact in facts.lookup(&key) {
        if let Some(subst) = unify(pattern, fact) {
            results.push(subst);
        }
    }
    results
}

/// Unify a pattern formula against a ground formula, returning variable bindings.
fn unify(pattern: &Formula, ground: &Formula) -> Option<Substitution> {
    match (pattern, ground) {
        (Formula::Lit(p_lit), Formula::Lit(g_lit)) => match (p_lit, g_lit) {
            (Literal::Pos(p_atom), Literal::Pos(g_atom))
            | (Literal::Neg(p_atom), Literal::Neg(g_atom)) => {
                if p_atom.predicate != g_atom.predicate || p_atom.terms.len() != g_atom.terms.len()
                {
                    return None;
                }
                let mut subst = HashMap::new();
                for (p_term, g_term) in p_atom.terms.iter().zip(g_atom.terms.iter()) {
                    match (p_term, g_term) {
                        (Term::Const(a), Term::Const(b)) if a == b => {}
                        (Term::Const(_), Term::Const(_)) => return None,
                        (Term::Var(v), Term::Const(c)) => {
                            if let Some(existing) = subst.get(v) {
                                if existing != c {
                                    return None;
                                }
                            } else {
                                subst.insert(v.clone(), c.clone());
                            }
                        }
                        _ => return None,
                    }
                }
                Some(subst)
            }
            _ => None,
        },
        _ => None,
    }
}

/// Merge two substitutions, returning None on conflicting bindings.
fn merge_substitutions(a: &Substitution, b: &Substitution) -> Option<Substitution> {
    let mut merged = a.clone();
    for (k, v) in b {
        if let Some(existing) = merged.get(k) {
            if existing != v {
                return None;
            }
        } else {
            merged.insert(k.clone(), v.clone());
        }
    }
    Some(merged)
}

/// Extract the substitution that transforms pattern antecedents into ground ones.
fn extract_substitution(pattern_ants: &[Formula], ground_ants: &[Formula]) -> Substitution {
    let mut subst = HashMap::new();
    for (p, g) in pattern_ants.iter().zip(ground_ants.iter()) {
        if let Some(s) = unify(p, g) {
            for (k, v) in s {
                subst.entry(k).or_insert(v);
            }
        }
    }
    subst
}

fn collect_vars(antecedents: &[Formula], consequent: &Formula) -> Vec<String> {
    let refs: Vec<&Formula> = antecedents
        .iter()
        .chain(std::iter::once(consequent))
        .collect();
    crate::language::grounding::collect_vars_from_formulae(&refs)
}

fn ground_formula_all(formula: &Formula, constants: &[String]) -> Vec<Formula> {
    let vars: Vec<String> = formula.variables().into_iter().cloned().collect();
    if vars.is_empty() {
        return vec![formula.clone()];
    }
    let mut results = vec![HashMap::new()];
    for var in &vars {
        let mut next = Vec::new();
        for subst in &results {
            for c in constants {
                let mut s = subst.clone();
                s.insert(var.clone(), c.clone());
                next.push(s);
            }
        }
        results = next;
    }
    results
        .into_iter()
        .map(|subst| apply_substitution(formula, &subst))
        .collect()
}

fn enumerate_substitutions_for(formula: &Formula, constants: &[String]) -> Vec<Substitution> {
    let vars: Vec<String> = formula.variables().into_iter().cloned().collect();
    if vars.is_empty() {
        return vec![HashMap::new()];
    }
    let mut results: Vec<Substitution> = vec![HashMap::new()];
    for var in &vars {
        let mut next = Vec::new();
        for subst in &results {
            for c in constants {
                let mut s = subst.clone();
                s.insert(var.clone(), c.clone());
                next.push(s);
            }
        }
        results = next;
    }
    results
}
