xlog/

main.rs

use clap::{Parser, Subcommand, ValueEnum};
use std::collections::HashMap;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use arrow::csv::WriterBuilder;
use arrow::util::pretty::pretty_format_batches;
use xlog_core::{symbol, MemoryBudget, Result, XlogError};
use xlog_cuda::{CudaDevice, CudaKernelProvider, GpuMemoryManager};
use xlog_gpu::logic::LogicProgram;
use xlog_ir::{EirBodyLiteral, EirTerm};
use xlog_logic::ast::{BodyLiteral, CompOp, ProbEngine, Program, Term};
use xlog_logic::compile::load_modules;
#[cfg(feature = "host-io")]
use xlog_logic::parse_program;
use xlog_logic::IncrementalParseResult;
use xlog_logic::{rewrite_magic_sets, MagicSetReport, MagicSetStatus, ParserSession};
use xlog_logic::{stratify, Compiler};
use xlog_logic::{QueryProofTrace, RuleProvenance};
#[cfg(feature = "host-io")]
use xlog_prob::exact::ExactDdnnfProgram;
#[cfg(feature = "host-io")]
use xlog_prob::exact::GpuConfig;
#[cfg(feature = "host-io")]
use xlog_prob::mc::{McEvalConfig, McProgram, McSamplingMethod};
use xlog_prob::provenance::{AggregateLiftReport, Value};

#[derive(Parser)]
#[command(author, version, about = "XLOG CLI")]
pub struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    Run(RunArgs),
    Prob(ProbArgs),
    Explain(ExplainArgs),
    Repl(ReplArgs),
    Watch(WatchArgs),
}

#[derive(Parser)]
struct RunArgs {
    source: PathBuf,
    #[arg(long, default_value = "0")]
    device: usize,
    #[arg(long, default_value = "1024")]
    memory_mb: u64,
    #[arg(long)]
    input: Vec<String>,
    #[arg(long, value_enum, default_value = "pretty")]
    output: OutputFormat,
    #[arg(long)]
    output_dir: Option<PathBuf>,
    /// Show execution statistics (timing, memory usage)
    #[arg(long)]
    stats: bool,
    /// Stats output format (human or json)
    #[arg(long, value_enum, default_value = "human")]
    stats_format: StatsFormat,
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
    /// Dump the compiled epistemic execution plan (EIR-derived GPU plan, world-view
    /// integrity constraints, and CPU-fallback counters) as JSON to this path.
    /// No-op for ordinary (non-epistemic) programs. This compiled
    /// epistemic-plan/EIR JSON dump exposes accepted `know`/`possible` literals
    /// and lets a caller assert `cpu_fallback == 0` off a real GPU run.
    #[arg(long)]
    epistemic_plan_json: Option<PathBuf>,
}

#[derive(Copy, Clone, ValueEnum, Default)]
enum StatsFormat {
    #[default]
    Human,
    Json,
}

#[derive(Parser)]
struct ProbArgs {
    source: PathBuf,
    #[arg(long, default_value = "0")]
    device: usize,
    #[arg(long, default_value = "1024")]
    memory_mb: u64,
    #[arg(long, value_enum)]
    prob_engine: Option<ProbEngineCli>,
    #[arg(long)]
    samples: Option<usize>,
    #[arg(long)]
    seed: Option<u64>,
    #[arg(long)]
    confidence: Option<f64>,
    #[arg(long, value_enum)]
    prob_method: Option<ProbMethodCli>,
    #[arg(long, alias = "max-nonmonotone-iterations")]
    prob_max_nonmonotone_iterations: Option<usize>,
    /// Allow the labeled CPU oracle when the resident GPU MC engine rejects
    /// the program (negation, aggregates, ...). Fail-closed when unset; the
    /// result is labeled `mc_engine: cpu-oracle` and is not GPU-native evidence.
    #[arg(long)]
    allow_cpu_oracle: bool,
    #[arg(long, value_enum, default_value = "pretty")]
    output: ProbOutputFormat,
    #[arg(long)]
    output_dir: Option<PathBuf>,
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
}

#[derive(Parser)]
struct ExplainArgs {
    source: PathBuf,
    #[arg(long, value_enum, default_value = "text")]
    format: ExplainFormat,
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
}

#[derive(Parser)]
struct ReplArgs {
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
}

#[derive(Parser)]
struct WatchArgs {
    source: PathBuf,
    #[arg(long, default_value = "250")]
    debounce_ms: u64,
    #[arg(long)]
    explain: bool,
    #[arg(long)]
    once: bool,
}

#[derive(Copy, Clone, ValueEnum)]
enum ExplainFormat {
    Text,
    Json,
    Dot,
}

#[derive(Copy, Clone, ValueEnum)]
enum OutputFormat {
    Pretty,
    Csv,
    Arrow,
}

#[derive(Copy, Clone, ValueEnum)]
enum ProbOutputFormat {
    Pretty,
    Csv,
    Arrow,
    Json,
}

#[derive(Copy, Clone, ValueEnum)]
enum ProbEngineCli {
    #[value(name = "exact_ddnnf")]
    ExactDdnnf,
    Mc,
}

#[derive(Copy, Clone, ValueEnum)]
enum ProbMethodCli {
    Rejection,
    #[value(name = "evidence_clamping")]
    EvidenceClamping,
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    match cli.command {
        Command::Run(args) => run_deterministic(args),
        Command::Prob(args) => run_probabilistic(args),
        Command::Explain(args) => explain(args),
        Command::Repl(args) => repl(args),
        Command::Watch(args) => watch(args),
    }
}

fn explain(args: ExplainArgs) -> Result<()> {
    let source = std::fs::read_to_string(&args.source).map_err(|e| {
        XlogError::Execution(format!("Failed to read {}: {}", args.source.display(), e))
    })?;
    let mut parser_session = ParserSession::new();
    let parsed = parser_session.parse_path(&args.source, &source)?;
    let parsed = resolve_explain_imports(parsed, &args.source, args.module_path)?;
    let report = build_explain_report(parsed, Some(&args.source))?;
    match args.format {
        ExplainFormat::Text => print_explain_text(&report),
        ExplainFormat::Json => print_explain_json(&report),
        ExplainFormat::Dot => print_magic_dot(&report.magic_sets),
    }
    Ok(())
}

fn resolve_explain_imports(
    mut parsed: IncrementalParseResult,
    source_path: &Path,
    module_path: Vec<PathBuf>,
) -> Result<IncrementalParseResult> {
    if parsed.program.imports.is_empty() {
        return Ok(parsed);
    }
    let resolver = load_modules(source_path, module_path)
        .map_err(|e| XlogError::Execution(format!("Module resolution failed: {}", e)))?;
    parsed.program = resolver
        .merge_imports(parsed.program)
        .map_err(|e| XlogError::Execution(format!("Module resolution failed: {}", e)))?;
    Ok(parsed)
}

fn repl(args: ReplArgs) -> Result<()> {
    let _ = args.module_path;
    let mut input = String::new();
    std::io::stdin()
        .read_to_string(&mut input)
        .map_err(|e| XlogError::Execution(format!("Failed to read stdin: {}", e)))?;
    let mut parser_session = ParserSession::new();
    let parsed = parser_session.parse_path("<repl>", &input)?;
    println!(
        "repl: statements={} cache_hits={} cache_misses={}",
        parsed.stats.statement_count, parsed.stats.hits, parsed.stats.misses
    );
    println!(
        "state: rules={} queries={} prob_queries={}",
        parsed.program.rules.len(),
        parsed.program.queries.len(),
        parsed.program.prob_queries.len()
    );
    Ok(())
}

fn watch(args: WatchArgs) -> Result<()> {
    let mut parser_session = ParserSession::new();
    loop {
        let source = std::fs::read_to_string(&args.source).map_err(|e| {
            XlogError::Execution(format!("Failed to read {}: {}", args.source.display(), e))
        })?;
        let parsed = parser_session.parse_path(&args.source, &source)?;
        println!(
            "watch: statements={} cache_hits={} cache_misses={}",
            parsed.stats.statement_count, parsed.stats.hits, parsed.stats.misses
        );
        if args.explain {
            let report = build_explain_report(parsed, Some(&args.source))?;
            print_explain_text(&report);
        }
        if args.once {
            break;
        }
        std::thread::sleep(Duration::from_millis(args.debounce_ms));
    }
    Ok(())
}

