core/core-rs/src/lib.rs

//! core-rs: Rust extension for CORE-AI
//!
//! Exposes hot-path operations to Python via PyO3:
//!   - geometric_product   (Cl(4,1) full product via precomputed table)
//!   - versor_apply        (sandwich product V*F*rev(V))
//!   - versor_condition    (||F*rev(F) - 1||_F)
//!   - cga_inner           (symmetric inner product)
//!   - vault_recall        (parallel top-k scan)
//!   - diffusion_step      (zero-copy graph diffusion step)
//!
//! All multivectors are f32 arrays of length 32, passed as numpy arrays.

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;

pub mod cga;
pub mod cl41;
pub mod diffusion;
pub mod vault;
pub mod versor;

use cga::cga_inner_raw;
use cl41::geometric_product_raw;
use diffusion::{graph_diffusion_step, unitize_f32};
use versor::{
    normalize_to_versor_raw, versor_apply_closed, versor_apply_closed_f64, versor_apply_raw,
    versor_condition_raw,
};

/// Geometric product in Cl(4,1). Accepts two contiguous float32 arrays of length 32.
///
/// Inputs are read via ``PyReadonlyArray1`` zero-copy views into the NumPy
/// buffer.  Wrong shape, dtype, or non-contiguous layout fails loudly — no
/// silent coercion.
#[pyfunction]
fn geometric_product(
    py: Python<'_>,
    a: numpy::PyReadonlyArray1<'_, f32>,
    b: numpy::PyReadonlyArray1<'_, f32>,
) -> PyResult<PyObject> {
    let a_slice = read_f32_cl41_mv(&a)?;
    let b_slice = read_f32_cl41_mv(&b)?;
    let result = geometric_product_raw(a_slice, b_slice)
        .map_err(|e| PyValueError::new_err(e.to_string()))?;
    f32_array_to_numpy(py, &result)
}

/// Sandwich product V*F*reverse(V).
#[pyfunction]
fn versor_apply(
    py: Python<'_>,
    v: &Bound<'_, pyo3::types::PyAny>,
    f: &Bound<'_, pyo3::types::PyAny>,
) -> PyResult<PyObject> {
    let v_slice = extract_f32_slice(v)?;
    let f_slice = extract_f32_slice(f)?;
    let result =
        versor_apply_raw(&v_slice, &f_slice).map_err(|e| PyValueError::new_err(e.to_string()))?;
    f32_array_to_numpy(py, &result)
}

/// Sandwich product V*F*reverse(V) with closure semantics.
/// Preserves null vectors, applies unit-versor closure with seed fallback.
#[pyfunction]
fn versor_apply_with_closure(
    py: Python<'_>,
    v: numpy::PyReadonlyArray1<'_, f32>,
    f: numpy::PyReadonlyArray1<'_, f32>,
) -> PyResult<PyObject> {
    let v_slice = read_f32_cl41_mv(&v)?;
    let f_slice = read_f32_cl41_mv(&f)?;
    let result =
        versor_apply_closed(v_slice, f_slice).map_err(|e| PyValueError::new_err(e.to_string()))?;
    f32_array_to_numpy(py, &result)
}

/// `versor_apply` f64 closure path — bit-identity port of Python
/// `algebra.versor.versor_apply` + `_close_applied_versor`.
/// Inputs and output are float64.  ADR-0020 parity surface.
#[pyfunction]
fn versor_apply_with_closure_f64(
    py: Python<'_>,
    v: numpy::PyReadonlyArray1<'_, f64>,
    f: numpy::PyReadonlyArray1<'_, f64>,
) -> PyResult<PyObject> {
    let v_slice = read_f64_cl41_mv(&v)?;
    let f_slice = read_f64_cl41_mv(&f)?;
    let result = versor_apply_closed_f64(v_slice, f_slice)
        .map_err(|e| PyValueError::new_err(e.to_string()))?;
    f64_array_to_numpy(py, &result)
}

/// ||F*reverse(F) - 1||_F. Returns scalar f32.
#[pyfunction]
fn versor_condition(f: numpy::PyReadonlyArray1<'_, f32>) -> PyResult<f32> {
    let f_slice = read_f32_cl41_mv(&f)?;
    versor_condition_raw(f_slice).map_err(|e| PyValueError::new_err(e.to_string()))
}

