xlog-logic 0.5.0

//! Abstract Syntax Tree for XLOG programs

use xlog_core::ScalarType;

/// A term in an atom
#[derive(Debug, Clone, PartialEq)]
pub enum Term {
    Variable(String),
    /// Anonymous wildcard `_` - each occurrence is a fresh unnamed variable
    Anonymous,
    Integer(i64),
    Float(f64),
    String(String),
    /// Interned symbol ID - use `xlog_core::symbol::resolve(id)` to get the string
    Symbol(u32),
    Aggregate(AggExpr),
}

impl Term {
    pub fn is_variable(&self) -> bool {
        matches!(self, Term::Variable(_))
    }

    /// Returns true if this is an anonymous wildcard `_`
    pub fn is_anonymous(&self) -> bool {
        matches!(self, Term::Anonymous)
    }

    /// Returns true if this is any kind of variable (named or anonymous)
    pub fn is_any_variable(&self) -> bool {
        matches!(self, Term::Variable(_) | Term::Anonymous)
    }

    pub fn is_constant(&self) -> bool {
        !self.is_any_variable() && !matches!(self, Term::Aggregate(_))
    }

    /// Returns the variable name, or None for anonymous/constants
    pub fn variable_name(&self) -> Option<&str> {
        match self {
            Term::Variable(name) => Some(name),
            _ => None,
        }
    }
}

/// Aggregate expression
#[derive(Debug, Clone, PartialEq)]
pub struct AggExpr {
    pub op: AggOp,
    pub variable: String,
}

/// Aggregation operator
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AggOp {
    Count,
    Sum,
    Min,
    Max,
    LogSumExp,
}

/// Arithmetic expression tree
#[derive(Debug, Clone, PartialEq)]
pub enum ArithExpr {
    Variable(String),
    Integer(i64),
    Float(f64),

    // Binary operations
    Add(Box<ArithExpr>, Box<ArithExpr>),
    Sub(Box<ArithExpr>, Box<ArithExpr>),
    Mul(Box<ArithExpr>, Box<ArithExpr>),
    Div(Box<ArithExpr>, Box<ArithExpr>),
    Mod(Box<ArithExpr>, Box<ArithExpr>),

    // Built-in functions
    Abs(Box<ArithExpr>),
    Min(Box<ArithExpr>, Box<ArithExpr>),
    Max(Box<ArithExpr>, Box<ArithExpr>),
    Pow(Box<ArithExpr>, Box<ArithExpr>),

    // Type cast
    Cast(Box<ArithExpr>, ScalarType),

    /// User-defined function call
    FuncCall {
        name: String,
        args: Vec<ArithExpr>,
    },

    /// Conditional expression (for expanded function bodies)
    Conditional {
        cond_left: Box<ArithExpr>,
        cond_op: CompOp,
        cond_right: Box<ArithExpr>,
        then_expr: Box<ArithExpr>,
        else_expr: Box<ArithExpr>,
    },
}

impl ArithExpr {
    /// Get all variable names used in this expression
    pub fn variables(&self) -> Vec<&str> {
        match self {
            ArithExpr::Variable(name) => vec![name.as_str()],
            ArithExpr::Integer(_) | ArithExpr::Float(_) => vec![],
            ArithExpr::Add(l, r)
            | ArithExpr::Sub(l, r)
            | ArithExpr::Mul(l, r)
            | ArithExpr::Div(l, r)
            | ArithExpr::Mod(l, r)
            | ArithExpr::Min(l, r)
            | ArithExpr::Max(l, r)
            | ArithExpr::Pow(l, r) => {
                let mut vars = l.variables();
                vars.extend(r.variables());
                vars
            }
            ArithExpr::Abs(e) | ArithExpr::Cast(e, _) => e.variables(),
            ArithExpr::FuncCall { args, .. } => args.iter().flat_map(|a| a.variables()).collect(),
            ArithExpr::Conditional {
                cond_left,
                cond_right,
                then_expr,
                else_expr,
                ..
            } => {
                let mut vars = cond_left.variables();
                vars.extend(cond_right.variables());
                vars.extend(then_expr.variables());
                vars.extend(else_expr.variables());
                vars
            }
        }
    }
}