struct ExplainReport {
    program: Program,
    parse_stats: xlog_logic::ParseCacheStats,
    epistemic: serde_json::Value,
    magic_sets: MagicSetReport,
    aggregate_lifting: Vec<AggregateLiftReport>,
    generated_rule_diagnostics: Vec<GeneratedRuleDiagnostic>,
    rule_provenance: Vec<RuleProvenance>,
    proof_traces: Vec<QueryProofTrace>,
    stratification_status: String,
    stratification_count: usize,
    rir_status: String,
    rir_sccs: usize,
    optimizer_status: String,
    optimizer_memory_peak: u64,
}

fn build_explain_report(
    parsed: xlog_logic::IncrementalParseResult,
    source_path: Option<&Path>,
) -> Result<ExplainReport> {
    let program = parsed.program;
    let magic_rewrite = rewrite_magic_sets(&program)?;
    let rule_provenance = xlog_logic::rule_provenance(&program, Some(&magic_rewrite.program));
    let proof_traces = xlog_logic::query_proof_traces(&program, &rule_provenance);
    let magic_sets = magic_rewrite.report;
    let aggregate_lifting = explain_aggregate_lifting(&program)?;
    let generated_rule_diagnostics = explain_generated_rule_diagnostics(&program, source_path);
    let epistemic = explain_epistemic(&program);
    let (stratification_status, stratification_count) = match stratify(&program) {
        Ok(strata) => ("ok".to_string(), strata.len()),
        Err(err) => (format!("error: {}", err), 0),
    };
    let mut compiler = Compiler::new();
    let (rir_status, rir_sccs, optimizer_status, optimizer_memory_peak) =
        match compiler.compile_program(&program) {
            Ok(plan) => (
                "ok".to_string(),
                plan.sccs.len(),
                "ok".to_string(),
                plan.est_memory_peak,
            ),
            Err(err) => (format!("error: {}", err), 0, "not_available".to_string(), 0),
        };
    Ok(ExplainReport {
        program,
        parse_stats: parsed.stats,
        epistemic,
        magic_sets,
        aggregate_lifting,
        generated_rule_diagnostics,
        rule_provenance,
        proof_traces,
        stratification_status,
        stratification_count,
        rir_status,
        rir_sccs,
        optimizer_status,
        optimizer_memory_peak,
    })
}

fn explain_epistemic(program: &Program) -> serde_json::Value {
    if !program_has_epistemic_literals(program) {
        let not_applicable = serde_json::json!({
            "status": "not_applicable",
            "reason": "program has no epistemic literals",
            "epistemic_literal_count": 0,
        });
        return serde_json::json!({
            "eir": not_applicable.clone(),
            "gpu_plan": not_applicable.clone(),
            "executable_plan": not_applicable,
        });
    }

    let eir = match xlog_logic::build_eir(program) {
        Ok(eir) => {
            let literals = eir
                .rules
                .iter()
                .enumerate()
                .flat_map(|(rule_index, rule)| {
                    rule.body.iter().filter_map(move |lit| match lit {
                        EirBodyLiteral::Epistemic(epistemic) => Some(serde_json::json!({
                            "rule_index": rule_index,
                            "literal": eir_epistemic_literal_json(epistemic),
                        })),
                        _ => None,
                    })
                })
                .collect::<Vec<_>>();
            let rule_summaries = eir
                .rules
                .iter()
                .enumerate()
                .map(|(rule_index, rule)| {
                    let epistemic_literal_count = rule
                        .body
                        .iter()
                        .filter(|lit| matches!(lit, EirBodyLiteral::Epistemic(_)))
                        .count();
                    let relational_body_atoms = rule
                        .body
                        .iter()
                        .filter(|lit| {
                            matches!(lit, EirBodyLiteral::Relational { negated: false, .. })
                        })
                        .count();
                    serde_json::json!({
                        "rule_index": rule_index,
                        "head": eir_atom_json(&rule.head),
                        "body_literal_count": rule.body.len(),
                        "epistemic_literal_count": epistemic_literal_count,
                        "relational_body_atoms": relational_body_atoms,
                    })
                })
                .collect::<Vec<_>>();
            serde_json::json!({
                "status": "ok",
                "mode": format!("{:?}", eir.mode),
                "rule_count": eir.rules.len(),
                "epistemic_literal_count": literals.len(),
                "literals": literals,
                "rules": rule_summaries,
            })
        }
        Err(err) => serde_json::json!({
            "status": "error",
            "error": err.to_string(),
        }),
    };

    let gpu_plan = match xlog_logic::epistemic::plan_epistemic_gpu_execution(program) {
        Ok(plan) => {
            let reductions = plan
                .reductions
                .iter()
                .map(|reduction| {
                    serde_json::json!({
                        "rule_index": reduction.rule_index,
                        "head_predicate": &reduction.head_predicate,
                        "relational_body_atoms": reduction.relational_body_atoms,
                        "wcoj_status": format!("{:?}", reduction.wcoj_status),
                    })
                })
                .collect::<Vec<_>>();
            let tuple_membership = plan
                .tuple_membership_bindings
                .iter()
                .map(|binding| {
                    serde_json::json!({
                        "literal_index": binding.literal_index,
                        "reduction_index": binding.reduction_index,
                        "predicate": &binding.predicate,
                        "arity": binding.arity,
                        "key_columns": binding.key_columns,
                        "key_terms": binding.key_terms.iter().map(eir_term_label).collect::<Vec<_>>(),
                        "bound_output_columns": binding.bound_output_columns,
                        "op": format!("{:?}", binding.op),
                        "negated": binding.negated,
                    })
                })
                .collect::<Vec<_>>();
            serde_json::json!({
                "status": "ok",
                "mode": format!("{:?}", plan.mode),
                "epistemic_literal_count": plan.epistemic_literals.len(),
                "required_phases": plan.required_phases.iter().map(|phase| format!("{:?}", phase)).collect::<Vec<_>>(),
                "required_kernel_phases": plan.required_kernel_phases.iter().map(|phase| format!("{:?}", phase)).collect::<Vec<_>>(),
                "required_buffers": plan.required_buffers.iter().map(|buffer| format!("{:?}", buffer)).collect::<Vec<_>>(),
                "reductions": reductions,
                "tuple_membership_bindings": tuple_membership,
                "solver_contract": {
                    "assumption_count": plan.solver_contract.assumption_bindings.len(),
                    "required_capabilities": plan.solver_contract.required_capabilities.iter().map(|cap| format!("{:?}", cap)).collect::<Vec<_>>(),
                    "required_statuses": plan.solver_contract.required_statuses.iter().map(|status| format!("{:?}", status)).collect::<Vec<_>>(),
                },
                "cpu_fallbacks": {
                    "candidate_enumeration": plan.cpu_fallbacks.candidate_enumeration,
                    "world_view_validation": plan.cpu_fallbacks.world_view_validation,
                    "solver_search": plan.cpu_fallbacks.solver_search,
                    "probabilistic_recompute": plan.cpu_fallbacks.probabilistic_recompute,
                    "is_zero": plan.cpu_fallbacks.is_zero(),
                }
            })
        }
        Err(err) => serde_json::json!({
            "status": "error",
            "error": err.to_string(),
        }),
    };

    let executable_plan = match xlog_logic::epistemic::compile_epistemic_gpu_execution(program) {
        Ok(plan) => serde_json::json!({
            "status": "ok",
            "relation_id_count": plan.relation_ids.len(),
            "reduced_runtime_sccs": plan.reduced_runtime_plan.sccs.len(),
            "reduced_runtime_est_memory_peak": plan.reduced_runtime_plan.est_memory_peak,
            "gpu_plan_literal_count": plan.gpu_plan.epistemic_literals.len(),
        }),
        Err(err) => serde_json::json!({
            "status": "error",
            "error": err.to_string(),
        }),
    };

    serde_json::json!({
        "eir": eir,
        "gpu_plan": gpu_plan,
        "executable_plan": executable_plan,
    })
}

