particles.rs

use bevy::prelude::Entity;
use processing::prelude::*;
use processing_render::geometry;
use pyo3::types::{PyDict, PyTuple};
use pyo3::{
    exceptions::{PyRuntimeError, PyTypeError, PyValueError},
    prelude::*,
};
use std::collections::HashMap;

use processing_render::particles::algebra::{
    combine as algebra_combine, extract as algebra_extract, generate as algebra_generate,
    lookup as algebra_lookup, map as algebra_map, mix as algebra_mix, pack as algebra_pack,
    reduce_components as algebra_reduce,
};
use processing_render::particles::compact::compact as compact_indices;

use crate::compute::{Buffer, Compute};
use crate::graphics::{Geometry, Image};

use processing::prelude::constants as c;
use processing_render::{
    COMBINE_ADD, COMBINE_DIV, COMBINE_MAX, COMBINE_MIN, COMBINE_MUL, COMBINE_POW, COMBINE_SUB,
};
use processing_render::{GEN_GAUSSIAN, GEN_SIGNED, GEN_UNIFORM};
use processing_render::{
    MAP_ABS, MAP_AFFINE, MAP_CLAMP, MAP_EQ, MAP_FLOOR, MAP_GEQ, MAP_GREATER, MAP_LEQ, MAP_LESS,
    MAP_LIMIT, MAP_NEGATE, MAP_NEQ, MAP_NORMALIZE, MAP_SQRT, MAP_SQUARE,
};
use processing_render::{
    REDUCE_LENGTH, REDUCE_MAX, REDUCE_MEAN, REDUCE_MIN, REDUCE_SUM, REDUCE_SUMSQ,
};

fn parse_map_op(s: &str) -> PyResult<u32> {
    match () {
        _ if s.eq_ignore_ascii_case(c::AFFINE) => Ok(MAP_AFFINE),
        _ if s.eq_ignore_ascii_case(c::ABS) => Ok(MAP_ABS),
        _ if s.eq_ignore_ascii_case(c::NEGATE) => Ok(MAP_NEGATE),
        _ if s.eq_ignore_ascii_case(c::CLAMP) => Ok(MAP_CLAMP),
        _ if s.eq_ignore_ascii_case(c::FLOOR) => Ok(MAP_FLOOR),
        _ if s.eq_ignore_ascii_case(c::SQUARE) => Ok(MAP_SQUARE),
        _ if s.eq_ignore_ascii_case(c::SQRT) => Ok(MAP_SQRT),
        _ if s.eq_ignore_ascii_case(c::GREATER) => Ok(MAP_GREATER),
        _ if s.eq_ignore_ascii_case(c::LESS) => Ok(MAP_LESS),
        _ if s.eq_ignore_ascii_case(c::GEQ) => Ok(MAP_GEQ),
        _ if s.eq_ignore_ascii_case(c::LEQ) => Ok(MAP_LEQ),
        _ if s.eq_ignore_ascii_case(c::EQ) => Ok(MAP_EQ),
        _ if s.eq_ignore_ascii_case(c::NEQ) => Ok(MAP_NEQ),
        _ if s.eq_ignore_ascii_case(c::NORMALIZE) => Ok(MAP_NORMALIZE),
        _ if s.eq_ignore_ascii_case(c::LIMIT) => Ok(MAP_LIMIT),
        _ => Err(PyValueError::new_err(format!("map: unknown op {s:?}"))),
    }
}

fn parse_combine_op(s: &str) -> PyResult<u32> {
    match () {
        _ if s.eq_ignore_ascii_case(c::ADD) => Ok(COMBINE_ADD),
        _ if s.eq_ignore_ascii_case(c::SUB) => Ok(COMBINE_SUB),
        _ if s.eq_ignore_ascii_case(c::MUL) => Ok(COMBINE_MUL),
        _ if s.eq_ignore_ascii_case(c::DIV) => Ok(COMBINE_DIV),
        _ if s.eq_ignore_ascii_case(c::MIN) => Ok(COMBINE_MIN),
        _ if s.eq_ignore_ascii_case(c::MAX) => Ok(COMBINE_MAX),
        _ if s.eq_ignore_ascii_case(c::POW) => Ok(COMBINE_POW),
        _ => Err(PyValueError::new_err(format!("combine: unknown op {s:?}"))),
    }
}

fn parse_reduce_op(s: &str) -> PyResult<u32> {
    match () {
        _ if s.eq_ignore_ascii_case(c::LENGTH) => Ok(REDUCE_LENGTH),
        _ if s.eq_ignore_ascii_case(c::SUM) => Ok(REDUCE_SUM),
        _ if s.eq_ignore_ascii_case(c::MIN) => Ok(REDUCE_MIN),
        _ if s.eq_ignore_ascii_case(c::MAX) => Ok(REDUCE_MAX),
        _ if s.eq_ignore_ascii_case(c::SUMSQ) => Ok(REDUCE_SUMSQ),
        _ if s.eq_ignore_ascii_case(c::MEAN) => Ok(REDUCE_MEAN),
        _ => Err(PyValueError::new_err(format!("reduce: unknown op {s:?}"))),
    }
}

