use laddu_runtime::{BinSpec, IntervalClosure, Predicate};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyAny};
use super::{error::to_py_err, expr::PyExpr, float_vec};
#[pyclass(name = "Predicate", module = "laddu", frozen, skip_from_py_object)]
#[derive(Clone, Debug)]
pub struct PyPredicate {
pub(crate) inner: Predicate,
}
#[pymethods]
impl PyPredicate {
fn __and__(&self, other: &Self) -> Self {
Self {
inner: self.inner.clone().and(other.inner.clone()),
}
}
fn __or__(&self, other: &Self) -> Self {
Self {
inner: self.inner.clone().or(other.inner.clone()),
}
}
fn __invert__(&self) -> Self {
Self {
inner: !self.inner.clone(),
}
}
#[staticmethod]
#[pyo3(signature = (value, *, low, high, closed="both"))]
fn between(value: &PyExpr, low: &PyExpr, high: &PyExpr, closed: &str) -> PyResult<Self> {
let closure = match closed {
"both" => IntervalClosure::Closed,
"neither" => IntervalClosure::Open,
"left" => IntervalClosure::LeftClosed,
"right" => IntervalClosure::RightClosed,
_ => {
return Err(PyValueError::new_err(
"closed must be 'both', 'neither', 'left', or 'right'",
));
}
};
Ok(Self {
inner: Predicate::between_with(
value.inner.clone(),
low.inner.clone(),
high.inner.clone(),
closure,
),
})
}
}
#[pyclass(name = "Bin", module = "laddu", frozen, skip_from_py_object)]
#[derive(Clone)]
pub struct PyBin {
pub(crate) inner: BinSpec,
}
#[pymethods]
impl PyBin {
#[new]
#[pyo3(signature = (
edges: "Sequence[float] | numpy.typing.NDArray[numpy.float32 | numpy.float64]"
))]
fn new(edges: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
inner: BinSpec::edges(float_vec(edges)?).map_err(to_py_err)?,
})
}
#[staticmethod]
#[pyo3(signature = (count, *, low, high))]
fn uniform(count: usize, low: f64, high: f64) -> PyResult<Self> {
Ok(Self {
inner: BinSpec::uniform(count, low, high).map_err(to_py_err)?,
})
}
#[getter]
fn edges(&self) -> Vec<f64> {
self.inner.edges_slice().to_vec()
}
}
impl_json_methods!(PyBin);