fn program_has_epistemic_literals(program: &Program) -> bool {
    program.rules.iter().any(|rule| {
        rule.body
            .iter()
            .any(|lit| matches!(lit, BodyLiteral::Epistemic(_)))
    }) || program.constraints.iter().any(|constraint| {
        constraint
            .body
            .iter()
            .any(|lit| matches!(lit, BodyLiteral::Epistemic(_)))
    })
}

fn eir_atom_json(atom: &xlog_ir::EirAtom) -> serde_json::Value {
    serde_json::json!({
        "predicate": &atom.predicate,
        "arity": atom.arity,
        "terms": atom.terms.iter().map(eir_term_label).collect::<Vec<_>>(),
    })
}

fn eir_epistemic_literal_json(lit: &xlog_ir::EirEpistemicLiteral) -> serde_json::Value {
    serde_json::json!({
        "op": format!("{:?}", lit.op),
        "negated": lit.negated,
        "atom": eir_atom_json(&lit.atom),
    })
}

fn eir_term_label(term: &EirTerm) -> String {
    match term {
        EirTerm::Variable(name) => name.clone(),
        EirTerm::Anonymous => "_".to_string(),
        EirTerm::Integer(value) => value.to_string(),
        EirTerm::FloatBits(bits) => f64::from_bits(*bits).to_string(),
        EirTerm::String(value) => value.clone(),
        EirTerm::Symbol(id) => symbol::resolve(*id),
        EirTerm::List(items) => format!(
            "[{}]",
            items
                .iter()
                .map(eir_term_label)
                .collect::<Vec<_>>()
                .join(", ")
        ),
        EirTerm::Cons { head, tail } => {
            format!("{}|{}", eir_term_label(head), eir_term_label(tail))
        }
        EirTerm::Compound { functor, args } => format!(
            "{}({})",
            functor,
            args.iter()
                .map(eir_term_label)
                .collect::<Vec<_>>()
                .join(", ")
        ),
        EirTerm::PredRef(name) => name.clone(),
        EirTerm::Aggregate { op, variable } => format!("{}({})", op, variable),
    }
}

fn explain_aggregate_lifting(program: &Program) -> Result<Vec<AggregateLiftReport>> {
    let has_probabilistic_source =
        !program.prob_facts.is_empty() || !program.annotated_disjunctions.is_empty();
    let has_aggregate_rule = program.proper_rules().any(|rule| rule.has_aggregation());
    if !(has_probabilistic_source && has_aggregate_rule) {
        return Ok(Vec::new());
    }
    Ok(xlog_prob::provenance::extract_from_program(program)?.aggregate_lifting)
}

struct GeneratedRuleDiagnostic {
    rule_head: String,
    source_relation: String,
    row_decisions: Vec<GeneratedRuleRowDecision>,
}

struct GeneratedRuleRowDecision {
    row_key: String,
    accepted: bool,
    failed_predicates: Vec<String>,
    threshold_comparisons: Vec<ThresholdComparison>,
    aggregate_inputs: Vec<String>,
}

struct ThresholdComparison {
    predicate: String,
    left: String,
    op: String,
    right: String,
    left_value: String,
    right_value: String,
    passed: bool,
}

fn explain_generated_rule_diagnostics(
    program: &Program,
    source_path: Option<&Path>,
) -> Vec<GeneratedRuleDiagnostic> {
    let external_rows = source_path
        .map(|path| load_external_relation_rows(program, path))
        .unwrap_or_default();
    let mut diagnostics = Vec::new();
    for rule in program
        .rules
        .iter()
        .filter(|rule| !rule.body.is_empty() && generated_rule_candidate(rule))
    {
        let Some(source_atom) = diagnostic_source_atom(rule) else {
            continue;
        };
        let mut row_decisions = Vec::new();
        for source_row in source_rows_for_atom(program, &external_rows, source_atom) {
            let Some(bindings) = bindings_for_source_row(source_atom, &source_row) else {
                continue;
            };

            let mut threshold_comparisons = Vec::new();
            for comparison in rule.body.iter().filter_map(|literal| match literal {
                BodyLiteral::Comparison(comparison) => Some(comparison),
                _ => None,
            }) {
                let left_value = bound_term(&comparison.left, &bindings);
                let right_value = bound_term(&comparison.right, &bindings);
                let passed = left_value
                    .as_ref()
                    .zip(right_value.as_ref())
                    .and_then(|(left, right)| compare_terms(left, comparison.op, right))
                    .unwrap_or(false);
                threshold_comparisons.push(ThresholdComparison {
                    predicate: format!(
                        "{} {} {}",
                        term_label(&comparison.left),
                        comp_op_label(comparison.op),
                        term_label(&comparison.right)
                    ),
                    left: term_label(&comparison.left),
                    op: comp_op_label(comparison.op).to_string(),
                    right: term_label(&comparison.right),
                    left_value: left_value
                        .as_ref()
                        .map(term_label)
                        .unwrap_or_else(|| "unbound".to_string()),
                    right_value: right_value
                        .as_ref()
                        .map(term_label)
                        .unwrap_or_else(|| "unbound".to_string()),
                    passed,
                });
            }
            let mut failed_predicates = predicate_failures(program, rule, source_atom, &bindings);
            failed_predicates.extend(
                threshold_comparisons
                    .iter()
                    .filter(|comparison| !comparison.passed)
                    .map(|comparison| comparison.predicate.clone()),
            );
            row_decisions.push(GeneratedRuleRowDecision {
                row_key: source_row
                    .first()
                    .map(term_label)
                    .unwrap_or_else(|| source_atom.predicate.clone()),
                accepted: failed_predicates.is_empty(),
                failed_predicates,
                threshold_comparisons,
                aggregate_inputs: vec![format!(
                    "{}({})",
                    source_atom.predicate,
                    source_row
                        .iter()
                        .map(term_label)
                        .collect::<Vec<_>>()
                        .join(", ")
                )],
            });
        }

        if !row_decisions.is_empty() {
            diagnostics.push(GeneratedRuleDiagnostic {
                rule_head: rule.head.predicate.clone(),
                source_relation: source_atom.predicate.clone(),
                row_decisions,
            });
        }
    }
    diagnostics
}

fn generated_rule_candidate(rule: &xlog_logic::ast::Rule) -> bool {
    rule.head.predicate.starts_with("generated_")
        || rule.head.predicate.starts_with("xlog_accepted_")
        || rule.head.predicate.starts_with("xlog_rejected_")
        || rule.body.iter().any(|literal| match literal {
            BodyLiteral::Positive(atom) | BodyLiteral::Negated(atom) => {
                diagnostic_source_predicate(&atom.predicate)
            }
            BodyLiteral::Epistemic(_) => false,
            BodyLiteral::Comparison(_) | BodyLiteral::IsExpr(_) | BodyLiteral::Univ(_) => false,
        })
}

fn diagnostic_source_atom(rule: &xlog_logic::ast::Rule) -> Option<&xlog_logic::ast::Atom> {
    rule.body.iter().find_map(|literal| match literal {
        BodyLiteral::Positive(atom) if diagnostic_source_predicate(&atom.predicate) => Some(atom),
        _ => None,
    })
}

fn diagnostic_source_predicate(predicate: &str) -> bool {
    predicate.starts_with("generated_")
        || predicate.ends_with("_candidate_input")
        || (predicate.contains("candidate") && predicate.ends_with("_input"))
}

fn source_rows_for_atom(
    program: &Program,
    external_rows: &HashMap<String, Vec<Vec<Term>>>,
    atom: &xlog_logic::ast::Atom,
) -> Vec<Vec<Term>> {
    let mut rows = program
        .rules
        .iter()
        .filter(|fact| {
            fact.body.is_empty()
                && fact.head.predicate == atom.predicate
                && fact.head.terms.len() == atom.terms.len()
        })
        .map(|fact| fact.head.terms.clone())
        .collect::<Vec<_>>();
    if let Some(external) = external_rows.get(&atom.predicate) {
        rows.extend(external.clone());
    }
    rows
}

