1#[allow(unused_imports)]
6use super::functions_2::*;
7use crate::Error;
8use crate::types::{PyContactResult, PyVec3};
9use serde::{Deserialize, Serialize};
10
11#[allow(unused_imports)]
12use super::functions::*;
13use super::functions::{PICKLE_MAGIC, PICKLE_VERSION};
14
15use std::collections::HashMap;
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct WorldState {
20 pub gravity: PyVec3,
22 pub time: f64,
24 pub positions: Vec<PyVec3>,
26 pub num_bodies: usize,
28}
29#[derive(Debug, Clone)]
31pub struct IncrementalExportConfig {
32 pub min_speed_threshold: f64,
34 pub max_batch_size: usize,
36 pub include_sleeping: bool,
38}
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
41pub struct SimBodyState {
42 pub handle: u32,
44 pub position: [f64; 3],
46 pub velocity: [f64; 3],
48 pub orientation: [f64; 4],
50 pub angular_velocity: [f64; 3],
52 pub is_sleeping: bool,
54 pub is_static: bool,
56 pub tag: Option<String>,
58}
59impl SimBodyState {
60 pub fn at_rest(handle: u32, position: [f64; 3]) -> Self {
62 Self {
63 handle,
64 position,
65 velocity: [0.0; 3],
66 orientation: [0.0, 0.0, 0.0, 1.0],
67 angular_velocity: [0.0; 3],
68 is_sleeping: false,
69 is_static: false,
70 tag: None,
71 }
72 }
73 pub fn speed(&self) -> f64 {
75 let v = &self.velocity;
76 (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt()
77 }
78 pub fn angular_speed(&self) -> f64 {
80 let w = &self.angular_velocity;
81 (w[0] * w[0] + w[1] * w[1] + w[2] * w[2]).sqrt()
82 }
83 pub fn kinetic_energy_proxy(&self) -> f64 {
85 let v = self.speed();
86 0.5 * v * v
87 }
88 pub fn distance_from_origin(&self) -> f64 {
90 let p = &self.position;
91 (p[0] * p[0] + p[1] * p[1] + p[2] * p[2]).sqrt()
92 }
93 pub fn is_at_rest(&self, linear_threshold: f64, angular_threshold: f64) -> bool {
95 self.speed() < linear_threshold && self.angular_speed() < angular_threshold
96 }
97}
98#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct SimulationSnapshot {
101 pub version: u32,
103 pub time: f64,
105 pub gravity: [f64; 3],
107 pub bodies: Vec<SimBodyState>,
109 pub contacts: Vec<PyContactResult>,
111 pub sleeping_count: usize,
113 pub description: Option<String>,
115 pub metadata: std::collections::HashMap<String, String>,
117}
118impl SimulationSnapshot {
119 pub const FORMAT_VERSION: u32 = 1;
121 pub fn empty() -> Self {
123 Self {
124 version: Self::FORMAT_VERSION,
125 time: 0.0,
126 gravity: [0.0, -9.81, 0.0],
127 bodies: Vec::new(),
128 contacts: Vec::new(),
129 sleeping_count: 0,
130 description: None,
131 metadata: std::collections::HashMap::new(),
132 }
133 }
134 pub fn body_count(&self) -> usize {
136 self.bodies.len()
137 }
138 pub fn sleeping_count(&self) -> usize {
140 self.sleeping_count
141 }
142 pub fn find_body(&self, handle: u32) -> Option<&SimBodyState> {
144 self.bodies.iter().find(|b| b.handle == handle)
145 }
146 pub fn find_by_tag(&self, tag: &str) -> Option<&SimBodyState> {
148 self.bodies.iter().find(|b| b.tag.as_deref() == Some(tag))
149 }
150 pub fn total_kinetic_energy_proxy(&self) -> f64 {
152 self.bodies
153 .iter()
154 .filter(|b| !b.is_sleeping)
155 .map(|b| b.kinetic_energy_proxy())
156 .sum()
157 }
158 pub fn to_pretty_json(&self) -> String {
160 serde_json::to_string_pretty(self).unwrap_or_else(|_| "{}".to_string())
161 }
162 pub fn to_json(&self) -> String {
164 serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
165 }
166 pub fn from_json(json: &str) -> Result<Self, crate::Error> {
168 serde_json::from_str(json)
169 .map_err(|e| crate::Error::General(format!("snapshot deserialization failed: {e}")))
170 }
171 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
173 self.metadata.insert(key.into(), value.into());
174 self
175 }
176 pub fn with_description(mut self, desc: impl Into<String>) -> Self {
178 self.description = Some(desc.into());
179 self
180 }
181 pub fn to_msgpack(&self) -> Vec<u8> {
183 let json = self.to_json();
184 let json_bytes = json.as_bytes();
185 let mut result = Vec::with_capacity(8 + json_bytes.len());
186 result.extend_from_slice(b"OXIP");
187 result.extend_from_slice(&(json_bytes.len() as u32).to_le_bytes());
188 result.extend_from_slice(json_bytes);
189 result
190 }
191 pub fn from_msgpack(data: &[u8]) -> Result<Self, Error> {
193 if data.len() < 8 {
194 return Err(Error::General("msgpack data too short".to_string()));
195 }
196 if &data[0..4] != b"OXIP" {
197 return Err(Error::General("invalid msgpack magic bytes".to_string()));
198 }
199 let len = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize;
200 if data.len() < 8 + len {
201 return Err(Error::General("msgpack data truncated".to_string()));
202 }
203 let json = std::str::from_utf8(&data[8..8 + len])
204 .map_err(|e| Error::General(format!("invalid UTF-8: {e}")))?;
205 Self::from_json(json)
206 }
207 pub fn static_body_count(&self) -> usize {
209 self.bodies.iter().filter(|b| b.is_static).count()
210 }
211 pub fn dynamic_body_count(&self) -> usize {
213 self.bodies.iter().filter(|b| !b.is_static).count()
214 }
215 pub fn handles(&self) -> Vec<u32> {
217 self.bodies.iter().map(|b| b.handle).collect()
218 }
219 pub fn find_by_tag_prefix(&self, prefix: &str) -> Vec<&SimBodyState> {
221 self.bodies
222 .iter()
223 .filter(|b| b.tag.as_deref().is_some_and(|t| t.starts_with(prefix)))
224 .collect()
225 }
226}
227impl SimulationSnapshot {
228 pub fn diff(&self, other: &SimulationSnapshot, position_threshold: f64) -> SnapshotDiff {
233 let a_map: HashMap<u32, &SimBodyState> =
234 self.bodies.iter().map(|b| (b.handle, b)).collect();
235 let b_map: HashMap<u32, &SimBodyState> =
236 other.bodies.iter().map(|b| (b.handle, b)).collect();
237 let removed: Vec<u32> = a_map
238 .keys()
239 .filter(|h| !b_map.contains_key(h))
240 .copied()
241 .collect();
242 let added: Vec<u32> = b_map
243 .keys()
244 .filter(|h| !a_map.contains_key(h))
245 .copied()
246 .collect();
247 let mut moved = Vec::new();
248 let mut max_disp = 0.0_f64;
249 for (handle, a_body) in &a_map {
250 if let Some(b_body) = b_map.get(handle) {
251 let dx = b_body.position[0] - a_body.position[0];
252 let dy = b_body.position[1] - a_body.position[1];
253 let dz = b_body.position[2] - a_body.position[2];
254 let disp = (dx * dx + dy * dy + dz * dz).sqrt();
255 if disp > max_disp {
256 max_disp = disp;
257 }
258 if disp > position_threshold {
259 moved.push(*handle);
260 }
261 }
262 }
263 SnapshotDiff {
264 removed,
265 added,
266 moved,
267 max_displacement: max_disp,
268 time_delta: other.time - self.time,
269 }
270 }
271}
272#[derive(Debug, Clone, Serialize, Deserialize)]
274pub struct BodyDict {
275 pub handle: u32,
277 pub pos: [f64; 3],
279 pub vel: [f64; 3],
281 pub quat: [f64; 4],
283 pub omega: [f64; 3],
285 pub sleeping: bool,
287 pub static_body: bool,
289 pub tag: Option<String>,
291}
292impl BodyDict {
293 #[allow(dead_code)]
295 pub fn from_sim_body(b: &SimBodyState) -> Self {
296 Self {
297 handle: b.handle,
298 pos: b.position,
299 vel: b.velocity,
300 quat: b.orientation,
301 omega: b.angular_velocity,
302 sleeping: b.is_sleeping,
303 static_body: b.is_static,
304 tag: b.tag.clone(),
305 }
306 }
307 #[allow(dead_code)]
309 pub fn to_sim_body(&self) -> SimBodyState {
310 SimBodyState {
311 handle: self.handle,
312 position: self.pos,
313 velocity: self.vel,
314 orientation: self.quat,
315 angular_velocity: self.omega,
316 is_sleeping: self.sleeping,
317 is_static: self.static_body,
318 tag: self.tag.clone(),
319 }
320 }
321}
322#[derive(Debug, Clone)]
335pub struct PickleEnvelope {
336 pub snapshot: SimulationSnapshot,
338}
339impl PickleEnvelope {
340 #[allow(dead_code)]
342 pub fn new(snapshot: SimulationSnapshot) -> Self {
343 Self { snapshot }
344 }
345 #[allow(dead_code)]
347 pub fn to_bytes(&self) -> Vec<u8> {
348 let json = self.snapshot.to_json();
349 let payload = json.as_bytes();
350 let mut buf = Vec::with_capacity(10 + payload.len());
351 buf.extend_from_slice(PICKLE_MAGIC);
352 buf.push(PICKLE_VERSION);
353 buf.extend_from_slice(&(payload.len() as u32).to_le_bytes());
354 buf.extend_from_slice(payload);
355 buf.push(b'.');
356 buf
357 }
358 #[allow(dead_code)]
360 pub fn from_bytes(data: &[u8]) -> Result<Self, Error> {
361 if data.len() < 10 {
362 return Err(Error::General("pickle envelope too short".to_string()));
363 }
364 if &data[0..4] != PICKLE_MAGIC {
365 return Err(Error::General("invalid pickle magic bytes".to_string()));
366 }
367 let _version = data[4];
368 let payload_len = u32::from_le_bytes([data[5], data[6], data[7], data[8]]) as usize;
369 if data.len() < 9 + payload_len + 1 {
370 return Err(Error::General("pickle envelope truncated".to_string()));
371 }
372 let json = std::str::from_utf8(&data[9..9 + payload_len])
373 .map_err(|e| Error::General(format!("invalid UTF-8 in pickle: {e}")))?;
374 let snapshot = SimulationSnapshot::from_json(json)?;
375 Ok(Self { snapshot })
376 }
377 #[allow(dead_code)]
379 pub fn to_hex(&self) -> String {
380 self.to_bytes()
381 .iter()
382 .map(|b| format!("{b:02x}"))
383 .collect::<Vec<_>>()
384 .join("")
385 }
386}
387#[derive(Debug, Clone)]
389pub struct SchemaValidationResult {
390 pub is_valid: bool,
392 pub errors: Vec<String>,
394}
395impl SchemaValidationResult {
396 #[allow(dead_code)]
398 pub fn ok() -> Self {
399 Self {
400 is_valid: true,
401 errors: Vec::new(),
402 }
403 }
404 #[allow(dead_code)]
406 pub fn err(msg: impl Into<String>) -> Self {
407 Self {
408 is_valid: false,
409 errors: vec![msg.into()],
410 }
411 }
412}
413#[derive(Debug, Clone, Serialize, Deserialize)]
418pub struct NumpyPositionArray {
419 pub data: Vec<f64>,
421 pub shape: [usize; 2],
423 pub dtype: String,
425 pub c_order: bool,
427}
428impl NumpyPositionArray {
429 #[allow(dead_code)]
431 pub fn from_snapshot(snap: &SimulationSnapshot) -> Self {
432 let n = snap.bodies.len();
433 let mut data = Vec::with_capacity(n * 3);
434 for b in &snap.bodies {
435 data.extend_from_slice(&b.position);
436 }
437 Self {
438 data,
439 shape: [n, 3],
440 dtype: "float64".to_string(),
441 c_order: true,
442 }
443 }
444 #[allow(dead_code)]
446 pub fn velocity_array(snap: &SimulationSnapshot) -> Self {
447 let n = snap.bodies.len();
448 let mut data = Vec::with_capacity(n * 3);
449 for b in &snap.bodies {
450 data.extend_from_slice(&b.velocity);
451 }
452 Self {
453 data,
454 shape: [n, 3],
455 dtype: "float64".to_string(),
456 c_order: true,
457 }
458 }
459 #[allow(dead_code)]
461 pub fn get_row(&self, i: usize) -> Option<[f64; 3]> {
462 if i >= self.shape[0] {
463 return None;
464 }
465 let base = i * 3;
466 Some([self.data[base], self.data[base + 1], self.data[base + 2]])
467 }
468 #[allow(dead_code)]
470 pub fn n_rows(&self) -> usize {
471 self.shape[0]
472 }
473 #[allow(dead_code)]
475 pub fn size(&self) -> usize {
476 self.data.len()
477 }
478 #[allow(dead_code)]
480 pub fn to_json(&self) -> String {
481 serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
482 }
483 #[allow(dead_code)]
485 pub fn to_raw_bytes(&self) -> Vec<u8> {
486 let mut buf = Vec::with_capacity(self.data.len() * 8);
487 for &v in &self.data {
488 buf.extend_from_slice(&v.to_le_bytes());
489 }
490 buf
491 }
492}
493#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
495pub struct BodyStateJson {
496 pub handle: u32,
498 pub position: [f64; 3],
500 pub velocity: [f64; 3],
502 pub orientation: [f64; 4],
504 pub angular_velocity: [f64; 3],
506 pub is_sleeping: bool,
508 pub is_static: bool,
510 pub tag: Option<String>,
512 pub schema_version: String,
514}
515impl BodyStateJson {
516 pub fn from_sim_body(body: &SimBodyState) -> Self {
518 Self {
519 handle: body.handle,
520 position: body.position,
521 velocity: body.velocity,
522 orientation: body.orientation,
523 angular_velocity: body.angular_velocity,
524 is_sleeping: body.is_sleeping,
525 is_static: body.is_static,
526 tag: body.tag.clone(),
527 schema_version: "1.0.0".to_string(),
528 }
529 }
530 pub fn to_sim_body(&self) -> SimBodyState {
532 SimBodyState {
533 handle: self.handle,
534 position: self.position,
535 velocity: self.velocity,
536 orientation: self.orientation,
537 angular_velocity: self.angular_velocity,
538 is_sleeping: self.is_sleeping,
539 is_static: self.is_static,
540 tag: self.tag.clone(),
541 }
542 }
543}
544#[derive(Debug, Clone, Serialize, Deserialize)]
546pub struct SimulationCheckpoint {
547 pub version: u32,
549 pub label: String,
551 pub timestamp: f64,
553 pub sim_time: f64,
555 pub step_count: u64,
557 pub gravity: [f64; 3],
559 pub bodies: Vec<BodyStateJson>,
561 pub metadata: std::collections::HashMap<String, String>,
563}
564impl SimulationCheckpoint {
565 pub const FORMAT_VERSION: u32 = 1;
567 pub fn empty(label: impl Into<String>) -> Self {
569 Self {
570 version: Self::FORMAT_VERSION,
571 label: label.into(),
572 timestamp: 0.0,
573 sim_time: 0.0,
574 step_count: 0,
575 gravity: [0.0, -9.81, 0.0],
576 bodies: Vec::new(),
577 metadata: std::collections::HashMap::new(),
578 }
579 }
580 pub fn from_snapshot(
582 snap: &SimulationSnapshot,
583 label: impl Into<String>,
584 step_count: u64,
585 timestamp: f64,
586 ) -> Self {
587 let bodies = snap
588 .bodies
589 .iter()
590 .map(BodyStateJson::from_sim_body)
591 .collect();
592 Self {
593 version: Self::FORMAT_VERSION,
594 label: label.into(),
595 timestamp,
596 sim_time: snap.time,
597 step_count,
598 gravity: snap.gravity,
599 bodies,
600 metadata: snap.metadata.clone(),
601 }
602 }
603 pub fn to_json(&self) -> String {
605 serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
606 }
607 pub fn from_json(json: &str) -> Result<Self, crate::Error> {
609 serde_json::from_str(json)
610 .map_err(|e| crate::Error::General(format!("checkpoint deserialization failed: {e}")))
611 }
612 pub fn to_snapshot(&self) -> SimulationSnapshot {
614 let bodies = self.bodies.iter().map(|b| b.to_sim_body()).collect();
615 SimulationSnapshot {
616 version: SimulationSnapshot::FORMAT_VERSION,
617 time: self.sim_time,
618 gravity: self.gravity,
619 bodies,
620 contacts: Vec::new(),
621 sleeping_count: 0,
622 description: Some(format!("Restored from checkpoint '{}'", self.label)),
623 metadata: self.metadata.clone(),
624 }
625 }
626 pub fn body_count(&self) -> usize {
628 self.bodies.len()
629 }
630}
631#[derive(Debug, Clone, Serialize, Deserialize)]
636pub struct IncrementalUpdate {
637 pub sequence: u64,
639 pub time: f64,
641 pub changed_bodies: Vec<SimBodyState>,
643 pub removed_handles: Vec<u32>,
645 pub added_handles: Vec<u32>,
647}
648impl IncrementalUpdate {
649 pub fn empty(sequence: u64, time: f64) -> Self {
651 Self {
652 sequence,
653 time,
654 changed_bodies: Vec::new(),
655 removed_handles: Vec::new(),
656 added_handles: Vec::new(),
657 }
658 }
659 pub fn to_json(&self) -> String {
661 serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
662 }
663 pub fn from_json(json: &str) -> Result<Self, Error> {
665 serde_json::from_str(json)
666 .map_err(|e| Error::General(format!("incremental update deserialization: {e}")))
667 }
668 pub fn is_empty(&self) -> bool {
670 self.changed_bodies.is_empty()
671 && self.removed_handles.is_empty()
672 && self.added_handles.is_empty()
673 }
674 pub fn change_count(&self) -> usize {
676 self.changed_bodies.len() + self.removed_handles.len() + self.added_handles.len()
677 }
678}
679#[derive(Debug, Clone, Serialize, Deserialize)]
681pub struct ExportBatch {
682 pub batch_index: usize,
684 pub is_last: bool,
686 pub total_batches: usize,
688 pub time: f64,
690 pub bodies: Vec<SimBodyState>,
692}
693impl ExportBatch {
694 #[allow(dead_code)]
696 pub fn to_json(&self) -> String {
697 serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
698 }
699 #[allow(dead_code)]
701 pub fn from_json(json: &str) -> Result<Self, Error> {
702 serde_json::from_str(json)
703 .map_err(|e| Error::General(format!("ExportBatch deserialization: {e}")))
704 }
705}
706#[derive(Debug, Clone, Serialize, Deserialize)]
708pub struct SchemaVersion {
709 pub major: u32,
711 pub minor: u32,
713 pub patch: u32,
715}
716impl SchemaVersion {
717 pub fn current() -> Self {
719 Self {
720 major: 1,
721 minor: 0,
722 patch: 0,
723 }
724 }
725 pub fn is_compatible_with(&self, other: &SchemaVersion) -> bool {
727 self.major == other.major
728 }
729 pub fn to_string_version(&self) -> String {
731 format!("{}.{}.{}", self.major, self.minor, self.patch)
732 }
733}
734#[derive(Debug, Clone, Serialize, Deserialize)]
739pub struct SnapshotDict {
740 pub version: u32,
742 pub time: f64,
744 pub gravity: [f64; 3],
746 pub bodies: Vec<BodyDict>,
748 pub n_contacts: usize,
750 pub description: Option<String>,
752}
753impl SnapshotDict {
754 #[allow(dead_code)]
756 pub fn from_snapshot(snap: &SimulationSnapshot) -> Self {
757 Self {
758 version: snap.version,
759 time: snap.time,
760 gravity: snap.gravity,
761 bodies: snap.bodies.iter().map(BodyDict::from_sim_body).collect(),
762 n_contacts: snap.contacts.len(),
763 description: snap.description.clone(),
764 }
765 }
766 #[allow(dead_code)]
768 pub fn to_snapshot(&self) -> SimulationSnapshot {
769 let bodies: Vec<SimBodyState> = self.bodies.iter().map(|b| b.to_sim_body()).collect();
770 let sleeping_count = bodies.iter().filter(|b| b.is_sleeping).count();
771 SimulationSnapshot {
772 version: self.version,
773 time: self.time,
774 gravity: self.gravity,
775 bodies,
776 contacts: Vec::new(),
777 sleeping_count,
778 description: self.description.clone(),
779 metadata: std::collections::HashMap::new(),
780 }
781 }
782 #[allow(dead_code)]
784 pub fn to_dict_json(&self) -> String {
785 serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
786 }
787 #[allow(dead_code)]
789 pub fn from_dict_json(json: &str) -> Result<Self, Error> {
790 serde_json::from_str(json)
791 .map_err(|e| Error::General(format!("SnapshotDict deserialization: {e}")))
792 }
793}
794#[derive(Debug, Clone)]
796#[allow(dead_code)]
797pub struct SnapshotDiff {
798 pub removed: Vec<u32>,
800 pub added: Vec<u32>,
802 pub moved: Vec<u32>,
804 pub max_displacement: f64,
806 pub time_delta: f64,
808}
809impl SnapshotDiff {
810 pub fn is_identical(&self) -> bool {
812 self.removed.is_empty() && self.added.is_empty() && self.moved.is_empty()
813 }
814 pub fn change_count(&self) -> usize {
816 self.removed.len() + self.added.len() + self.moved.len()
817 }
818}
819#[derive(Debug, Clone)]
821#[allow(dead_code)]
822pub struct ValidationResult {
823 pub is_valid: bool,
825 pub issues: Vec<String>,
827}