dirt_clump 0.1.4

//! Multisphere/clump rigid body composites for non-spherical DEM particles.
//!
//! A **clump** is a rigid body composed of multiple overlapping spheres. Each sphere
//! participates in normal contact detection, but forces are aggregated to its
//! [`MultisphereBody`]. The body integrates translational + rotational (Euler equation)
//! dynamics, then sub-sphere positions/velocities are derived from body state.
//!
//! # Architecture
//!
//! - **No phantom parent atom.** Body data lives in [`MultisphereBodyStore`], not in
//!   the atom arrays. Each sub-sphere atom references its body via `body_id` in
//!   [`ClumpAtom`].
//! - **Full inertia tensor.** Analytically via parallel axis theorem, or Monte Carlo
//!   for overlapping spheres. Diagonalized to principal moments + axes quaternion.
//! - **Euler equation integration.** Torque → principal frame, Euler α, half-kick ω.
//!
//! # Contact-exclusion contract
//!
//! Sub-spheres of the *same* body must never push on each other — they are held
//! rigid, so an intra-body "contact" would be a spurious internal force. This
//! crate does **not** silently filter the neighbor list; instead it exposes
//! [`same_body`], and every contact-force consumer is responsible for calling it
//! to skip same-body pairs:
//!
//! ```ignore
//! // inside a pair loop over (i, j) from the neighbor list:
//! if same_body(&clump_data, i, j) {
//!     continue; // sub-spheres of one rigid body do not interact
//! }
//! ```
//!
//! [`same_body`] returns `true` only when both atoms carry a non-zero `body_id`
//! and those ids match. The granular contact kernels in `dirt_granular` already
//! honor this contract; any custom force plugin that iterates the neighbor list
//! must do the same, or rigid bodies will self-repel and fly apart.
//!
//! # Rigid-body integration
//!
//! Each body is advanced with a velocity-Verlet scheme that follows LIGGGHTS's
//! `FixMultisphere` (helper names are mirrored: [`body::angmom_to_omega`],
//! `richardson`, `vecquat`). Two design choices are worth calling out:
//!
//! - **Angular momentum is the integrated state.** The body stores `angmom`
//!   (space-frame angular momentum) and treats it as the primary rotational
//!   variable, half-kicked by the torque each step (`L += ½ dt τ`). Angular
//!   velocity `omega` is *derived* from `angmom` on demand via
//!   [`body::angmom_to_omega`], never integrated directly. This keeps the update
//!   stable for the asymmetric inertia tensors of non-spherical clumps, where
//!   integrating `omega` through Euler's equation directly is fragile.
//! - **Richardson-extrapolated quaternion update.** The orientation quaternion
//!   is advanced by a private `richardson` routine: it takes one full step and
//!   two half-steps (recomputing `omega` at the half-step orientation) and
//!   extrapolates `q ← 2 q_half − q_full` for second-order accuracy, then
//!   renormalizes.
//!
//! ## Two quaternions
//!
//! A body carries **two** orientation quaternions, and the distinction matters:
//!
//! - `quaternion` — **body → space**: the live orientation of the body frame in
//!   the world, updated every step by the Richardson integrator.
//! - `principal_axes` — **body → principal**: a *fixed* rotation (set at
//!   creation) from the body frame to the frame in which the inertia tensor is
//!   diagonal (`principal_moments`).
//!
//! Inertia work is only simple in the principal frame, so wherever a
//! principal-frame mapping is needed the two are composed:
//! `q_principal→space = quaternion * principal_axes` (Hamilton product). That
//! composed quaternion is what [`body::angmom_to_omega`] expects — it rotates
//! `angmom` into the principal frame, divides by `principal_moments`, and
//! rotates the resulting `omega` back to space.
//!
//! # Inertia computation
//!
//! At body creation the inertia tensor is built one of two ways, auto-selected
//! by [`body::has_overlap`]:
//!
//! - **Non-overlapping spheres** → [`compute_inertia_tensor_analytical`]
//!   (exact, parallel-axis theorem).
//! - **Overlapping spheres** → [`compute_inertia_tensor_montecarlo`] with a
//!   **hardcoded 100 000 samples**. Monte Carlo integration over the bounding
//!   box double-counts overlap volume correctly. The default sampler derives a
//!   deterministic seed from the clump definition, density, and sample count, so
//!   repeated runs of the same clump definition get the same estimate; callers
//!   that need their own bounded stream can use
//!   [`compute_inertia_tensor_montecarlo_seeded`]. Sampling error remains the
//!   usual Monte Carlo error (roughly a few percent at this sample count). Both
//!   paths are then diagonalized into `principal_moments` + `principal_axes`.
//!
//! The scalar helper [`compute_clump_inertia`] is **legacy** — it returns only a
//! single averaged moment (the trace / 3) and is kept for backward
//! compatibility. New code should use the full-tensor functions above.
//!
//! # Configuration
//!
//! Clump definitions and insertion live under the `[clump]` TOML section
//! (separate from `[dem]` which uses `deny_unknown_fields`).
//!
//! ```toml
//! [[clump.definitions]]
//! name = "dimer"
//! spheres = [
//!     { offset = [-0.0003, 0.0, 0.0], radius = 0.001 },
//!     { offset = [0.0003, 0.0, 0.0], radius = 0.001 },
//! ]
//!
//! [[clump.insert]]
//! definition = "dimer"
//! count = 100
//! density = 2500.0
//! material = "glass"
//! velocity = 0.5
//! region = { type = "block", min = [0.001, 0.001, 0.001], max = [0.019, 0.019, 0.019] }
//! ```

// Public API documentation-completeness gate: every public item in this crate
// must carry a doc comment. Enforced on both `cargo build` (rustc) and
// `cargo doc` (rustdoc; e.g. `RUSTDOCFLAGS="-D missing_docs"`). Document real
// API intent here — do not add empty doc comments just to satisfy the lint.
#![deny(missing_docs)]

use std::collections::HashMap;
use std::f64::consts::PI;

use grass_app::prelude::*;
use grass_scheduler::prelude::*;
use rand::{rngs::StdRng, Rng, SeedableRng};
use serde::Deserialize;
use soil_derive::AtomData;

use soil_core::{
    register_atom_data, Accum, Atom, AtomData, AtomDataRegistry, CommResource, Config, Domain,
    Optional, ParticleSimScheduleSet, ParticleStore, ParticleStoreError, ParticlesWith, Read, Real,
    Region, RunState, ScheduleSetupSet, Write, EXCHANGE, REVERSE_SEND_FORCE,
};

#[cfg(feature = "mpi_backend")]
use soil_core::CommTopology;

use dirt_atom::DemAtom;
use dirt_schedule::{
    CLUMP_EXCHANGE, CLUMP_FINAL_INTEGRATION, CLUMP_FORCE_AGGREGATION, CLUMP_GHOST_CUTOFF,
    CLUMP_INITIAL_INTEGRATION, CLUMP_INSERT, CLUMP_LOST_ATOM_CHECK, CLUMP_PBC,
    CLUMP_POSITION_UPDATE, CLUMP_PRE_EXCHANGE_UPDATE, CLUMP_REMAP, CLUMP_RESTORE, CLUMP_SNAP,
    CONTACT_FORCE,
};

pub mod body;
pub use body::{
    compute_inertia_tensor_analytical, compute_inertia_tensor_montecarlo,
    compute_inertia_tensor_montecarlo_seeded, diagonalize_inertia, has_overlap,
    jacobi_eigendecomposition, rotation_matrix_to_quaternion, MultisphereBody,
    MultisphereBodyStore,
};

// ── Configuration ────────────────────────────────────────────────────────────

/// A single sphere within a clump definition.
#[derive(Deserialize, Clone, Debug)]
#[serde(deny_unknown_fields)]
pub struct ClumpSphereConfig {
    /// Offset from clump center of mass in body frame [x, y, z].
    pub offset: [f64; 3],
    /// Sphere radius.
    pub radius: f64,
}

/// A clump type definition from `[[clump.definitions]]`.
#[derive(Deserialize, Clone, Debug)]
#[serde(deny_unknown_fields)]
pub struct ClumpDef {
    /// Name of this clump type, referenced by insertion config.
    pub name: String,
    /// Spheres composing this clump (positions relative to COM).
    pub spheres: Vec<ClumpSphereConfig>,
}

/// A clump insertion block from `[[clump.insert]]`.
#[derive(Deserialize, Clone, Debug)]
pub struct ClumpInsertConfig {
    /// Name of the clump definition to insert.
    pub definition: String,
    /// Number of clumps to insert.
    pub count: u32,
    /// Particle density (kg/m³).
    pub density: f64,
    /// Material name (must match a `[[dem.materials]]` entry).
    pub material: String,
    /// Random velocity magnitude (m/s). Each component is uniform in [-v, +v].
    #[serde(default)]
    pub velocity: Option<f64>,
    /// Insertion region. Defaults to domain bounds inset by effective clump radius.
    #[serde(default)]
    pub region: Option<Region>,
    /// When `true`, each inserted clump is given an independent, uniformly random
    /// 3-D orientation (its definition offsets are rotated by a random unit
    /// quaternion). Default `false` — all clumps share the definition orientation.
    /// Randomising is essential for shape-anisotropic bodies (e.g. hooking
    /// "monkeys"): identically-oriented anisotropic clumps interlock coherently
    /// and jam/overlap on compaction, whereas random orientations pack naturally.
    #[serde(default)]
    pub random_orientation: bool,
    /// Seed for deterministic clump insertion. Defaults to 0.
    ///
    /// The seed drives candidate positions, optional random orientations, and
    /// random velocities so repeated runs with the same config produce the same
    /// inserted clump state.
    #[serde(default)]
    pub seed: Option<u64>,
}