fn bindings_for_source_row(
    atom: &xlog_logic::ast::Atom,
    row: &[Term],
) -> Option<HashMap<String, Term>> {
    if atom.terms.len() != row.len() {
        return None;
    }
    let mut bindings = HashMap::new();
    for (pattern, value) in atom.terms.iter().zip(row) {
        match pattern {
            Term::Variable(name) => {
                if let Some(existing) = bindings.get(name) {
                    if existing != value {
                        return None;
                    }
                } else {
                    bindings.insert(name.clone(), value.clone());
                }
            }
            Term::Anonymous => {}
            _ if pattern == value => {}
            _ => return None,
        }
    }
    Some(bindings)
}

fn load_external_relation_rows(
    program: &Program,
    source_path: &Path,
) -> HashMap<String, Vec<Vec<Term>>> {
    let mut loaded = HashMap::new();
    for decl in &program.predicates {
        let Some((relation_path, columns)) = external_relation_source(source_path, decl) else {
            continue;
        };
        if columns.len() != decl.columns.len() {
            continue;
        }
        let Ok(source) = std::fs::read_to_string(&relation_path) else {
            continue;
        };
        let Ok(json) = serde_json::from_str::<serde_json::Value>(&source) else {
            continue;
        };
        let Some(rows) = json.get("rows").and_then(serde_json::Value::as_array) else {
            continue;
        };
        let mut relation_rows = Vec::new();
        for row in rows {
            let Some(object) = row.as_object() else {
                continue;
            };
            let mut terms = Vec::with_capacity(columns.len());
            let mut complete = true;
            for column in &columns {
                match object.get(column).and_then(json_value_to_term) {
                    Some(term) => terms.push(term),
                    None => {
                        complete = false;
                        break;
                    }
                }
            }
            if complete {
                relation_rows.push(terms);
            }
        }
        if !relation_rows.is_empty() {
            loaded.insert(decl.name.clone(), relation_rows);
        }
    }
    loaded
}

fn external_relation_source(
    source_path: &Path,
    decl: &xlog_logic::ast::PredDecl,
) -> Option<(PathBuf, Vec<String>)> {
    if let Some(from_manifest) = external_relation_source_from_manifest(source_path, decl) {
        return Some(from_manifest);
    }
    let columns = declared_column_names(decl)?;
    let source_dir = source_path.parent()?;
    for candidate in relation_json_candidates(source_dir, &decl.name) {
        if candidate.exists() {
            return Some((candidate, columns));
        }
    }
    None
}

fn external_relation_source_from_manifest(
    source_path: &Path,
    decl: &xlog_logic::ast::PredDecl,
) -> Option<(PathBuf, Vec<String>)> {
    let source_dir = source_path.parent()?;
    let mut manifests = vec![source_dir.join("xlog_hypothesis_execution.json")];
    if let Some(parent) = source_dir.parent() {
        manifests.push(parent.join("xlog_hypothesis_execution.json"));
    }
    for manifest_path in manifests {
        let Ok(source) = std::fs::read_to_string(&manifest_path) else {
            continue;
        };
        let Ok(json) = serde_json::from_str::<serde_json::Value>(&source) else {
            continue;
        };
        let Some(columns) = json
            .get("relation_input_columns")
            .and_then(serde_json::Value::as_array)
            .map(|items| {
                items
                    .iter()
                    .filter_map(serde_json::Value::as_str)
                    .map(ToString::to_string)
                    .collect::<Vec<_>>()
            })
        else {
            continue;
        };
        if columns.len() != decl.columns.len() {
            continue;
        }
        let Some(path_value) = json
            .get("relation_input_path")
            .and_then(serde_json::Value::as_str)
        else {
            continue;
        };
        let relation_path = PathBuf::from(path_value);
        let relation_path = if relation_path.is_absolute() {
            relation_path
        } else {
            manifest_path
                .parent()
                .unwrap_or_else(|| Path::new("."))
                .join(relation_path)
        };
        if relation_path.exists() {
            return Some((relation_path, columns));
        }
    }
    None
}

fn declared_column_names(decl: &xlog_logic::ast::PredDecl) -> Option<Vec<String>> {
    decl.columns
        .iter()
        .map(|column| column.name.clone())
        .collect()
}

fn relation_json_candidates(source_dir: &Path, predicate: &str) -> Vec<PathBuf> {
    let mut candidates = vec![source_dir.join(format!("{predicate}.json"))];
    if let Some(stem) = predicate.strip_suffix("_input") {
        candidates.push(source_dir.join(format!("{stem}_relation.json")));
    }
    candidates
}

fn json_value_to_term(value: &serde_json::Value) -> Option<Term> {
    if let Some(value) = value.as_i64() {
        Some(Term::Integer(value))
    } else if let Some(value) = value.as_u64() {
        i64::try_from(value).ok().map(Term::Integer)
    } else if let Some(value) = value.as_str() {
        Some(Term::String(value.to_string()))
    } else {
        value
            .as_bool()
            .map(|value| Term::Integer(if value { 1 } else { 0 }))
    }
}

fn predicate_failures(
    program: &Program,
    rule: &xlog_logic::ast::Rule,
    source_atom: &xlog_logic::ast::Atom,
    bindings: &HashMap<String, Term>,
) -> Vec<String> {
    let mut failures = Vec::new();
    for literal in &rule.body {
        match literal {
            BodyLiteral::Positive(atom) if atom.predicate != source_atom.predicate => {
                if !matching_fact_exists(program, atom, bindings) {
                    failures.push(atom_label(atom));
                }
            }
            BodyLiteral::Negated(atom) => {
                if matching_fact_exists(program, atom, bindings) {
                    failures.push(format!("not {}", atom_label(atom)));
                }
            }
            BodyLiteral::Positive(_)
            | BodyLiteral::Epistemic(_)
            | BodyLiteral::Comparison(_)
            | BodyLiteral::IsExpr(_)
            | BodyLiteral::Univ(_) => {}
        }
    }
    failures
}

fn matching_fact_exists(
    program: &Program,
    atom: &xlog_logic::ast::Atom,
    bindings: &HashMap<String, Term>,
) -> bool {
    program.rules.iter().any(|fact| {
        fact.body.is_empty()
            && fact.head.predicate == atom.predicate
            && fact.head.terms.len() == atom.terms.len()
            && atom
                .terms
                .iter()
                .zip(&fact.head.terms)
                .all(|(pattern, value)| match bound_term(pattern, bindings) {
                    Some(bound) => bound == *value,
                    None => false,
                })
    })
}

fn atom_label(atom: &xlog_logic::ast::Atom) -> String {
    format!(
        "{}({})",
        atom.predicate,
        atom.terms
            .iter()
            .map(term_label)
            .collect::<Vec<_>>()
            .join(", ")
    )
}

fn bound_term(term: &Term, bindings: &HashMap<String, Term>) -> Option<Term> {
    match term {
        Term::Variable(name) => bindings.get(name).cloned(),
        _ => Some(term.clone()),
    }
}

fn compare_terms(left: &Term, op: CompOp, right: &Term) -> Option<bool> {
    match (left, right) {
        (Term::Integer(left), Term::Integer(right)) => Some(compare_i64(*left, op, *right)),
        (Term::String(left), Term::String(right)) => match op {
            CompOp::Eq => Some(left == right),
            CompOp::Ne => Some(left != right),
            _ => None,
        },
        (Term::Symbol(left), Term::Symbol(right)) => match op {
            CompOp::Eq => Some(left == right),
            CompOp::Ne => Some(left != right),
            _ => None,
        },
        _ => None,
    }
}

fn compare_i64(left: i64, op: CompOp, right: i64) -> bool {
    match op {
        CompOp::Eq => left == right,
        CompOp::Ne => left != right,
        CompOp::Lt => left < right,
        CompOp::Le => left <= right,
        CompOp::Gt => left > right,
        CompOp::Ge => left >= right,
    }
}

fn comp_op_label(op: CompOp) -> &'static str {
    match op {
        CompOp::Eq => "==",
        CompOp::Ne => "!=",
        CompOp::Lt => "<",
        CompOp::Le => "<=",
        CompOp::Gt => ">",
        CompOp::Ge => ">=",
    }
}

