#[allow(unused_imports)]
use super::functions_2::*;
use crate::Error;
use crate::types::{PyContactResult, PyVec3};
use serde::{Deserialize, Serialize};
#[allow(unused_imports)]
use super::functions::*;
use super::functions::{PICKLE_MAGIC, PICKLE_VERSION};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorldState {
pub gravity: PyVec3,
pub time: f64,
pub positions: Vec<PyVec3>,
pub num_bodies: usize,
}
#[derive(Debug, Clone)]
pub struct IncrementalExportConfig {
pub min_speed_threshold: f64,
pub max_batch_size: usize,
pub include_sleeping: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SimBodyState {
pub handle: u32,
pub position: [f64; 3],
pub velocity: [f64; 3],
pub orientation: [f64; 4],
pub angular_velocity: [f64; 3],
pub is_sleeping: bool,
pub is_static: bool,
pub tag: Option<String>,
}
impl SimBodyState {
pub fn at_rest(handle: u32, position: [f64; 3]) -> Self {
Self {
handle,
position,
velocity: [0.0; 3],
orientation: [0.0, 0.0, 0.0, 1.0],
angular_velocity: [0.0; 3],
is_sleeping: false,
is_static: false,
tag: None,
}
}
pub fn speed(&self) -> f64 {
let v = &self.velocity;
(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt()
}
pub fn angular_speed(&self) -> f64 {
let w = &self.angular_velocity;
(w[0] * w[0] + w[1] * w[1] + w[2] * w[2]).sqrt()
}
pub fn kinetic_energy_proxy(&self) -> f64 {
let v = self.speed();
0.5 * v * v
}
pub fn distance_from_origin(&self) -> f64 {
let p = &self.position;
(p[0] * p[0] + p[1] * p[1] + p[2] * p[2]).sqrt()
}
pub fn is_at_rest(&self, linear_threshold: f64, angular_threshold: f64) -> bool {
self.speed() < linear_threshold && self.angular_speed() < angular_threshold
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimulationSnapshot {
pub version: u32,
pub time: f64,
pub gravity: [f64; 3],
pub bodies: Vec<SimBodyState>,
pub contacts: Vec<PyContactResult>,
pub sleeping_count: usize,
pub description: Option<String>,
pub metadata: std::collections::HashMap<String, String>,
}
impl SimulationSnapshot {
pub const FORMAT_VERSION: u32 = 1;
pub fn empty() -> Self {
Self {
version: Self::FORMAT_VERSION,
time: 0.0,
gravity: [0.0, -9.81, 0.0],
bodies: Vec::new(),
contacts: Vec::new(),
sleeping_count: 0,
description: None,
metadata: std::collections::HashMap::new(),
}
}
pub fn body_count(&self) -> usize {
self.bodies.len()
}
pub fn sleeping_count(&self) -> usize {
self.sleeping_count
}
pub fn find_body(&self, handle: u32) -> Option<&SimBodyState> {
self.bodies.iter().find(|b| b.handle == handle)
}
pub fn find_by_tag(&self, tag: &str) -> Option<&SimBodyState> {
self.bodies.iter().find(|b| b.tag.as_deref() == Some(tag))
}
pub fn total_kinetic_energy_proxy(&self) -> f64 {
self.bodies
.iter()
.filter(|b| !b.is_sleeping)
.map(|b| b.kinetic_energy_proxy())
.sum()
}
pub fn to_pretty_json(&self) -> String {
serde_json::to_string_pretty(self).unwrap_or_else(|_| "{}".to_string())
}
pub fn to_json(&self) -> String {
serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
}
pub fn from_json(json: &str) -> Result<Self, crate::Error> {
serde_json::from_str(json)
.map_err(|e| crate::Error::General(format!("snapshot deserialization failed: {e}")))
}
pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.metadata.insert(key.into(), value.into());
self
}
pub fn with_description(mut self, desc: impl Into<String>) -> Self {
self.description = Some(desc.into());
self
}
pub fn to_msgpack(&self) -> Vec<u8> {
let json = self.to_json();
let json_bytes = json.as_bytes();
let mut result = Vec::with_capacity(8 + json_bytes.len());
result.extend_from_slice(b"OXIP");
result.extend_from_slice(&(json_bytes.len() as u32).to_le_bytes());
result.extend_from_slice(json_bytes);
result
}
pub fn from_msgpack(data: &[u8]) -> Result<Self, Error> {
if data.len() < 8 {
return Err(Error::General("msgpack data too short".to_string()));
}
if &data[0..4] != b"OXIP" {
return Err(Error::General("invalid msgpack magic bytes".to_string()));
}
let len = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize;
if data.len() < 8 + len {
return Err(Error::General("msgpack data truncated".to_string()));
}
let json = std::str::from_utf8(&data[8..8 + len])
.map_err(|e| Error::General(format!("invalid UTF-8: {e}")))?;
Self::from_json(json)
}
pub fn static_body_count(&self) -> usize {
self.bodies.iter().filter(|b| b.is_static).count()
}
pub fn dynamic_body_count(&self) -> usize {
self.bodies.iter().filter(|b| !b.is_static).count()
}
pub fn handles(&self) -> Vec<u32> {
self.bodies.iter().map(|b| b.handle).collect()
}
pub fn find_by_tag_prefix(&self, prefix: &str) -> Vec<&SimBodyState> {
self.bodies
.iter()
.filter(|b| b.tag.as_deref().is_some_and(|t| t.starts_with(prefix)))
.collect()
}
}
impl SimulationSnapshot {
pub fn diff(&self, other: &SimulationSnapshot, position_threshold: f64) -> SnapshotDiff {
let a_map: HashMap<u32, &SimBodyState> =
self.bodies.iter().map(|b| (b.handle, b)).collect();
let b_map: HashMap<u32, &SimBodyState> =
other.bodies.iter().map(|b| (b.handle, b)).collect();
let removed: Vec<u32> = a_map
.keys()
.filter(|h| !b_map.contains_key(h))
.copied()
.collect();
let added: Vec<u32> = b_map
.keys()
.filter(|h| !a_map.contains_key(h))
.copied()
.collect();
let mut moved = Vec::new();
let mut max_disp = 0.0_f64;
for (handle, a_body) in &a_map {
if let Some(b_body) = b_map.get(handle) {
let dx = b_body.position[0] - a_body.position[0];
let dy = b_body.position[1] - a_body.position[1];
let dz = b_body.position[2] - a_body.position[2];
let disp = (dx * dx + dy * dy + dz * dz).sqrt();
if disp > max_disp {
max_disp = disp;
}
if disp > position_threshold {
moved.push(*handle);
}
}
}
SnapshotDiff {
removed,
added,
moved,
max_displacement: max_disp,
time_delta: other.time - self.time,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BodyDict {
pub handle: u32,
pub pos: [f64; 3],
pub vel: [f64; 3],
pub quat: [f64; 4],
pub omega: [f64; 3],
pub sleeping: bool,
pub static_body: bool,
pub tag: Option<String>,
}
impl BodyDict {
#[allow(dead_code)]
pub fn from_sim_body(b: &SimBodyState) -> Self {
Self {
handle: b.handle,
pos: b.position,
vel: b.velocity,
quat: b.orientation,
omega: b.angular_velocity,
sleeping: b.is_sleeping,
static_body: b.is_static,
tag: b.tag.clone(),
}
}
#[allow(dead_code)]
pub fn to_sim_body(&self) -> SimBodyState {
SimBodyState {
handle: self.handle,
position: self.pos,
velocity: self.vel,
orientation: self.quat,
angular_velocity: self.omega,
is_sleeping: self.sleeping,
is_static: self.static_body,
tag: self.tag.clone(),
}
}
}
#[derive(Debug, Clone)]
pub struct PickleEnvelope {
pub snapshot: SimulationSnapshot,
}
impl PickleEnvelope {
#[allow(dead_code)]
pub fn new(snapshot: SimulationSnapshot) -> Self {
Self { snapshot }
}
#[allow(dead_code)]
pub fn to_bytes(&self) -> Vec<u8> {
let json = self.snapshot.to_json();
let payload = json.as_bytes();
let mut buf = Vec::with_capacity(10 + payload.len());
buf.extend_from_slice(PICKLE_MAGIC);
buf.push(PICKLE_VERSION);
buf.extend_from_slice(&(payload.len() as u32).to_le_bytes());
buf.extend_from_slice(payload);
buf.push(b'.');
buf
}
#[allow(dead_code)]
pub fn from_bytes(data: &[u8]) -> Result<Self, Error> {
if data.len() < 10 {
return Err(Error::General("pickle envelope too short".to_string()));
}
if &data[0..4] != PICKLE_MAGIC {
return Err(Error::General("invalid pickle magic bytes".to_string()));
}
let _version = data[4];
let payload_len = u32::from_le_bytes([data[5], data[6], data[7], data[8]]) as usize;
if data.len() < 9 + payload_len + 1 {
return Err(Error::General("pickle envelope truncated".to_string()));
}
let json = std::str::from_utf8(&data[9..9 + payload_len])
.map_err(|e| Error::General(format!("invalid UTF-8 in pickle: {e}")))?;
let snapshot = SimulationSnapshot::from_json(json)?;
Ok(Self { snapshot })
}
#[allow(dead_code)]
pub fn to_hex(&self) -> String {
self.to_bytes()
.iter()
.map(|b| format!("{b:02x}"))
.collect::<Vec<_>>()
.join("")
}
}
#[derive(Debug, Clone)]
pub struct SchemaValidationResult {
pub is_valid: bool,
pub errors: Vec<String>,
}
impl SchemaValidationResult {
#[allow(dead_code)]
pub fn ok() -> Self {
Self {
is_valid: true,
errors: Vec::new(),
}
}
#[allow(dead_code)]
pub fn err(msg: impl Into<String>) -> Self {
Self {
is_valid: false,
errors: vec![msg.into()],
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NumpyPositionArray {
pub data: Vec<f64>,
pub shape: [usize; 2],
pub dtype: String,
pub c_order: bool,
}
impl NumpyPositionArray {
#[allow(dead_code)]
pub fn from_snapshot(snap: &SimulationSnapshot) -> Self {
let n = snap.bodies.len();
let mut data = Vec::with_capacity(n * 3);
for b in &snap.bodies {
data.extend_from_slice(&b.position);
}
Self {
data,
shape: [n, 3],
dtype: "float64".to_string(),
c_order: true,
}
}
#[allow(dead_code)]
pub fn velocity_array(snap: &SimulationSnapshot) -> Self {
let n = snap.bodies.len();
let mut data = Vec::with_capacity(n * 3);
for b in &snap.bodies {
data.extend_from_slice(&b.velocity);
}
Self {
data,
shape: [n, 3],
dtype: "float64".to_string(),
c_order: true,
}
}
#[allow(dead_code)]
pub fn get_row(&self, i: usize) -> Option<[f64; 3]> {
if i >= self.shape[0] {
return None;
}
let base = i * 3;
Some([self.data[base], self.data[base + 1], self.data[base + 2]])
}
#[allow(dead_code)]
pub fn n_rows(&self) -> usize {
self.shape[0]
}
#[allow(dead_code)]
pub fn size(&self) -> usize {
self.data.len()
}
#[allow(dead_code)]
pub fn to_json(&self) -> String {
serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
}
#[allow(dead_code)]
pub fn to_raw_bytes(&self) -> Vec<u8> {
let mut buf = Vec::with_capacity(self.data.len() * 8);
for &v in &self.data {
buf.extend_from_slice(&v.to_le_bytes());
}
buf
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct BodyStateJson {
pub handle: u32,
pub position: [f64; 3],
pub velocity: [f64; 3],
pub orientation: [f64; 4],
pub angular_velocity: [f64; 3],
pub is_sleeping: bool,
pub is_static: bool,
pub tag: Option<String>,
pub schema_version: String,
}
impl BodyStateJson {
pub fn from_sim_body(body: &SimBodyState) -> Self {
Self {
handle: body.handle,
position: body.position,
velocity: body.velocity,
orientation: body.orientation,
angular_velocity: body.angular_velocity,
is_sleeping: body.is_sleeping,
is_static: body.is_static,
tag: body.tag.clone(),
schema_version: "1.0.0".to_string(),
}
}
pub fn to_sim_body(&self) -> SimBodyState {
SimBodyState {
handle: self.handle,
position: self.position,
velocity: self.velocity,
orientation: self.orientation,
angular_velocity: self.angular_velocity,
is_sleeping: self.is_sleeping,
is_static: self.is_static,
tag: self.tag.clone(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimulationCheckpoint {
pub version: u32,
pub label: String,
pub timestamp: f64,
pub sim_time: f64,
pub step_count: u64,
pub gravity: [f64; 3],
pub bodies: Vec<BodyStateJson>,
pub metadata: std::collections::HashMap<String, String>,
}
impl SimulationCheckpoint {
pub const FORMAT_VERSION: u32 = 1;
pub fn empty(label: impl Into<String>) -> Self {
Self {
version: Self::FORMAT_VERSION,
label: label.into(),
timestamp: 0.0,
sim_time: 0.0,
step_count: 0,
gravity: [0.0, -9.81, 0.0],
bodies: Vec::new(),
metadata: std::collections::HashMap::new(),
}
}
pub fn from_snapshot(
snap: &SimulationSnapshot,
label: impl Into<String>,
step_count: u64,
timestamp: f64,
) -> Self {
let bodies = snap
.bodies
.iter()
.map(BodyStateJson::from_sim_body)
.collect();
Self {
version: Self::FORMAT_VERSION,
label: label.into(),
timestamp,
sim_time: snap.time,
step_count,
gravity: snap.gravity,
bodies,
metadata: snap.metadata.clone(),
}
}
pub fn to_json(&self) -> String {
serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
}
pub fn from_json(json: &str) -> Result<Self, crate::Error> {
serde_json::from_str(json)
.map_err(|e| crate::Error::General(format!("checkpoint deserialization failed: {e}")))
}
pub fn to_snapshot(&self) -> SimulationSnapshot {
let bodies = self.bodies.iter().map(|b| b.to_sim_body()).collect();
SimulationSnapshot {
version: SimulationSnapshot::FORMAT_VERSION,
time: self.sim_time,
gravity: self.gravity,
bodies,
contacts: Vec::new(),
sleeping_count: 0,
description: Some(format!("Restored from checkpoint '{}'", self.label)),
metadata: self.metadata.clone(),
}
}
pub fn body_count(&self) -> usize {
self.bodies.len()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IncrementalUpdate {
pub sequence: u64,
pub time: f64,
pub changed_bodies: Vec<SimBodyState>,
pub removed_handles: Vec<u32>,
pub added_handles: Vec<u32>,
}
impl IncrementalUpdate {
pub fn empty(sequence: u64, time: f64) -> Self {
Self {
sequence,
time,
changed_bodies: Vec::new(),
removed_handles: Vec::new(),
added_handles: Vec::new(),
}
}
pub fn to_json(&self) -> String {
serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
}
pub fn from_json(json: &str) -> Result<Self, Error> {
serde_json::from_str(json)
.map_err(|e| Error::General(format!("incremental update deserialization: {e}")))
}
pub fn is_empty(&self) -> bool {
self.changed_bodies.is_empty()
&& self.removed_handles.is_empty()
&& self.added_handles.is_empty()
}
pub fn change_count(&self) -> usize {
self.changed_bodies.len() + self.removed_handles.len() + self.added_handles.len()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExportBatch {
pub batch_index: usize,
pub is_last: bool,
pub total_batches: usize,
pub time: f64,
pub bodies: Vec<SimBodyState>,
}
impl ExportBatch {
#[allow(dead_code)]
pub fn to_json(&self) -> String {
serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
}
#[allow(dead_code)]
pub fn from_json(json: &str) -> Result<Self, Error> {
serde_json::from_str(json)
.map_err(|e| Error::General(format!("ExportBatch deserialization: {e}")))
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SchemaVersion {
pub major: u32,
pub minor: u32,
pub patch: u32,
}
impl SchemaVersion {
pub fn current() -> Self {
Self {
major: 1,
minor: 0,
patch: 0,
}
}
pub fn is_compatible_with(&self, other: &SchemaVersion) -> bool {
self.major == other.major
}
pub fn to_string_version(&self) -> String {
format!("{}.{}.{}", self.major, self.minor, self.patch)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotDict {
pub version: u32,
pub time: f64,
pub gravity: [f64; 3],
pub bodies: Vec<BodyDict>,
pub n_contacts: usize,
pub description: Option<String>,
}
impl SnapshotDict {
#[allow(dead_code)]
pub fn from_snapshot(snap: &SimulationSnapshot) -> Self {
Self {
version: snap.version,
time: snap.time,
gravity: snap.gravity,
bodies: snap.bodies.iter().map(BodyDict::from_sim_body).collect(),
n_contacts: snap.contacts.len(),
description: snap.description.clone(),
}
}
#[allow(dead_code)]
pub fn to_snapshot(&self) -> SimulationSnapshot {
let bodies: Vec<SimBodyState> = self.bodies.iter().map(|b| b.to_sim_body()).collect();
let sleeping_count = bodies.iter().filter(|b| b.is_sleeping).count();
SimulationSnapshot {
version: self.version,
time: self.time,
gravity: self.gravity,
bodies,
contacts: Vec::new(),
sleeping_count,
description: self.description.clone(),
metadata: std::collections::HashMap::new(),
}
}
#[allow(dead_code)]
pub fn to_dict_json(&self) -> String {
serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
}
#[allow(dead_code)]
pub fn from_dict_json(json: &str) -> Result<Self, Error> {
serde_json::from_str(json)
.map_err(|e| Error::General(format!("SnapshotDict deserialization: {e}")))
}
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct SnapshotDiff {
pub removed: Vec<u32>,
pub added: Vec<u32>,
pub moved: Vec<u32>,
pub max_displacement: f64,
pub time_delta: f64,
}
impl SnapshotDiff {
pub fn is_identical(&self) -> bool {
self.removed.is_empty() && self.added.is_empty() && self.moved.is_empty()
}
pub fn change_count(&self) -> usize {
self.removed.len() + self.added.len() + self.moved.len()
}
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct ValidationResult {
pub is_valid: bool,
pub issues: Vec<String>,
}