//! Grounding: eliminating variables by substituting all constant combinations.
//!
//! An ASPIC+ system may contain schematic (non-ground) rules and premises whose
//! formulae include variables. Before argument construction can proceed, every
//! variable must be replaced by concrete constants — a process called
//! *grounding*.
//!
//! This module provides:
//!
//! * [`GroundedSystem`] — the fully-ground output, consumed by argument
//!   constructors.
//! * `GroundingContext` — the internal engine that enumerates substitutions
//!   and applies them to rules, premises, and contrariness mappings.
//! * [`Substitution`] — a variable-to-constant mapping.
//! * Helper functions ([`apply_substitution`], [`collect_vars_from_formulae`])
//!   reused by both the naïve and semi-naïve grounders in
//!   [`crate::pipeline::grounding`].

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

use crate::error::AspicError;
use crate::language::formula::{Formula, Term};
use crate::system::contrariness::ContrarinessFunction;
use crate::system::knowledge::{KnowledgeBase, Premise};
use crate::system::rules::{DefeasibleRule, StrictRule};
use crate::system::AspicSystem;

/// A variable-to-constant mapping used during grounding.
pub type Substitution = HashMap<String, String>;

/// Apply a [`Substitution`] to a formula, replacing each variable with its
/// corresponding constant.
pub fn apply_substitution(formula: &Formula, subst: &Substitution) -> Formula {
    let mut result = formula.clone();
    for (var, constant) in subst {
        result = result.substitute(var, &Term::Const(constant.clone()));
    }
    result
}

/// Collect all distinct variables appearing in a slice of formulae,
/// preserving first-occurrence order.
pub fn collect_vars_from_formulae(formulae: &[&Formula]) -> Vec<String> {
    let mut seen = HashSet::new();
    let mut vars = Vec::new();
    for f in formulae {
        for v in f.variables() {
            if seen.insert(v.clone()) {
                vars.push(v.clone());
            }
        }
    }
    vars
}

/// A fully grounded ASPIC+ system, ready for argument construction.
///
/// All variables have been eliminated; every rule and premise is ground.
/// Produced by a [`crate::pipeline::traits::Grounder`] implementation such as
/// [`crate::pipeline::grounding::NaiveGrounder`] or
/// [`crate::pipeline::grounding::SemiNaiveGrounder`].
#[derive(Clone, Debug)]
pub struct GroundedSystem {
    /// The knowledge base with all premise formulae fully ground.
    pub knowledge_base: KnowledgeBase,
    /// Strict rules after variable elimination; each ground instance gets a
    /// unique [`RuleId`](crate::system::rules::RuleId) derived from the
    /// original rule's ID.
    pub strict_rules: Vec<StrictRule>,
    /// Defeasible rules after variable elimination, with names (if present)
    /// also ground.
    pub defeasible_rules: Vec<DefeasibleRule>,
    /// The contrariness function with all entries ground against the same
    /// Herbrand universe.
    pub contrariness: ContrarinessFunction,
}

/// Internal grounding engine that enumerates constant substitutions and
/// applies them to an [`AspicSystem`].
pub(crate) struct GroundingContext {
    pub(crate) constants: BTreeSet<String>,
}

impl GroundingContext {
    pub(crate) fn from_system(system: &AspicSystem) -> Self {
        GroundingContext {
            constants: system.collect_constants(),
        }
    }

    pub(crate) fn ground_formula(&self, formula: &Formula) -> Vec<Formula> {
        let vars: Vec<String> = formula.variables().into_iter().cloned().collect();
        if vars.is_empty() {
            return vec![formula.clone()];
        }

        let substitutions = self.all_substitutions(&vars);
        substitutions
            .into_iter()
            .map(|subst| {
                let mut f = formula.clone();
                for (var, constant) in &subst {
                    f = f.substitute(var, &Term::Const(constant.clone()));
                }
                f
            })
            .collect()
    }

    #[allow(clippy::unnecessary_wraps)]
    pub(crate) fn ground_system(&self, system: &AspicSystem) -> Result<GroundedSystem, AspicError> {
        let mut grounded_premises = Vec::new();
        for premise in &system.knowledge_base.premises {
            let ground_formulae = self.ground_formula(&premise.formula);
            for gf in ground_formulae {
                grounded_premises.push(Premise {
                    formula: gf,
                    kind: premise.kind.clone(),
                });
            }
        }

        let mut grounded_strict = Vec::new();
        for rule in &system.strict_rules {
            let all_formulae: Vec<&Formula> = rule
                .antecedents
                .iter()
                .chain(std::iter::once(&rule.consequent))
                .collect();
            let vars = collect_vars_from_formulae(&all_formulae);
            if vars.is_empty() {
                grounded_strict.push(rule.clone());
            } else {
                let substitutions = self.all_substitutions(&vars);
                for (i, subst) in substitutions.into_iter().enumerate() {
                    let mut antecedents = rule.antecedents.clone();
                    let mut consequent = rule.consequent.clone();
                    for (var, constant) in &subst {
                        let term = Term::Const(constant.clone());
                        antecedents = antecedents
                            .iter()
                            .map(|a| a.substitute(var, &term))
                            .collect();
                        consequent = consequent.substitute(var, &term);
                    }
                    grounded_strict.push(StrictRule {
                        id: crate::system::rules::RuleId(format!("{}_{}", rule.id.0, i)),
                        antecedents,
                        consequent,
                    });
                }
            }
        }

        let mut grounded_defeasible = Vec::new();
        for rule in &system.defeasible_rules {
            let mut all_formulae: Vec<&Formula> = rule
                .antecedents
                .iter()
                .chain(std::iter::once(&rule.consequent))
                .collect();
            if let Some(name) = &rule.name {
                all_formulae.push(name);
            }
            let vars = collect_vars_from_formulae(&all_formulae);
            if vars.is_empty() {
                grounded_defeasible.push(rule.clone());
            } else {
                let substitutions = self.all_substitutions(&vars);
                for (i, subst) in substitutions.into_iter().enumerate() {
                    let mut antecedents = rule.antecedents.clone();
                    let mut consequent = rule.consequent.clone();
                    let mut name = rule.name.clone();
                    for (var, constant) in &subst {
                        let term = Term::Const(constant.clone());
                        antecedents = antecedents
                            .iter()
                            .map(|a| a.substitute(var, &term))
                            .collect();
                        consequent = consequent.substitute(var, &term);
                        if let Some(n) = &name {
                            name = Some(n.substitute(var, &term));
                        }
                    }
                    grounded_defeasible.push(DefeasibleRule {
                        id: crate::system::rules::RuleId(format!("{}_{}", rule.id.0, i)),
                        antecedents,
                        consequent,
                        name,
                    });
                }
            }
        }

        let mut grounded_contrariness = system.contrariness.clone();
        grounded_contrariness.ground(self);

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

    pub(crate) fn all_substitutions(&self, vars: &[String]) -> Vec<Vec<(String, String)>> {
        if vars.is_empty() {
            return vec![vec![]];
        }
        let constants: Vec<&String> = self.constants.iter().collect();
        let mut result = vec![vec![]];
        for var in vars {
            let mut new_result = Vec::new();
            for subst in &result {
                for constant in &constants {
                    let mut new_subst = subst.clone();
                    new_subst.push((var.clone(), (*constant).clone()));
                    new_result.push(new_subst);
                }
            }
            result = new_result;
        }
        result
    }
}