fn term_label(term: &Term) -> String {
    match term {
        Term::Variable(name) => name.clone(),
        Term::Anonymous => "_".to_string(),
        Term::Integer(value) => value.to_string(),
        Term::Float(value) => value.to_string(),
        Term::String(value) => value.clone(),
        Term::Symbol(id) => symbol::resolve(*id),
        Term::List(items) => format!(
            "[{}]",
            items.iter().map(term_label).collect::<Vec<_>>().join(", ")
        ),
        Term::Cons { head, tail } => format!("{}|{}", term_label(head), term_label(tail)),
        Term::Compound { functor, args } => format!(
            "{}({})",
            functor,
            args.iter().map(term_label).collect::<Vec<_>>().join(", ")
        ),
        Term::PredRef(name) => name.clone(),
        Term::Aggregate(agg) => format!("{:?}({})", agg.op, agg.variable),
    }
}

fn print_explain_text(report: &ExplainReport) {
    println!("parse:");
    println!("  statements: {}", report.parse_stats.statement_count);
    println!("ast:");
    println!("  rules: {}", report.program.rules.len());
    println!("  queries: {}", report.program.queries.len());
    println!("stratification:");
    println!("  status: {}", report.stratification_status);
    println!("  strata: {}", report.stratification_count);
    println!("rir:");
    println!("  status: {}", report.rir_status);
    println!("  sccs: {}", report.rir_sccs);
    println!("optimizer:");
    println!("  status: {}", report.optimizer_status);
    println!("  est_memory_peak: {}", report.optimizer_memory_peak);
    println!("wcoj:");
    println!("  status: reported");
    println!("epistemic:");
    let eir_status = report
        .epistemic
        .get("eir")
        .and_then(|value| value.get("status"))
        .and_then(serde_json::Value::as_str)
        .unwrap_or("not_available");
    let gpu_plan_status = report
        .epistemic
        .get("gpu_plan")
        .and_then(|value| value.get("status"))
        .and_then(serde_json::Value::as_str)
        .unwrap_or("not_available");
    let executable_status = report
        .epistemic
        .get("executable_plan")
        .and_then(|value| value.get("status"))
        .and_then(serde_json::Value::as_str)
        .unwrap_or("not_available");
    println!("  eir: {}", eir_status);
    println!("  gpu_plan: {}", gpu_plan_status);
    println!("  executable_plan: {}", executable_status);
    print_magic_text(&report.magic_sets);
    if !report.aggregate_lifting.is_empty() {
        println!("aggregate_lifting:");
        for entry in &report.aggregate_lifting {
            println!(
                "  - predicate: {} operator: {} status: {} domain: {} uncertain: {} cap: {}",
                entry.predicate,
                entry.operator,
                entry.status.as_str(),
                entry.domain_size,
                entry.uncertain_rows,
                entry.cap
            );
        }
    }
    if !report.rule_provenance.is_empty() {
        println!("rule_provenance:");
        for entry in &report.rule_provenance {
            println!(
                "  - id: {} source_kind: {} head: {}",
                entry.rule_id,
                entry.source_kind.as_str(),
                entry.head
            );
        }
    }
    if !report.proof_traces.is_empty() {
        println!("proof_traces:");
        for entry in &report.proof_traces {
            println!(
                "  - query: {} rules: {} source_facts: {}",
                entry.query,
                entry.rule_ids.len(),
                entry.source_facts.len()
            );
        }
    }
}

fn print_magic_text(report: &MagicSetReport) {
    println!("magic_sets:");
    println!("  status: {}", magic_status_label(report.status));
    if !report.adorned_predicates.is_empty() {
        println!("  adorned_predicates:");
        for pred in &report.adorned_predicates {
            println!("    - {}", pred);
        }
    }
    if !report.generated_predicates.is_empty() {
        println!("  generated_predicates:");
        for pred in &report.generated_predicates {
            println!("    - {}", pred);
        }
    }
    if !report.declined_reasons.is_empty() {
        println!("  declined_reasons:");
        for reason in &report.declined_reasons {
            println!("    - {}", reason);
        }
    }
}

fn print_explain_json(report: &ExplainReport) {
    println!("{{");
    println!("  \"parse\": {{");
    println!(
        "    \"statements\": {},",
        report.parse_stats.statement_count
    );
    println!("    \"cache_hits\": {},", report.parse_stats.hits);
    println!("    \"cache_misses\": {}", report.parse_stats.misses);
    println!("  }},");
    println!("  \"ast\": {{");
    println!("    \"rules\": {},", report.program.rules.len());
    println!("    \"queries\": {},", report.program.queries.len());
    println!(
        "    \"prob_queries\": {}",
        report.program.prob_queries.len()
    );
    println!("  }},");
    println!("  \"stratification\": {{");
    println!(
        "    \"status\": \"{}\",",
        json_escape(&report.stratification_status)
    );
    println!("    \"strata\": {}", report.stratification_count);
    println!("  }},");
    println!("  \"rir\": {{");
    println!("    \"status\": \"{}\",", json_escape(&report.rir_status));
    println!("    \"sccs\": {}", report.rir_sccs);
    println!("  }},");
    println!("  \"optimizer\": {{");
    println!(
        "    \"status\": \"{}\",",
        json_escape(&report.optimizer_status)
    );
    println!("    \"est_memory_peak\": {}", report.optimizer_memory_peak);
    println!("  }},");
    println!("  \"wcoj\": {{");
    println!("    \"status\": \"reported\"");
    println!("  }},");
    println!("  \"epistemic\": {},", report.epistemic);
    println!("  \"magic_sets\": {{");
    println!(
        "    \"status\": \"{}\",",
        json_escape(magic_status_label(report.magic_sets.status))
    );
    println!(
        "    \"adorned_predicates\": {},",
        json_string_array(&report.magic_sets.adorned_predicates)
    );
    println!(
        "    \"generated_predicates\": {},",
        json_string_array(&report.magic_sets.generated_predicates)
    );
    println!(
        "    \"declined_reasons\": {}",
        json_string_array(&report.magic_sets.declined_reasons)
    );
    println!("  }},");
    println!("  \"probability\": {{");
    println!(
        "    \"engine\": \"{}\",",
        match report.program.prob_engine() {
            ProbEngine::ExactDdnnf => "exact_ddnnf",
            ProbEngine::Mc => "mc",
        }
    );
    println!(
        "    \"aggregate_lifting_count\": {}",
        report.aggregate_lifting.len()
    );
    println!("  }},");
    println!("  \"aggregate_lifting\": [");
    for (idx, entry) in report.aggregate_lifting.iter().enumerate() {
        let suffix = if idx + 1 == report.aggregate_lifting.len() {
            ""
        } else {
            ","
        };
        println!("    {{");
        println!(
            "      \"predicate\": \"{}\",",
            json_escape(&entry.predicate)
        );
        println!(
            "      \"group_key\": {},",
            json_value_array(&entry.group_key)
        );
        println!("      \"operator\": \"{}\",", json_escape(&entry.operator));
        println!(
            "      \"finite_domain_source\": \"{}\",",
            json_escape(&entry.finite_domain_source)
        );
        println!(
            "      \"deterministic_rows\": {},",
            entry.deterministic_rows
        );
        println!("      \"uncertain_rows\": {},", entry.uncertain_rows);
        println!("      \"domain_size\": {},", entry.domain_size);
        println!("      \"cap\": {},", entry.cap);
        println!("      \"status\": \"{}\",", entry.status.as_str());
        println!("      \"reason\": \"{}\",", json_escape(&entry.reason));
        println!("      \"naive_outcomes\": {},", entry.naive_outcomes);
        println!(
            "      \"dynamic_programming_states\": {}",
            entry.dynamic_programming_states
        );
        println!("    }}{}", suffix);
    }
    println!("  ],");
    print_rule_provenance_json(&report.rule_provenance);
    println!(",");
    print_proof_traces_json(&report.proof_traces);
    println!(",");
    print_generated_rule_diagnostics_json(&report.generated_rule_diagnostics);
    println!("}}");
}