/// Is-expression for variable binding: Z is X + Y
#[derive(Debug, Clone, PartialEq)]
pub struct IsExpr {
    pub target: String, // Must be fresh (unbound) variable
    pub expr: ArithExpr,
}

/// An atom (predicate applied to terms)
#[derive(Debug, Clone, PartialEq)]
pub struct Atom {
    pub predicate: String,
    pub terms: Vec<Term>,
}

impl Atom {
    pub fn arity(&self) -> usize {
        self.terms.len()
    }

    pub fn variables(&self) -> Vec<&str> {
        self.terms
            .iter()
            .filter_map(|t| t.variable_name())
            .collect()
    }
}

/// Comparison operator
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompOp {
    Eq,
    Ne,
    Lt,
    Le,
    Gt,
    Ge,
}

/// A comparison expression
#[derive(Debug, Clone, PartialEq)]
pub struct Comparison {
    pub left: Term,
    pub op: CompOp,
    pub right: Term,
}

/// A literal in the body of a rule
#[derive(Debug, Clone, PartialEq)]
pub enum BodyLiteral {
    Positive(Atom),
    Negated(Atom),
    Comparison(Comparison),
    IsExpr(IsExpr),
}

impl BodyLiteral {
    pub fn is_positive(&self) -> bool {
        matches!(self, BodyLiteral::Positive(_))
    }

    pub fn is_negated(&self) -> bool {
        matches!(self, BodyLiteral::Negated(_))
    }

    pub fn atom(&self) -> Option<&Atom> {
        match self {
            BodyLiteral::Positive(a) | BodyLiteral::Negated(a) => Some(a),
            BodyLiteral::Comparison(_) | BodyLiteral::IsExpr(_) => None,
        }
    }

    pub fn variables(&self) -> Vec<&str> {
        match self {
            BodyLiteral::Positive(a) | BodyLiteral::Negated(a) => a.variables(),
            BodyLiteral::Comparison(c) => {
                let mut vars = vec![];
                if let Some(v) = c.left.variable_name() {
                    vars.push(v);
                }
                if let Some(v) = c.right.variable_name() {
                    vars.push(v);
                }
                vars
            }
            BodyLiteral::IsExpr(is_expr) => {
                let mut vars = is_expr.expr.variables();
                vars.push(is_expr.target.as_str());
                vars
            }
        }
    }
}

/// A rule (head :- body)
#[derive(Debug, Clone, PartialEq)]
pub struct Rule {
    pub head: Atom,
    pub body: Vec<BodyLiteral>,
}

impl Rule {
    pub fn is_fact(&self) -> bool {
        self.body.is_empty()
    }

    pub fn has_negation(&self) -> bool {
        self.body.iter().any(|l| l.is_negated())
    }

    pub fn has_aggregation(&self) -> bool {
        self.head
            .terms
            .iter()
            .any(|t| matches!(t, Term::Aggregate(_)))
    }

    pub fn body_predicates(&self) -> Vec<&str> {
        self.body
            .iter()
            .filter_map(|l| l.atom().map(|a| a.predicate.as_str()))
            .collect()
    }

    pub fn head_variables(&self) -> Vec<&str> {
        self.head.variables()
    }

    pub fn body_variables(&self) -> Vec<&str> {
        self.body.iter().flat_map(|l| l.variables()).collect()
    }
}

/// A constraint (:- body)
#[derive(Debug, Clone, PartialEq)]
pub struct Constraint {
    pub body: Vec<BodyLiteral>,
}

/// A query (?- atom)
#[derive(Debug, Clone, PartialEq)]
pub struct Query {
    pub atom: Atom,
}

/// Probabilistic engine selection
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProbEngine {
    ExactDdnnf,
    Mc,
}

/// Probabilistic compilation caching
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProbCache {
    On,
    Off,
}

