pyxlog/

lib.rs

//! Python bindings for XLOG via PyO3.
#![allow(missing_docs)] // PyO3 #[pyclass] / #[pymethods] generate pub items without docs
#![allow(
    clippy::large_enum_variant,
    clippy::needless_range_loop,
    clippy::too_many_arguments,
    clippy::type_complexity
)]

use std::collections::{HashMap, HashSet};
use std::os::raw::{c_char, c_void};
use std::sync::Arc;

use pyo3::exceptions::{PyBufferError, PyMemoryError, PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};

use xlog_core::{MemoryBudget, Schema};
use xlog_cuda::{
    device_runtime::{
        AsyncCudaResource, DeviceMemoryResource, GlobalDeviceBudget, StreamPool, XlogDeviceRuntime,
    },
    CudaBuffer, CudaDevice, CudaKernelProvider, DlpackManagedTensor, GpuMemoryManager,
};
#[cfg(feature = "arrow-device-import")]
use xlog_cuda::{ArrowDeviceArray, ArrowDeviceArrayOwned};
use xlog_gpu::logic as gpu_logic;
use xlog_logic::ast::ProbEngine;
use xlog_neural::{NetworkRegistry, TensorSourceRegistry};
use xlog_prob::exact::GpuConfig;

use xlog_core::RelId;
use xlog_ir::ExecutionPlan;
use xlog_logic::ast::Program as AstProgram;
use xlog_runtime::{Executor, RelationStore};

mod neural_registry;
use neural_registry::NeuralPredicateRegistry;
mod diagnostic_serialization;
mod dlpack;
mod epistemic;
mod ilp;
mod ilp_exact;
mod ilp_gpu;
mod joint_carrier;
mod logic;
mod neural;
mod program;
mod relation_metadata;
mod training;
mod types;
pub(crate) use diagnostic_serialization::{pack_query_proof_traces, pack_rule_provenance};
pub(crate) use program::{
    CachedCircuit, CompiledProbProgram, HardFilter, InputSource, JoinPlan, NeuralGroup,
    QuerySignature,
};
use relation_metadata::RelationMetadataStore;

const DLPACK_CAPSULE_NAME: &[u8] = b"dltensor\0";
const USED_DLPACK_CAPSULE_NAME: &[u8] = b"used_dltensor\0";
// Every pyxlog DLPack import is consumed on CudaDevice's legacy default
// stream. The Python Array API reserves integer 1 for that CUDA stream.
const DLPACK_CUDA_LEGACY_DEFAULT_STREAM: i64 = 1;

#[cfg(feature = "arrow-device-import")]
const ARROW_DEVICE_ARRAY_CAPSULE_NAME: &[u8] = b"arrow_device_array\0";
#[cfg(feature = "arrow-device-import")]
const USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME: &[u8] = b"used_arrow_device_array\0";

unsafe extern "C" fn dlpack_capsule_destructor(capsule: *mut pyo3::ffi::PyObject) {
    if capsule.is_null() {
        return;
    }

    let valid =
        pyo3::ffi::PyCapsule_IsValid(capsule, DLPACK_CAPSULE_NAME.as_ptr() as *const c_char);
    if valid == 0 {
        return;
    }

    let ptr =
        pyo3::ffi::PyCapsule_GetPointer(capsule, DLPACK_CAPSULE_NAME.as_ptr() as *const c_char);
    if ptr.is_null() {
        pyo3::ffi::PyErr_Clear();
        return;
    }

    let managed = ptr as *mut xlog_cuda::DLManagedTensor;
    drop(DlpackManagedTensor::from_raw(managed));
}

pub(crate) fn dlpack_capsule_from_tensor(
    py: Python<'_>,
    tensor: DlpackManagedTensor,
) -> PyResult<Py<PyAny>> {
    let raw = tensor.into_raw();
    let ptr = raw as *mut c_void;
    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let capsule = unsafe {
        pyo3::ffi::PyCapsule_New(
            ptr,
            DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
            Some(dlpack_capsule_destructor),
        )
    };
    if capsule.is_null() {
        // SAFETY: the pointer is a valid owned Python object pointer returned by the C API
        unsafe {
            drop(DlpackManagedTensor::from_raw(raw));
        }
        return Err(PyRuntimeError::new_err("Failed to create DLPack capsule"));
    }
    // SAFETY: capsule is a non-null owned pointer returned by PyCapsule_New; PyO3 takes ownership
    let obj: Py<PyAny> = unsafe { Bound::from_owned_ptr(py, capsule) }.unbind();
    Ok(obj)
}

