1use crate::*;
2use ::ndarray;
3use boxcars;
4pub use derive_new;
5use lazy_static::lazy_static;
6pub use paste;
7use serde::Serialize;
8use std::sync::Arc;
9
10#[derive(Debug, Clone, PartialEq, Serialize)]
21pub struct NDArrayColumnHeaders {
22 pub global_headers: Vec<String>,
23 pub player_headers: Vec<String>,
24}
25
26impl NDArrayColumnHeaders {
27 pub fn new(global_headers: Vec<String>, player_headers: Vec<String>) -> Self {
28 Self {
29 global_headers,
30 player_headers,
31 }
32 }
33}
34
35#[derive(Debug, Clone, PartialEq, Serialize)]
44pub struct ReplayMetaWithHeaders {
45 pub replay_meta: ReplayMeta,
46 pub column_headers: NDArrayColumnHeaders,
47}
48
49impl ReplayMetaWithHeaders {
50 pub fn headers_vec(&self) -> Vec<String> {
51 self.headers_vec_from(|_, _info, index| format!("Player {index} - "))
52 }
53
54 pub fn headers_vec_from<F>(&self, player_prefix_getter: F) -> Vec<String>
55 where
56 F: Fn(&Self, &PlayerInfo, usize) -> String,
57 {
58 self.column_headers
59 .global_headers
60 .iter()
61 .cloned()
62 .chain(self.replay_meta.player_order().enumerate().flat_map(
63 move |(player_index, info)| {
64 let player_prefix = player_prefix_getter(self, info, player_index);
65 self.column_headers
66 .player_headers
67 .iter()
68 .map(move |header| format!("{player_prefix}{header}"))
69 },
70 ))
71 .collect()
72 }
73}
74
75pub struct NDArrayCollector<F> {
86 feature_adders: FeatureAdders<F>,
87 player_feature_adders: PlayerFeatureAdders<F>,
88 data: Vec<F>,
89 replay_meta: Option<ReplayMeta>,
90 frames_added: usize,
91}
92
93impl<F> NDArrayCollector<F> {
94 pub fn new(
112 feature_adders: FeatureAdders<F>,
113 player_feature_adders: PlayerFeatureAdders<F>,
114 ) -> Self {
115 Self {
116 feature_adders,
117 player_feature_adders,
118 data: Vec::new(),
119 replay_meta: None,
120 frames_added: 0,
121 }
122 }
123
124 pub fn get_column_headers(&self) -> NDArrayColumnHeaders {
132 let global_headers = self
133 .feature_adders
134 .iter()
135 .flat_map(move |fa| {
136 fa.get_column_headers()
137 .iter()
138 .map(move |column_name| column_name.to_string())
139 })
140 .collect();
141 let player_headers = self
142 .player_feature_adders
143 .iter()
144 .flat_map(move |pfa| {
145 pfa.get_column_headers()
146 .iter()
147 .map(move |base_name| base_name.to_string())
148 })
149 .collect();
150 NDArrayColumnHeaders::new(global_headers, player_headers)
151 }
152
153 pub fn get_ndarray(self) -> SubtrActorResult<ndarray::Array2<F>> {
164 self.get_meta_and_ndarray().map(|a| a.1)
165 }
166
167 pub fn get_meta_and_ndarray(
177 self,
178 ) -> SubtrActorResult<(ReplayMetaWithHeaders, ndarray::Array2<F>)> {
179 let features_per_row = self.try_get_frame_feature_count()?;
180 let expected_length = features_per_row * self.frames_added;
181 assert!(self.data.len() == expected_length);
182 let column_headers = self.get_column_headers();
183 Ok((
184 ReplayMetaWithHeaders {
185 replay_meta: self.replay_meta.ok_or(SubtrActorError::new(
186 SubtrActorErrorVariant::CouldNotBuildReplayMeta,
187 ))?,
188 column_headers,
189 },
190 ndarray::Array2::from_shape_vec((self.frames_added, features_per_row), self.data)
191 .map_err(SubtrActorErrorVariant::NDArrayShapeError)
192 .map_err(SubtrActorError::new)?,
193 ))
194 }
195
196 pub fn process_and_get_meta_and_headers(
212 &mut self,
213 replay: &boxcars::Replay,
214 ) -> SubtrActorResult<ReplayMetaWithHeaders> {
215 let mut processor = ReplayProcessor::new(replay)?;
216 processor.process_long_enough_to_get_actor_ids()?;
217 self.maybe_set_replay_meta(&processor)?;
218 Ok(ReplayMetaWithHeaders {
219 replay_meta: self
220 .replay_meta
221 .as_ref()
222 .ok_or(SubtrActorError::new(
223 SubtrActorErrorVariant::CouldNotBuildReplayMeta,
224 ))?
225 .clone(),
226 column_headers: self.get_column_headers(),
227 })
228 }
229
230 fn try_get_frame_feature_count(&self) -> SubtrActorResult<usize> {
231 let player_count = self
232 .replay_meta
233 .as_ref()
234 .ok_or(SubtrActorError::new(
235 SubtrActorErrorVariant::CouldNotBuildReplayMeta,
236 ))?
237 .player_count();
238 let global_feature_count: usize = self
239 .feature_adders
240 .iter()
241 .map(|fa| fa.features_added())
242 .sum();
243 let player_feature_count: usize = self
244 .player_feature_adders
245 .iter() .map(|pfa| pfa.features_added() * player_count)
247 .sum();
248 Ok(global_feature_count + player_feature_count)
249 }
250
251 fn maybe_set_replay_meta(&mut self, processor: &ReplayProcessor) -> SubtrActorResult<()> {
252 if self.replay_meta.is_none() {
253 self.replay_meta = Some(processor.get_replay_meta()?);
254 }
255 Ok(())
256 }
257}
258
259impl<F> Collector for NDArrayCollector<F> {
260 fn process_frame(
261 &mut self,
262 processor: &ReplayProcessor,
263 frame: &boxcars::Frame,
264 frame_number: usize,
265 current_time: f32,
266 ) -> SubtrActorResult<collector::TimeAdvance> {
267 self.maybe_set_replay_meta(processor)?;
268
269 if !processor.ball_rigid_body_exists()? {
270 return Ok(collector::TimeAdvance::NextFrame);
271 }
272
273 for feature_adder in self.feature_adders.iter() {
274 feature_adder.add_features(
275 processor,
276 frame,
277 frame_number,
278 current_time,
279 &mut self.data,
280 )?;
281 }
282
283 for player_id in processor.iter_player_ids_in_order() {
284 for player_feature_adder in self.player_feature_adders.iter() {
285 player_feature_adder.add_features(
286 player_id,
287 processor,
288 frame,
289 frame_number,
290 current_time,
291 &mut self.data,
292 )?;
293 }
294 }
295
296 self.frames_added += 1;
297
298 Ok(collector::TimeAdvance::NextFrame)
299 }
300}
301
302impl NDArrayCollector<f32> {
303 pub fn from_strings(fa_names: &[&str], pfa_names: &[&str]) -> SubtrActorResult<Self> {
304 let feature_adders: Vec<Arc<dyn FeatureAdder<f32> + Send + Sync>> = fa_names
305 .iter()
306 .map(|name| {
307 Ok(NAME_TO_GLOBAL_FEATURE_ADDER
308 .get(name)
309 .ok_or_else(|| {
310 SubtrActorError::new(SubtrActorErrorVariant::UnknownFeatureAdderName(
311 name.to_string(),
312 ))
313 })?
314 .clone())
315 })
316 .collect::<SubtrActorResult<Vec<_>>>()?;
317 let player_feature_adders: Vec<Arc<dyn PlayerFeatureAdder<f32> + Send + Sync>> = pfa_names
318 .iter()
319 .map(|name| {
320 Ok(NAME_TO_PLAYER_FEATURE_ADDER
321 .get(name)
322 .ok_or_else(|| {
323 SubtrActorError::new(SubtrActorErrorVariant::UnknownFeatureAdderName(
324 name.to_string(),
325 ))
326 })?
327 .clone())
328 })
329 .collect::<SubtrActorResult<Vec<_>>>()?;
330 Ok(Self::new(feature_adders, player_feature_adders))
331 }
332}
333
334impl<F: TryFrom<f32> + Send + Sync + 'static> Default for NDArrayCollector<F>
335where
336 <F as TryFrom<f32>>::Error: std::fmt::Debug,
337{
338 fn default() -> Self {
339 NDArrayCollector::new(
340 vec![BallRigidBody::arc_new()],
341 vec![
342 PlayerRigidBody::arc_new(),
343 PlayerBoost::arc_new(),
344 PlayerAnyJump::arc_new(),
345 ],
346 )
347 }
348}
349
350pub trait FeatureAdder<F> {
357 fn features_added(&self) -> usize {
358 self.get_column_headers().len()
359 }
360
361 fn get_column_headers(&self) -> &[&str];
362
363 fn add_features(
364 &self,
365 processor: &ReplayProcessor,
366 frame: &boxcars::Frame,
367 frame_count: usize,
368 current_time: f32,
369 vector: &mut Vec<F>,
370 ) -> SubtrActorResult<()>;
371}
372
373pub type FeatureAdders<F> = Vec<Arc<dyn FeatureAdder<F> + Send + Sync>>;
374
375pub trait LengthCheckedFeatureAdder<F, const N: usize> {
381 fn get_column_headers_array(&self) -> &[&str; N];
382
383 fn get_features(
384 &self,
385 processor: &ReplayProcessor,
386 frame: &boxcars::Frame,
387 frame_count: usize,
388 current_time: f32,
389 ) -> SubtrActorResult<[F; N]>;
390}
391
392#[macro_export]
401macro_rules! impl_feature_adder {
402 ($struct_name:ident) => {
403 impl<F: TryFrom<f32>> FeatureAdder<F> for $struct_name<F>
404 where
405 <F as TryFrom<f32>>::Error: std::fmt::Debug,
406 {
407 fn add_features(
408 &self,
409 processor: &ReplayProcessor,
410 frame: &boxcars::Frame,
411 frame_count: usize,
412 current_time: f32,
413 vector: &mut Vec<F>,
414 ) -> SubtrActorResult<()> {
415 Ok(
416 vector.extend(self.get_features(
417 processor,
418 frame,
419 frame_count,
420 current_time,
421 )?),
422 )
423 }
424
425 fn get_column_headers(&self) -> &[&str] {
426 self.get_column_headers_array()
427 }
428 }
429 };
430}
431
432pub trait PlayerFeatureAdder<F> {
440 fn features_added(&self) -> usize {
441 self.get_column_headers().len()
442 }
443
444 fn get_column_headers(&self) -> &[&str];
445
446 fn add_features(
447 &self,
448 player_id: &PlayerId,
449 processor: &ReplayProcessor,
450 frame: &boxcars::Frame,
451 frame_count: usize,
452 current_time: f32,
453 vector: &mut Vec<F>,
454 ) -> SubtrActorResult<()>;
455}
456
457pub type PlayerFeatureAdders<F> = Vec<Arc<dyn PlayerFeatureAdder<F> + Send + Sync>>;
458
459pub trait LengthCheckedPlayerFeatureAdder<F, const N: usize> {
465 fn get_column_headers_array(&self) -> &[&str; N];
466
467 fn get_features(
468 &self,
469 player_id: &PlayerId,
470 processor: &ReplayProcessor,
471 frame: &boxcars::Frame,
472 frame_count: usize,
473 current_time: f32,
474 ) -> SubtrActorResult<[F; N]>;
475}
476
477#[macro_export]
487macro_rules! impl_player_feature_adder {
488 ($struct_name:ident) => {
489 impl<F: TryFrom<f32>> PlayerFeatureAdder<F> for $struct_name<F>
490 where
491 <F as TryFrom<f32>>::Error: std::fmt::Debug,
492 {
493 fn add_features(
494 &self,
495 player_id: &PlayerId,
496 processor: &ReplayProcessor,
497 frame: &boxcars::Frame,
498 frame_count: usize,
499 current_time: f32,
500 vector: &mut Vec<F>,
501 ) -> SubtrActorResult<()> {
502 Ok(vector.extend(self.get_features(
503 player_id,
504 processor,
505 frame,
506 frame_count,
507 current_time,
508 )?))
509 }
510
511 fn get_column_headers(&self) -> &[&str] {
512 self.get_column_headers_array()
513 }
514 }
515 };
516}
517
518impl<G, F, const N: usize> FeatureAdder<F> for (G, &[&str; N])
519where
520 G: Fn(&ReplayProcessor, &boxcars::Frame, usize, f32) -> SubtrActorResult<[F; N]>,
521{
522 fn add_features(
523 &self,
524 processor: &ReplayProcessor,
525 frame: &boxcars::Frame,
526 frame_count: usize,
527 current_time: f32,
528 vector: &mut Vec<F>,
529 ) -> SubtrActorResult<()> {
530 vector.extend(self.0(processor, frame, frame_count, current_time)?);
531 Ok(())
532 }
533
534 fn get_column_headers(&self) -> &[&str] {
535 self.1.as_slice()
536 }
537}
538
539impl<G, F, const N: usize> PlayerFeatureAdder<F> for (G, &[&str; N])
540where
541 G: Fn(&PlayerId, &ReplayProcessor, &boxcars::Frame, usize, f32) -> SubtrActorResult<[F; N]>,
542{
543 fn add_features(
544 &self,
545 player_id: &PlayerId,
546 processor: &ReplayProcessor,
547 frame: &boxcars::Frame,
548 frame_count: usize,
549 current_time: f32,
550 vector: &mut Vec<F>,
551 ) -> SubtrActorResult<()> {
552 vector.extend(self.0(
553 player_id,
554 processor,
555 frame,
556 frame_count,
557 current_time,
558 )?);
559 Ok(())
560 }
561
562 fn get_column_headers(&self) -> &[&str] {
563 self.1.as_slice()
564 }
565}
566
567#[macro_export]
602macro_rules! build_global_feature_adder {
603 ($struct_name:ident, $prop_getter:expr, $( $column_names:expr ),* $(,)?) => {
604
605 #[derive(derive_new::new)]
606 pub struct $struct_name<F> {
607 _zero: std::marker::PhantomData<F>,
608 }
609
610 impl<F: Sync + Send + TryFrom<f32> + 'static> $struct_name<F> where
611 <F as TryFrom<f32>>::Error: std::fmt::Debug,
612 {
613 pub fn arc_new() -> std::sync::Arc<dyn FeatureAdder<F> + Send + Sync + 'static> {
614 std::sync::Arc::new(Self::new())
615 }
616 }
617
618 global_feature_adder!(
619 $struct_name,
620 $prop_getter,
621 $( $column_names ),*
622 );
623 }
624}
625
626#[macro_export]
638macro_rules! global_feature_adder {
639 ($struct_name:ident, $prop_getter:expr, $( $column_names:expr ),* $(,)?) => {
640 macro_rules! _global_feature_adder {
641 ($count:ident) => {
642 impl<F: TryFrom<f32>> LengthCheckedFeatureAdder<F, $count> for $struct_name<F>
643 where
644 <F as TryFrom<f32>>::Error: std::fmt::Debug,
645 {
646 fn get_column_headers_array(&self) -> &[&str; $count] {
647 &[$( $column_names ),*]
648 }
649
650 fn get_features(
651 &self,
652 processor: &ReplayProcessor,
653 frame: &boxcars::Frame,
654 frame_count: usize,
655 current_time: f32,
656 ) -> SubtrActorResult<[F; $count]> {
657 $prop_getter(self, processor, frame, frame_count, current_time)
658 }
659 }
660
661 impl_feature_adder!($struct_name);
662 };
663 }
664 paste::paste! {
665 const [<$struct_name:snake:upper _LENGTH>]: usize = [$($column_names),*].len();
666 _global_feature_adder!([<$struct_name:snake:upper _LENGTH>]);
667 }
668 }
669}
670
671#[macro_export]
730macro_rules! build_player_feature_adder {
731 ($struct_name:ident, $prop_getter:expr, $( $column_names:expr ),* $(,)?) => {
732 #[derive(derive_new::new)]
733 pub struct $struct_name<F> {
734 _zero: std::marker::PhantomData<F>,
735 }
736
737 impl<F: Sync + Send + TryFrom<f32> + 'static> $struct_name<F> where
738 <F as TryFrom<f32>>::Error: std::fmt::Debug,
739 {
740 pub fn arc_new() -> std::sync::Arc<dyn PlayerFeatureAdder<F> + Send + Sync + 'static> {
741 std::sync::Arc::new(Self::new())
742 }
743 }
744
745 player_feature_adder!(
746 $struct_name,
747 $prop_getter,
748 $( $column_names ),*
749 );
750 }
751}
752
753#[macro_export]
766macro_rules! player_feature_adder {
767 ($struct_name:ident, $prop_getter:expr, $( $column_names:expr ),* $(,)?) => {
768 macro_rules! _player_feature_adder {
769 ($count:ident) => {
770 impl<F: TryFrom<f32>> LengthCheckedPlayerFeatureAdder<F, $count> for $struct_name<F>
771 where
772 <F as TryFrom<f32>>::Error: std::fmt::Debug,
773 {
774 fn get_column_headers_array(&self) -> &[&str; $count] {
775 &[$( $column_names ),*]
776 }
777
778 fn get_features(
779 &self,
780 player_id: &PlayerId,
781 processor: &ReplayProcessor,
782 frame: &boxcars::Frame,
783 frame_count: usize,
784 current_time: f32,
785 ) -> SubtrActorResult<[F; $count]> {
786 $prop_getter(self, player_id, processor, frame, frame_count, current_time)
787 }
788 }
789
790 impl_player_feature_adder!($struct_name);
791 };
792 }
793 paste::paste! {
794 const [<$struct_name:snake:upper _LENGTH>]: usize = [$($column_names),*].len();
795 _player_feature_adder!([<$struct_name:snake:upper _LENGTH>]);
796 }
797 }
798}
799
800pub fn convert_float_conversion_error<T>(_: T) -> SubtrActorError {
803 SubtrActorError::new(SubtrActorErrorVariant::FloatConversionError)
804}
805
806#[macro_export]
817macro_rules! convert_all {
818 ($err:expr, $( $item:expr ),* $(,)?) => {{
819 Ok([
820 $( $item.try_into().map_err($err)? ),*
821 ])
822 }};
823}
824
825#[macro_export]
847macro_rules! convert_all_floats {
848 ($( $item:expr ),* $(,)?) => {{
849 convert_all!(convert_float_conversion_error, $( $item ),*)
850 }};
851}
852
853fn or_zero_boxcars_3f() -> boxcars::Vector3f {
854 boxcars::Vector3f {
855 x: 0.0,
856 y: 0.0,
857 z: 0.0,
858 }
859}
860
861type RigidBodyArrayResult<F> = SubtrActorResult<[F; 12]>;
862
863pub fn get_rigid_body_properties<F: TryFrom<f32>>(
873 rigid_body: &boxcars::RigidBody,
874) -> RigidBodyArrayResult<F>
875where
876 <F as TryFrom<f32>>::Error: std::fmt::Debug,
877{
878 let linear_velocity = rigid_body
879 .linear_velocity
880 .unwrap_or_else(or_zero_boxcars_3f);
881 let angular_velocity = rigid_body
882 .angular_velocity
883 .unwrap_or_else(or_zero_boxcars_3f);
884 let rotation = rigid_body.rotation;
885 let location = rigid_body.location;
886 let (rx, ry, rz) =
887 glam::quat(rotation.x, rotation.y, rotation.z, rotation.w).to_euler(glam::EulerRot::XYZ);
888 convert_all_floats!(
889 location.x,
890 location.y,
891 location.z,
892 rx,
893 ry,
894 rz,
895 linear_velocity.x,
896 linear_velocity.y,
897 linear_velocity.z,
898 angular_velocity.x,
899 angular_velocity.y,
900 angular_velocity.z,
901 )
902}
903
904pub fn get_rigid_body_properties_no_velocities<F: TryFrom<f32>>(
913 rigid_body: &boxcars::RigidBody,
914) -> SubtrActorResult<[F; 7]>
915where
916 <F as TryFrom<f32>>::Error: std::fmt::Debug,
917{
918 let rotation = rigid_body.rotation;
919 let location = rigid_body.location;
920 convert_all_floats!(
921 location.x, location.y, location.z, rotation.x, rotation.y, rotation.z, rotation.w
922 )
923}
924
925fn default_rb_state<F: TryFrom<f32>>() -> RigidBodyArrayResult<F>
926where
927 <F as TryFrom<f32>>::Error: std::fmt::Debug,
928{
929 convert_all!(
930 convert_float_conversion_error,
931 0.0,
935 0.0,
936 0.0,
937 0.0,
938 0.0,
939 0.0,
940 0.0,
941 0.0,
942 0.0,
943 0.0,
944 0.0,
945 0.0,
946 )
947}
948
949fn default_rb_state_no_velocities<F: TryFrom<f32>>() -> SubtrActorResult<[F; 7]>
950where
951 <F as TryFrom<f32>>::Error: std::fmt::Debug,
952{
953 convert_all_floats!(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,)
954}
955
956build_global_feature_adder!(
957 SecondsRemaining,
958 |_, processor: &ReplayProcessor, _frame, _index, _current_time| {
959 convert_all_floats!(processor.get_seconds_remaining()?.clone() as f32)
960 },
961 "seconds remaining"
962);
963
964build_global_feature_adder!(
965 CurrentTime,
966 |_, _processor, _frame, _index, current_time: f32| { convert_all_floats!(current_time) },
967 "current time"
968);
969
970build_global_feature_adder!(
972 ReplicatedStateName,
973 |_, processor: &ReplayProcessor, _frame, _index, _current_time| {
974 convert_all_floats!(processor.get_replicated_state_name().unwrap_or(0) as f32)
975 },
976 "game state"
977);
978
979build_global_feature_adder!(
981 ReplicatedGameStateTimeRemaining,
982 |_, processor: &ReplayProcessor, _frame, _index, _current_time| {
983 convert_all_floats!(processor
984 .get_replicated_game_state_time_remaining()
985 .unwrap_or(0) as f32)
986 },
987 "kickoff countdown"
988);
989
990build_global_feature_adder!(
992 BallHasBeenHit,
993 |_, processor: &ReplayProcessor, _frame, _index, _current_time| {
994 convert_all_floats!(if processor.get_ball_has_been_hit().unwrap_or(false) {
995 1.0
996 } else {
997 0.0
998 })
999 },
1000 "ball has been hit"
1001);
1002
1003build_global_feature_adder!(
1004 FrameTime,
1005 |_, _processor, frame: &boxcars::Frame, _index, _current_time| {
1006 convert_all_floats!(frame.time)
1007 },
1008 "frame time"
1009);
1010
1011build_global_feature_adder!(
1012 BallRigidBody,
1013 |_, processor: &ReplayProcessor, _frame, _index, _current_time| {
1014 get_rigid_body_properties(processor.get_ball_rigid_body()?)
1015 },
1016 "Ball - position x",
1017 "Ball - position y",
1018 "Ball - position z",
1019 "Ball - rotation x",
1020 "Ball - rotation y",
1021 "Ball - rotation z",
1022 "Ball - linear velocity x",
1023 "Ball - linear velocity y",
1024 "Ball - linear velocity z",
1025 "Ball - angular velocity x",
1026 "Ball - angular velocity y",
1027 "Ball - angular velocity z",
1028);
1029
1030build_global_feature_adder!(
1031 BallRigidBodyNoVelocities,
1032 |_, processor: &ReplayProcessor, _frame, _index, _current_time| {
1033 get_rigid_body_properties_no_velocities(processor.get_ball_rigid_body()?)
1034 },
1035 "Ball - position x",
1036 "Ball - position y",
1037 "Ball - position z",
1038 "Ball - rotation x",
1039 "Ball - rotation y",
1040 "Ball - rotation z",
1041 "Ball - rotation w",
1042);
1043
1044build_global_feature_adder!(
1047 VelocityAddedBallRigidBodyNoVelocities,
1048 |_, processor: &ReplayProcessor, _frame, _index, current_time: f32| {
1049 get_rigid_body_properties_no_velocities(
1050 &processor.get_velocity_applied_ball_rigid_body(current_time)?,
1051 )
1052 },
1053 "Ball - position x",
1054 "Ball - position y",
1055 "Ball - position z",
1056 "Ball - rotation x",
1057 "Ball - rotation y",
1058 "Ball - rotation z",
1059 "Ball - rotation w",
1060);
1061
1062#[derive(derive_new::new)]
1063pub struct InterpolatedBallRigidBodyNoVelocities<F> {
1064 close_enough_to_frame_time: f32,
1065 _zero: std::marker::PhantomData<F>,
1066}
1067
1068impl<F> InterpolatedBallRigidBodyNoVelocities<F> {
1069 pub fn arc_new(close_enough_to_frame_time: f32) -> Arc<Self> {
1070 Arc::new(Self::new(close_enough_to_frame_time))
1071 }
1072}
1073
1074global_feature_adder!(
1075 InterpolatedBallRigidBodyNoVelocities,
1076 |s: &InterpolatedBallRigidBodyNoVelocities<F>,
1077 processor: &ReplayProcessor,
1078 _frame: &boxcars::Frame,
1079 _index,
1080 current_time: f32| {
1081 processor
1082 .get_interpolated_ball_rigid_body(current_time, s.close_enough_to_frame_time)
1083 .map(|v| get_rigid_body_properties_no_velocities(&v))
1084 .unwrap_or_else(|_| default_rb_state_no_velocities())
1085 },
1086 "Ball - position x",
1087 "Ball - position y",
1088 "Ball - position z",
1089 "Ball - rotation x",
1090 "Ball - rotation y",
1091 "Ball - rotation z",
1092 "Ball - rotation w",
1093);
1094
1095build_player_feature_adder!(
1096 PlayerRigidBody,
1097 |_, player_id: &PlayerId, processor: &ReplayProcessor, _frame, _index, _current_time: f32| {
1098 if let Ok(rb) = processor.get_player_rigid_body(player_id) {
1099 get_rigid_body_properties(rb)
1100 } else {
1101 default_rb_state()
1102 }
1103 },
1104 "position x",
1105 "position y",
1106 "position z",
1107 "rotation x",
1108 "rotation y",
1109 "rotation z",
1110 "linear velocity x",
1111 "linear velocity y",
1112 "linear velocity z",
1113 "angular velocity x",
1114 "angular velocity y",
1115 "angular velocity z",
1116);
1117
1118build_player_feature_adder!(
1119 PlayerRigidBodyNoVelocities,
1120 |_, player_id: &PlayerId, processor: &ReplayProcessor, _frame, _index, _current_time: f32| {
1121 if let Ok(rb) = processor.get_player_rigid_body(player_id) {
1122 get_rigid_body_properties_no_velocities(rb)
1123 } else {
1124 default_rb_state_no_velocities()
1125 }
1126 },
1127 "position x",
1128 "position y",
1129 "position z",
1130 "rotation x",
1131 "rotation y",
1132 "rotation z",
1133 "rotation w"
1134);
1135
1136build_player_feature_adder!(
1139 VelocityAddedPlayerRigidBodyNoVelocities,
1140 |_, player_id: &PlayerId, processor: &ReplayProcessor, _frame, _index, current_time: f32| {
1141 if let Ok(rb) = processor.get_velocity_applied_player_rigid_body(player_id, current_time) {
1142 get_rigid_body_properties_no_velocities(&rb)
1143 } else {
1144 default_rb_state_no_velocities()
1145 }
1146 },
1147 "position x",
1148 "position y",
1149 "position z",
1150 "rotation x",
1151 "rotation y",
1152 "rotation z",
1153 "rotation w"
1154);
1155
1156#[derive(derive_new::new)]
1157pub struct InterpolatedPlayerRigidBodyNoVelocities<F> {
1158 close_enough_to_frame_time: f32,
1159 _zero: std::marker::PhantomData<F>,
1160}
1161
1162impl<F> InterpolatedPlayerRigidBodyNoVelocities<F> {
1163 pub fn arc_new(close_enough_to_frame_time: f32) -> Arc<Self> {
1164 Arc::new(Self::new(close_enough_to_frame_time))
1165 }
1166}
1167
1168player_feature_adder!(
1169 InterpolatedPlayerRigidBodyNoVelocities,
1170 |s: &InterpolatedPlayerRigidBodyNoVelocities<F>,
1171 player_id: &PlayerId,
1172 processor: &ReplayProcessor,
1173 _frame: &boxcars::Frame,
1174 _index,
1175 current_time: f32| {
1176 processor
1177 .get_interpolated_player_rigid_body(
1178 player_id,
1179 current_time,
1180 s.close_enough_to_frame_time,
1181 )
1182 .map(|v| get_rigid_body_properties_no_velocities(&v))
1183 .unwrap_or_else(|_| default_rb_state_no_velocities())
1184 },
1185 "i position x",
1186 "i position y",
1187 "i position z",
1188 "i rotation x",
1189 "i rotation y",
1190 "i rotation z",
1191 "i rotation w"
1192);
1193
1194build_player_feature_adder!(
1195 PlayerBoost,
1196 |_, player_id: &PlayerId, processor: &ReplayProcessor, _frame, _index, _current_time: f32| {
1197 convert_all_floats!(processor.get_player_boost_level(player_id).unwrap_or(0.0))
1198 },
1199 "boost level"
1200);
1201
1202fn u8_get_f32(v: u8) -> SubtrActorResult<f32> {
1203 Ok(v.into())
1204}
1205
1206build_player_feature_adder!(
1207 PlayerJump,
1208 |_,
1209 player_id: &PlayerId,
1210 processor: &ReplayProcessor,
1211 _frame,
1212 _frame_number,
1213 _current_time: f32| {
1214 convert_all_floats!(
1215 processor
1216 .get_dodge_active(player_id)
1217 .and_then(u8_get_f32)
1218 .unwrap_or(0.0),
1219 processor
1220 .get_jump_active(player_id)
1221 .and_then(u8_get_f32)
1222 .unwrap_or(0.0),
1223 processor
1224 .get_double_jump_active(player_id)
1225 .and_then(u8_get_f32)
1226 .unwrap_or(0.0),
1227 )
1228 },
1229 "dodge active",
1230 "jump active",
1231 "double jump active"
1232);
1233
1234build_player_feature_adder!(
1235 PlayerAnyJump,
1236 |_,
1237 player_id: &PlayerId,
1238 processor: &ReplayProcessor,
1239 _frame,
1240 _frame_number,
1241 _current_time: f32| {
1242 let dodge_is_active = processor.get_dodge_active(player_id).unwrap_or(0) % 2;
1243 let jump_is_active = processor.get_jump_active(player_id).unwrap_or(0) % 2;
1244 let double_jump_is_active = processor.get_double_jump_active(player_id).unwrap_or(0) % 2;
1245 let value: f32 = [dodge_is_active, jump_is_active, double_jump_is_active]
1246 .into_iter()
1247 .enumerate()
1248 .map(|(index, is_active)| (1 << index) * is_active)
1249 .sum::<u8>() as f32;
1250 convert_all_floats!(value)
1251 },
1252 "any_jump_active"
1253);
1254
1255const DEMOLISH_APPEARANCE_FRAME_COUNT: usize = 30;
1256
1257build_player_feature_adder!(
1258 PlayerDemolishedBy,
1259 |_,
1260 player_id: &PlayerId,
1261 processor: &ReplayProcessor,
1262 _frame,
1263 frame_number,
1264 _current_time: f32| {
1265 let demolisher_index = processor
1266 .demolishes
1267 .iter()
1268 .find(|demolish_info| {
1269 &demolish_info.victim == player_id
1270 && frame_number - demolish_info.frame < DEMOLISH_APPEARANCE_FRAME_COUNT
1271 })
1272 .map(|demolish_info| {
1273 processor
1274 .iter_player_ids_in_order()
1275 .position(|player_id| player_id == &demolish_info.attacker)
1276 .unwrap_or_else(|| processor.iter_player_ids_in_order().count())
1277 })
1278 .and_then(|v| i32::try_from(v).ok())
1279 .unwrap_or(-1);
1280 convert_all_floats!(demolisher_index as f32)
1281 },
1282 "player demolished by"
1283);
1284
1285build_player_feature_adder!(
1286 PlayerRigidBodyQuaternions,
1287 |_, player_id: &PlayerId, processor: &ReplayProcessor, _frame, _index, _current_time: f32| {
1288 if let Ok(rb) = processor.get_player_rigid_body(player_id) {
1289 let rotation = rb.rotation;
1290 let location = rb.location;
1291 convert_all_floats!(
1292 location.x, location.y, location.z, rotation.x, rotation.y, rotation.z, rotation.w
1293 )
1294 } else {
1295 convert_all_floats!(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0)
1296 }
1297 },
1298 "position x",
1299 "position y",
1300 "position z",
1301 "quaternion x",
1302 "quaternion y",
1303 "quaternion z",
1304 "quaternion w"
1305);
1306
1307build_global_feature_adder!(
1308 BallRigidBodyQuaternions,
1309 |_, processor: &ReplayProcessor, _frame, _index, _current_time| {
1310 let rb = processor.get_ball_rigid_body()?;
1311 let rotation = rb.rotation;
1312 let location = rb.location;
1313 convert_all_floats!(
1314 location.x, location.y, location.z, rotation.x, rotation.y, rotation.z, rotation.w
1315 )
1316 },
1317 "Ball - position x",
1318 "Ball - position y",
1319 "Ball - position z",
1320 "Ball - quaternion x",
1321 "Ball - quaternion y",
1322 "Ball - quaternion z",
1323 "Ball - quaternion w"
1324);
1325
1326lazy_static! {
1327 static ref NAME_TO_GLOBAL_FEATURE_ADDER: std::collections::HashMap<&'static str, Arc<dyn FeatureAdder<f32> + Send + Sync + 'static>> = {
1328 let mut m: std::collections::HashMap<
1329 &'static str,
1330 Arc<dyn FeatureAdder<f32> + Send + Sync + 'static>,
1331 > = std::collections::HashMap::new();
1332 macro_rules! insert_adder {
1333 ($adder_name:ident, $( $arguments:expr ),*) => {
1334 m.insert(stringify!($adder_name), $adder_name::<f32>::arc_new($ ( $arguments ),*));
1335 };
1336 ($adder_name:ident) => {
1337 insert_adder!($adder_name,)
1338 }
1339 }
1340 insert_adder!(BallRigidBody);
1341 insert_adder!(BallRigidBodyNoVelocities);
1342 insert_adder!(BallRigidBodyQuaternions);
1343 insert_adder!(VelocityAddedBallRigidBodyNoVelocities);
1344 insert_adder!(InterpolatedBallRigidBodyNoVelocities, 0.0);
1345 insert_adder!(SecondsRemaining);
1346 insert_adder!(CurrentTime);
1347 insert_adder!(FrameTime);
1348 insert_adder!(ReplicatedStateName);
1349 insert_adder!(ReplicatedGameStateTimeRemaining);
1350 insert_adder!(BallHasBeenHit);
1351 m
1352 };
1353 static ref NAME_TO_PLAYER_FEATURE_ADDER: std::collections::HashMap<
1354 &'static str,
1355 Arc<dyn PlayerFeatureAdder<f32> + Send + Sync + 'static>,
1356 > = {
1357 let mut m: std::collections::HashMap<
1358 &'static str,
1359 Arc<dyn PlayerFeatureAdder<f32> + Send + Sync + 'static>,
1360 > = std::collections::HashMap::new();
1361 macro_rules! insert_adder {
1362 ($adder_name:ident, $( $arguments:expr ),*) => {
1363 m.insert(stringify!($adder_name), $adder_name::<f32>::arc_new($ ( $arguments ),*));
1364 };
1365 ($adder_name:ident) => {
1366 insert_adder!($adder_name,)
1367 };
1368 }
1369 insert_adder!(PlayerRigidBody);
1370 insert_adder!(PlayerRigidBodyNoVelocities);
1371 insert_adder!(PlayerRigidBodyQuaternions);
1372 insert_adder!(VelocityAddedPlayerRigidBodyNoVelocities);
1373 insert_adder!(InterpolatedPlayerRigidBodyNoVelocities, 0.003);
1374 insert_adder!(PlayerBoost);
1375 insert_adder!(PlayerJump);
1376 insert_adder!(PlayerAnyJump);
1377 insert_adder!(PlayerDemolishedBy);
1378 m
1379 };
1380}