/// TOML `[clump]` — top-level clump configuration.
///
/// Separate from `[dem]` because `DemConfig` uses `deny_unknown_fields`.
#[derive(Deserialize, Clone, Default)]
pub struct ClumpTopConfig {
    /// Clump shape definitions (`[[clump.definitions]]`), each a named
    /// multisphere template referenced by insertion commands.
    #[serde(default)]
    pub definitions: Option<Vec<ClumpDef>>,
    /// Clump insertion commands (`[[clump.insert]]`) that place instances of a
    /// named definition into the domain.
    #[serde(default)]
    pub insert: Option<Vec<ClumpInsertConfig>>,
}

// ── Per-atom clump data ─────────────────────────────────────────────────────

/// Per-atom clump membership and body-frame offset data.
///
/// Every atom gets these fields. For atoms not in a clump, `body_id` is 0.
/// Sub-spheres store their body-frame offset for position reconstruction.
#[derive(AtomData)]
pub struct ClumpAtom {
    /// Body ID this atom belongs to (0 = not in a clump).
    /// Encoded as f64 for AtomData compatibility; use as u32.
    #[forward]
    pub body_id: Vec<f64>,

    /// Local offset from body COM in body frame [x, y, z].
    #[forward]
    pub body_offset: Vec<[f64; 3]>,
}

impl Default for ClumpAtom {
    fn default() -> Self {
        Self::new()
    }
}

impl ClumpAtom {
    /// Creates an empty `ClumpAtom` column with no per-atom entries.
    pub fn new() -> Self {
        ClumpAtom {
            body_id: Vec::new(),
            body_offset: Vec::new(),
        }
    }
}

// ── Clump registry (runtime data) ──────────────────────────────────────────

/// Runtime storage for clump definitions, looked up during insertion.
pub struct ClumpRegistry {
    /// All registered clump definitions, indexed by insertion-time name lookup.
    pub defs: Vec<ClumpDef>,
}

impl ClumpRegistry {
    /// Creates an empty registry with no clump definitions.
    pub fn new() -> Self {
        ClumpRegistry { defs: Vec::new() }
    }

    /// Returns the clump definition with the given `name`, or `None` if no
    /// definition by that name has been registered.
    pub fn find(&self, name: &str) -> Option<&ClumpDef> {
        self.defs.iter().find(|d| d.name == name)
    }
}

impl Default for ClumpRegistry {
    fn default() -> Self {
        Self::new()
    }
}

// ── Quaternion utilities ─────────────────────────────────────────────────────

/// Rotate a vector by a quaternion q = [w, x, y, z].
#[inline]
pub fn quat_rotate(q: [f64; 4], v: [f64; 3]) -> [f64; 3] {
    let w = q[0];
    let qx = q[1];
    let qy = q[2];
    let qz = q[3];

    let cx = qy * v[2] - qz * v[1];
    let cy = qz * v[0] - qx * v[2];
    let cz = qx * v[1] - qy * v[0];

    [
        v[0] + 2.0 * (w * cx + qy * cz - qz * cy),
        v[1] + 2.0 * (w * cy + qz * cx - qx * cz),
        v[2] + 2.0 * (w * cz + qx * cy - qy * cx),
    ]
}

/// Cross product of two 3-vectors.
#[inline]
pub fn cross(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
    [
        a[1] * b[2] - a[2] * b[1],
        a[2] * b[0] - a[0] * b[2],
        a[0] * b[1] - a[1] * b[0],
    ]
}

// ── Legacy scalar inertia (kept for backward compatibility) ─────────────────

/// **Legacy.** Compute a single scalar moment of inertia for a clump (the
/// average of the diagonal, i.e. trace / 3 — a spherical approximation).
///
/// Kept only for backward compatibility. It discards the off-diagonal coupling
/// and anisotropy that non-spherical clumps actually have, so the integrator
/// does **not** use it; the body-creation path uses the full tensor from
/// [`compute_inertia_tensor_analytical`] / [`compute_inertia_tensor_montecarlo`]
/// instead. Prefer those for any new code.
pub fn compute_clump_inertia(spheres: &[ClumpSphereConfig], density: f64) -> (f64, f64) {
    let (mass, tensor) = compute_inertia_tensor_analytical(spheres, density);
    let avg = (tensor[0][0] + tensor[1][1] + tensor[2][2]) / 3.0;
    (mass, avg)
}

// ── Plugin ──────────────────────────────────────────────────────────────────

/// Plugin that adds multisphere/clump rigid body support to DIRT.
///
/// Registers:
/// - [`ClumpAtom`] per-atom data (body_id, body_offset)
/// - [`ClumpRegistry`] resource with clump definitions from config
/// - [`MultisphereBodyStore`] resource for rigid body state
/// - Body integration systems (Euler equations)
/// - Force aggregation + position update systems
pub struct ClumpPlugin;

impl Plugin for ClumpPlugin {
    fn dependencies(&self) -> Vec<std::any::TypeId> {
        grass_app::type_ids![dirt_atom::DemAtomPlugin]
    }

    fn build(&self, app: &mut App) {
        register_atom_data!(app, ClumpAtom::new());

        let mut registry = ClumpRegistry::new();
        let clump_config = Config::load::<ClumpTopConfig>(app, "clump");
        if let Some(defs) = clump_config.definitions {
            for def in defs {
                assert!(
                    !def.spheres.is_empty(),
                    "Clump '{}' must have at least one sphere",
                    def.name
                );
                registry.defs.push(def);
            }
        }

        app.add_resource(registry);
        app.add_resource(MultisphereBodyStore::new());

        // Clump insertion from [[clump.insert]] config (runs after normal particle insertion)
        app.add_setup_system(
            clump_insert_atoms.label(CLUMP_INSERT),
            ScheduleSetupSet::Setup,
        );

        // Set minimum ghost cutoff for clumps BEFORE neighbor_setup computes bins.
        // neighbor_setup will use max(its_value, domain.ghost_cutoff).
        app.add_setup_system(
            extend_ghost_cutoff_for_clumps.label(CLUMP_GHOST_CUTOFF),
            ScheduleSetupSet::Setup,
        );

        // Before any exchange, snap sub-sphere positions to body COM so they
        // always migrate to the same rank as their body. Must run before
        // exchange_bodies so all bodies are still local for the lookup.
        app.add_update_system(
            snap_subspheres_to_body_com
                .label(CLUMP_SNAP)
                .before(CLUMP_EXCHANGE)
                .before(EXCHANGE),
            ParticleSimScheduleSet::Exchange,
        );

        // Body exchange: migrate bodies whose COM left the local subdomain.
        app.add_update_system(
            exchange_bodies.label(CLUMP_EXCHANGE).before(EXCHANGE),
            ParticleSimScheduleSet::Exchange,
        );

        // After atom exchange, restore sub-sphere positions from body state
        // (undo the snap-to-COM done before exchange).
        app.add_update_system(
            restore_subsphere_positions
                .label(CLUMP_RESTORE)
                .after(EXCHANGE),
            ParticleSimScheduleSet::Exchange,
        );

        // Affinely remap body COMs when the box RESIZES (e.g. `[deform]` erate
        // compression). `apply_deform` remaps the sub-sphere *atoms* and the box,
        // but a rigid body is driven by its COM (the sub-spheres are re-rigidified
        // from it each step), so without this the clumps ignore the compaction and
        // the shrinking box overruns them into an overlap blow-up. A no-op for pure
        // Lees–Edwards shear (tilt only, size unchanged) and for undeformed runs.
        app.add_resource(ClumpBoxState::default());
        app.add_update_system(
            remap_bodies_on_box_resize
                .label(CLUMP_REMAP)
                .before(CLUMP_INITIAL_INTEGRATION),
            ParticleSimScheduleSet::InitialIntegration,
        );

        // Body initial integration (half-kick + drift + quaternion update)
        app.add_update_system(
            integrate_bodies_initial.label(CLUMP_INITIAL_INTEGRATION),
            ParticleSimScheduleSet::InitialIntegration,
        );

        // PBC for body COM
        app.add_update_system(
            pbc_multisphere_bodies.label(CLUMP_PBC),
            ParticleSimScheduleSet::PostInitialIntegration,
        );

        // Force aggregation: sub-sphere forces → body force/torque
        // Must run after reverse_send_force so ghost sub-sphere forces are
        // accumulated back to their owning atoms before aggregation to bodies.
        app.add_update_system(
            aggregate_clump_forces
                .label(CLUMP_FORCE_AGGREGATION)
                .after(CONTACT_FORCE)
                .after(REVERSE_SEND_FORCE),
            ParticleSimScheduleSet::PostForce,
        );

        // Body final integration (half-kick after new forces)
        app.add_update_system(
            integrate_bodies_final.label(CLUMP_FINAL_INTEGRATION),
            ParticleSimScheduleSet::FinalIntegration,
        );

        // Update sub-sphere positions before exchange so atoms migrate to the
        // same rank as their body (prevents orphan sub-spheres whose forces
        // would have no local body to aggregate to).
        app.add_update_system(
            update_clump_positions
                .label(CLUMP_PRE_EXCHANGE_UPDATE)
                .after(CLUMP_PBC),
            ParticleSimScheduleSet::PostInitialIntegration,
        );

        // Derive sub-sphere pos/vel from body state (end of step)
        app.add_update_system(
            update_clump_positions.label(CLUMP_POSITION_UPDATE),
            ParticleSimScheduleSet::PostFinalIntegration,
        );

        // Lost atom detection (lightweight, every 1000 steps)
        app.add_update_system(
            check_lost_clump_atoms
                .label(CLUMP_LOST_ATOM_CHECK)
                .after(CLUMP_POSITION_UPDATE),
            ParticleSimScheduleSet::PostFinalIntegration,
        );
    }