fn print_rule_provenance_json(entries: &[RuleProvenance]) {
    println!("  \"rule_provenance\": [");
    for (idx, entry) in entries.iter().enumerate() {
        let suffix = if idx + 1 == entries.len() { "" } else { "," };
        println!("    {{");
        println!("      \"rule_id\": \"{}\",", json_escape(&entry.rule_id));
        println!("      \"head\": \"{}\",", json_escape(&entry.head));
        println!(
            "      \"source_kind\": \"{}\",",
            json_escape(entry.source_kind.as_str())
        );
        println!(
            "      \"source_span\": {},",
            json_optional_string(entry.source_span.as_deref())
        );
        println!(
            "      \"generation_trace_hash\": {},",
            json_optional_string(entry.generation_trace_hash.as_deref())
        );
        println!(
            "      \"support_relation_ids\": {},",
            json_string_array(&entry.support_relation_ids)
        );
        println!(
            "      \"counterexample_relation_ids\": {}",
            json_string_array(&entry.counterexample_relation_ids)
        );
        println!("    }}{}", suffix);
    }
    println!("  ]");
}

fn print_proof_traces_json(entries: &[QueryProofTrace]) {
    println!("  \"proof_traces\": [");
    for (idx, entry) in entries.iter().enumerate() {
        let suffix = if idx + 1 == entries.len() { "" } else { "," };
        println!("    {{");
        println!("      \"query_id\": \"{}\",", json_escape(&entry.query_id));
        println!("      \"query\": \"{}\",", json_escape(&entry.query));
        println!(
            "      \"answer_relation\": \"{}\",",
            json_escape(&entry.answer_relation)
        );
        println!(
            "      \"rule_ids\": {},",
            json_string_array(&entry.rule_ids)
        );
        println!(
            "      \"source_facts\": {},",
            json_string_array(&entry.source_facts)
        );
        println!(
            "      \"rejected_alternatives\": {}",
            json_string_array(&entry.rejected_alternatives)
        );
        println!("    }}{}", suffix);
    }
    println!("  ]");
}

fn print_generated_rule_diagnostics_json(entries: &[GeneratedRuleDiagnostic]) {
    println!("  \"generated_rule_diagnostics\": [");
    for (idx, entry) in entries.iter().enumerate() {
        let suffix = if idx + 1 == entries.len() { "" } else { "," };
        println!("    {{");
        println!(
            "      \"rule_head\": \"{}\",",
            json_escape(&entry.rule_head)
        );
        println!(
            "      \"source_relation\": \"{}\",",
            json_escape(&entry.source_relation)
        );
        println!("      \"row_decisions\": [");
        for (row_idx, row) in entry.row_decisions.iter().enumerate() {
            let row_suffix = if row_idx + 1 == entry.row_decisions.len() {
                ""
            } else {
                ","
            };
            println!("        {{");
            println!("          \"row_key\": \"{}\",", json_escape(&row.row_key));
            println!("          \"accepted\": {},", row.accepted);
            println!(
                "          \"failed_predicates\": {},",
                json_string_array(&row.failed_predicates)
            );
            println!("          \"threshold_comparisons\": [");
            for (comparison_idx, comparison) in row.threshold_comparisons.iter().enumerate() {
                let comparison_suffix = if comparison_idx + 1 == row.threshold_comparisons.len() {
                    ""
                } else {
                    ","
                };
                println!("            {{");
                println!(
                    "              \"predicate\": \"{}\",",
                    json_escape(&comparison.predicate)
                );
                println!(
                    "              \"left\": \"{}\",",
                    json_escape(&comparison.left)
                );
                println!("              \"op\": \"{}\",", json_escape(&comparison.op));
                println!(
                    "              \"right\": \"{}\",",
                    json_escape(&comparison.right)
                );
                println!(
                    "              \"left_value\": \"{}\",",
                    json_escape(&comparison.left_value)
                );
                println!(
                    "              \"right_value\": \"{}\",",
                    json_escape(&comparison.right_value)
                );
                println!("              \"passed\": {}", comparison.passed);
                println!("            }}{}", comparison_suffix);
            }
            println!("          ],");
            println!(
                "          \"aggregate_inputs\": {}",
                json_string_array(&row.aggregate_inputs)
            );
            println!("        }}{}", row_suffix);
        }
        println!("      ]");
        println!("    }}{}", suffix);
    }
    println!("  ]");
}

fn print_magic_dot(report: &MagicSetReport) {
    println!("digraph xlog_magic_sets {{");
    println!(
        "  status [label=\"status: {}\"];",
        magic_status_label(report.status)
    );
    for pred in &report.generated_predicates {
        println!("  \"{}\" [shape=box];", dot_escape(pred));
    }
    for pred in &report.adorned_predicates {
        println!("  \"{}\" [shape=ellipse];", dot_escape(pred));
    }
    println!("}}");
}

fn magic_status_label(status: MagicSetStatus) -> &'static str {
    match status {
        MagicSetStatus::Disabled => "disabled",
        MagicSetStatus::Applied => "applied",
        MagicSetStatus::Declined => "declined",
    }
}

fn json_string_array(items: &[String]) -> String {
    let values = items
        .iter()
        .map(|item| format!("\"{}\"", json_escape(item)))
        .collect::<Vec<_>>()
        .join(", ");
    format!("[{}]", values)
}

fn json_value_array(items: &[Value]) -> String {
    let values = items.iter().map(json_value).collect::<Vec<_>>().join(", ");
    format!("[{}]", values)
}

fn json_value(value: &Value) -> String {
    match value {
        Value::I64(v) => v.to_string(),
        Value::F64(bits) => {
            let v = f64::from_bits(*bits);
            if v.is_finite() {
                v.to_string()
            } else {
                format!("\"{}\"", json_escape(&v.to_string()))
            }
        }
        Value::Symbol(id) => format!("\"{}\"", json_escape(&symbol::resolve(*id))),
        Value::String(s) => format!("\"{}\"", json_escape(s)),
    }
}

fn json_optional_string(value: Option<&str>) -> String {
    match value {
        Some(value) => format!("\"{}\"", json_escape(value)),
        None => "null".to_string(),
    }
}

fn json_escape(value: &str) -> String {
    value.replace('\\', "\\\\").replace('"', "\\\"")
}

fn dot_escape(value: &str) -> String {
    value.replace('\\', "\\\\").replace('"', "\\\"")
}

fn make_provider(device: usize, memory_mb: u64) -> Result<Arc<CudaKernelProvider>> {
    let device = Arc::new(CudaDevice::new(device)?);
    let memory = Arc::new(GpuMemoryManager::new(
        device.clone(),
        MemoryBudget::with_limit(memory_mb * 1024 * 1024),
    ));
    Ok(Arc::new(CudaKernelProvider::new(device, memory)?))
}

fn parse_inputs(inputs: &[String]) -> Result<HashMap<String, PathBuf>> {
    let mut out = HashMap::new();
    for entry in inputs {
        let (name, path) = entry.split_once('=').ok_or_else(|| {
            XlogError::Execution(format!("Invalid --input '{}', expected rel=path", entry))
        })?;
        out.insert(name.to_string(), PathBuf::from(path));
    }
    Ok(out)
}