#[cfg(feature = "arrow-device-import")]
unsafe extern "C" fn arrow_device_array_capsule_destructor(capsule: *mut pyo3::ffi::PyObject) {
    if capsule.is_null() {
        return;
    }

    let valid = pyo3::ffi::PyCapsule_IsValid(
        capsule,
        ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
    );
    if valid == 0 {
        return;
    }

    let ptr = pyo3::ffi::PyCapsule_GetPointer(
        capsule,
        ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
    );
    if ptr.is_null() {
        pyo3::ffi::PyErr_Clear();
        return;
    }

    drop(ArrowDeviceArrayOwned::from_raw(
        ptr as *mut ArrowDeviceArray,
    ));
}

#[cfg(feature = "arrow-device-import")]
pub(crate) fn arrow_device_capsule_from_device_array(
    py: Python<'_>,
    device_array: ArrowDeviceArrayOwned,
) -> PyResult<Py<PyAny>> {
    let raw = device_array.into_raw();
    let ptr = raw as *mut c_void;
    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let capsule = unsafe {
        pyo3::ffi::PyCapsule_New(
            ptr,
            ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
            Some(arrow_device_array_capsule_destructor),
        )
    };
    if capsule.is_null() {
        // SAFETY: the pointer is a valid owned Python object pointer returned by the C API
        unsafe {
            drop(ArrowDeviceArrayOwned::from_raw(raw));
        }
        return Err(PyRuntimeError::new_err(
            "Failed to create Arrow device array capsule",
        ));
    }
    // SAFETY: capsule is a non-null owned pointer returned by PyCapsule_New; PyO3 takes ownership
    let obj: Py<PyAny> = unsafe { Bound::from_owned_ptr(py, capsule) }.unbind();
    Ok(obj)
}

