subtr-actor 1.0.0

Rocket League replay transformer
Documentation
use super::*;

const SOCCAR_CEILING_Z: f32 = 2044.0;
const CEILING_CONTACT_MAX_GAP: f32 = 90.0;
const CEILING_CONTACT_MIN_ROOF_ALIGNMENT: f32 = 0.72;
const CEILING_SHOT_MAX_TOUCH_AFTER_CONTACT_SECONDS: f32 = 1.35;
const CEILING_SHOT_MIN_TOUCH_SEPARATION: f32 = 120.0;
const CEILING_SHOT_MIN_PLAYER_HEIGHT: f32 = 260.0;
const CEILING_SHOT_MIN_BALL_HEIGHT: f32 = 220.0;
const CEILING_SHOT_MIN_FORWARD_ALIGNMENT: f32 = 0.12;
const CEILING_SHOT_MIN_FORWARD_APPROACH_SPEED: f32 = 90.0;
const CEILING_SHOT_MIN_BALL_SPEED_CHANGE: f32 = 120.0;
const CEILING_SHOT_MIN_CONFIDENCE: f32 = 0.54;

#[derive(Debug, Clone, PartialEq, Serialize, ts_rs::TS)]
#[ts(export)]
pub struct CeilingShotEvent {
    pub time: f32,
    pub frame: usize,
    #[ts(as = "crate::interop::ts_bindings::RemoteIdTs")]
    pub player: PlayerId,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub player_position: Option<[f32; 3]>,
    pub is_team_0: bool,
    pub ceiling_contact_time: f32,
    pub ceiling_contact_frame: usize,
    pub time_since_ceiling_contact: f32,
    pub ceiling_contact_position: [f32; 3],
    pub touch_position: [f32; 3],
    pub local_ball_position: [f32; 3],
    pub separation_from_ceiling: f32,
    pub roof_alignment: f32,
    pub forward_alignment: f32,
    pub forward_approach_speed: f32,
    pub ball_speed_change: f32,
    pub confidence: f32,
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct RecentCeilingContact {
    time: f32,
    frame: usize,
    position: [f32; 3],
    roof_alignment: f32,
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct CeilingContactObservation {
    position: glam::Vec3,
    roof_alignment: f32,
}

#[derive(Debug, Clone, Default, PartialEq)]
pub struct CeilingShotCalculator {
    events: EventStream<CeilingShotEvent>,
    recent_ceiling_contacts: HashMap<PlayerId, RecentCeilingContact>,
    previous_ball_velocity: Option<glam::Vec3>,
}

impl CeilingShotCalculator {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn events(&self) -> &[CeilingShotEvent] {
        self.events.all()
    }

    pub fn new_events(&self) -> &[CeilingShotEvent] {
        self.events.new_events()
    }

    fn normalize_score(value: f32, min_value: f32, max_value: f32) -> f32 {
        if max_value <= min_value {
            return 0.0;
        }

        ((value - min_value) / (max_value - min_value)).clamp(0.0, 1.0)
    }

    fn ball_speed_change(
        frame: &FrameInfo,
        ball: &BallFrameState,
        previous_ball_velocity: Option<glam::Vec3>,
    ) -> f32 {
        const BALL_GRAVITY_Z: f32 = -650.0;

        let Some(ball) = ball.sample() else {
            return 0.0;
        };
        let Some(previous_ball_velocity) = previous_ball_velocity else {
            return 0.0;
        };

        let expected_linear_delta = glam::Vec3::new(0.0, 0.0, BALL_GRAVITY_Z * frame.dt.max(0.0));
        let residual_linear_impulse =
            ball.velocity() - previous_ball_velocity - expected_linear_delta;
        residual_linear_impulse.length()
    }

    fn ceiling_contact_observation(player: &PlayerSample) -> Option<CeilingContactObservation> {
        let rigid_body = player.rigid_body.as_ref()?;
        let position = player.position()?;
        let gap_to_ceiling = SOCCAR_CEILING_Z - position.z;
        if !(0.0..=CEILING_CONTACT_MAX_GAP).contains(&gap_to_ceiling) {
            return None;
        }

        let up = quat_to_glam(&rigid_body.rotation) * glam::Vec3::Z;
        let roof_alignment = (-up).dot(glam::Vec3::Z);
        if roof_alignment < CEILING_CONTACT_MIN_ROOF_ALIGNMENT {
            return None;
        }

        Some(CeilingContactObservation {
            position,
            roof_alignment,
        })
    }

    fn update_recent_ceiling_contacts(&mut self, frame: &FrameInfo, players: &PlayerFrameState) {
        for player in &players.players {
            let observation = Self::ceiling_contact_observation(player);
            let Some(observation) = observation else {
                continue;
            };

            self.recent_ceiling_contacts.insert(
                player.player_id.clone(),
                RecentCeilingContact {
                    time: frame.time,
                    frame: frame.frame_number,
                    position: observation.position.to_array(),
                    roof_alignment: observation.roof_alignment,
                },
            );
        }
    }

    fn prune_recent_ceiling_contacts(&mut self, current_time: f32) {
        self.recent_ceiling_contacts.retain(|_, contact| {
            current_time - contact.time <= CEILING_SHOT_MAX_TOUCH_AFTER_CONTACT_SECONDS
        });
    }

    fn candidate_event(
        &self,
        ball: &BallFrameState,
        player: &PlayerSample,
        touch_event: &TouchEvent,
        recent_contact: RecentCeilingContact,
        ball_speed_change: f32,
    ) -> Option<CeilingShotEvent> {
        let ball = ball.sample()?;
        let player_position = player.position()?;
        let player_rigid_body = player.rigid_body.as_ref()?;
        let ball_position = ball.position();

        if player_position.z < CEILING_SHOT_MIN_PLAYER_HEIGHT
            || ball_position.z < CEILING_SHOT_MIN_BALL_HEIGHT
        {
            return None;
        }

        let time_since_ceiling_contact = touch_event.time - recent_contact.time;
        if !(0.0..=CEILING_SHOT_MAX_TOUCH_AFTER_CONTACT_SECONDS)
            .contains(&time_since_ceiling_contact)
        {
            return None;
        }

        let separation_from_ceiling = SOCCAR_CEILING_Z - player_position.z;
        if separation_from_ceiling < CEILING_SHOT_MIN_TOUCH_SEPARATION {
            return None;
        }

        let relative_ball_position = ball_position - player_position;
        if relative_ball_position.length_squared() <= f32::EPSILON {
            return None;
        }

        let player_rotation = quat_to_glam(&player_rigid_body.rotation);
        let local_ball_position = player_rotation.inverse() * relative_ball_position;
        if local_ball_position.x < -120.0
            || local_ball_position.y.abs() > 260.0
            || local_ball_position.z.abs() > 240.0
        {
            return None;
        }

        let to_ball = relative_ball_position.normalize_or_zero();
        let forward = player_rotation * glam::Vec3::X;
        let forward_alignment = forward.dot(to_ball);
        if forward_alignment < CEILING_SHOT_MIN_FORWARD_ALIGNMENT {
            return None;
        }

        let forward_approach_speed = player.velocity().unwrap_or(glam::Vec3::ZERO).dot(to_ball);
        if forward_approach_speed < CEILING_SHOT_MIN_FORWARD_APPROACH_SPEED {
            return None;
        }
        if ball_speed_change < CEILING_SHOT_MIN_BALL_SPEED_CHANGE {
            return None;
        }

        let timing_score = 1.0
            - Self::normalize_score(
                time_since_ceiling_contact,
                0.10,
                CEILING_SHOT_MAX_TOUCH_AFTER_CONTACT_SECONDS,
            );
        let separation_score = Self::normalize_score(separation_from_ceiling, 140.0, 520.0);
        let height_score = Self::normalize_score(
            player_position.z.max(ball_position.z),
            CEILING_SHOT_MIN_BALL_HEIGHT,
            900.0,
        );
        let alignment_score =
            Self::normalize_score(forward_alignment, CEILING_SHOT_MIN_FORWARD_ALIGNMENT, 0.92);
        let approach_score = Self::normalize_score(
            forward_approach_speed,
            CEILING_SHOT_MIN_FORWARD_APPROACH_SPEED,
            900.0,
        );
        let impulse_score =
            Self::normalize_score(ball_speed_change, CEILING_SHOT_MIN_BALL_SPEED_CHANGE, 900.0);
        let contact_score = Self::normalize_score(
            recent_contact.roof_alignment,
            CEILING_CONTACT_MIN_ROOF_ALIGNMENT,
            0.98,
        );

        let confidence = 0.20 * timing_score
            + 0.15 * separation_score
            + 0.12 * height_score
            + 0.17 * alignment_score
            + 0.16 * approach_score
            + 0.10 * impulse_score
            + 0.10 * contact_score;
        if confidence < CEILING_SHOT_MIN_CONFIDENCE {
            return None;
        }

        Some(CeilingShotEvent {
            time: touch_event.time,
            frame: touch_event.frame,
            player: player.player_id.clone(),
            player_position: Some(player_position.to_array()),
            is_team_0: player.is_team_0,
            ceiling_contact_time: recent_contact.time,
            ceiling_contact_frame: recent_contact.frame,
            time_since_ceiling_contact,
            ceiling_contact_position: recent_contact.position,
            touch_position: ball_position.to_array(),
            local_ball_position: local_ball_position.to_array(),
            separation_from_ceiling,
            roof_alignment: recent_contact.roof_alignment,
            forward_alignment,
            forward_approach_speed,
            ball_speed_change,
            confidence,
        })
    }

    fn apply_touch_events(
        &mut self,
        frame: &FrameInfo,
        ball: &BallFrameState,
        players: &PlayerFrameState,
        touch_events: &[TouchEvent],
    ) {
        let ball_speed_change = Self::ball_speed_change(frame, ball, self.previous_ball_velocity);

        for touch_event in touch_events {
            let Some(player_id) = touch_event.player.as_ref() else {
                continue;
            };
            let Some(player) = players
                .players
                .iter()
                .find(|player| &player.player_id == player_id)
            else {
                continue;
            };
            let Some(recent_contact) = self.recent_ceiling_contacts.get(player_id).copied() else {
                continue;
            };
            let Some(event) =
                self.candidate_event(ball, player, touch_event, recent_contact, ball_speed_change)
            else {
                continue;
            };
            self.events.push(event);
        }
    }

    fn reset_live_play_state(&mut self, ball: &BallFrameState) {
        self.recent_ceiling_contacts.clear();
        self.previous_ball_velocity = ball.velocity();
    }

    pub fn update_parts(
        &mut self,
        frame: &FrameInfo,
        ball: &BallFrameState,
        players: &PlayerFrameState,
        touch_events: &[TouchEvent],
        live_play_state: &LivePlayState,
    ) -> SubtrActorResult<()> {
        self.events.begin_update();
        if !live_play_state.is_live_play {
            self.reset_live_play_state(ball);
            return Ok(());
        }
        self.prune_recent_ceiling_contacts(frame.time);
        self.apply_touch_events(frame, ball, players, touch_events);
        self.update_recent_ceiling_contacts(frame, players);
        self.previous_ball_velocity = ball.velocity();
        Ok(())
    }
}