fn run_deterministic(args: RunArgs) -> Result<()> {
    let provider = make_provider(args.device, args.memory_mb)?;
    let source = std::fs::read_to_string(&args.source).map_err(|e| {
        XlogError::Execution(format!("Failed to read {}: {}", args.source.display(), e))
    })?;

    // Check if the source has any imports that need resolution
    let has_imports = source.contains("use ");

    // Load and merge modules if there are imports
    let program = if has_imports {
        let resolver = load_modules(&args.source, args.module_path.clone())
            .map_err(|e| XlogError::Execution(format!("Module resolution failed: {}", e)))?;
        LogicProgram::compile_with_resolver(&source, &resolver)?
    } else {
        LogicProgram::compile(&source)?
    };
    let mut inputs = HashMap::new();
    for (name, path) in parse_inputs(&args.input)? {
        let buf = provider.read_arrow_ipc_stream_file(&path)?;
        inputs.insert(name, buf);
    }

    let result = program.evaluate_with_options(provider.clone(), inputs, args.stats)?;

    // Dump the compiled epistemic execution plan after a successful GPU run, so
    // the JSON corresponds to a real accepted hot-path execution.
    if let Some(plan_path) = &args.epistemic_plan_json {
        match program.epistemic_plan_json() {
            Some(json) => {
                std::fs::write(plan_path, json).map_err(|e| {
                    XlogError::Execution(format!(
                        "Failed to write epistemic plan JSON {}: {}",
                        plan_path.display(),
                        e
                    ))
                })?;
                eprintln!("epistemic plan dumped to {}", plan_path.display());
            }
            None => {
                eprintln!(
                    "note: --epistemic-plan-json given but program has no epistemic literals; no plan dumped"
                );
            }
        }
    }

    // Emit query results
    emit_logic_results(
        provider.as_ref(),
        &result.queries,
        args.output,
        args.output_dir.as_deref(),
    )?;

    // Emit stats if requested
    if args.stats {
        if let Some(stats) = result.stats {
            let stats_output = match args.stats_format {
                StatsFormat::Human => stats.format_human(),
                StatsFormat::Json => stats.format_json(),
            };
            eprintln!("{}", stats_output);
        }
        // Symbol table statistics
        eprintln!(
            "Symbols: {} interned ({} bytes)",
            symbol::count(),
            symbol::memory_usage()
        );
    }

    Ok(())
}

fn run_probabilistic(args: ProbArgs) -> Result<()> {
    #[cfg(not(feature = "host-io"))]
    {
        let _ = args;
        return Err(XlogError::Execution(
            "Host output is disabled (feature \"host-io\" is OFF). Use device-resident APIs (DLPack) or rebuild with --features host-io.".to_string(),
        ));
    }

    #[cfg(feature = "host-io")]
    {
        let source = std::fs::read_to_string(&args.source).map_err(|e| {
            XlogError::Execution(format!("Failed to read {}: {}", args.source.display(), e))
        })?;
        let parsed_program = parse_program(&source)?;

        // Validate module imports if any search paths are provided
        if !args.module_path.is_empty() {
            let _ = load_modules(&args.source, args.module_path.clone())
                .map_err(|e| XlogError::Execution(format!("Module resolution failed: {}", e)))?;
        }

        let mut config = GpuConfig::default();
        config.device_ordinal = args.device;
        config.memory_bytes = args.memory_mb * 1024 * 1024;

        match resolve_prob_engine(&args, &parsed_program) {
            ProbEngineCli::ExactDdnnf => {
                let prog = ExactDdnnfProgram::compile_source_with_gpu(&source, config)?;
                let result = prog.evaluate()?;
                emit_prob_exact(result, args.output, args.output_dir.as_deref())
            }
            ProbEngineCli::Mc => {
                let prog = McProgram::compile_source_with_gpu(&source, config)?;
                let mut cfg = McEvalConfig::from_directives(&parsed_program.directives)?;
                apply_mc_cli_overrides(&args, &mut cfg)?;
                // `evaluate` runs the GPU-native device hot loop and then
                // materializes the result on the host (downloads the final
                // query/evidence counts after the loop) so the CLI can print
                // probabilities and confidence intervals. The hot loop itself is
                // zero-host; this final download is host-result materialization,
                // not a hot-loop transfer. Device-resident consumers that want to
                // keep counts on the GPU use `McProgram::evaluate_gpu_device`.
                let result = prog.evaluate(cfg)?;
                emit_prob_mc(result, args.output, args.output_dir.as_deref())
            }
        }
    }
}

#[cfg(feature = "host-io")]
fn resolve_prob_engine(args: &ProbArgs, program: &Program) -> ProbEngineCli {
    args.prob_engine
        .unwrap_or_else(|| match program.directives.prob_engine_or_default() {
            ProbEngine::ExactDdnnf => ProbEngineCli::ExactDdnnf,
            ProbEngine::Mc => ProbEngineCli::Mc,
        })
}

#[cfg(feature = "host-io")]
fn apply_mc_cli_overrides(args: &ProbArgs, cfg: &mut McEvalConfig) -> Result<()> {
    if let Some(samples) = args.samples {
        cfg.samples = samples;
    }
    if let Some(seed) = args.seed {
        cfg.seed = seed;
    }
    if let Some(confidence) = args.confidence {
        cfg.confidence = confidence;
    }
    if let Some(iterations) = args.prob_max_nonmonotone_iterations {
        cfg.max_nonmonotone_iterations = iterations;
    }
    if let Some(method) = args.prob_method {
        cfg.sampling_method = Some(match method {
            ProbMethodCli::Rejection => McSamplingMethod::Rejection,
            ProbMethodCli::EvidenceClamping => McSamplingMethod::EvidenceClamping,
        });
    }
    cfg.allow_cpu_oracle_fallback = args.allow_cpu_oracle;
    cfg.validate()
}

fn emit_logic_results(
    provider: &CudaKernelProvider,
    queries: &[xlog_gpu::logic::LogicQueryResult],
    format: OutputFormat,
    output_dir: Option<&Path>,
) -> Result<()> {
    for (i, q) in queries.iter().enumerate() {
        if q.buffer.schema().arity() == 0 && matches!(format, OutputFormat::Pretty) {
            println!(
                "{}\nrows: {}",
                q.relation_name,
                provider.device_row_count(&q.buffer)?
            );
            continue;
        }
        let batch = provider.to_arrow_record_batch(&q.buffer)?;
        match format {
            OutputFormat::Pretty => {
                let formatted = pretty_format_batches(&[batch])
                    .map_err(|e| XlogError::Execution(format!("Pretty print failed: {}", e)))?;
                println!("{}\n{}", q.relation_name, formatted);
            }
            OutputFormat::Csv => {
                let mut out = Vec::new();
                {
                    let mut writer = WriterBuilder::new().build(&mut out);
                    writer
                        .write(&batch)
                        .map_err(|e| XlogError::Execution(format!("CSV write failed: {}", e)))?;
                }
                println!("{}\n{}", q.relation_name, String::from_utf8_lossy(&out));
            }
            OutputFormat::Arrow => {
                let dir = output_dir.unwrap_or_else(|| Path::new("."));
                let path = dir.join(format!("query_{}.arrow", i));
                provider.write_arrow_ipc_stream_file(&q.buffer, &path)?;
                println!("wrote {}", path.display());
            }
        }
    }
    Ok(())
}