/// Project F onto versor manifold: F / sqrt(|F*rev(F)|).
#[pyfunction]
fn normalize_to_versor(py: Python<'_>, f: &Bound<'_, pyo3::types::PyAny>) -> PyResult<PyObject> {
    let f_slice = extract_f32_slice(f)?;
    let result =
        normalize_to_versor_raw(&f_slice).map_err(|e| PyValueError::new_err(e.to_string()))?;
    f32_array_to_numpy(py, &result)
}

/// Symmetric CGA inner product: 0.5 * scalar(X*Y + Y*X).
#[pyfunction]
fn cga_inner(
    x: numpy::PyReadonlyArray1<'_, f32>,
    y: numpy::PyReadonlyArray1<'_, f32>,
) -> PyResult<f32> {
    let x_slice = read_f32_cl41_mv(&x)?;
    let y_slice = read_f32_cl41_mv(&y)?;
    cga_inner_raw(x_slice, y_slice).map_err(|e| PyValueError::new_err(e.to_string()))
}

/// Embed a Euclidean point [x, y, z] into the CGA null cone.
#[pyfunction]
fn embed_point(
    py: Python<'_>,
    p: numpy::PyReadonlyArray1<'_, f32>,
) -> PyResult<PyObject> {
    let p_slice = read_f32_xyz(&p)?;
    let result = crate::cga::embed_point_raw(p_slice);
    f32_array_to_numpy(py, &result)
}

/// Re-project a multivector onto the null cone by Euclidean read-back + re-embed.
#[pyfunction]
fn null_project(
    py: Python<'_>,
    x: numpy::PyReadonlyArray1<'_, f32>,
) -> PyResult<PyObject> {
    let x_slice = read_f32_cl41_mv(&x)?;
    let result = crate::cga::null_project_raw(x_slice);
    f32_array_to_numpy(py, &result)
}

/// Check whether a multivector lies on the null cone.
#[pyfunction]
fn is_null(
    x: numpy::PyReadonlyArray1<'_, f32>,
    tol: f32,
) -> PyResult<bool> {
    let x_slice = read_f32_cl41_mv(&x)?;
    crate::cga::is_null_raw(x_slice, tol)
        .map_err(|e| PyValueError::new_err(e.to_string()))
}

/// Parallel top-k vault recall by CGA inner product (zero-copy).
///
/// Per ADR-0020 follow-on (task #35): accepts a 2D numpy
/// (N, 32) float32 array via `PyReadonlyArray2`, which exposes a
/// view *directly into the numpy buffer*.  No Python→Rust copy,
/// no re-chunking — Rayon scores straight off the source slice.
/// This is the load-bearing reason for the Rust path: NumPy
/// already SIMD-vectorises the same kernel; the only win Rust
/// can offer is *avoiding the marshalling tax* and adding
/// thread-parallel scoring on top.
#[pyfunction]
fn vault_recall(
    versors: numpy::PyReadonlyArray2<'_, f32>,
    query: numpy::PyReadonlyArray1<'_, f32>,
    top_k: usize,
) -> PyResult<Vec<(usize, f32)>> {
    let view = versors.as_array();
    let shape = view.shape();
    if shape.len() != 2 || shape[1] != 32 {
        return Err(PyValueError::new_err(format!(
            "versors must be shape (N, 32), got {:?}",
            shape
        )));
    }
    let n = shape[0];
    let q_slice = query
        .as_slice()
        .map_err(|e| PyValueError::new_err(format!("query must be contiguous f32 (32,): {}", e)))?;
    if q_slice.len() != 32 {
        return Err(PyValueError::new_err(format!(
            "query must have length 32, got {}",
            q_slice.len()
        )));
    }
    let v_slice = versors.as_slice().map_err(|e| {
        PyValueError::new_err(format!("versors must be C-contiguous f32 (N, 32): {}", e))
    })?;
    let mut q_arr = [0f32; 32];
    q_arr.copy_from_slice(q_slice);

    crate::vault::vault_recall_flat(v_slice, n, &q_arr, top_k)
        .map_err(|e| PyValueError::new_err(e.to_string()))
}