fn parse_generate_mode(s: &str) -> PyResult<u32> {
    match () {
        _ if s.eq_ignore_ascii_case(c::UNIFORM) => Ok(GEN_UNIFORM),
        _ if s.eq_ignore_ascii_case(c::SIGNED) => Ok(GEN_SIGNED),
        _ if s.eq_ignore_ascii_case(c::GAUSSIAN) => Ok(GEN_GAUSSIAN),
        _ => Err(PyValueError::new_err(format!(
            "generate: unknown mode {s:?}"
        ))),
    }
}

const NEIGHBOR_SUM: u32 = 0;
const NEIGHBOR_MEAN: u32 = 1;
const NEIGHBOR_COUNT: u32 = 2;

fn parse_neighbor_op(s: &str) -> PyResult<u32> {
    match () {
        _ if s.eq_ignore_ascii_case(c::SUM) => Ok(NEIGHBOR_SUM),
        _ if s.eq_ignore_ascii_case(c::MEAN) => Ok(NEIGHBOR_MEAN),
        _ if s.eq_ignore_ascii_case(c::COUNT) || s.eq_ignore_ascii_case(c::DENSITY) => {
            Ok(NEIGHBOR_COUNT)
        }
        _ => Err(PyValueError::new_err(format!("neighbor: unknown op {s:?}"))),
    }
}

fn parse_falloff(s: &str) -> PyResult<u32> {
    match () {
        _ if s.eq_ignore_ascii_case(c::CONSTANT) => Ok(FALLOFF_CONST),
        _ if s.eq_ignore_ascii_case(c::LINEAR) => Ok(FALLOFF_LINEAR),
        _ if s.eq_ignore_ascii_case(c::SMOOTHSTEP) => Ok(FALLOFF_SMOOTHSTEP),
        _ if s.eq_ignore_ascii_case(c::QUADRATIC) => Ok(FALLOFF_QUADRATIC),
        _ if s.eq_ignore_ascii_case(c::CUBIC) => Ok(FALLOFF_CUBIC),
        _ if s.eq_ignore_ascii_case(c::INVERSE) => Ok(FALLOFF_INVERSE),
        _ if s.eq_ignore_ascii_case(c::INVERSE_SQUARE) => Ok(FALLOFF_INVERSE_SQUARE),
        _ => Err(PyValueError::new_err(format!(
            "neighbor: unknown falloff {s:?}"
        ))),
    }
}

/// `emit()` data: flat numbers, or one value per particle (`[x, y, z]`, a `Color`, ...)
/// converted like `Buffer.write`.
fn emit_bytes(name: &str, value: &Bound<'_, PyAny>) -> PyResult<Vec<u8>> {
    let mut bytes = Vec::new();
    for item in value.try_iter()? {
        let item = item?;
        // a bare int would become an i32 shader value; attributes are floats
        if let Ok(f) = item.extract::<f32>() {
            bytes.extend_from_slice(&f.to_le_bytes());
            continue;
        }
        let sv = crate::material::py_to_shader_value(&item)?;
        let b = sv.to_bytes().ok_or_else(|| {
            PyTypeError::new_err(format!("attribute '{name}': unsupported value {sv:?}"))
        })?;
        bytes.extend_from_slice(&b);
    }
    Ok(bytes)
}

#[pyclass(unsendable)]
pub struct Grid {
    pub(crate) entity: Entity,
}

#[pymethods]
impl Grid {
    /// Opaque id for this object.
    pub fn id(&self) -> u64 {
        self.entity.to_bits()
    }

    pub fn build(&self, position: &Buffer) -> PyResult<()> {
        grid_build(self.entity, position.entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))
    }

    #[getter]
    pub fn cell_size(&self) -> PyResult<f32> {
        Ok(grid_get(self.entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?
            .params
            .cell_size)
    }
}

impl Drop for Grid {
    fn drop(&mut self) {
        let _ = grid_destroy(self.entity);
    }
}

static PHYSICS_COMPUTES: std::sync::Mutex<Option<HashMap<String, Entity>>> =
    std::sync::Mutex::new(None);

fn physics_compute(name: &str) -> PyResult<Option<Entity>> {
    let lower = name.to_ascii_lowercase();
    if let Some(cache) = PHYSICS_COMPUTES.lock().unwrap().as_ref()
        && let Some(&e) = cache.get(&lower)
    {
        return Ok(Some(e));
    }
    let created = if lower == c::NOISE {
        particles_kernel_noise()
    } else if lower == c::TRANSFORM {
        particles_kernel_transform()
    } else if lower == c::ATTRACT {
        particles_kernel_attract()
    } else if lower == c::DRAG {
        particles_kernel_drag()
    } else if lower == c::VORTEX {
        particles_kernel_vortex()
    } else if lower == c::FORCE {
        particles_kernel_force()
    } else if lower == c::INTEGRATE {
        particles_kernel_integrate()
    } else if lower == c::AGE {
        particles_kernel_age()
    } else if lower == c::IMPULSE {
        particles_kernel_impulse()
    } else if lower == c::ORIENT {
        particles_kernel_orient()
    } else if lower == c::FIELD {
        particles_kernel_field()
    } else if lower == c::BOUNDS_SPHERE {
        particles_kernel_bounds_sphere()
    } else if lower == c::BOUNDS_BOX {
        particles_kernel_bounds_box()
    } else {
        return Ok(None);
    };
    let entity = created.map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
    PHYSICS_COMPUTES
        .lock()
        .unwrap()
        .get_or_insert_with(HashMap::new)
        .insert(lower, entity);
    Ok(Some(entity))
}