    fn try_build(&self, app: &mut App) -> Result<(), AppError> {
        validate_clump_config(app)?;
        self.build(app);
        Ok(())
    }
}

fn validate_clump_config(app: &mut App) -> Result<(), AppError> {
    let config = Config::try_load::<ClumpTopConfig>(app, "clump")
        .map_err(|error| AppError::message(error.to_string()))?;
    let materials = app
        .get_resource_ref::<dirt_atom::MaterialTable>()
        .ok_or_else(|| AppError::message("ClumpPlugin requires DemAtomPlugin"))?;
    let defs = config.definitions.as_deref().unwrap_or_default();
    for def in defs {
        if def.spheres.is_empty() {
            return Err(AppError::message(format!(
                "Clump '{}' must have at least one sphere",
                def.name
            )));
        }
    }
    for insert in config.insert.as_deref().unwrap_or_default() {
        if !defs.iter().any(|def| def.name == insert.definition) {
            return Err(AppError::message(format!(
                "Clump definition '{}' not found",
                insert.definition
            )));
        }
        if !materials.names.iter().any(|name| name == &insert.material) {
            return Err(AppError::message(format!(
                "Material '{}' not found in [[dem.materials]]",
                insert.material
            )));
        }
        if let Some(region) = &insert.region {
            validate_clump_insertion_region(region)?;
        }
    }
    Ok(())
}

fn validate_clump_insertion_region(region: &Region) -> Result<(), AppError> {
    // Do not duplicate SOIL's geometry rules here.  The insertion system uses
    // this exact sampler, whose fallible API validates finite dimensions,
    // positive radii, bounded planes, empty boolean regions, and unsuccessful
    // rejection sampling.  A private, seeded RNG makes this preflight
    // deterministic and deliberately leaves the insertion RNG untouched.
    let mut rng = StdRng::seed_from_u64(0xC1A0_5EED);
    region.random_point_inside(&mut rng).map_err(|error| {
        AppError::message(format!(
            "[[clump.insert]] region cannot be sampled: {error}"
        ))
    })?;
    Ok(())
}

// ── Systems ─────────────────────────────────────────────────────────────────

/// Extend `Domain::ghost_cutoff` by the maximum clump bounding radius.
///
/// Ensures all sub-spheres of bodies near subdomain boundaries are visible as
/// ghosts on the body-owning rank. Mirrors LIGGGHTS's `extend_cut_ghost()`.
fn extend_ghost_cutoff_for_clumps(
    clump_registry: Res<ClumpRegistry>,
    mut domain: ResMut<Domain>,
    comm: Res<CommResource>,
) {
    let mut max_r_bound: f64 = 0.0;
    for def in &clump_registry.defs {
        for sphere in &def.spheres {
            let r =
                (sphere.offset[0].powi(2) + sphere.offset[1].powi(2) + sphere.offset[2].powi(2))
                    .sqrt()
                    + sphere.radius;
            max_r_bound = max_r_bound.max(r);
        }
    }
    if max_r_bound > 0.0 {
        let extension = 2.0 * max_r_bound;
        domain.ghost_cutoff += extension;
        if comm.rank() == 0 {
            println!(
                "ClumpPlugin: extended ghost_cutoff by {:.6} (2 * R_bound={:.6}) → {:.6}",
                extension, max_r_bound, domain.ghost_cutoff
            );
        }
    }
}

/// Snap sub-sphere positions to their body COM before atom exchange.
///
/// This ensures sub-spheres always migrate to the same rank as their body.
/// Without this, a sub-sphere near a subdomain boundary could exchange to a
/// different rank than its body, causing orphaned forces.
fn snap_subspheres_to_body_com(
    mut atoms: ResMut<Atom>,
    bodies: Res<MultisphereBodyStore>,
    particles: ParticlesWith<'_, Read<ClumpAtom>>,
) {
    particles.with(|clump| {
        let nlocal = atoms.nlocal as usize;
        for i in 0..nlocal {
            if i >= clump.body_id.len() {
                break;
            }
            let bid = clump.body_id[i] as u32;
            if bid == 0 {
                continue;
            }
            if let Some(body_idx) = bodies.map(bid) {
                let com = bodies.bodies[body_idx].com_pos;
                atoms.pos[i] = [com[0] as Real, com[1] as Real, com[2] as Real];
            }
        }
    });
}

/// Restore sub-sphere positions from body state after atom exchange.
///
/// Undoes the snap-to-COM and sets correct offset positions, velocities, and angular velocities.
/// Also regenerates the body ID→index map after body exchange.
fn restore_subsphere_positions(
    mut atoms: ResMut<Atom>,
    mut bodies: ResMut<MultisphereBodyStore>,
    particles: ParticlesWith<'_, (Read<ClumpAtom>, Write<DemAtom>)>,
) {
    bodies.generate_map();
    particles.with(|(clump, mut dem)| {
        let nlocal = atoms.nlocal as usize;
        for i in 0..nlocal {
            if i >= clump.body_id.len() {
                break;
            }
            let bid = clump.body_id[i] as u32;
            if bid == 0 {
                continue;
            }
            if let Some(body_idx) = bodies.map(bid) {
                let body = &bodies.bodies[body_idx];
                let rotated = quat_rotate(body.quaternion, clump.body_offset[i]);
                atoms.pos[i] = [
                    (body.com_pos[0] + rotated[0]) as Real,
                    (body.com_pos[1] + rotated[1]) as Real,
                    (body.com_pos[2] + rotated[2]) as Real,
                ];
                let omega_cross_r = cross(body.omega, rotated);
                atoms.vel[i] = [
                    (body.com_vel[0] + omega_cross_r[0]) as Real,
                    (body.com_vel[1] + omega_cross_r[1]) as Real,
                    (body.com_vel[2] + omega_cross_r[2]) as Real,
                ];
                dem.omega[i] = body.omega;
            }
        }
    });
}

/// Initial half-step: integrate all rigid bodies (Euler equations).
fn integrate_bodies_initial(atoms: Res<Atom>, mut bodies: ResMut<MultisphereBodyStore>) {
    let dt = atoms.dt;
    for body in &mut bodies.bodies {
        body::integrate_body_initial(body, dt);
    }
}

/// Final half-step: integrate all rigid bodies after new forces.
fn integrate_bodies_final(atoms: Res<Atom>, mut bodies: ResMut<MultisphereBodyStore>) {
    let dt = atoms.dt;
    for body in &mut bodies.bodies {
        body::integrate_body_final(body, dt);
    }
}

/// Tracks the box bounds seen on the previous step so [`remap_bodies_on_box_resize`]
/// can detect a resize and apply the matching affine COM transform.
#[derive(Default)]
pub struct ClumpBoxState {
    prev_low: [f64; 3],
    prev_size: [f64; 3],
    initialized: bool,
}

/// Affinely remap rigid-body COMs when the simulation box is resized by
/// `[deform]` (erate/vel/final). `apply_deform` moves the sub-sphere atoms and
/// the box each step, but a clump's motion is carried by its COM (the sub-spheres
/// are re-rigidified from `COM + R·offset` afterwards), so the box change must be
/// applied to the COM too — using the *same* centered-affine map as the atom
/// remap: `com' = new_center + (com − old_center)·(new_size/old_size)`.
///
/// Size-preserving deformation (Lees–Edwards xy shear tilts the box but keeps each
/// edge length) yields `scale = 1` on every axis → exact no-op, so shear and
/// undeformed runs are untouched (no numerical drift).
fn remap_bodies_on_box_resize(
    mut bodies: ResMut<MultisphereBodyStore>,
    domain: Res<Domain>,
    mut state: ResMut<ClumpBoxState>,
) {
    let new_low = domain.boundaries_low;
    let new_size = domain.size;
    if !state.initialized {
        state.prev_low = new_low;
        state.prev_size = new_size;
        state.initialized = true;
        return;
    }
    for d in 0..3 {
        let old_size = state.prev_size[d];
        if old_size <= 0.0 || (new_size[d] - old_size).abs() <= 1e-15 * old_size {
            continue; // axis unchanged (shear / no deform)
        }
        let old_center = state.prev_low[d] + 0.5 * old_size;
        let new_center = new_low[d] + 0.5 * new_size[d];
        let scale = new_size[d] / old_size;
        for body in &mut bodies.bodies {
            body.com_pos[d] = new_center + (body.com_pos[d] - old_center) * scale;
        }
    }
    state.prev_low = new_low;
    state.prev_size = new_size;
}