/// Compilation/evaluation directives (e.g., `#pragma ...`)
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Directives {
    pub prob_engine: Option<ProbEngine>,
    pub prob_cache: Option<ProbCache>,
    pub max_recursion_depth: Option<u32>,
}

impl Directives {
    pub fn prob_engine_or_default(&self) -> ProbEngine {
        self.prob_engine.unwrap_or(ProbEngine::ExactDdnnf)
    }

    pub fn max_recursion_depth_or_default(&self) -> u32 {
        self.max_recursion_depth.unwrap_or(1000)
    }
}

/// A probabilistic fact (`p::atom.`)
#[derive(Debug, Clone, PartialEq)]
pub struct ProbFact {
    pub prob: f64,
    pub atom: Atom,
}

/// Neural predicate declaration
///
/// Neural predicates connect neural networks to probabilistic logic.
/// Syntax: `nn(network, [inputs], output, [labels]) :: pred(args).`
///
/// The neural network produces probability distributions over labels,
/// which become probabilistic facts in the logic program.
///
/// # Examples
/// ```text
/// nn(mnist_net, [X], Y, [0,1,2,3,4,5,6,7,8,9]) :: digit(X, Y).
/// nn(encoder, [Text], Embedding) :: encode(Text, Embedding).
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct NeuralPredDecl {
    /// Name of the registered neural network
    pub network: String,
    /// Input variable names (bind to tensor sources)
    pub inputs: Vec<String>,
    /// Output variable name
    pub output: String,
    /// Optional classification labels (for classification networks)
    /// If None, the network produces embeddings
    pub labels: Option<Vec<NeuralLabel>>,
    /// The predicate this neural network defines
    pub predicate: Atom,
}

/// A label in a neural predicate classification
///
/// Labels can be integers or symbols (identifiers).
#[derive(Debug, Clone, PartialEq)]
pub enum NeuralLabel {
    Integer(i64),
    Symbol(String),
}

/// A learnable rule template parameterized by a named tensor mask.
/// Used for differentiable ILP — the mask selects which (body1, body2, head)
/// combinations are active during execution.
#[derive(Debug, Clone)]
pub struct LearnableRule {
    pub mask_name: String,
    pub head: Atom,
    pub body: Vec<BodyLiteral>,
}

/// Annotated disjunction (`p1::a1; p2::a2.`)
#[derive(Debug, Clone, PartialEq)]
pub struct AnnotatedDisjunction {
    pub choices: Vec<ProbFact>,
}

/// Evidence statement (`evidence(atom, true|false).`)
#[derive(Debug, Clone, PartialEq)]
pub struct Evidence {
    pub atom: Atom,
    pub value: bool,
}

/// Probabilistic query statement (`query(atom).`)
#[derive(Debug, Clone, PartialEq)]
pub struct ProbQuery {
    pub atom: Atom,
}

/// Import statement: use module. or use module::{pred1, pred2}.
#[derive(Debug, Clone, PartialEq)]
pub struct UseDecl {
    /// Module path segments, e.g., ["utils", "math"]
    pub module_path: Vec<String>,
    /// Specific imports (None = import all public)
    pub imports: Option<Vec<String>>,
}

/// Domain declaration
#[derive(Debug, Clone, PartialEq)]
pub struct DomainDecl {
    pub name: String,
    pub typ: ScalarType,
}

/// Predicate declaration
#[derive(Debug, Clone, PartialEq)]
pub struct PredDecl {
    pub name: String,
    pub types: Vec<ScalarType>,
    pub is_private: bool,
}

/// Function parameter with optional type annotation
#[derive(Debug, Clone, PartialEq)]
pub struct FuncParam {
    pub name: String,
    pub typ: Option<ScalarType>,
}

/// Conditional expression: if X < 0 then A else B
#[derive(Debug, Clone, PartialEq)]
pub struct CondExpr {
    /// Left side of condition
    pub cond_left: ArithExpr,
    /// Comparison operator
    pub cond_op: CompOp,
    /// Right side of condition
    pub cond_right: ArithExpr,
    /// Value if condition is true
    pub then_branch: Box<FuncBody>,
    /// Value if condition is false
    pub else_branch: Box<FuncBody>,
}