#[cfg(feature = "host-io")]
fn emit_prob_exact(
    result: xlog_prob::exact::ExactResult,
    format: ProbOutputFormat,
    output_dir: Option<&Path>,
) -> Result<()> {
    if matches!(format, ProbOutputFormat::Json) {
        print_prob_exact_json(result);
        return Ok(());
    }

    let mut atoms = Vec::new();
    let mut probs = Vec::new();
    let mut log_probs = Vec::new();
    for q in result.query_probs {
        atoms.push(atom_to_string(&q.atom));
        probs.push(q.prob);
        log_probs.push(q.log_prob);
    }

    let batch = arrow::record_batch::RecordBatch::try_from_iter(vec![
        (
            "atom",
            Arc::new(arrow::array::StringArray::from(atoms)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "prob",
            Arc::new(arrow::array::Float64Array::from(probs)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "log_prob",
            Arc::new(arrow::array::Float64Array::from(log_probs)) as Arc<dyn arrow::array::Array>,
        ),
    ])
    .map_err(|e| XlogError::Execution(format!("Failed to build prob batch: {}", e)))?;

    emit_batch(
        "prob",
        &batch,
        prob_output_as_batch_format(format),
        output_dir,
    )
}

#[cfg(feature = "host-io")]
fn emit_prob_mc(
    result: xlog_prob::mc::McResult,
    format: ProbOutputFormat,
    output_dir: Option<&Path>,
) -> Result<()> {
    if matches!(format, ProbOutputFormat::Json) {
        print_prob_mc_json(result);
        return Ok(());
    }

    let total_samples = result.total_samples as u64;
    let evidence_samples = result.evidence_samples as u64;
    let seed = result.seed;
    let confidence = result.confidence;
    let sampling_method = result.sampling_method.as_str().to_string();
    let mc_engine = result.engine.as_str().to_string();

    let mut atoms = Vec::new();
    let mut probs = Vec::new();
    let mut log_probs = Vec::new();
    let mut stderr = Vec::new();
    let mut ci_low = Vec::new();
    let mut ci_high = Vec::new();
    let mut total_samples_col = Vec::new();
    let mut evidence_samples_col = Vec::new();
    let mut seed_col = Vec::new();
    let mut confidence_col = Vec::new();
    let mut sampling_method_col = Vec::new();
    let mut mc_engine_col = Vec::new();
    for q in result.query_estimates {
        atoms.push(atom_to_string(&q.atom));
        probs.push(q.prob);
        log_probs.push(q.log_prob);
        stderr.push(q.stderr);
        ci_low.push(q.ci_low);
        ci_high.push(q.ci_high);
        total_samples_col.push(total_samples);
        evidence_samples_col.push(evidence_samples);
        seed_col.push(seed);
        confidence_col.push(confidence);
        sampling_method_col.push(sampling_method.clone());
        mc_engine_col.push(mc_engine.clone());
    }

    let batch = arrow::record_batch::RecordBatch::try_from_iter(vec![
        (
            "atom",
            Arc::new(arrow::array::StringArray::from(atoms)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "prob",
            Arc::new(arrow::array::Float64Array::from(probs)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "log_prob",
            Arc::new(arrow::array::Float64Array::from(log_probs)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "stderr",
            Arc::new(arrow::array::Float64Array::from(stderr)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "ci_low",
            Arc::new(arrow::array::Float64Array::from(ci_low)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "ci_high",
            Arc::new(arrow::array::Float64Array::from(ci_high)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "total_samples",
            Arc::new(arrow::array::UInt64Array::from(total_samples_col))
                as Arc<dyn arrow::array::Array>,
        ),
        (
            "evidence_samples",
            Arc::new(arrow::array::UInt64Array::from(evidence_samples_col))
                as Arc<dyn arrow::array::Array>,
        ),
        (
            "seed",
            Arc::new(arrow::array::UInt64Array::from(seed_col)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "confidence",
            Arc::new(arrow::array::Float64Array::from(confidence_col))
                as Arc<dyn arrow::array::Array>,
        ),
        (
            "sampling_method",
            Arc::new(arrow::array::StringArray::from(sampling_method_col))
                as Arc<dyn arrow::array::Array>,
        ),
        (
            "mc_engine",
            Arc::new(arrow::array::StringArray::from(mc_engine_col))
                as Arc<dyn arrow::array::Array>,
        ),
    ])
    .map_err(|e| XlogError::Execution(format!("Failed to build mc batch: {}", e)))?;

    emit_batch(
        "prob",
        &batch,
        prob_output_as_batch_format(format),
        output_dir,
    )
}

#[cfg(feature = "host-io")]
fn prob_output_as_batch_format(format: ProbOutputFormat) -> OutputFormat {
    match format {
        ProbOutputFormat::Pretty => OutputFormat::Pretty,
        ProbOutputFormat::Csv => OutputFormat::Csv,
        ProbOutputFormat::Arrow => OutputFormat::Arrow,
        ProbOutputFormat::Json => unreachable!("json output is handled before batch emission"),
    }
}

#[cfg(feature = "host-io")]
fn print_prob_exact_json(result: xlog_prob::exact::ExactResult) {
    println!("{{");
    println!("  \"engine\": \"exact_ddnnf\",");
    println!("  \"queries\": [");
    let len = result.query_probs.len();
    for (idx, q) in result.query_probs.into_iter().enumerate() {
        let suffix = if idx + 1 == len { "" } else { "," };
        println!("    {{");
        println!(
            "      \"atom\": \"{}\",",
            json_escape(&atom_to_string(&q.atom))
        );
        println!("      \"prob\": {},", q.prob);
        println!("      \"log_prob\": {}", q.log_prob);
        println!("    }}{}", suffix);
    }
    println!("  ]");
    println!("}}");
}

#[cfg(feature = "host-io")]
fn print_prob_mc_json(result: xlog_prob::mc::McResult) {
    let total_samples = result.total_samples;
    let evidence_samples = result.evidence_samples;
    let seed = result.seed;
    let confidence = result.confidence;
    let sampling_method = result.sampling_method.as_str();
    let mc_engine = result.engine.as_str();
    println!("{{");
    println!("  \"engine\": \"mc\",");
    println!("  \"mc_engine\": \"{}\",", mc_engine);
    println!("  \"total_samples\": {},", total_samples);
    println!("  \"evidence_samples\": {},", evidence_samples);
    println!("  \"seed\": {},", seed);
    println!("  \"confidence\": {},", confidence);
    println!("  \"sampling_method\": \"{}\",", sampling_method);
    println!("  \"queries\": [");
    let len = result.query_estimates.len();
    for (idx, q) in result.query_estimates.into_iter().enumerate() {
        let suffix = if idx + 1 == len { "" } else { "," };
        println!("    {{");
        println!(
            "      \"atom\": \"{}\",",
            json_escape(&atom_to_string(&q.atom))
        );
        println!("      \"prob\": {},", q.prob);
        println!("      \"log_prob\": {},", q.log_prob);
        println!("      \"stderr\": {},", q.stderr);
        println!("      \"ci_low\": {},", q.ci_low);
        println!("      \"ci_high\": {},", q.ci_high);
        println!("      \"total_samples\": {},", total_samples);
        println!("      \"evidence_samples\": {}", evidence_samples);
        println!("    }}{}", suffix);
    }
    println!("  ]");
    println!("}}");
}

#[cfg(feature = "host-io")]
fn emit_batch(
    name: &str,
    batch: &arrow::record_batch::RecordBatch,
    format: OutputFormat,
    output_dir: Option<&Path>,
) -> Result<()> {
    match format {
        OutputFormat::Pretty => {
            let formatted = pretty_format_batches(std::slice::from_ref(batch))
                .map_err(|e| XlogError::Execution(format!("Pretty print failed: {}", e)))?;
            println!("{}\n{}", name, formatted);
        }
        OutputFormat::Csv => {
            let mut out = Vec::new();
            {
                let mut writer = WriterBuilder::new().build(&mut out);
                writer
                    .write(batch)
                    .map_err(|e| XlogError::Execution(format!("CSV write failed: {}", e)))?;
            }
            println!("{}\n{}", name, String::from_utf8_lossy(&out));
        }
        OutputFormat::Arrow => {
            let dir = output_dir.unwrap_or_else(|| Path::new("."));
            let path = dir.join(format!("{}_prob.arrow", name));
            let mut out = Vec::new();
            let mut writer =
                arrow::ipc::writer::StreamWriter::try_new(&mut out, &batch.schema())
                    .map_err(|e| XlogError::Execution(format!("Arrow writer failed: {}", e)))?;
            writer
                .write(batch)
                .map_err(|e| XlogError::Execution(format!("Arrow write failed: {}", e)))?;
            writer
                .finish()
                .map_err(|e| XlogError::Execution(format!("Arrow finish failed: {}", e)))?;
            std::fs::write(&path, out)
                .map_err(|e| XlogError::Execution(format!("Arrow write file failed: {}", e)))?;
            println!("wrote {}", path.display());
        }
    }
    Ok(())
}

#[cfg(feature = "host-io")]
fn atom_to_string(atom: &xlog_prob::provenance::GroundAtom) -> String {
    use xlog_prob::provenance::Value;

    if atom.args.is_empty() {
        return format!("{}()", atom.predicate);
    }

    let mut out = String::new();
    out.push_str(&atom.predicate);
    out.push('(');
    for (i, arg) in atom.args.iter().enumerate() {
        if i != 0 {
            out.push_str(", ");
        }
        match arg {
            Value::I64(v) => out.push_str(&v.to_string()),
            Value::F64(bits) => out.push_str(&f64::from_bits(*bits).to_string()),
            Value::Symbol(sym) => out.push_str(&symbol::resolve(*sym)),
            Value::String(v) => out.push_str(v),
        }
    }
    out.push(')');
    out
}