/// Unitize a multivector via the Cl(4,1) exponential map.
/// Distinguishes boost planes (cosh/sinh) from rotation planes (cos/sin).
#[pyfunction]
fn unitize_expmap(py: Python<'_>, v: &Bound<'_, pyo3::types::PyAny>) -> PyResult<PyObject> {
    let v_slice = extract_f32_slice(v)?;
    let result = unitize_f32(&v_slice);
    f32_array_to_numpy(py, &result)
}

/// One forward step of graph diffusion.
///
/// Takes ``fields`` (N x 32 float32 numpy) and ``edges`` (E x 2 int32
/// numpy) as zero-copy ``PyReadonlyArray2`` views; returns the new
/// fields as an owned ``PyArray2<f32>`` plus the scalar L2 delta.
///
/// The previous signature took ``Vec<f32>`` + ``Vec<i32>``, which forced
/// PyO3 to box-unbox every element through Python's float/int object
/// representation on the way in, and required a ``numpy.array(...)
/// .reshape(...)`` round-trip on the way out.  For a 200-step pulse
/// over a small graph this was the dominant cost — Rust-vs-Python
/// parity (0.99x) on the speedup bench was paying for marshalling,
/// not algorithm.  Zero-copy ``PyReadonlyArray2`` + ``bytemuck`` slice
/// reinterpretation removes both ends of that tax; the inner kernel
/// (``diffusion::graph_diffusion_step``) is unchanged.
#[pyfunction]
fn diffusion_step<'py>(
    py: Python<'py>,
    fields: numpy::PyReadonlyArray2<'py, f32>,
    edges: numpy::PyReadonlyArray2<'py, i32>,
    damping: f64,
) -> PyResult<(Bound<'py, numpy::PyArray2<f32>>, f64)> {
    // ``shape()`` lives on the ndarray view, not directly on
    // ``PyReadonlyArray2`` — go through ``as_array()`` to get the view.
    let fields_view = fields.as_array();
    let fields_shape = fields_view.shape();
    if fields_shape.len() != 2 || fields_shape[1] != 32 {
        return Err(PyValueError::new_err(format!(
            "fields must be shape (N, 32), got {:?}",
            fields_shape
        )));
    }
    let n_nodes = fields_shape[0];

    let edges_view = edges.as_array();
    let edges_shape = edges_view.shape();
    if edges_shape.len() != 2 || edges_shape[1] != 2 {
        return Err(PyValueError::new_err(format!(
            "edges must be shape (E, 2), got {:?}",
            edges_shape
        )));
    }

    let fields_slice = fields.as_slice().map_err(|e| {
        PyValueError::new_err(format!("fields must be C-contiguous f32 (N, 32): {}", e))
    })?;
    let edges_slice = edges.as_slice().map_err(|e| {
        PyValueError::new_err(format!("edges must be C-contiguous i32 (E, 2): {}", e))
    })?;

    // ``[f32; 32]`` and ``[i32; 2]`` are both ``Pod`` (arrays of POD
    // primitives), so reinterpretation of the contiguous numpy buffer
    // into the kernel's expected slice types is zero-copy.
    let fields_blocks: &[[f32; 32]] = bytemuck::cast_slice(fields_slice);
    let edges_blocks: &[[i32; 2]] = bytemuck::cast_slice(edges_slice);

    let (new_fields, delta) = graph_diffusion_step(fields_blocks, edges_blocks, damping);

    // ``Vec<[f32; 32]>`` → ``Vec<f32>`` is a zero-copy reinterpretation
    // of the allocation (requires the ``extern_crate_alloc`` bytemuck
    // feature; see Cargo.toml).
    //
    // We use ``numpy::ndarray::Array2`` (numpy 0.21's re-export of
    // ndarray 0.15) rather than ``ndarray::Array2`` to keep crate
    // versions aligned — the workspace pulls ndarray 0.16 for the
    // ``diffusion`` module but ``numpy::IntoPyArray`` is implemented
    // for ndarray 0.15's types only.
    let flat: Vec<f32> = bytemuck::allocation::cast_vec(new_fields);
    let arr = numpy::ndarray::Array2::from_shape_vec((n_nodes, 32), flat)
        .map_err(|e| PyValueError::new_err(e.to_string()))?;
    Ok((numpy::IntoPyArray::into_pyarray_bound(arr, py), delta))
}