/// Function body - arithmetic, conditional, or predicate-based
#[derive(Debug, Clone, PartialEq)]
pub enum FuncBody {
    /// Pure arithmetic expression: X * X
    Arithmetic(ArithExpr),
    /// Conditional expression: if X < 0 then ...
    Conditional(CondExpr),
    /// Predicate-based: P :- parent(X, P)
    Predicate {
        /// Result variable
        result: String,
        /// Body literals
        body: Vec<BodyLiteral>,
    },
}

/// User-defined function
#[derive(Debug, Clone, PartialEq)]
pub struct FuncDef {
    /// Function name
    pub name: String,
    /// Parameters
    pub params: Vec<FuncParam>,
    /// Optional return type annotation
    pub return_type: Option<ScalarType>,
    /// Function body
    pub body: FuncBody,
    /// Is this function private?
    pub is_private: bool,
}

/// A complete XLOG program
#[derive(Debug, Clone, Default)]
pub struct Program {
    pub imports: Vec<UseDecl>,
    pub functions: Vec<FuncDef>,
    pub domains: Vec<DomainDecl>,
    pub predicates: Vec<PredDecl>,
    pub rules: Vec<Rule>,
    pub constraints: Vec<Constraint>,
    pub queries: Vec<Query>,
    pub prob_facts: Vec<ProbFact>,
    pub annotated_disjunctions: Vec<AnnotatedDisjunction>,
    pub evidence: Vec<Evidence>,
    pub prob_queries: Vec<ProbQuery>,
    pub neural_predicates: Vec<NeuralPredDecl>,
    pub learnable_rules: Vec<LearnableRule>,
    pub directives: Directives,
}

impl Program {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn facts(&self) -> impl Iterator<Item = &Rule> {
        self.rules.iter().filter(|r| r.is_fact())
    }

    pub fn proper_rules(&self) -> impl Iterator<Item = &Rule> {
        self.rules.iter().filter(|r| !r.is_fact())
    }

    pub fn defined_predicates(&self) -> Vec<&str> {
        self.rules
            .iter()
            .map(|r| r.head.predicate.as_str())
            .collect::<std::collections::HashSet<_>>()
            .into_iter()
            .collect()
    }

    pub fn is_probabilistic_profile(&self) -> bool {
        !self.prob_facts.is_empty()
            || !self.annotated_disjunctions.is_empty()
            || !self.evidence.is_empty()
            || !self.prob_queries.is_empty()
            || self.directives.prob_engine.is_some()
            || self.directives.prob_cache.is_some()
    }

    pub fn prob_engine(&self) -> ProbEngine {
        self.directives.prob_engine_or_default()
    }