/// Wrap body COM through periodic boundaries and update body image flags.
///
/// Triclinic (Lees–Edwards) boxes wrap the COM in fractional (lamda) coordinates —
/// so a gradient (y) crossing picks up the box x-tilt automatically — and apply the
/// streaming-velocity remap to the COM velocity only (a uniform translation of a
/// rigid body, adding no spurious spin; the body wraps as a unit so it never tears).
fn pbc_multisphere_bodies(mut bodies: ResMut<MultisphereBodyStore>, domain: Res<Domain>) {
    if domain.triclinic {
        let periodic = domain.periodic_flags();
        let bvel = domain.boundary_vel;
        for body in &mut bodies.bodies {
            let mut lam = domain.x2lamda(body.com_pos);
            let mut dy = 0i32;
            for d in 0..3 {
                if periodic[d] {
                    if lam[d] < 0.0 {
                        lam[d] += 1.0;
                        body.image[d] -= 1;
                        if d == 1 {
                            dy -= 1;
                        }
                    } else if lam[d] >= 1.0 {
                        lam[d] -= 1.0;
                        body.image[d] += 1;
                        if d == 1 {
                            dy += 1;
                        }
                    }
                }
            }
            body.com_pos = domain.lamda2x(lam);
            if dy != 0 {
                let s = dy as f64;
                body.com_vel[0] -= s * bvel[0];
                body.com_vel[1] -= s * bvel[1];
                body.com_vel[2] -= s * bvel[2];
            }
        }
        return;
    }

    for body in &mut bodies.bodies {
        for d in 0..3 {
            if domain.is_periodic(d) {
                let low = domain.boundaries_low[d];
                let size = domain.size[d];
                let high = low + size;
                if body.com_pos[d] < low {
                    body.com_pos[d] += size;
                    body.image[d] -= 1;
                } else if body.com_pos[d] >= high {
                    body.com_pos[d] -= size;
                    body.image[d] += 1;
                }
            }
        }
    }
}

/// Exchange bodies between processors when their COM leaves the local subdomain.
///
/// Mirrors the atom exchange pattern in `comm.rs`: for each dimension, scan bodies
/// whose COM is outside `[sub_domain_low, sub_domain_high)`, pack them, send to the
/// neighbor processor, receive incoming bodies, and rebuild the ID→index map.
#[cfg(feature = "mpi_backend")]
fn exchange_bodies(
    comm: Res<CommResource>,
    topo: Res<CommTopology>,
    mut bodies: ResMut<MultisphereBodyStore>,
    domain: Res<Domain>,
) {
    let decomp = comm.processor_decomposition();

    let mut lo_buf: Vec<f64> = Vec::new();
    let mut hi_buf: Vec<f64> = Vec::new();

    for dim in 0..3usize {
        if decomp[dim] == 1 {
            continue;
        }

        let lo_proc = topo.swap_directions[0][dim];
        let hi_proc = topo.swap_directions[1][dim];

        lo_buf.clear();
        hi_buf.clear();
        let mut lo_count = 0u32;
        let mut hi_count = 0u32;

        // Scan bodies in reverse, pack those with COM outside subdomain.
        // Triclinic: classify the COM in fractional (lamda) coordinates against the
        // lamda subdomain bounds (the box is a unit cube there).
        let triclinic = domain.triclinic;
        let (sub_lo, sub_hi) = if triclinic {
            (domain.sub_lamda_low[dim], domain.sub_lamda_high[dim])
        } else {
            (domain.sub_domain_low[dim], domain.sub_domain_high[dim])
        };
        for i in (0..bodies.bodies.len()).rev() {
            let pos = if triclinic {
                domain.x2lamda(bodies.bodies[i].com_pos)[dim]
            } else {
                bodies.bodies[i].com_pos[dim]
            };
            if pos < sub_lo {
                lo_count += 1;
                bodies.bodies[i].pack(&mut lo_buf);
                bodies.bodies.swap_remove(i);
            } else if pos >= sub_hi {
                hi_count += 1;
                bodies.bodies[i].pack(&mut hi_buf);
                bodies.bodies.swap_remove(i);
            }
        }

        lo_buf.push(lo_count as f64);
        hi_buf.push(hi_count as f64);

        // Send lo, receive from hi
        if lo_proc != -1 && hi_proc != -1 {
            let msg = comm.sendrecv_f64(lo_proc, &lo_buf, hi_proc);
            unpack_bodies_from_msg(&msg, &mut bodies.bodies);
        } else if lo_proc != -1 {
            comm.send_f64(lo_proc, &lo_buf);
        } else if hi_proc != -1 {
            let msg = comm.recv_f64(hi_proc);
            unpack_bodies_from_msg(&msg, &mut bodies.bodies);
        }

        // Send hi, receive from lo
        if hi_proc != -1 && lo_proc != -1 {
            let msg = comm.sendrecv_f64(hi_proc, &hi_buf, lo_proc);
            unpack_bodies_from_msg(&msg, &mut bodies.bodies);
        } else if hi_proc != -1 {
            comm.send_f64(hi_proc, &hi_buf);
        } else if lo_proc != -1 {
            let msg = comm.recv_f64(lo_proc);
            unpack_bodies_from_msg(&msg, &mut bodies.bodies);
        }
    }

    bodies.generate_map();
}

/// Unpack bodies from a received message buffer.
#[cfg(feature = "mpi_backend")]
fn unpack_bodies_from_msg(msg: &[f64], bodies: &mut Vec<MultisphereBody>) {
    let count = msg[msg.len() - 1] as usize;
    let data = &msg[..msg.len() - 1];
    let mut pos = 0;
    for _ in 0..count {
        let (body, consumed) = MultisphereBody::unpack(&data[pos..]);
        bodies.push(body);
        pos += consumed;
    }
}

/// No-op body exchange for single-process builds.
#[cfg(not(feature = "mpi_backend"))]
fn exchange_bodies() {}

/// Aggregate forces from sub-sphere atoms onto their parent body.
///
/// For each sub-sphere with `body_id > 0`:
/// - Accumulate force onto body
/// - Compute torque: `r × F` where `r` is the rotated body offset
/// - Accumulate sub-sphere contact torque onto body
/// - Zero sub-sphere force and torque
pub fn aggregate_clump_forces(
    mut atoms: ResMut<Atom>,
    mut bodies: ResMut<MultisphereBodyStore>,
    particles: ParticlesWith<'_, (Write<DemAtom>, Optional<Read<ClumpAtom>>)>,
) {
    particles.with(|(mut dem, clump)| {
        let clump = match clump {
            Some(c) => c,
            None => return,
        };

        // Zero body accumulators
        for body in &mut bodies.bodies {
            body.zero_accumulators();
        }

        let nlocal = atoms.nlocal as usize;

        // Collect contributions to avoid borrow conflicts
        struct Contrib {
            body_idx: usize,
            force: [f64; 3],
            torque: [f64; 3],
            atom_idx: usize,
        }

        let mut contribs = Vec::new();

        for i in 0..nlocal {
            if i >= clump.body_id.len() {
                break;
            }
            let bid = clump.body_id[i] as u32;
            if bid == 0 {
                continue;
            }

            let body_idx = match bodies.map(bid) {
                Some(idx) => idx,
                None => continue,
            };

            let body = &bodies.bodies[body_idx];

            // r = rotated body_offset (current space-frame displacement)
            let rotated = quat_rotate(body.quaternion, clump.body_offset[i]);

            let f_raw = atoms.force[i];
            let f = [f_raw[0] as f64, f_raw[1] as f64, f_raw[2] as f64];
            let torque_from_force = cross(rotated, f);

            let sub_torque = if i < dem.torque.len() {
                dem.torque[i]
            } else {
                [0.0; 3]
            };

            contribs.push(Contrib {
                body_idx,
                force: f,
                torque: [
                    torque_from_force[0] + sub_torque[0],
                    torque_from_force[1] + sub_torque[1],
                    torque_from_force[2] + sub_torque[2],
                ],
                atom_idx: i,
            });
        }

        // Apply contributions to bodies
        for c in &contribs {
            let body = &mut bodies.bodies[c.body_idx];
            for d in 0..3 {
                body.force[d] += c.force[d];
                body.torque[d] += c.torque[d];
            }
        }

        // Zero sub-sphere forces and torques
        for c in &contribs {
            atoms.force[c.atom_idx] = [0.0; 3];
            if c.atom_idx < dem.torque.len() {
                dem.torque[c.atom_idx] = [0.0; 3];
            }
        }
    });
}

/// Derive sub-sphere positions, velocities, and angular velocities from body state.
///
/// For each sub-sphere: `pos = COM + q * body_offset`, `vel = COM_vel + omega × (q * offset)`,
/// `omega = body.omega` (rigid body constraint — all sub-spheres share the body angular velocity).
pub fn update_clump_positions(
    mut atoms: ResMut<Atom>,
    bodies: Res<MultisphereBodyStore>,
    particles: ParticlesWith<'_, (Write<DemAtom>, Optional<Read<ClumpAtom>>)>,
) {
    particles.with(|(mut dem, clump)| {
        let clump = match clump {
            Some(c) => c,
            None => return,
        };

        let nlocal = atoms.nlocal as usize;

        struct SubUpdate {
            idx: usize,
            pos: [f64; 3],
            vel: [f64; 3],
            omega: [f64; 3],
        }

        let mut updates: Vec<SubUpdate> = Vec::new();

        for i in 0..nlocal {
            if i >= clump.body_id.len() {
                break;
            }
            let bid = clump.body_id[i] as u32;
            if bid == 0 {
                continue;
            }

            let body_idx = match bodies.map(bid) {
                Some(idx) => idx,
                None => continue,
            };

            let body = &bodies.bodies[body_idx];
            let rotated = quat_rotate(body.quaternion, clump.body_offset[i]);

            let new_pos = [
                body.com_pos[0] + rotated[0],
                body.com_pos[1] + rotated[1],
                body.com_pos[2] + rotated[2],
            ];

            let omega_cross_r = cross(body.omega, rotated);
            let new_vel = [
                body.com_vel[0] + omega_cross_r[0],
                body.com_vel[1] + omega_cross_r[1],
                body.com_vel[2] + omega_cross_r[2],
            ];

            updates.push(SubUpdate {
                idx: i,
                pos: new_pos,
                vel: new_vel,
                omega: body.omega,
            });
        }

        for u in updates {
            atoms.pos[u.idx] = [u.pos[0] as Real, u.pos[1] as Real, u.pos[2] as Real];
            atoms.vel[u.idx] = [u.vel[0] as Real, u.vel[1] as Real, u.vel[2] as Real];
            dem.omega[u.idx] = u.omega;
        }
    });
}