fn kw<'a>(kwargs: Option<&Bound<'a, PyDict>>, key: &str) -> Option<Bound<'a, PyAny>> {
    kwargs.and_then(|d| d.get_item(key).ok().flatten())
}

fn kw_f32(kwargs: Option<&Bound<'_, PyDict>>, key: &str, default: f32) -> PyResult<f32> {
    match kw(kwargs, key) {
        Some(v) => v.extract(),
        None => Ok(default),
    }
}

fn kw_u32(kwargs: Option<&Bound<'_, PyDict>>, key: &str, default: u32) -> PyResult<u32> {
    match kw(kwargs, key) {
        Some(v) => v.extract(),
        None => Ok(default),
    }
}

fn kw_bool(kwargs: Option<&Bound<'_, PyDict>>, key: &str, default: bool) -> PyResult<bool> {
    match kw(kwargs, key) {
        Some(v) => v.extract(),
        None => Ok(default),
    }
}

fn map_params(kwargs: Option<&Bound<'_, PyDict>>, op: u32) -> PyResult<(f32, f32)> {
    Ok(match op {
        MAP_AFFINE => (
            kw_f32(kwargs, "scale", 1.0)?,
            kw_f32(kwargs, "offset", 0.0)?,
        ),
        MAP_CLAMP => (kw_f32(kwargs, "lo", 0.0)?, kw_f32(kwargs, "hi", 1.0)?),
        MAP_NORMALIZE => (kw_f32(kwargs, "length", 1.0)?, 0.0),
        MAP_LIMIT => (
            kw_f32(kwargs, "min_length", 0.0)?,
            kw_f32(kwargs, "max_length", f32::MAX)?,
        ),
        MAP_GREATER | MAP_LESS | MAP_GEQ | MAP_LEQ | MAP_EQ | MAP_NEQ => (
            kw_f32(kwargs, "threshold", 0.0)?,
            kw_f32(kwargs, "epsilon", 1.0e-6)?,
        ),
        _ => (0.0, 0.0),
    })
}

fn map_param_keys(op: u32) -> &'static [&'static str] {
    match op {
        MAP_AFFINE => &["scale", "offset"],
        MAP_CLAMP => &["lo", "hi"],
        MAP_NORMALIZE => &["length"],
        MAP_LIMIT => &["min_length", "max_length"],
        MAP_GREATER | MAP_LESS | MAP_GEQ | MAP_LEQ | MAP_EQ | MAP_NEQ => &["threshold", "epsilon"],
        _ => &[],
    }
}

fn operand_names(verb: &str) -> &'static [&'static str] {
    let is = |c: &str| verb.eq_ignore_ascii_case(c);
    if is(c::MAP) || is(c::REDUCE) || is(c::EXTRACT) || is(c::LOOKUP) || is(c::NEIGHBOR) {
        &["a"]
    } else if is(c::COMBINE) {
        &["a", "b"]
    } else if is(c::MIX) {
        &["a", "b", "t"]
    } else {
        &[]
    }
}

fn with_operands<'py>(
    verb: &str,
    operands: &Bound<'py, PyTuple>,
    kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Option<Bound<'py, PyDict>>> {
    let py = operands.py();
    let merged = match kwargs {
        Some(k) => k.copy()?,
        None => PyDict::new(py),
    };
    let packs = verb.eq_ignore_ascii_case(c::PACK);
    let names = operand_names(verb);
    for name in names.iter().chain(packs.then_some(&"sources")) {
        if merged.contains(*name)? {
            return Err(PyTypeError::new_err(format!(
                "apply({verb}): pass operands positionally, not as `{name}=`"
            )));
        }
    }
    if packs {
        merged.set_item("sources", operands.to_list())?;
    } else if operands.len() > names.len() {
        return Err(PyTypeError::new_err(format!(
            "apply({verb}): takes at most {} operand(s), got {}",
            names.len(),
            operands.len()
        )));
    } else {
        for (name, operand) in names.iter().zip(operands.iter()) {
            merged.set_item(*name, operand)?;
        }
    }
    Ok(Some(merged))
}

fn reject_unknown_kwargs(kwargs: Option<&Bound<'_, PyDict>>, valid: &[&str]) -> PyResult<()> {
    let Some(kwargs) = kwargs else {
        return Ok(());
    };
    for key in kwargs.keys() {
        let name: String = key.extract()?;
        if !valid.iter().any(|v| *v == name) {
            return Err(PyValueError::new_err(format!(
                "apply(): unknown parameter {name:?} (valid: {})",
                valid.join(", ")
            )));
        }
    }
    Ok(())
}