fn read_f32_cl41_mv<'a>(arr: &'a numpy::PyReadonlyArray1<'a, f32>) -> PyResult<&'a [f32; 32]> {
    let len = arr.len()?;
    if len != 32 {
        return Err(PyValueError::new_err(format!(
            "expected contiguous float32 array of length 32, got length {}",
            len
        )));
    }
    let slice = arr.as_slice().map_err(|e| {
        PyValueError::new_err(format!("input must be C-contiguous float32 (32,): {}", e))
    })?;
    slice
        .try_into()
        .map_err(|_| PyValueError::new_err("expected contiguous float32 array of length 32"))
}

fn read_f64_cl41_mv<'a>(arr: &'a numpy::PyReadonlyArray1<'a, f64>) -> PyResult<&'a [f64; 32]> {
    let len = arr.len()?;
    if len != 32 {
        return Err(PyValueError::new_err(format!(
            "expected contiguous float64 array of length 32, got length {}",
            len
        )));
    }
    let slice = arr.as_slice().map_err(|e| {
        PyValueError::new_err(format!("input must be C-contiguous float64 (32,): {}", e))
    })?;
    slice
        .try_into()
        .map_err(|_| PyValueError::new_err("expected contiguous float64 array of length 32"))
}

fn read_f32_xyz<'a>(arr: &'a numpy::PyReadonlyArray1<'a, f32>) -> PyResult<&'a [f32; 3]> {
    let len = arr.len()?;
    if len != 3 {
        return Err(PyValueError::new_err(format!(
            "expected contiguous float32 array of length 3, got length {}",
            len
        )));
    }
    let slice = arr.as_slice().map_err(|e| {
        PyValueError::new_err(format!(
            "input must be C-contiguous float32 (3,): {}",
            e
        ))
    })?;
    slice.try_into().map_err(|_| {
        PyValueError::new_err("expected contiguous float32 array of length 3")
    })
}

fn extract_f32_slice(obj: &Bound<'_, pyo3::types::PyAny>) -> PyResult<[f32; 32]> {
    let np = obj.py().import_bound("numpy")?;
    let arr = np.call_method1("asarray", (obj, "float32"))?;
    let flat = arr.call_method0("flatten")?;
    let list: Vec<f32> = flat.extract()?;
    if list.len() != 32 {
        return Err(PyValueError::new_err(format!(
            "Expected array of length 32, got {}",
            list.len()
        )));
    }
    let mut out = [0f32; 32];
    out.copy_from_slice(&list);
    Ok(out)
}

fn f32_array_to_numpy(py: Python<'_>, data: &[f32; 32]) -> PyResult<PyObject> {
    let np = py.import_bound("numpy")?;
    let list: Vec<f32> = data.to_vec();
    let arr = np.call_method1("array", (list, "float32"))?;
    Ok(arr.into_py(py))
}

fn f64_array_to_numpy(py: Python<'_>, data: &[f64; 32]) -> PyResult<PyObject> {
    let np = py.import_bound("numpy")?;
    let list: Vec<f64> = data.to_vec();
    let arr = np.call_method1("array", (list, "float64"))?;
    Ok(arr.into_py(py))
}

#[pymodule]
fn core_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(geometric_product, m)?)?;
    m.add_function(wrap_pyfunction!(versor_apply, m)?)?;
    m.add_function(wrap_pyfunction!(versor_apply_with_closure, m)?)?;
    m.add_function(wrap_pyfunction!(versor_apply_with_closure_f64, m)?)?;
    m.add_function(wrap_pyfunction!(versor_condition, m)?)?;
    m.add_function(wrap_pyfunction!(normalize_to_versor, m)?)?;
    m.add_function(wrap_pyfunction!(cga_inner, m)?)?;
    m.add_function(wrap_pyfunction!(embed_point, m)?)?;
    m.add_function(wrap_pyfunction!(null_project, m)?)?;
    m.add_function(wrap_pyfunction!(is_null, m)?)?;
    m.add_function(wrap_pyfunction!(vault_recall, m)?)?;
    m.add_function(wrap_pyfunction!(unitize_expmap, m)?)?;
    m.add_function(wrap_pyfunction!(diffusion_step, m)?)?;
    Ok(())
}