/// Diagnostic: check that each body has the expected number of local sub-sphere atoms.
///
/// Runs every 1000 steps. Warns on mismatch but does not delete atoms.
fn check_lost_clump_atoms(
    atoms: Res<Atom>,
    bodies: Res<MultisphereBodyStore>,
    particles: ParticlesWith<'_, Optional<Read<ClumpAtom>>>,
    comm: Res<CommResource>,
    run_state: Res<RunState>,
) {
    if run_state.total_cycle % 1000 != 0 {
        return;
    }

    particles.with(|clump| {
        let Some(clump) = clump else {
            return;
        };

        let nlocal = atoms.nlocal as usize;
        let mut counts: HashMap<u32, usize> = HashMap::new();

        for i in 0..nlocal {
            if i >= clump.body_id.len() {
                break;
            }
            let bid = clump.body_id[i] as u32;
            if bid > 0 {
                *counts.entry(bid).or_default() += 1;
            }
        }

        for body in &bodies.bodies {
            let expected = body.sub_sphere_tags.len();
            let actual = counts.get(&body.id).copied().unwrap_or(0);
            if actual != expected {
                eprintln!(
                    "WARNING: Body {} has {}/{} atoms on rank {}",
                    body.id,
                    actual,
                    expected,
                    comm.rank()
                );
            }
        }
    });
}

// ── Clump insertion from config ──────────────────────────────────────────────

/// Setup system: insert clumps from `[[clump.insert]]` config blocks.
///
/// For each insertion block, looks up the named clump definition from the
/// [`ClumpRegistry`], then inserts `count` clumps at random non-overlapping
/// positions within the specified region.
fn clump_insert_atoms(
    comm: Res<CommResource>,
    domain: Res<Domain>,
    mut atoms: ResMut<Atom>,
    registry: Res<AtomDataRegistry>,
    clump_registry: Res<ClumpRegistry>,
    mut body_store: ResMut<MultisphereBodyStore>,
    clump_config: Res<ClumpTopConfig>,
    material_table: Res<dirt_atom::MaterialTable>,
    scheduler_manager: Res<SchedulerManager>,
) {
    // Setup systems re-run at the start of every `[[run]]` stage. Clumps must be
    // inserted only once, at the first stage — otherwise each subsequent stage
    // (e.g. compress, shear) would insert another `count` clumps (and, with a
    // now-tight box, spin in the non-overlapping-placement retry loop). This
    // mirrors the stage-0 guard in `dem_insert_atoms` for `[[particles.insert]]`.
    if scheduler_manager.index != 0 {
        return;
    }

    let inserts = match clump_config.insert {
        Some(ref v) => v,
        None => return,
    };

    if comm.rank() != 0 {
        return;
    }

    for insert in inserts {
        let def = clump_registry.find(&insert.definition).unwrap_or_else(|| {
            panic!(
                "Clump definition '{}' not found. Available: {:?}",
                insert.definition,
                clump_registry
                    .defs
                    .iter()
                    .map(|d| &d.name)
                    .collect::<Vec<_>>()
            );
        });

        // Resolve material index
        let mat_idx = material_table
            .names
            .iter()
            .position(|n| n == &insert.material)
            .unwrap_or_else(|| {
                panic!(
                    "Material '{}' not found in [[dem.materials]]",
                    insert.material
                );
            }) as u32;
        let cutoff_padding = material_table.liquid_bridge_cutoff_padding(mat_idx);

        // Compute effective radius for overlap checks (max sub-sphere extent from COM)
        let eff_radius = def
            .spheres
            .iter()
            .map(|s| {
                let d = (s.offset[0].powi(2) + s.offset[1].powi(2) + s.offset[2].powi(2)).sqrt();
                d + s.radius
            })
            .fold(0.0_f64, f64::max);

        // Determine insertion region
        let region = insert.region.clone().unwrap_or_else(|| Region::Block {
            min: [
                domain.boundaries_low[0] + eff_radius,
                domain.boundaries_low[1] + eff_radius,
                domain.boundaries_low[2] + eff_radius,
            ],
            max: [
                domain.boundaries_high[0] - eff_radius,
                domain.boundaries_high[1] - eff_radius,
                domain.boundaries_high[2] - eff_radius,
            ],
        });

        println!(
            "ClumpInsert: inserting {} '{}' clumps (eff_r={:.4}mm, rho={}, mat='{}')",
            insert.count,
            insert.definition,
            eff_radius * 1000.0,
            insert.density,
            insert.material,
        );

        let mut rng = clump_insert_rng(insert);
        let inserted = insert_clumps_with_rng(
            &mut atoms,
            &registry,
            &mut body_store,
            def,
            insert,
            mat_idx,
            cutoff_padding,
            eff_radius,
            &region,
            &mut rng,
        );

        if inserted < insert.count {
            eprintln!(
                "WARNING: Could only insert {}/{} clumps after {} attempts.",
                inserted,
                insert.count,
                insert.count as u64 * 1_000_000
            );
        }
    }
}

fn clump_insert_rng(insert: &ClumpInsertConfig) -> StdRng {
    StdRng::seed_from_u64(insert.seed.unwrap_or(0))
}

#[allow(clippy::too_many_arguments)]
fn insert_clumps_with_rng<R: Rng>(
    atoms: &mut Atom,
    registry: &AtomDataRegistry,
    body_store: &mut MultisphereBodyStore,
    def: &ClumpDef,
    insert: &ClumpInsertConfig,
    mat_idx: u32,
    cutoff_padding: f64,
    eff_radius: f64,
    region: &Region,
    rng: &mut R,
) -> u32 {
    // Track inserted COM positions for overlap avoidance
    let mut com_positions: Vec<[f64; 3]> = Vec::new();
    let mut inserted = 0u32;
    let mut attempts = 0u64;
    let max_attempts = insert.count as u64 * 1_000_000;
    let mut next_clump_id = body_store.bodies.len() as u32 + 1;

    while inserted < insert.count && attempts < max_attempts {
        attempts += 1;
        let pos = region.random_point_inside(rng).unwrap_or_else(|e| {
            panic!("ClumpPlugin preflight should reject invalid insertion regions: {e}")
        });

        // Check overlap with existing clump COMs
        let min_sep = 2.0 * eff_radius * 1.05; // 5% margin
        let mut overlaps = false;

        // Check against existing atoms
        for i in 0..atoms.len() {
            let dx = pos[0] - atoms.pos[i][0] as f64;
            let dy = pos[1] - atoms.pos[i][1] as f64;
            let dz = pos[2] - atoms.pos[i][2] as f64;
            let dist_sq = dx * dx + dy * dy + dz * dz;
            let min_d = eff_radius + atoms.cutoff_radius[i] as f64;
            if dist_sq < min_d * min_d {
                overlaps = true;
                break;
            }
        }

        // Check against already-inserted clump COMs in this batch
        if !overlaps {
            for existing in &com_positions {
                let dx = pos[0] - existing[0];
                let dy = pos[1] - existing[1];
                let dz = pos[2] - existing[2];
                let dist_sq = dx * dx + dy * dy + dz * dz;
                if dist_sq < min_sep * min_sep {
                    overlaps = true;
                    break;
                }
            }
        }

        if overlaps {
            continue;
        }

        // Optional per-clump random orientation: rotate the definition's offsets
        // by a uniformly random unit quaternion (Shoemake).
        let rotated_def;
        let def: &ClumpDef = if insert.random_orientation {
            let u1 = rng.random_range(0.0..1.0f64);
            let u2 = rng.random_range(0.0..1.0f64);
            let u3 = rng.random_range(0.0..1.0f64);
            let two_pi = std::f64::consts::TAU;
            let q = [
                (1.0 - u1).sqrt() * (two_pi * u2).sin(),
                (1.0 - u1).sqrt() * (two_pi * u2).cos(),
                u1.sqrt() * (two_pi * u3).sin(),
                u1.sqrt() * (two_pi * u3).cos(),
            ];
            rotated_def = ClumpDef {
                name: def.name.clone(),
                spheres: def
                    .spheres
                    .iter()
                    .map(|s| ClumpSphereConfig {
                        offset: quat_rotate(q, s.offset),
                        radius: s.radius,
                    })
                    .collect(),
            };
            &rotated_def
        } else {
            def
        };

        // Generate velocity
        let vel = if let Some(v_mag) = insert.velocity {
            [
                rng.random_range(-v_mag..v_mag),
                rng.random_range(-v_mag..v_mag),
                rng.random_range(-v_mag..v_mag),
            ]
        } else {
            [0.0; 3]
        };

        try_insert_clump_with_cutoff_padding(
            atoms,
            registry,
            body_store,
            def,
            pos,
            vel,
            insert.density,
            mat_idx,
            cutoff_padding,
            next_clump_id,
        )
        .expect("validated clump configuration must accept transactional rows");

        com_positions.push(pos);
        next_clump_id += 1;
        inserted += 1;
    }

    inserted
}