fn kw_op(
    kwargs: Option<&Bound<'_, PyDict>>,
    default: u32,
    parse: fn(&str) -> PyResult<u32>,
) -> PyResult<u32> {
    match kw(kwargs, "op") {
        Some(v) => parse(&v.extract::<String>()?),
        None => Ok(default),
    }
}

#[pyclass(eq, eq_int, from_py_object)]
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum AttributeFormat {
    Float = 1,
    Float2 = 2,
    Float3 = 3,
    Float4 = 4,
}

impl AttributeFormat {
    pub(crate) fn to_inner(self) -> geometry::AttributeFormat {
        match self {
            Self::Float => geometry::AttributeFormat::Float,
            Self::Float2 => geometry::AttributeFormat::Float2,
            Self::Float3 => geometry::AttributeFormat::Float3,
            Self::Float4 => geometry::AttributeFormat::Float4,
        }
    }

    pub(crate) fn from_inner(inner: geometry::AttributeFormat) -> Self {
        match inner {
            geometry::AttributeFormat::Float => Self::Float,
            geometry::AttributeFormat::Float2 => Self::Float2,
            geometry::AttributeFormat::Float3 => Self::Float3,
            geometry::AttributeFormat::Float4 => Self::Float4,
        }
    }

    pub(crate) fn float_count(self) -> usize {
        match self {
            Self::Float => 1,
            Self::Float2 => 2,
            Self::Float3 => 3,
            Self::Float4 => 4,
        }
    }
}

#[pyclass(unsendable, frozen, hash, eq, from_py_object)]
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct Attribute {
    pub(crate) entity: Entity,
}

#[pymethods]
impl Attribute {
    /// Opaque id for this object.
    pub fn id(&self) -> u64 {
        self.entity.to_bits()
    }

    #[new]
    pub fn new(name: &str, format: AttributeFormat) -> PyResult<Self> {
        let entity = geometry_attribute_create(name, format.to_inner())
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        Ok(Self { entity })
    }

    #[staticmethod]
    pub fn position() -> Self {
        Self {
            entity: geometry_attribute_position(),
        }
    }
    #[staticmethod]
    pub fn normal() -> Self {
        Self {
            entity: geometry_attribute_normal(),
        }
    }
    #[staticmethod]
    pub fn color() -> Self {
        Self {
            entity: geometry_attribute_color(),
        }
    }
    #[staticmethod]
    pub fn uv() -> Self {
        Self {
            entity: geometry_attribute_uv(),
        }
    }
    #[staticmethod]
    pub fn rotation() -> Self {
        Self {
            entity: geometry_attribute_rotation(),
        }
    }
    #[staticmethod]
    pub fn scale() -> Self {
        Self {
            entity: geometry_attribute_scale(),
        }
    }
    #[staticmethod]
    pub fn life() -> Self {
        Self {
            entity: geometry_attribute_life(),
        }
    }
    #[staticmethod]
    pub fn velocity() -> Self {
        Self {
            entity: geometry_attribute_velocity(),
        }
    }
    #[staticmethod]
    pub fn age() -> Self {
        Self {
            entity: geometry_attribute_age(),
        }
    }