#[cfg(feature = "arrow-device-import")]
pub(crate) fn arrow_device_from_py(obj: &Bound<'_, PyAny>) -> PyResult<ArrowDeviceArrayOwned> {
    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    if unsafe {
        pyo3::ffi::PyCapsule_IsValid(
            obj.as_ptr(),
            ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    } == 0
    {
        return Err(PyValueError::new_err(
            "Expected an Arrow device array capsule (arrow_device_array)",
        ));
    }

    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let ptr = unsafe {
        pyo3::ffi::PyCapsule_GetPointer(
            obj.as_ptr(),
            ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    };
    if ptr.is_null() {
        return Err(PyRuntimeError::new_err(
            "Failed to get Arrow device array pointer",
        ));
    }

    // Mark consumed so the capsule destructor doesn't free the pointer we now own.
    // SAFETY: capsule is valid (checked above); renaming marks it consumed so the destructor skips cleanup
    let rc = unsafe {
        pyo3::ffi::PyCapsule_SetName(
            obj.as_ptr(),
            USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    };
    if rc != 0 {
        return Err(PyRuntimeError::new_err(
            "Failed to mark Arrow device array capsule as consumed",
        ));
    }

    // SAFETY: ptr is non-null (checked above) and points to an ArrowDeviceArray matching the Arrow C Data Interface layout
    Ok(unsafe { ArrowDeviceArrayOwned::from_raw(ptr as *mut ArrowDeviceArray) })
}

pub(crate) fn provider_from_config(config: GpuConfig) -> xlog_core::Result<CudaKernelProvider> {
    let device = Arc::new(CudaDevice::new(config.device_ordinal)?);
    let stream_pool = Arc::new(StreamPool::with_defaults(Arc::clone(&device)));
    let async_resource: Box<dyn DeviceMemoryResource + Send + Sync> =
        Box::new(AsyncCudaResource::new(
            Arc::clone(&device),
            config.device_ordinal as u32,
            Arc::clone(&stream_pool),
        ));
    let budget_limit = usize::try_from(config.memory_bytes).unwrap_or(usize::MAX);
    let budgeted: Box<dyn DeviceMemoryResource + Send + Sync> =
        Box::new(GlobalDeviceBudget::new(async_resource, budget_limit));
    let runtime = Arc::new(XlogDeviceRuntime::with_resource(
        Arc::clone(&device),
        config.device_ordinal as u32,
        stream_pool,
        budgeted,
    ));
    let memory = Arc::new(GpuMemoryManager::with_runtime(
        device.clone(),
        MemoryBudget::with_limit(config.memory_bytes),
        runtime,
    ));
    CudaKernelProvider::with_runtime(device, memory)
}

pub(crate) fn enforce_call_memory_limit(
    provider: &Arc<CudaKernelProvider>,
    memory_mb: Option<u64>,
) -> PyResult<()> {
    let Some(memory_mb) = memory_mb else {
        return Ok(());
    };
    if memory_mb == 0 {
        return Err(PyValueError::new_err("memory_mb must be > 0"));
    }
    let memory_limit_bytes = memory_mb.saturating_mul(1024 * 1024);
    let allocated_bytes = provider.memory().allocated_bytes();
    if allocated_bytes > memory_limit_bytes {
        return Err(PyMemoryError::new_err(format!(
            "per-call memory limit exceeded before evaluation: allocated_bytes={} memory_limit_bytes={}",
            allocated_bytes, memory_limit_bytes
        )));
    }
    Ok(())
}

pub(crate) fn provider_memory_stats(
    py: Python<'_>,
    provider: &Arc<CudaKernelProvider>,
) -> PyResult<Py<PyAny>> {
    let dict = PyDict::new(py);
    let memory = provider.memory();
    dict.set_item("allocated_bytes", memory.allocated_bytes())?;
    dict.set_item("memory_limit_bytes", memory.budget().device_bytes)?;
    dict.set_item("peak_memory_bytes", memory.peak_bytes())?;
    dict.set_item("status", "available")?;
    Ok(dict.into())
}

pub(crate) fn parse_prob_engine_override(s: &str) -> PyResult<ProbEngine> {
    let v = s.trim().to_ascii_lowercase();
    match v.as_str() {
        "exact_ddnnf" | "exact" | "ddnnf" => Ok(ProbEngine::ExactDdnnf),
        "mc" => Ok(ProbEngine::Mc),
        other => Err(PyValueError::new_err(format!(
            "Unknown prob_engine '{}'; expected 'exact_ddnnf' or 'mc'",
            other
        ))),
    }
}

pub(crate) fn dlpack_from_py(obj: &Bound<'_, PyAny>) -> PyResult<DlpackManagedTensor> {
    let py = obj.py();

    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let capsule_obj: Bound<'_, PyAny> = if unsafe {
        pyo3::ffi::PyCapsule_IsValid(obj.as_ptr(), DLPACK_CAPSULE_NAME.as_ptr() as *const c_char)
    } != 0
    {
        obj.clone()
    } else if obj.hasattr("__dlpack__")? {
        let (device_type, device_id): (i32, i32) =
            obj.call_method0("__dlpack_device__")?.extract()?;
        if device_type != xlog_cuda::dlpack::K_DLCUDA {
            return Err(PyBufferError::new_err(format!(
                "Unsupported DLPack producer device type {device_type} (device {device_id}); \
                 XLOG requires CUDA device memory (kDLCUDA=2)"
            )));
        }
        // Passing the consumer stream makes the producer order any pending
        // non-default-stream writes before XLOG reads the tensor. A raw
        // capsule cannot negotiate synchronization and must already be ready
        // for the legacy default stream when supplied by the caller.
        let kwargs = PyDict::new(py);
        kwargs.set_item("stream", DLPACK_CUDA_LEGACY_DEFAULT_STREAM)?;
        obj.call_method("__dlpack__", (), Some(&kwargs))?
    } else {
        return Err(PyValueError::new_err(
            "Expected a DLPack capsule or an object with __dlpack__",
        ));
    };

    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    if unsafe {
        pyo3::ffi::PyCapsule_IsValid(
            capsule_obj.as_ptr(),
            DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    } == 0
    {
        return Err(PyValueError::new_err("Invalid DLPack capsule"));
    }

    // SAFETY: capsule validity was checked immediately before this call; pointer lifetime is managed by the capsule
    let ptr = unsafe {
        pyo3::ffi::PyCapsule_GetPointer(
            capsule_obj.as_ptr(),
            DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    };
    if ptr.is_null() {
        return Err(PyRuntimeError::new_err("Failed to get DLPack pointer"));
    }

    // SAFETY: capsule is valid (checked above); renaming marks it consumed so the destructor skips cleanup
    let rc = unsafe {
        pyo3::ffi::PyCapsule_SetName(
            capsule_obj.as_ptr(),
            USED_DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
        )
    };
    if rc != 0 {
        return Err(PyRuntimeError::new_err(
            "Failed to mark DLPack capsule as consumed",
        ));
    }

    // SAFETY: ptr is non-null (checked above) and points to a DLManagedTensor matching the DLPack specification layout
    Ok(unsafe { DlpackManagedTensor::from_raw(ptr as *mut xlog_cuda::DLManagedTensor) })
}

#[pyfunction]
fn dlpack_is_cuda(obj: &Bound<'_, PyAny>) -> PyResult<bool> {
    // SAFETY: capsule validity is checked before reading the DLPack header. This
    // does not consume the capsule; ownership remains with its destructor.
    if unsafe {
        pyo3::ffi::PyCapsule_IsValid(obj.as_ptr(), DLPACK_CAPSULE_NAME.as_ptr() as *const c_char)
    } == 0
    {
        return Err(PyValueError::new_err(
            "Expected a DLPack capsule (dltensor)",
        ));
    }

    // SAFETY: capsule validity was checked immediately before this call.
    let ptr = unsafe {
        pyo3::ffi::PyCapsule_GetPointer(obj.as_ptr(), DLPACK_CAPSULE_NAME.as_ptr() as *const c_char)
    };
    if ptr.is_null() {
        return Err(PyRuntimeError::new_err("Failed to get DLPack pointer"));
    }

    // SAFETY: ptr is non-null and points to a DLManagedTensor owned by the capsule.
    let managed = unsafe { &*(ptr as *const xlog_cuda::DLManagedTensor) };
    Ok(managed.dl_tensor.device.device_type == xlog_cuda::dlpack::K_DLCUDA)
}

#[pyclass(name = "DifferentiableProofTraceMap")]
pub struct PyDifferentiableProofTraceMap {
    inner: xlog_logic::DifferentiableProofTraceMap,
}

fn pack_differentiable_proof_trace(
    py: Python<'_>,
    trace: &xlog_logic::ProofTrace,
) -> PyResult<Py<PyAny>> {
    let dict = PyDict::new(py);
    dict.set_item("proof_id", trace.proof_id)?;
    dict.set_item("answer_key", &trace.answer_key)?;
    dict.set_item("clause_id", &trace.clause_id)?;
    dict.set_item("support_atoms", &trace.support_atoms)?;
    dict.set_item("weight", trace.weight)?;
    dict.set_item("gradient", trace.gradient)?;
    Ok(dict.into())
}

#[pymethods]
impl PyDifferentiableProofTraceMap {
    #[new]
    fn new() -> Self {
        Self {
            inner: xlog_logic::DifferentiableProofTraceMap::new(),
        }
    }

    fn insert(
        &mut self,
        answer_key: String,
        clause_id: String,
        support_atoms: Vec<String>,
        initial_weight: f64,
    ) -> PyResult<u64> {
        if !initial_weight.is_finite() {
            return Err(PyValueError::new_err(
                "initial_weight must be a finite float",
            ));
        }
        Ok(self.inner.insert(xlog_logic::ProofTraceSpec {
            answer_key,
            clause_id,
            support_atoms,
            initial_weight,
        }))
    }

    fn trace(&self, py: Python<'_>, proof_id: u64) -> PyResult<Option<Py<PyAny>>> {
        self.inner
            .trace(proof_id)
            .map(|trace| pack_differentiable_proof_trace(py, trace))
            .transpose()
    }

    fn traces(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let list = PyList::empty(py);
        for trace in self.inner.traces() {
            list.append(pack_differentiable_proof_trace(py, trace)?)?;
        }
        Ok(list.into())
    }

    fn accumulate_binary_logistic_gradients(
        &mut self,
        targets: Vec<(String, f64)>,
    ) -> PyResult<f64> {
        if targets.iter().any(|(_, target)| !target.is_finite()) {
            return Err(PyValueError::new_err("targets must be finite floats"));
        }
        Ok(self.inner.accumulate_binary_logistic_gradients(&targets))
    }

    fn apply_gradients(&mut self, learning_rate: f64) -> PyResult<()> {
        if !learning_rate.is_finite() || learning_rate < 0.0 {
            return Err(PyValueError::new_err(
                "learning_rate must be a finite non-negative float",
            ));
        }
        self.inner.apply_gradients(learning_rate);
        Ok(())
    }
}

#[pyclass]
pub struct Program;

#[pyclass]
pub struct CompiledProgram {
    pub(crate) program: CompiledProbProgram,
    pub(crate) output_provider: Arc<CudaKernelProvider>,
    /// Registry for neural networks
    pub(crate) network_registry: NetworkRegistry,
    /// Registry for neural predicate metadata (predicate -> network/labels)
    pub(crate) neural_registry: NeuralPredicateRegistry,
    /// Names of neural networks declared in the program (from nn() declarations)
    pub(crate) declared_networks: HashSet<String>,
    /// Map from network name to form: true = embedding, false = classification
    pub(crate) declared_network_forms: HashMap<String, bool>,
    /// Registry for tensor data sources (images, embeddings, etc.)
    pub(crate) tensor_sources: TensorSourceRegistry,
    /// Name of the Stage-B existential-join domain tensor source, as supplied by
    /// the Python driver (single source of truth). `None` until a domain source is
    /// registered; read by the join forward to resolve `DomainRow`/`ConstDummy`
    /// groups instead of an engine-side hardcoded name.
    pub(crate) domain_source: Option<String>,
    /// Which ROW of the join-domain tensor holds which domain CONSTANT, as stated by
    /// the Python driver: `domain_ids[j]` is the constant whose feature vector is row
    /// `j`. This is the SINGLE source of truth for "constant -> feature row" — the
    /// torch-side mixture and this circuit look the row up in the same list, so they
    /// cannot drift apart.
    ///
    /// The ids arrive WITH the tensor (`register_domain_tensor_source` requires them), so
    /// there is no "no ids registered" state to fall back from. The fallback that used to
    /// live here — row = the constant's position in the relation's materialized domain —
    /// is the defect this map exists to close: it reads a different row from the torch
    /// path for any domain that is not exactly `0..D-1`, silently, and no test that
    /// stays inside that coincidence can see it. Keeping it as a reachable branch would
    /// leave the same silent-wrong mode one call away.
    ///
    /// Empty until a domain source is registered — and that state IS reachable from a
    /// caller mistake, not only from an internal bug: a join signature is compiled for
    /// every rule defining the train head (`joint_candidate_eligibility`), so a program
    /// whose driver forgot `domain_inputs` reaches this map while it is still empty. The
    /// lookup then fails with "constant N is not in domain_ids", which is true but points
    /// at the wrong parameter, so the Python driver checks `domain_inputs` for every join
    /// candidate BEFORE it asks for eligibility. A lookup miss here is therefore a real
    /// error to surface (a joined constant with no id), never something to paper over.
    pub(crate) domain_ids: Vec<i64>,
    /// Original program source (for dynamic query compilation)
    pub(crate) _source: String,
    /// Parsed program AST (for signature analysis)
    pub(crate) ast: xlog_logic::ast::Program,
    /// GPU configuration
    pub(crate) _gpu_config: GpuConfig,
    /// Probabilistic inference engine
    pub(crate) _prob_engine: ProbEngine,
    /// Cache of analyzed query signatures.
    pub(crate) query_signature_cache: HashMap<String, QuerySignature>,
    /// Cache of compiled circuits by template signature
    pub(crate) circuit_cache: HashMap<String, CachedCircuit>,
    /// Number of circuit-template cache hits observed by neural training paths.
    pub(crate) circuit_cache_hits: usize,
    /// Number of circuit-template cache misses observed by neural training paths.
    pub(crate) circuit_cache_misses: usize,
    /// Number of times the template compilation path executed.
    pub(crate) template_compile_count: usize,
    /// When true, batch queries sharing the same circuit template in training.
    pub(crate) batch_queries: bool,
    /// Latest circuit compilation profile (populated on cache miss when profiling).
    pub(crate) last_compile_profile: Option<xlog_prob::compilation::CircuitCompileProfile>,
}

#[pyclass]
pub struct LogicProgram;

#[pyclass]
pub struct CompiledLogicProgram {
    pub(crate) program: Arc<gpu_logic::LogicProgram>,
    pub(crate) provider: Arc<CudaKernelProvider>,
}

/// A fixed accepted-evidence exact circuit with mutable independent fact priors.
///
/// Clones share one serialized native state. Updating or evaluating one clone
/// excludes concurrent work on the same circuit and is visible to every clone.
#[pyclass(skip_from_py_object)]
#[derive(Clone)]
pub struct CompiledConditionedProgram {
    #[cfg(feature = "host-io")]
    pub(crate) program: xlog_prob::epistemic_production::PreparedConditionedProgram,
    #[cfg(feature = "host-io")]
    pub(crate) result_provider: Arc<CudaKernelProvider>,
}

#[pyclass]
pub struct LogicRelationSession {
    pub(crate) program: Arc<gpu_logic::LogicProgram>,
    pub(crate) provider: Arc<CudaKernelProvider>,
    pub(crate) relation_store: RelationStore,
    pub(crate) evaluation_store: Option<gpu_logic::LogicMaterializedStore>,
    pub(crate) session_runtime: Option<gpu_logic::LogicSessionRuntime>,
    pub(crate) last_delta_stats: Option<LogicDeltaStats>,
    pub(crate) relation_callbacks: Vec<RelationChangeCallback>,
    pub(crate) next_relation_callback_id: u64,
    pub(crate) relation_generations: HashMap<String, u64>,
    pub(crate) relation_metadata: RelationMetadataStore,
}

pub(crate) struct RelationChangeCallback {
    pub id: u64,
    pub callback: Py<PyAny>,
}

#[derive(Clone, Debug)]
pub(crate) struct LogicDeltaStats {
    pub input_delta_count: usize,
    pub changed_relations: usize,
    pub changed_relation_names: Vec<String>,
    pub insert_rows: u64,
    pub delete_rows: u64,
    pub has_deletes: bool,
    pub affected_sccs: usize,
    pub recomputed_sccs: usize,
    pub incremental_sccs: usize,
    pub coalesced_insert_rows: u64,
    pub coalesced_delete_rows: u64,
    pub canceled_rows: u64,
    pub equivalent_to_full_recompute: Option<bool>,
    pub planner_telemetry: gpu_logic::DeltaPlannerTelemetry,
    pub debug_trace: Vec<String>,
}

#[pyclass]
pub struct LogicQueryResult {
    #[pyo3(get)]
    pub relation_name: String,
    #[pyo3(get)]
    pub columns: Vec<String>,
    #[pyo3(get)]
    pub sort_labels: Vec<String>,
    #[pyo3(get)]
    pub tensors: Vec<Py<PyAny>>,
    #[pyo3(get)]
    pub num_rows: usize,
    #[pyo3(get)]
    pub is_true: bool,
}

#[pyclass]
pub struct LogicEvalResult {
    #[pyo3(get)]
    pub queries: Vec<Py<LogicQueryResult>>,
}

#[pyclass]
pub struct IlpTaggedCreditDeviceResult {
    #[pyo3(get)]
    pub fact_row_offsets: Py<PyAny>,
    #[pyo3(get)]
    pub entry_indices: Py<PyAny>,
    #[pyo3(get)]
    pub entry_i: Py<PyAny>,
    #[pyo3(get)]
    pub entry_j: Py<PyAny>,
    #[pyo3(get)]
    pub entry_k: Py<PyAny>,
}

#[pyclass]
pub struct McDeviceEvalResult {
    /// Per-query satisfying-sample counts. DLPack int32 tensor on CUDA.
    #[pyo3(get)]
    pub query_counts: Py<PyAny>,
    /// Evidence satisfying-sample count. DLPack int32 tensor with shape [1] on CUDA.
    #[pyo3(get)]
    pub evidence_count: Py<PyAny>,
    #[pyo3(get)]
    pub total_samples: usize,
    #[pyo3(get)]
    pub seed: u64,
    #[pyo3(get)]
    pub confidence: f64,
    #[pyo3(get)]
    pub nonmonotone_semantics: String,
    #[pyo3(get)]
    pub nonmonotone_sccs: usize,
    #[pyo3(get)]
    pub nonmonotone_cycles: usize,
    #[pyo3(get)]
    pub nonmonotone_iteration_limit_hits: usize,
    #[pyo3(get)]
    pub sampling_method: String,
    #[pyo3(get)]
    pub resident_no_host_certified: bool,
    #[pyo3(get)]
    pub resident_no_host_policy_result: String,
    #[pyo3(get)]
    pub resident_no_host_tracked_dtoh_calls: u64,
    #[pyo3(get)]
    pub resident_no_host_tracked_htod_calls: u64,
    #[pyo3(get)]
    pub resident_no_host_host_loop_iterations: u64,
    #[pyo3(get)]
    pub resident_no_host_per_sample_host_launches: u64,
    #[pyo3(get)]
    pub resident_no_host_untracked_metadata_reads: u64,
    #[pyo3(get)]
    pub resident_no_host_engine_launches: u64,
    #[pyo3(get)]
    pub resident_no_host_host_fixpoint_iterations: u64,
    #[pyo3(get)]
    pub resident_no_host_per_operator_host_allocations: u64,
}

#[pyclass]
pub struct EvalResult {
    #[pyo3(get)]
    pub atoms: Vec<String>,
    #[pyo3(get)]
    pub prob: Py<PyAny>,
    #[pyo3(get)]
    pub log_prob: Py<PyAny>,
    #[pyo3(get)]
    pub num_vars: usize,
    /// Exact log-evidence `log Z_E` (natural log). `None` for Monte Carlo results.
    #[pyo3(get)]
    pub log_z_e: Option<f64>,
    #[pyo3(get)]
    pub grad_true: Option<Vec<Py<PyAny>>>,
    #[pyo3(get)]
    pub grad_false: Option<Vec<Py<PyAny>>>,
    #[pyo3(get)]
    pub approx: bool,
    #[pyo3(get)]
    pub stderr: Option<Py<PyAny>>,
    #[pyo3(get)]
    pub ci_low: Option<Py<PyAny>>,
    #[pyo3(get)]
    pub ci_high: Option<Py<PyAny>>,
    #[pyo3(get)]
    pub samples: Option<usize>,
    #[pyo3(get)]
    pub evidence_samples: Option<usize>,
    #[pyo3(get)]
    pub seed: Option<u64>,
    #[pyo3(get)]
    pub confidence: Option<f64>,
    #[pyo3(get)]
    pub nonmonotone_semantics: Option<String>,
    #[pyo3(get)]
    pub nonmonotone_sccs: Option<usize>,
    #[pyo3(get)]
    pub nonmonotone_cycles: Option<usize>,
    #[pyo3(get)]
    pub nonmonotone_iteration_limit_hits: Option<usize>,
    #[pyo3(get)]
    pub sampling_method: Option<String>,
    /// MC only: which engine produced the result — `"gpu-resident"` for the
    /// resident megakernel engine, `"cpu-oracle"` for the explicitly opted-in
    /// CPU oracle. `None` for exact inference.
    #[pyo3(get)]
    pub mc_engine: Option<String>,
}

/// Exact probabilities conditioned on an accepted epistemic world view.
///
/// `prob`/`log_prob` are DLPack capsules over device memory, like `EvalResult`.
/// `trace` carries the production-path counters of the epistemic->probability
/// adapter: they are the evidence that conditioning actually happened on the GPU.
///
/// `log_z_e` is log P(evidence): the exact log-probability of the conditioned
/// evidence under the probabilistic program's distribution, computed by weighted
/// model counting over the compiled circuit. Query probabilities are
/// `exp(log_z_eq - log_z_e)`. When the conditioned atoms are independent root
/// facts it coincides with the log of the product of their priors, but that is a
/// special case, not the definition: evidence on a derived atom, on atoms sharing
/// an ancestor, or negated evidence all diverge from the product form.
#[pyclass]
pub struct EpistemicEvalResult {
    #[pyo3(get)]
    pub atoms: Vec<String>,
    #[pyo3(get)]
    pub prob: Py<PyAny>,
    #[pyo3(get)]
    pub log_prob: Py<PyAny>,
    #[pyo3(get)]
    pub log_z_e: f64,
    #[pyo3(get)]
    pub trace: Py<PyAny>,
}

/// Summary of one accepted epistemic GPU execution.
///
/// `accepted_world_views == 0` means the program ran but nothing was accepted.
/// `evaluate_conditioned` on that same program RAISES `RuntimeError` rather than
/// returning an unconditioned result — this method is the non-raising way to detect
/// the state. `know_operator_count`/`possible_operator_count` are plan-level
/// censuses and stay non-zero even when nothing is accepted.
#[pyclass]
pub struct EpistemicEvidence {
    #[pyo3(get)]
    pub epistemic_mode: String,
    #[pyo3(get)]
    pub know_operator_count: usize,
    #[pyo3(get)]
    pub possible_operator_count: usize,
    #[pyo3(get)]
    pub accepted_candidates: usize,
    #[pyo3(get)]
    pub rejected_candidates: usize,
    #[pyo3(get)]
    pub accepted_world_views: usize,
    #[pyo3(get)]
    pub final_output_rows: usize,
}

// =========================================================================
// Training Infrastructure
// =========================================================================

/// Statistics for a single training epoch.
// Result-only class: nothing extracts it back from Python, so the
// Clone-derived automatic `FromPyObject` is explicitly skipped.
#[pyclass(skip_from_py_object)]
#[derive(Clone)]
pub struct EpochStats {
    /// Average loss across all batches in the epoch
    #[pyo3(get)]
    pub avg_loss: f64,
    /// Number of batches processed
    #[pyo3(get)]
    pub num_batches: usize,
    /// Total number of queries processed
    #[pyo3(get)]
    pub total_queries: usize,
}

/// Training history tracking loss over epochs and batches.
// Result-only class: nothing extracts it back from Python, so the
// Clone-derived automatic `FromPyObject` is explicitly skipped.
#[pyclass(skip_from_py_object)]
#[derive(Clone)]
pub struct TrainingHistory {
    /// Loss at the end of each epoch
    #[pyo3(get)]
    pub epoch_losses: Vec<f64>,
    /// Wall-clock time (seconds) for each epoch
    #[pyo3(get)]
    pub epoch_times: Vec<f64>,
    /// Loss for each batch across all epochs
    #[pyo3(get)]
    pub batch_losses: Vec<f64>,
    /// True if training was stopped early due to validation loss plateau.
    #[pyo3(get)]
    pub stopped_early: bool,
}

#[pyclass]
pub struct IlpProgramFactory;

#[pyclass]
pub struct CompiledIlpProgram {
    pub(crate) base_source: String,
    pub(crate) _learnable_source: String,
    pub(crate) ast: AstProgram,
    pub(crate) executor: Executor,
    pub(crate) provider: Arc<CudaKernelProvider>,
    pub(crate) plan: ExecutionPlan,
    pub(crate) rel_index: Vec<(RelId, String)>,
    pub(crate) schemas: HashMap<String, Schema>,
    pub(crate) left_keys: Vec<usize>,
    pub(crate) right_keys: Vec<usize>,
    pub(crate) head_projection: Vec<usize>,
    pub(crate) compiled_schema_size: usize,
    pub(crate) head_rel_name: String,
    pub(crate) max_active_rules: usize,
    pub(crate) candidate_map: Option<HashMap<(u32, u32, u32), u32>>,
    pub(crate) candidate_order: Option<Vec<(u32, u32, u32)>>,
    pub(crate) relation_overrides: HashMap<String, CudaBuffer>,
    /// Maximum bytes for per-chunk temp allocations (masks, prefix sums,
    /// chunk-local COO scratch). The final merged COO buffer is exact-NNZ
    /// sized and may exceed this budget. Default: 16 MB.
    pub(crate) coo_chunk_budget: u64,
    /// When true, raise instead of falling back to chunked COO path.
    /// Use in zero-D2H benchmarks and CI gates. Default: false.
    pub(crate) strict_zero_dtoh: bool,
    /// Per-phase wall-clock of the compile that built this program (ms).
    pub(crate) compile_timing: Vec<(&'static str, f64)>,
}

#[pymodule]
#[pyo3(name = "_native")]
fn pyxlog(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    m.add_class::<Program>()?;
    m.add_class::<CompiledProgram>()?;
    m.add_class::<LogicProgram>()?;
    m.add_class::<CompiledLogicProgram>()?;
    m.add_class::<CompiledConditionedProgram>()?;
    m.add_class::<LogicRelationSession>()?;
    m.add_class::<relation_metadata::RelationEvidence>()?;
    m.add_class::<LogicQueryResult>()?;
    m.add_class::<LogicEvalResult>()?;
    m.add_class::<McDeviceEvalResult>()?;
    m.add_class::<EvalResult>()?;
    m.add_class::<EpistemicEvalResult>()?;
    m.add_class::<EpistemicEvidence>()?;
    // Training infrastructure
    m.add_class::<PyDifferentiableProofTraceMap>()?;
    m.add_class::<EpochStats>()?;
    m.add_class::<TrainingHistory>()?;
    // ILP bindings
    m.add_class::<IlpProgramFactory>()?;
    m.add_class::<CompiledIlpProgram>()?;
    m.add_class::<IlpTaggedCreditDeviceResult>()?;
    m.add_function(wrap_pyfunction!(training::train_model, m)?)?;
    m.add_function(wrap_pyfunction!(training::train_model_tensor, m)?)?;
    m.add_function(wrap_pyfunction!(dlpack::dlpack_roundtrip, m)?)?;
    m.add_function(wrap_pyfunction!(dlpack_is_cuda, m)?)?;
    #[cfg(feature = "arrow-device-import")]
    m.add_function(wrap_pyfunction!(dlpack::export_arrow_device, m)?)?;
    #[cfg(feature = "arrow-device-import")]
    m.add_function(wrap_pyfunction!(dlpack::import_arrow_device, m)?)?;
    // Joint constraint carrier bindings
    m.add_class::<joint_carrier::JointConstraintCarrier>()?;
    m.add(
        "CarrierRefused",
        _py.get_type::<joint_carrier::CarrierRefused>(),
    )?;
    m.add(
        "SolverResourceExhausted",
        _py.get_type::<joint_carrier::SolverResourceExhausted>(),
    )?;
    m.add(
        "RelationMetadataError",
        _py.get_type::<relation_metadata::RelationMetadataError>(),
    )?;
    m.add("SOLVER_ABI_IDENTITY", xlog_cuda::SOLVER_ABI_IDENTITY)?;
    Ok(())
}