    /// Merge another program's exports into this program.
    /// Used for importing modules - adds predicates, functions, rules from the imported module.
    /// Only merges public items (private items are not exported).
    ///
    /// # Arguments
    /// * `other` - The program to merge from
    /// * `imported_items` - Optional set of specific items to import. If None, imports all public items.
    pub fn merge_from(
        &mut self,
        other: &Program,
        imported_items: Option<&std::collections::HashSet<String>>,
    ) {
        use std::collections::HashSet;

        // Track which predicates are private in the source
        let private_preds: HashSet<&str> = other
            .predicates
            .iter()
            .filter(|p| p.is_private)
            .map(|p| p.name.as_str())
            .collect();

        let _private_funcs: HashSet<&str> = other
            .functions
            .iter()
            .filter(|f| f.is_private)
            .map(|f| f.name.as_str())
            .collect();

        // Merge predicate declarations (only public ones)
        for pred in &other.predicates {
            if pred.is_private {
                continue;
            }
            // Check if this is in the import list (if specified)
            if let Some(items) = imported_items {
                if !items.contains(&pred.name) {
                    continue;
                }
            }
            // Avoid duplicate declarations
            if !self.predicates.iter().any(|p| p.name == pred.name) {
                self.predicates.push(pred.clone());
            }
        }

        // Merge functions (only public ones)
        for func in &other.functions {
            if func.is_private {
                continue;
            }
            if let Some(items) = imported_items {
                if !items.contains(&func.name) {
                    continue;
                }
            }
            // Avoid duplicate functions
            if !self.functions.iter().any(|f| f.name == func.name) {
                self.functions.push(func.clone());
            }
        }

        // Merge rules (facts and rules for public predicates)
        for rule in &other.rules {
            // Skip if the head predicate is private
            if private_preds.contains(rule.head.predicate.as_str()) {
                continue;
            }
            // Check import list for facts/rules
            if let Some(items) = imported_items {
                if !items.contains(&rule.head.predicate) {
                    continue;
                }
            }
            self.rules.push(rule.clone());
        }

        // Merge domains
        for domain in &other.domains {
            if !self.domains.iter().any(|d| d.name == domain.name) {
                self.domains.push(domain.clone());
            }
        }
    }
}

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

    #[test]
    fn test_term_variable() {
        let term = Term::Variable("X".to_string());
        assert!(term.is_variable());
        assert!(!term.is_constant());
    }

    #[test]
    fn test_term_constant() {
        let term = Term::Integer(42);
        assert!(!term.is_variable());
        assert!(term.is_constant());
    }

    #[test]
    fn test_atom_arity() {
        let atom = Atom {
            predicate: "edge".to_string(),
            terms: vec![Term::Integer(1), Term::Integer(2)],
        };
        assert_eq!(atom.arity(), 2);
    }

    #[test]
    fn test_atom_variables() {
        let atom = Atom {
            predicate: "edge".to_string(),
            terms: vec![Term::Variable("X".to_string()), Term::Integer(2)],
        };
        let vars = atom.variables();
        assert_eq!(vars, vec!["X"]);
    }

    #[test]
    fn test_rule_is_fact() {
        let fact = Rule {
            head: Atom {
                predicate: "edge".to_string(),
                terms: vec![Term::Integer(1), Term::Integer(2)],
            },
            body: vec![],
        };
        assert!(fact.is_fact());
    }

    #[test]
    fn test_rule_has_negation() {
        let rule = Rule {
            head: Atom {
                predicate: "isolated".to_string(),
                terms: vec![Term::Variable("X".to_string())],
            },
            body: vec![
                BodyLiteral::Positive(Atom {
                    predicate: "node".to_string(),
                    terms: vec![Term::Variable("X".to_string())],
                }),
                BodyLiteral::Negated(Atom {
                    predicate: "edge".to_string(),
                    terms: vec![
                        Term::Variable("X".to_string()),
                        Term::Variable("Y".to_string()),
                    ],
                }),
            ],
        };
        assert!(rule.has_negation());
    }

    #[test]
    fn test_program_facts() {
        let mut program = Program::new();
        program.rules.push(Rule {
            head: Atom {
                predicate: "edge".to_string(),
                terms: vec![Term::Integer(1), Term::Integer(2)],
            },
            body: vec![],
        });
        program.rules.push(Rule {
            head: Atom {
                predicate: "reach".to_string(),
                terms: vec![
                    Term::Variable("X".to_string()),
                    Term::Variable("Y".to_string()),
                ],
            },
            body: vec![BodyLiteral::Positive(Atom {
                predicate: "edge".to_string(),
                terms: vec![
                    Term::Variable("X".to_string()),
                    Term::Variable("Y".to_string()),
                ],
            })],
        });
        assert_eq!(program.facts().count(), 1);
        assert_eq!(program.proper_rules().count(), 1);
    }

    #[test]
    fn test_arith_expr_structure() {
        let expr = ArithExpr::Add(
            Box::new(ArithExpr::Variable("X".to_string())),
            Box::new(ArithExpr::Integer(1)),
        );
        assert!(matches!(expr, ArithExpr::Add(_, _)));
    }

    #[test]
    fn test_is_expr_structure() {
        let is_expr = IsExpr {
            target: "Z".to_string(),
            expr: ArithExpr::Variable("Y".to_string()),
        };
        assert_eq!(is_expr.target, "Z");
    }
}