// ── Clump insertion helper ──────────────────────────────────────────────────

/// Insert a single clump at the given COM position.
///
/// Creates N sub-sphere atoms and one [`MultisphereBody`] entry.
/// No parent atom is created — the body resource holds COM state.
///
/// Returns the number of atoms inserted (N sub-spheres).
pub fn insert_clump(
    atoms: &mut Atom,
    registry: &AtomDataRegistry,
    body_store: &mut MultisphereBodyStore,
    def: &ClumpDef,
    com_pos: [f64; 3],
    com_vel: [f64; 3],
    density: f64,
    atom_type: u32,
    clump_id: u32,
) -> usize {
    try_insert_clump_with_cutoff_padding(
        atoms, registry, body_store, def, com_pos, com_vel, density, atom_type, 0.0, clump_id,
    )
    .expect("registered clump rows must accept transactional insertion")
}

#[allow(clippy::too_many_arguments)]
fn try_insert_clump_with_cutoff_padding(
    atoms: &mut Atom,
    registry: &AtomDataRegistry,
    body_store: &mut MultisphereBodyStore,
    def: &ClumpDef,
    com_pos: [f64; 3],
    com_vel: [f64; 3],
    density: f64,
    atom_type: u32,
    cutoff_padding: f64,
    clump_id: u32,
) -> Result<usize, ParticleStoreError> {
    // Compute inertia tensor (auto-detect overlap)
    let (total_mass, tensor) = if has_overlap(&def.spheres) {
        compute_inertia_tensor_montecarlo(&def.spheres, density, 100_000)
    } else {
        compute_inertia_tensor_analytical(&def.spheres, density)
    };

    let (principal_moments, principal_axes) = diagonalize_inertia(tensor);

    let base_tag = atoms.get_max_tag() + 1;

    // Create MultisphereBody
    let mut body_offsets = Vec::with_capacity(def.spheres.len());
    let mut sub_sphere_radii = Vec::with_capacity(def.spheres.len());
    let mut sub_sphere_tags = Vec::with_capacity(def.spheres.len());

    for (si, sphere) in def.spheres.iter().enumerate() {
        let sub_tag = base_tag + si as u32;
        body_offsets.push(sphere.offset);
        sub_sphere_radii.push(sphere.radius);
        sub_sphere_tags.push(sub_tag);
    }

    let body = MultisphereBody {
        id: clump_id,
        com_pos,
        com_vel,
        quaternion: [1.0, 0.0, 0.0, 0.0],
        omega: [0.0; 3],
        angmom: [0.0; 3],
        principal_moments,
        principal_axes,
        total_mass,
        inv_mass: if total_mass > 0.0 {
            1.0 / total_mass
        } else {
            0.0
        },
        force: [0.0; 3],
        torque: [0.0; 3],
        image: [0; 3],
        body_offsets,
        sub_sphere_radii,
        sub_sphere_tags,
    };
    // Insert sub-sphere atoms first.  The body is deliberately committed only
    // after every ParticleStore row has accepted its default: an extension
    // rejection must not leave an orphan body or a partially materialized
    // clump.  Roll back already accepted rows through the facade as well.
    let original_natoms = atoms.natoms;
    let original_nlocal = atoms.nlocal;
    for (si, sphere) in def.spheres.iter().enumerate() {
        let sub_tag = base_tag + si as u32;
        let sub_pos = [
            com_pos[0] + sphere.offset[0],
            com_pos[1] + sphere.offset[1],
            com_pos[2] + sphere.offset[2],
        ];

        let sub_mass = density * (4.0 / 3.0) * PI * sphere.radius.powi(3);

        let global_natoms = atoms.natoms + 1;
        if let Err(error) = ParticleStore::new(atoms, registry).push_default_local(global_natoms) {
            while atoms.nlocal > original_nlocal {
                let last = atoms.nlocal as usize - 1;
                ParticleStore::new(atoms, registry)
                    .swap_remove(last)
                    .expect("previously accepted clump rows must remain removable");
            }
            atoms.natoms = original_natoms;
            return Err(error);
        }
        let i = atoms.len() - 1;
        atoms.tag[i] = sub_tag;
        atoms.atom_type[i] = atom_type;
        atoms.origin_index[i] = 0;
        atoms.pos[i] = [sub_pos[0] as Real, sub_pos[1] as Real, sub_pos[2] as Real];
        atoms.vel[i] = [com_vel[0] as Real, com_vel[1] as Real, com_vel[2] as Real];
        atoms.force[i] = [0.0 as Accum; 3];
        atoms.mass[i] = sub_mass as Real;
        atoms.inv_mass[i] = 0.0 as Real;
        atoms.cutoff_radius[i] = (sphere.radius + cutoff_padding.max(0.0)) as Real;
        atoms.image[i] = [0, 0, 0];
        atoms.is_ghost[i] = false;
        let mut dem = registry.expect_mut::<DemAtom>("insert_clump");
        dem.radius[i] = sphere.radius;
        dem.density[i] = density;
        dem.inv_inertia[i] = 0.0;
        dem.quaternion[i] = [1.0, 0.0, 0.0, 0.0];
        dem.omega[i] = [0.0; 3];
        dem.ang_mom[i] = [0.0; 3];
        dem.torque[i] = [0.0; 3];
        dem.body_id[i] = clump_id as f64;
        drop(dem);
        let mut clump_data = registry.expect_mut::<ClumpAtom>("insert_clump");
        clump_data.body_id[i] = clump_id as f64;
        clump_data.body_offset[i] = sphere.offset;

        let _ = si; // suppress unused warning
    }

    body_store.bodies.push(body);
    body_store.generate_map();
    Ok(def.spheres.len())
}

/// Check if two atoms belong to the same rigid body (for contact exclusion).
///
/// Returns `true` only when both `i` and `j` carry a non-zero `body_id` and
/// those ids match (out-of-range indices return `false`). This is the
/// **contact-exclusion contract**: every force plugin that iterates the
/// neighbor list must call `same_body` and `continue` on a `true` result, so
/// that sub-spheres of one rigid body never push on each other. The
/// `dirt_granular` contact kernels already do this; custom force plugins must
/// too. See the crate-level "Contact-exclusion contract" section.
#[inline]
pub fn same_body(clump_data: &ClumpAtom, i: usize, j: usize) -> bool {
    if i >= clump_data.body_id.len() || j >= clump_data.body_id.len() {
        return false;
    }
    let ci = clump_data.body_id[i];
    let cj = clump_data.body_id[j];
    ci > 0.0 && cj > 0.0 && (ci - cj).abs() < 0.5
}

/// Check if atom i is a rigid body sub-sphere.
#[inline]
pub fn is_body_atom(clump_data: &ClumpAtom, i: usize) -> bool {
    i < clump_data.body_id.len() && clump_data.body_id[i] > 0.0
}

// ── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use dirt_atom::{DemAtom, DemAtomPlugin, Elastic, Friction, Material};
    use dirt_test_utils::{ParticleFixture, ParticleSpec};
    use soil_core::{Atom, AtomData, AtomDataRegistry, ParticleStoreError, SingleProcessComm};

    /// A deliberately malformed extension used to prove that clump construction
    /// never commits its body resource before all particle rows are accepted.
    #[derive(Default)]
    struct RejectDefaultRow;

    impl AtomData for RejectDefaultRow {
        fn as_any(&self) -> &dyn std::any::Any {
            self
        }
        fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
            self
        }
        fn snapshot(&self) -> Box<dyn AtomData> {
            Box::new(Self)
        }
        fn len(&self) -> usize {
            0
        }
        unsafe fn push_default(&mut self) {}
        unsafe fn truncate(&mut self, _: usize) {}
        unsafe fn swap_remove(&mut self, _: usize) {}
        fn pack(&self, _: usize, _: &mut Vec<f64>) {}
        unsafe fn unpack(&mut self, _: &[f64]) -> usize {
            0
        }
        unsafe fn apply_permutation(&mut self, _: &[usize], _: usize) {}
    }

    /// Non-overlapping dimer (center distance > r1 + r2) for deterministic tests.
    fn make_dimer_def() -> ClumpDef {
        ClumpDef {
            name: "dimer".to_string(),
            spheres: vec![
                ClumpSphereConfig {
                    offset: [-0.0015, 0.0, 0.0],
                    radius: 0.001,
                },
                ClumpSphereConfig {
                    offset: [0.0015, 0.0, 0.0],
                    radius: 0.001,
                },
            ],
        }
    }

    fn setup_clump_test() -> (Atom, AtomDataRegistry, MultisphereBodyStore) {
        let mut registry = AtomDataRegistry::new();
        registry.try_register(DemAtom::new(), 0).unwrap();
        registry.try_register(ClumpAtom::new(), 0).unwrap();
        (Atom::new(), registry, MultisphereBodyStore::new())
    }

    #[test]
    fn fixture_registers_clump_extension_with_matching_rows() {
        let mut fixture = ParticleFixture::single(ParticleSpec::new(7, [0.0; 3], 0.001)).build();
        let mut clump = ClumpAtom::new();
        clump.body_id.push(3.0);
        clump.body_offset.push([0.0; 3]);
        fixture.register_atom_data(clump);
        let clump = fixture.registry.expect::<ClumpAtom>("fixture clump");
        assert!(is_body_atom(&clump, 0));
    }

    #[test]
    fn degenerate_clump_region_is_a_typed_plugin_error() {
        let mut app = App::new();
        app.add_resource(Config::from_str(
            r#"
[[dem.materials]]
name = "glass"
youngs_mod = 8.7e9
poisson_ratio = 0.3
restitution = 0.9
friction = 0.5

[[clump.definitions]]
name = "dimer"
spheres = [{ offset = [0.0, 0.0, 0.0], radius = 0.001 }]

[[clump.insert]]
definition = "dimer"
count = 1
density = 2500.0
material = "glass"
region = { type = "block", min = [1.0, 1.0, 1.0], max = [1.0, 2.0, 2.0] }
"#,
        ));
        app.try_add_plugins(DemAtomPlugin)
            .expect("valid material setup must satisfy ClumpPlugin dependency");

        let error = match app.try_add_plugins(ClumpPlugin) {
            Err(error) => error,
            Ok(_) => panic!("degenerate clump insertion region must fail preflight"),
        };
        assert!(error
            .to_string()
            .contains("min[0] must be less than max[0]"));
    }

    fn clump_state_bits(atoms: &Atom, bodies: &MultisphereBodyStore) -> Vec<u64> {
        let mut bits = Vec::new();
        for pos in atoms.pos.iter() {
            bits.extend(pos.iter().map(|x| (*x as f64).to_bits()));
        }
        for vel in atoms.vel.iter() {
            bits.extend(vel.iter().map(|x| (*x as f64).to_bits()));
        }
        for body in &bodies.bodies {
            bits.extend(body.com_pos.iter().map(|x| x.to_bits()));
            bits.extend(body.com_vel.iter().map(|x| x.to_bits()));
            for sphere_offset in &body.body_offsets {
                bits.extend(sphere_offset.iter().map(|x| x.to_bits()));
            }
        }
        bits
    }

    fn seeded_insert_snapshot(seed: Option<u64>) -> Vec<u64> {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();
        let insert = ClumpInsertConfig {
            definition: "dimer".to_string(),
            count: 6,
            density: 2500.0,
            material: "glass".to_string(),
            velocity: Some(0.25),
            region: Some(Region::Block {
                min: [-0.02, -0.02, -0.02],
                max: [0.02, 0.02, 0.02],
            }),
            random_orientation: true,
            seed,
        };
        let eff_radius = def
            .spheres
            .iter()
            .map(|s| {
                let d = (s.offset[0].powi(2) + s.offset[1].powi(2) + s.offset[2].powi(2)).sqrt();
                d + s.radius
            })
            .fold(0.0_f64, f64::max);
        let region = insert.region.clone().expect("test region");
        let mut rng = clump_insert_rng(&insert);

        let inserted = insert_clumps_with_rng(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            &insert,
            0,
            0.0,
            eff_radius,
            &region,
            &mut rng,
        );

        assert_eq!(inserted, insert.count);
        clump_state_bits(&atoms, &bodies)
    }

    fn config_insert_snapshot(seed: Option<u64>) -> Vec<u64> {
        let mut app = App::new();
        let mut registry = AtomDataRegistry::new();
        registry.try_register(DemAtom::new(), 0).unwrap();
        registry.try_register(ClumpAtom::new(), 0).unwrap();

        let mut domain = Domain::new();
        domain.boundaries_low = [-0.03; 3];
        domain.boundaries_high = [0.03; 3];
        domain.sub_domain_low = domain.boundaries_low;
        domain.sub_domain_high = domain.boundaries_high;
        domain.size = [0.06; 3];
        domain.sub_length = domain.size;
        domain.volume = 0.06_f64.powi(3);

        let mut clump_registry = ClumpRegistry::new();
        clump_registry.defs.push(make_dimer_def());

        let insert = ClumpInsertConfig {
            definition: "dimer".to_string(),
            count: 6,
            density: 2500.0,
            material: "glass".to_string(),
            velocity: Some(0.25),
            region: Some(Region::Block {
                min: [-0.02, -0.02, -0.02],
                max: [0.02, 0.02, 0.02],
            }),
            random_orientation: true,
            seed,
        };

        let mut materials = dirt_atom::MaterialTable::new();
        materials
            .add(
                Material::new("glass", Elastic::new(8.7e9, 0.3, 0.9)).with_friction(Friction {
                    sliding: 0.5,
                    ..Friction::default()
                }),
            )
            .unwrap();
        materials.build_pair_tables();

        app.add_resource(Atom::new());
        app.add_resource(registry);
        app.add_resource(CommResource(Box::new(SingleProcessComm::new())));
        app.add_resource(domain);
        app.add_resource(clump_registry);
        app.add_resource(MultisphereBodyStore::new());
        app.add_resource(ClumpTopConfig {
            definitions: None,
            insert: Some(vec![insert]),
        });
        app.add_resource(materials);
        app.add_resource(SchedulerManager::default());
        app.add_setup_system(
            clump_insert_atoms.label(CLUMP_INSERT),
            ScheduleSetupSet::Setup,
        );
        app.organize_systems();
        app.setup();

        let atoms = app.get_resource_ref::<Atom>().unwrap();
        let bodies = app.get_resource_ref::<MultisphereBodyStore>().unwrap();
        clump_state_bits(&atoms, &bodies)
    }

    #[test]
    fn test_quat_rotate_identity() {
        let q = [1.0, 0.0, 0.0, 0.0];
        let v = [1.0, 2.0, 3.0];
        let result = quat_rotate(q, v);
        assert!((result[0] - 1.0).abs() < 1e-12);
        assert!((result[1] - 2.0).abs() < 1e-12);
        assert!((result[2] - 3.0).abs() < 1e-12);
    }

    #[test]
    fn test_quat_rotate_90_degrees_z() {
        let angle = std::f64::consts::FRAC_PI_2;
        let half = angle * 0.5;
        let q = [half.cos(), 0.0, 0.0, half.sin()];
        let v = [1.0, 0.0, 0.0];
        let result = quat_rotate(q, v);
        assert!((result[0]).abs() < 1e-12);
        assert!((result[1] - 1.0).abs() < 1e-12);
        assert!((result[2]).abs() < 1e-12);
    }

    #[test]
    fn test_compute_clump_inertia_single_sphere() {
        let spheres = vec![ClumpSphereConfig {
            offset: [0.0, 0.0, 0.0],
            radius: 0.001,
        }];
        let density = 2500.0;
        let (mass, inertia) = compute_clump_inertia(&spheres, density);

        let expected_mass = density * (4.0 / 3.0) * PI * 0.001_f64.powi(3);
        let expected_inertia = 0.4 * expected_mass * 0.001 * 0.001;

        assert!((mass - expected_mass).abs() < 1e-15);
        assert!(
            (inertia - expected_inertia).abs() / expected_inertia < 1e-12,
            "got {}, expected {}",
            inertia,
            expected_inertia
        );
    }

    #[test]
    fn test_insert_clump_creates_correct_atoms() {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();

        let count = insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            [0.0, 0.0, 0.0],
            [0.0; 3],
            2500.0,
            0,
            1,
        );

        assert_eq!(count, 2, "Should insert 2 sub-spheres (no parent atom)");
        assert_eq!(atoms.nlocal, 2);
        assert_eq!(atoms.natoms, 2);
        assert_eq!(bodies.bodies.len(), 1);

        // Both atoms are sub-spheres with real radii
        let dem = registry.expect::<DemAtom>("test_insert_clump_creates_correct_atoms");
        assert!((dem.radius[0] - 0.001).abs() < 1e-10);
        assert!((dem.radius[1] - 0.001).abs() < 1e-10);

        // Sub-sphere positions offset from COM
        assert!((atoms.pos[0][0] - (-0.0015)).abs() < 1e-10);
        assert!((atoms.pos[1][0] - 0.0015).abs() < 1e-10);

        // Sub-spheres have zero inv_mass
        assert_eq!(atoms.inv_mass[0], 0.0);
        assert_eq!(atoms.inv_mass[1], 0.0);

        // Body has correct mass
        let r = 0.001;
        let m_sphere = 2500.0 * (4.0 / 3.0) * PI * r * r * r;
        assert!(
            (bodies.bodies[0].total_mass - 2.0 * m_sphere).abs() / (2.0 * m_sphere) < 1e-12,
            "mass: got {}, expected {}",
            bodies.bodies[0].total_mass,
            2.0 * m_sphere
        );

        // Body has principal moments (diagonalized)
        assert!(bodies.bodies[0].principal_moments[0] > 0.0);
    }

    #[test]
    fn clump_row_rejection_rolls_back_atoms_before_body_commit() {
        let mut registry = AtomDataRegistry::new();
        registry.try_register(DemAtom::new(), 0).unwrap();
        registry.try_register(ClumpAtom::new(), 0).unwrap();
        registry.try_register(RejectDefaultRow, 0).unwrap();
        let mut atoms = Atom::new();
        let mut bodies = MultisphereBodyStore::new();
        let error = try_insert_clump_with_cutoff_padding(
            &mut atoms,
            &registry,
            &mut bodies,
            &make_dimer_def(),
            [0.0; 3],
            [0.0; 3],
            2500.0,
            0,
            0.0,
            7,
        )
        .unwrap_err();
        assert_eq!(error, ParticleStoreError::MalformedExtensionRecord);
        assert!(atoms.is_empty());
        assert_eq!((atoms.nlocal, atoms.nghost, atoms.natoms), (0, 0, 0));
        assert!(registry.validate_rows(0));
        assert!(bodies.bodies.is_empty());
        assert_eq!(bodies.find_by_id(7), None);
    }

    #[test]
    fn test_seeded_clump_insertion_is_byte_stable() {
        let first = seeded_insert_snapshot(Some(20260705));
        let second = seeded_insert_snapshot(Some(20260705));
        assert_eq!(
            first, second,
            "same [[clump.insert]] seed must reproduce positions, velocities, and orientations"
        );

        let different_seed = seeded_insert_snapshot(Some(20260706));
        assert_ne!(
            first, different_seed,
            "changing [[clump.insert]] seed should change the insertion stream"
        );

        let default_a = seeded_insert_snapshot(None);
        let default_b = seeded_insert_snapshot(None);
        assert_eq!(
            default_a, default_b,
            "omitting [[clump.insert]] seed should still use the deterministic default"
        );
    }

    #[test]
    fn test_config_clump_insert_system_is_byte_stable() {
        let first = config_insert_snapshot(Some(20260705));
        let second = config_insert_snapshot(Some(20260705));
        assert_eq!(
            first, second,
            "the clump_insert_atoms setup system must honor [[clump.insert]] seed"
        );

        let different_seed = config_insert_snapshot(Some(20260706));
        assert_ne!(
            first, different_seed,
            "changing [[clump.insert]] seed should change the config insertion stream"
        );

        let default_a = config_insert_snapshot(None);
        let default_b = config_insert_snapshot(None);
        assert_eq!(
            default_a, default_b,
            "the clump_insert_atoms setup system must use the deterministic default seed"
        );
    }

    #[test]
    fn test_same_body_exclusion() {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();

        insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            [0.0, 0.0, 0.0],
            [0.0; 3],
            2500.0,
            0,
            1,
        );

        // Atoms 0 and 1 are in same body
        assert!(same_body(
            &registry.expect::<ClumpAtom>("test_same_body_exclusion"),
            0,
            1
        ));
        // Backward compat
        assert!(same_body(
            &registry.expect::<ClumpAtom>("test_same_body_exclusion"),
            0,
            1
        ));
    }

    #[test]
    fn test_different_bodies_not_excluded() {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();

        insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            [0.0, 0.0, 0.0],
            [0.0; 3],
            2500.0,
            0,
            1,
        );
        insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            [0.01, 0.0, 0.0],
            [0.0; 3],
            2500.0,
            0,
            2,
        );

        // Sub-spheres from different bodies not excluded
        let clump = registry.expect::<ClumpAtom>("test_different_bodies_not_excluded");
        assert!(!same_body(&clump, 0, 2)); // body 1 sub vs body 2 sub
        assert!(!same_body(&clump, 1, 3));
    }

    #[test]
    fn test_force_aggregation() {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();

        insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            [0.0, 0.0, 0.0],
            [0.0; 3],
            2500.0,
            0,
            1,
        );

        // Apply force to sub-sphere 0 (at x = -0.0015)
        atoms.force[0] = [0.0, 0.0, 10.0];

        let mut app = App::new();
        app.add_resource(atoms);
        app.add_resource(registry);
        app.add_resource(bodies);
        app.add_update_system(aggregate_clump_forces, ParticleSimScheduleSet::PostForce);
        app.organize_systems();
        app.run();

        let atoms = app.get_resource_ref::<Atom>().unwrap();
        let bodies = app.get_resource_ref::<MultisphereBodyStore>().unwrap();

        // Force transferred to body
        assert!(
            (bodies.bodies[0].force[2] - 10.0).abs() < 1e-10,
            "Body z-force should be 10.0, got {}",
            bodies.bodies[0].force[2]
        );

        // Sub-sphere force zeroed
        assert!(
            atoms.force[0][2].abs() < 1e-10,
            "Sub-sphere force should be zeroed"
        );

        // Torque: r × F where r = q * offset = [-0.0015, 0, 0] (identity q)
        // [-0.0015, 0, 0] × [0, 0, 10] = [0, 0.015, 0]
        assert!(
            (bodies.bodies[0].torque[1] - 0.015).abs() < 1e-10,
            "Body y-torque should be 0.015, got {}",
            bodies.bodies[0].torque[1]
        );
    }

    #[test]
    fn test_position_update_after_rotation() {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();

        insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            [0.0, 0.0, 0.0],
            [0.0; 3],
            2500.0,
            0,
            1,
        );

        // Rotate body 90° about z-axis
        let angle = std::f64::consts::FRAC_PI_2;
        let half = angle * 0.5;
        bodies.bodies[0].quaternion = [half.cos(), 0.0, 0.0, half.sin()];

        let mut app = App::new();
        app.add_resource(atoms);
        app.add_resource(registry);
        app.add_resource(bodies);
        app.add_update_system(
            update_clump_positions,
            ParticleSimScheduleSet::PostFinalIntegration,
        );
        app.organize_systems();
        app.run();

        let atoms = app.get_resource_ref::<Atom>().unwrap();

        // After 90° z rotation:
        // Sub 0 offset [-0.0015, 0, 0] -> [0, -0.0015, 0]
        assert!((atoms.pos[0][0]).abs() < 1e-10);
        assert!((atoms.pos[0][1] - (-0.0015)).abs() < 1e-10);

        // Sub 1 offset [0.0015, 0, 0] -> [0, 0.0015, 0]
        assert!((atoms.pos[1][0]).abs() < 1e-10);
        assert!((atoms.pos[1][1] - 0.0015).abs() < 1e-10);
    }

    #[test]
    fn test_dimer_free_fall() {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();

        let com_pos = [0.0, 0.0, 0.1];
        insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            com_pos,
            [0.0; 3],
            2500.0,
            0,
            1,
        );
        atoms.dt = 1e-6;

        let gravity_z = -9.81;
        let total_mass = bodies.bodies[0].total_mass;

        let nsteps = 100;
        let dt = atoms.dt;
        let mut expected_vel_z = 0.0;
        let mut expected_pos_z = com_pos[2];

        for _ in 0..nsteps {
            // Apply gravity to body
            bodies.bodies[0].force = [0.0, 0.0, total_mass * gravity_z];

            // Initial half-step
            body::integrate_body_initial(&mut bodies.bodies[0], dt);

            // Expected trajectory
            expected_vel_z += 0.5 * dt * gravity_z;
            expected_pos_z += expected_vel_z * dt;

            // Final half-step (same force)
            bodies.bodies[0].force = [0.0, 0.0, total_mass * gravity_z];
            body::integrate_body_final(&mut bodies.bodies[0], dt);
            expected_vel_z += 0.5 * dt * gravity_z;
        }

        assert!(
            (bodies.bodies[0].com_pos[2] - expected_pos_z).abs() < 1e-14,
            "COM z: got {}, expected {}",
            bodies.bodies[0].com_pos[2],
            expected_pos_z
        );
    }

    #[test]
    fn test_subsphere_velocity_from_rotation() {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();

        insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            2500.0,
            0,
            1,
        );

        bodies.bodies[0].omega = [0.0, 0.0, 100.0];

        let mut app = App::new();
        app.add_resource(atoms);
        app.add_resource(registry);
        app.add_resource(bodies);
        app.add_update_system(
            update_clump_positions,
            ParticleSimScheduleSet::PostFinalIntegration,
        );
        app.organize_systems();
        app.run();

        let atoms = app.get_resource_ref::<Atom>().unwrap();

        // Sub 1 at offset [0.0015, 0, 0]:
        // vel = [1,0,0] + [0,0,100]×[0.0015,0,0] = [1, 0.15, 0]
        assert!((atoms.vel[1][0] - 1.0).abs() < 1e-10);
        assert!((atoms.vel[1][1] - 0.15).abs() < 1e-10);
    }

    #[test]
    fn test_contact_on_one_sphere_creates_torque() {
        let (mut atoms, registry, mut bodies) = setup_clump_test();
        let def = make_dimer_def();

        insert_clump(
            &mut atoms,
            &registry,
            &mut bodies,
            &def,
            [0.0, 0.0, 0.0],
            [0.0; 3],
            2500.0,
            0,
            1,
        );

        // Force in y on right sub-sphere (at x = +0.0015)
        atoms.force[1] = [0.0, 5.0, 0.0];

        let mut app = App::new();
        app.add_resource(atoms);
        app.add_resource(registry);
        app.add_resource(bodies);
        app.add_update_system(aggregate_clump_forces, ParticleSimScheduleSet::PostForce);
        app.organize_systems();
        app.run();

        let bodies = app.get_resource_ref::<MultisphereBodyStore>().unwrap();

        // Force on body
        assert!((bodies.bodies[0].force[1] - 5.0).abs() < 1e-10);

        // Torque: [0.0015, 0, 0] × [0, 5, 0] = [0, 0, 0.0075]
        assert!(
            (bodies.bodies[0].torque[2] - 0.0075).abs() < 1e-10,
            "z-torque should be 0.0075, got {}",
            bodies.bodies[0].torque[2]
        );
    }
}