    #[getter]
    pub fn name(&self) -> PyResult<String> {
        let (name, _) = geometry_attribute_info(self.entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        Ok(name)
    }

    #[getter]
    pub fn format(&self) -> PyResult<AttributeFormat> {
        let (_, fmt) = geometry_attribute_info(self.entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        Ok(AttributeFormat::from_inner(fmt))
    }
}

#[pyclass(unsendable)]
pub struct Particles {
    pub(crate) entity: Entity,
    name_to_attr: HashMap<String, (Entity, AttributeFormat)>,
}

impl Particles {
    fn build_name_index(
        attrs: &[Attribute],
    ) -> PyResult<HashMap<String, (Entity, AttributeFormat)>> {
        let mut map = HashMap::with_capacity(attrs.len());
        for attr in attrs {
            let (name, fmt) = geometry_attribute_info(attr.entity)
                .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
            map.insert(name, (attr.entity, AttributeFormat::from_inner(fmt)));
        }
        Ok(map)
    }

    /// Resolve a built-in attribute name to its `Attribute`.
    fn builtin_attribute(name: &str) -> Option<Attribute> {
        Some(match name {
            "position" => Attribute::position(),
            "velocity" => Attribute::velocity(),
            "normal" => Attribute::normal(),
            "color" => Attribute::color(),
            "uv" => Attribute::uv(),
            "rotation" => Attribute::rotation(),
            "scale" => Attribute::scale(),
            "life" => Attribute::life(),
            "age" => Attribute::age(),
            _ => return None,
        })
    }

    /// Resolve an attribute name or `Attribute` to its entity. Unknown names must have been
    /// written first, e.g. via `out=`.
    fn resolve_attribute(&self, attribute: &Bound<'_, PyAny>) -> PyResult<Entity> {
        if let Ok(attr) = attribute.extract::<Attribute>() {
            return Ok(attr.entity);
        }
        if let Ok(name) = attribute.extract::<String>() {
            if let Some(attr) = Self::builtin_attribute(&name) {
                return Ok(attr.entity);
            }
            if let Some((entity, _)) = self.name_to_attr.get(&name) {
                return Ok(*entity);
            }
            if let Some(entity) = geometry_attribute_find(&name)
                .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?
            {
                return Ok(entity);
            }
            return Err(PyValueError::new_err(format!(
                "no attribute \"{name}\" yet; write to it first, e.g. out=\"{name}\""
            )));
        }
        Err(PyTypeError::new_err(
            "buffer() expects an attribute name or an Attribute",
        ))
    }

    fn resolve_operand(&self, val: &Bound<'_, PyAny>) -> PyResult<(Entity, u32)> {
        if let Ok(b) = val.extract::<PyRef<Buffer>>() {
            let comp = b.components().ok_or_else(|| {
                PyRuntimeError::new_err(
                    "operand buffer has no element type; pass a particle attribute \
                     or a buffer created with typed data",
                )
            })?;
            return Ok((b.entity, comp));
        }
        let attr_entity = self.resolve_attribute(val)?;
        let buf = particles_ensure_attribute(self.entity, attr_entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        let (_, fmt) = geometry_attribute_info(attr_entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        let comp = match AttributeFormat::from_inner(fmt) {
            AttributeFormat::Float => 1,
            AttributeFormat::Float2 => 2,
            AttributeFormat::Float3 => 3,
            AttributeFormat::Float4 => 4,
        };
        Ok((buf, comp))
    }

    fn operand(&self, kwargs: Option<&Bound<'_, PyDict>>, key: &str) -> PyResult<(Entity, u32)> {
        let val = kw(kwargs, key).ok_or_else(|| {
            PyTypeError::new_err(match key {
                "out" => "apply(): missing `out=`",
                _ => "apply(): missing an operand",
            })
        })?;
        self.resolve_operand(&val)
    }

    /// Resolve `out=`, creating it with `components` if new. An existing `out` must match.
    fn output(&self, kwargs: Option<&Bound<'_, PyDict>>, components: u32) -> PyResult<Entity> {
        let rt = |e: error::ProcessingError| PyRuntimeError::new_err(format!("{e}"));
        let val =
            kw(kwargs, "out").ok_or_else(|| PyTypeError::new_err("apply(): missing `out=`"))?;
        if let Ok(name) = val.extract::<String>()
            && Self::builtin_attribute(&name).is_none()
            && !self.name_to_attr.contains_key(&name)
            && geometry_attribute_find(&name).map_err(rt)?.is_none()
        {
            use processing_render::geometry::AttributeFormat as F;
            let format = match components {
                1 => F::Float,
                2 => F::Float2,
                3 => F::Float3,
                4 => F::Float4,
                n => {
                    return Err(PyValueError::new_err(format!(
                        "apply(): can't create \"{name}\" with {n} components"
                    )));
                }
            };
            let attr = geometry_attribute_create(name, format).map_err(rt)?;
            return particles_ensure_attribute(self.entity, attr).map_err(rt);
        }
        let (out, comp) = self.resolve_operand(&val)?;
        if comp != components {
            return Err(PyValueError::new_err(format!(
                "apply(): `out` has {comp} components but this writes {components}"
            )));
        }
        Ok(out)
    }

    /// `out=` if given, otherwise in place.
    fn dest(
        &self,
        kwargs: Option<&Bound<'_, PyDict>>,
        in_place: Entity,
        components: u32,
    ) -> PyResult<Entity> {
        match kw(kwargs, "out") {
            Some(_) => self.output(kwargs, components),
            None => Ok(in_place),
        }
    }

    /// Build a particle system (backs `create_particles`) from a count or a `Geometry`.
    pub(crate) fn create(
        source: &Bound<'_, PyAny>,
        attributes: Option<Vec<PyRef<Attribute>>>,
    ) -> PyResult<Self> {
        let attrs: Vec<Attribute> = match attributes {
            Some(list) => list.iter().map(|a| (**a).clone()).collect(),
            None => vec![Attribute::position()],
        };
        let attr_entities: Vec<Entity> = attrs.iter().map(|a| a.entity).collect();

        let entity = if let Ok(geometry) = source.extract::<PyRef<Geometry>>() {
            particles_create_from_geometry(geometry.entity, attr_entities)
        } else if let Ok(capacity) = source.extract::<u32>() {
            particles_create(capacity, attr_entities)
        } else {
            return Err(PyTypeError::new_err(
                "create_particles(): pass a particle count or a Geometry",
            ));
        }
        .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;

        Ok(Self {
            entity,
            name_to_attr: Particles::build_name_index(&attrs)?,
        })
    }

    fn apply_named(&self, name: &str, kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<()> {
        fn rt(e: impl std::fmt::Display) -> PyErr {
            PyRuntimeError::new_err(format!("{e}"))
        }

        if let Some(entity) = physics_compute(name)? {
            if let Some(kwargs) = kwargs {
                crate::compute::set_compute_kwargs(entity, kwargs)?;
            }
            return particles_apply(self.entity, entity).map_err(rt);
        }

        if name.eq_ignore_ascii_case(c::MAP) {
            let (a, comp) = self.operand(kwargs, "a")?;
            let out = self.dest(kwargs, a, comp)?;
            let op = kw_op(kwargs, MAP_AFFINE, parse_map_op)?;
            let mut valid = vec!["a", "out", "op"];
            valid.extend_from_slice(map_param_keys(op));
            reject_unknown_kwargs(kwargs, &valid)?;
            let (p0, p1) = map_params(kwargs, op)?;
            algebra_map(out, a, comp, op, p0, p1).map_err(rt)
        } else if name.eq_ignore_ascii_case(c::COMBINE) {
            reject_unknown_kwargs(kwargs, &["a", "b", "out", "op", "b_scale", "b_offset"])?;
            let (a, comp) = self.operand(kwargs, "a")?;
            let (b, b_comp) = self.operand(kwargs, "b")?;
            if b_comp != comp && b_comp != 1 {
                return Err(PyValueError::new_err(format!(
                    "apply(combine): the first operand has {comp} components, so the second needs {comp} or 1, not {b_comp}"
                )));
            }
            let out = self.dest(kwargs, a, comp)?;
            let op = kw_op(kwargs, COMBINE_ADD, parse_combine_op)?;
            let b_scale = kw_f32(kwargs, "b_scale", 1.0)?;
            let b_offset = kw_f32(kwargs, "b_offset", 0.0)?;
            algebra_combine(out, a, b, comp, op, b_scale, b_offset).map_err(rt)
        } else if name.eq_ignore_ascii_case(c::MIX) {
            reject_unknown_kwargs(
                kwargs,
                &["a", "b", "t", "out", "t_scale", "t_offset", "t_clamp"],
            )?;
            let (a, comp) = self.operand(kwargs, "a")?;
            let (b, b_comp) = self.operand(kwargs, "b")?;
            let (t, t_comp) = self.operand(kwargs, "t")?;
            if b_comp != comp {
                return Err(PyValueError::new_err(format!(
                    "apply(mix): the first operand has {comp} components but the second has {b_comp}"
                )));
            }
            if t_comp != 1 {
                return Err(PyValueError::new_err(format!(
                    "apply(mix): the third operand must be one value per particle, not {t_comp} components"
                )));
            }
            let out = self.dest(kwargs, a, comp)?;
            let t_scale = kw_f32(kwargs, "t_scale", 1.0)?;
            let t_offset = kw_f32(kwargs, "t_offset", 0.0)?;
            let t_clamp = kw_bool(kwargs, "t_clamp", true)?;
            algebra_mix(out, a, b, t, comp, t_scale, t_offset, t_clamp).map_err(rt)
        } else if name.eq_ignore_ascii_case(c::LOOKUP) {
            reject_unknown_kwargs(
                kwargs,
                &[
                    "a",
                    "out",
                    "tex",
                    "u_scale",
                    "u_offset",
                    "v_scale",
                    "v_offset",
                    "color_scale",
                ],
            )?;
            let (a, in_comp) = self.operand(kwargs, "a")?;
            let out = self.output(kwargs, 4)?;
            let tex = kw(kwargs, "tex")
                .ok_or_else(|| PyRuntimeError::new_err("apply(lookup): missing 'tex' Image"))?
                .extract::<PyRef<Image>>()
                .map_err(|_| PyRuntimeError::new_err("apply(lookup): 'tex' must be an Image"))?
                .entity;
            let u_scale = kw_f32(kwargs, "u_scale", 1.0)?;
            let u_offset = kw_f32(kwargs, "u_offset", 0.0)?;
            let v_scale = kw_f32(kwargs, "v_scale", 1.0)?;
            let v_offset = kw_f32(kwargs, "v_offset", 0.0)?;
            let color_scale = kw_f32(kwargs, "color_scale", 1.0)?;
            algebra_lookup(
                out,
                a,
                tex,
                in_comp,
                u_scale,
                u_offset,
                v_scale,
                v_offset,
                color_scale,
            )
            .map_err(rt)
        } else if name.eq_ignore_ascii_case(c::REDUCE) {
            reject_unknown_kwargs(kwargs, &["a", "out", "op"])?;
            let (a, comp) = self.operand(kwargs, "a")?;
            let out = self.output(kwargs, 1)?;
            let op = kw_op(kwargs, REDUCE_LENGTH, parse_reduce_op)?;
            algebra_reduce(out, a, comp, op).map_err(rt)
        } else if name.eq_ignore_ascii_case(c::EXTRACT) {
            reject_unknown_kwargs(kwargs, &["a", "out", "index"])?;
            let (a, comp) = self.operand(kwargs, "a")?;
            let out = self.output(kwargs, 1)?;
            let index = kw_u32(kwargs, "index", 0)?;
            algebra_extract(out, a, comp, index).map_err(rt)
        } else if name.eq_ignore_ascii_case(c::PACK) {
            reject_unknown_kwargs(kwargs, &["out", "sources"])?;
            let sources = kw(kwargs, "sources")
                .ok_or_else(|| PyRuntimeError::new_err("apply(pack): missing 'sources' list"))?;
            let items: Vec<Bound<'_, PyAny>> = sources.extract()?;
            let out = self.output(kwargs, items.len() as u32)?;
            let mut entities = Vec::with_capacity(items.len());
            for item in &items {
                entities.push(self.resolve_operand(item)?.0);
            }
            algebra_pack(out, &entities).map_err(rt)
        } else if name.eq_ignore_ascii_case(c::GENERATE) {
            reject_unknown_kwargs(
                kwargs,
                &["out", "mode", "seed", "scale", "offset", "components"],
            )?;
            // nothing to infer from, so a new `out` needs `components=`
            let (out, comp) = match kw(kwargs, "components") {
                Some(c) => {
                    let comp = c.extract::<u32>()?;
                    (self.output(kwargs, comp)?, comp)
                }
                None => match kw(kwargs, "out").map(|v| self.resolve_operand(&v)) {
                    Some(Ok(resolved)) => resolved,
                    Some(Err(_)) => {
                        return Err(PyValueError::new_err(
                            "apply(generate): pass components= to create a new `out`",
                        ));
                    }
                    None => return Err(PyTypeError::new_err("apply(): missing `out=`")),
                },
            };
            let mode = match kw(kwargs, "mode") {
                Some(v) => parse_generate_mode(&v.extract::<String>()?)?,
                None => GEN_UNIFORM,
            };
            let seed = kw_u32(kwargs, "seed", 0)?;
            let scale = kw_f32(kwargs, "scale", 1.0)?;
            let offset = kw_f32(kwargs, "offset", 0.0)?;
            algebra_generate(out, comp, mode, seed, scale, offset).map_err(rt)
        } else if name.eq_ignore_ascii_case(c::FIND_NEIGHBORS) {
            reject_unknown_kwargs(kwargs, &["grid", "radius", "max"])?;
            let grid = kw(kwargs, "grid")
                .ok_or_else(|| PyTypeError::new_err("apply(find_neighbors): missing `grid=`"))?
                .extract::<PyRef<Grid>>()
                .map_err(|_| PyTypeError::new_err("apply(find_neighbors): `grid` must be a Grid"))?
                .entity;
            let cell = grid_get(grid).map_err(rt)?.params.cell_size;
            let radius = kw_f32(kwargs, "radius", cell)?;
            let max = kw_u32(kwargs, "max", 64)?;
            particles_find_neighbors(self.entity, grid, radius, max).map_err(rt)
        } else if name.eq_ignore_ascii_case(c::NEIGHBOR) {
            reject_unknown_kwargs(
                kwargs,
                &["a", "out", "grid", "op", "radius", "falloff", "relative"],
            )?;
            let op = kw_op(kwargs, NEIGHBOR_MEAN, parse_neighbor_op)?;
            let relative = match kw(kwargs, "relative") {
                Some(v) => v.extract::<bool>()?,
                None => false,
            };
            // `grid=` searches exactly; otherwise read the lists from find_neighbors
            let grid = match kw(kwargs, "grid") {
                Some(g) => {
                    if relative {
                        return Err(PyValueError::new_err(
                            "apply(neighbor): relative=True reads the neighbor lists, so drop grid=",
                        ));
                    }
                    let g = g.extract::<PyRef<Grid>>().map_err(|_| {
                        PyTypeError::new_err("apply(neighbor): `grid` must be a Grid")
                    })?;
                    Some(g.entity)
                }
                None => None,
            };
            let (falloff_default, radius) = match grid {
                Some(g) => {
                    let cell = grid_get(g).map_err(rt)?.params.cell_size;
                    (
                        FALLOFF_SMOOTHSTEP,
                        kw_f32(kwargs, "radius", cell)?.min(cell),
                    )
                }
                None => (FALLOFF_CONST, kw_f32(kwargs, "radius", f32::MAX)?),
            };
            let falloff = match kw(kwargs, "falloff") {
                Some(v) => parse_falloff(&v.extract::<String>()?)?,
                None => falloff_default,
            };

            let (a, components) = if op == NEIGHBOR_COUNT {
                let a = match kw(kwargs, "a") {
                    Some(_) => self.operand(kwargs, "a")?.0,
                    None => {
                        let pos = Self::builtin_attribute("position")
                            .expect("position is a built-in")
                            .entity;
                        particles_ensure_attribute(self.entity, pos).map_err(rt)?
                    }
                };
                (a, 1u32)
            } else {
                self.operand(kwargs, "a")?
            };
            let out = self.output(kwargs, components)?;
            match grid {
                Some(grid) => {
                    particles_gather(self.entity, grid, a, out, op, radius, falloff, components)
                }
                None => particles_neighbor_reduce(
                    self.entity,
                    a,
                    out,
                    op,
                    radius,
                    falloff,
                    components,
                    relative,
                ),
            }
            .map_err(rt)
        } else {
            Err(PyValueError::new_err(format!(
                "apply(): unknown operation {name:?}"
            )))
        }
    }
}

#[pymethods]
impl Particles {
    /// Opaque id for this object.
    pub fn id(&self) -> u64 {
        self.entity.to_bits()
    }

    #[getter]
    pub fn capacity(&self) -> PyResult<u32> {
        particles_capacity(self.entity).map_err(|e| PyRuntimeError::new_err(format!("{e}")))
    }

    #[pyo3(signature = (attribute, default=None))]
    pub fn add_attribute(
        &mut self,
        attribute: PyRef<Attribute>,
        default: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<()> {
        let default_value = default
            .map(crate::material::py_to_shader_value)
            .transpose()?;
        particles_attribute_add(self.entity, attribute.entity, default_value)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        let (name, fmt) = geometry_attribute_info(attribute.entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        self.name_to_attr
            .insert(name, (attribute.entity, AttributeFormat::from_inner(fmt)));
        Ok(())
    }

    /// The GPU buffer for an attribute, materialized on demand. `attribute` is a
    /// built-in name (`"position"`, `"velocity"`, `"color"`, `"scale"`, `"life"`,
    /// `"age"`, `"normal"`, `"uv"`, `"rotation"`), the name of any created
    /// attribute, or an `Attribute`.
    pub fn buffer(slf: &Bound<'_, Self>, attribute: &Bound<'_, PyAny>) -> PyResult<Buffer> {
        let this = slf.borrow();
        let attr_entity = this.resolve_attribute(attribute)?;
        let buf = particles_ensure_attribute(this.entity, attr_entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        let (_, fmt) = geometry_attribute_info(attr_entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        let element_type = match AttributeFormat::from_inner(fmt) {
            AttributeFormat::Float => shader_value::ShaderValue::Float(0.0),
            AttributeFormat::Float2 => shader_value::ShaderValue::Float2([0.0; 2]),
            AttributeFormat::Float3 => shader_value::ShaderValue::Float3([0.0; 3]),
            AttributeFormat::Float4 => shader_value::ShaderValue::Float4([0.0; 4]),
        };
        Ok(
            Buffer::from_entity(buf, Some(element_type))
                .with_owner(slf.clone().into_any().unbind()),
        )
    }

    pub fn index_buffer(slf: &Bound<'_, Self>, index_count: u32) -> PyResult<Buffer> {
        let entity = particles_set_connectivity(slf.borrow().entity, index_count)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        Ok(
            Buffer::from_entity(entity, Some(shader_value::ShaderValue::UInt(0)))
                .with_owner(slf.clone().into_any().unbind()),
        )
    }

    pub fn draw_args(slf: &Bound<'_, Self>) -> PyResult<Buffer> {
        let entity = particles_connectivity_indirect(slf.borrow().entity)
            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
        Ok(
            Buffer::from_entity(entity, Some(shader_value::ShaderValue::UInt(0)))
                .with_owner(slf.clone().into_any().unbind()),
        )
    }

    pub fn reset_indices(&self) -> PyResult<()> {
        particles_reset_indices(self.entity).map_err(|e| PyRuntimeError::new_err(format!("{e}")))
    }

    /// Run an operation: operands positionally, `out=` and options as keywords, e.g.
    /// `p.apply(COMBINE, "velocity", "force", op="add")`.
    #[pyo3(signature = (kind, *operands, **kwargs))]
    pub fn apply(
        &self,
        kind: &Bound<'_, PyAny>,
        operands: &Bound<'_, PyTuple>,
        kwargs: Option<&Bound<'_, PyDict>>,
    ) -> PyResult<()> {
        if let Ok(compute) = kind.extract::<PyRef<Compute>>() {
            if !operands.is_empty() {
                return Err(PyTypeError::new_err(
                    "apply(compute): pass kernel parameters as keywords",
                ));
            }
            if let Some(kwargs) = kwargs {
                compute.set(Some(kwargs))?;
            }
            return particles_apply(self.entity, compute.entity)
                .map_err(|e| PyRuntimeError::new_err(format!("{e}")));
        }
        let name: String = kind.extract().map_err(|_| {
            PyTypeError::new_err(
                "apply(): first argument must be an operation constant or a Compute",
            )
        })?;
        let kwargs = with_operands(&name, operands, kwargs)?;
        self.apply_named(&name, kwargs.as_ref())
    }

    #[pyo3(signature = (n, **kwargs))]
    pub fn emit(&self, n: u32, kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<()> {
        let Some(kwargs) = kwargs else {
            return particles_emit(self.entity, n, vec![])
                .map_err(|e| PyRuntimeError::new_err(format!("{e}")));
        };
        let mut data: Vec<(Entity, Vec<u8>)> = Vec::new();