tdkpin/tdkpin-rs/src/simulation.rs

use std::{
    fs,
    path::{Path, PathBuf},
    str::FromStr,
};

use serde::Serialize;

use crate::game::{ClawSpriteBank, Controls, Event, Game};

pub const SIMULATION_HZ: u32 = 120;
const SIMULATION_HZ_F64: f64 = 120.0;
const SIMULATION_DT: f32 = 1.0 / 120.0;
const MAX_SIMULATION_STEPS: u64 = 72_000;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Scenario {
    Loading,
    Attract,
    HighScores,
    NameEntry,
    Effect7,
    Autoplay,
    Launcher,
    Flippers,
    Panel,
    Targets,
    Claw1,
    Claw6,
    Claw7,
    Claw18,
}

impl Scenario {
    pub const fn name(self) -> &'static str {
        match self {
            Self::Loading => "loading",
            Self::Attract => "attract",
            Self::HighScores => "highscores",
            Self::NameEntry => "name-entry",
            Self::Effect7 => "effect-7",
            Self::Autoplay => "autoplay",
            Self::Launcher => "launcher",
            Self::Flippers => "flippers",
            Self::Panel => "panel",
            Self::Targets => "targets",
            Self::Claw1 => "claw-1",
            Self::Claw6 => "claw-6",
            Self::Claw7 => "claw-7",
            Self::Claw18 => "claw-18",
        }
    }

    const fn claw_terminal(self) -> Option<u8> {
        match self {
            Self::Claw1 => Some(1),
            Self::Claw6 => Some(6),
            Self::Claw7 => Some(7),
            Self::Claw18 => Some(18),
            Self::Loading
            | Self::Attract
            | Self::HighScores
            | Self::NameEntry
            | Self::Effect7
            | Self::Autoplay
            | Self::Launcher
            | Self::Flippers
            | Self::Panel
            | Self::Targets => None,
        }
    }
}

impl FromStr for Scenario {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "loading" => Ok(Self::Loading),
            "attract" => Ok(Self::Attract),
            "highscores" => Ok(Self::HighScores),
            "name-entry" => Ok(Self::NameEntry),
            "effect-7" => Ok(Self::Effect7),
            "autoplay" => Ok(Self::Autoplay),
            "launcher" => Ok(Self::Launcher),
            "flippers" => Ok(Self::Flippers),
            "panel" => Ok(Self::Panel),
            "targets" => Ok(Self::Targets),
            "claw-1" => Ok(Self::Claw1),
            "claw-6" => Ok(Self::Claw6),
            "claw-7" => Ok(Self::Claw7),
            "claw-18" => Ok(Self::Claw18),
            _ => Err(format!(
                "unknown scenario {value:?}; expected loading, attract, highscores, name-entry, effect-7, autoplay, launcher, flippers, panel, targets, claw-1, claw-6, claw-7, or claw-18"
            )),
        }
    }
}

#[derive(Debug, PartialEq)]
pub struct Request {
    pub scenario: Scenario,
    pub target_step: u64,
    pub screenshot: Option<PathBuf>,
    pub trace: Option<PathBuf>,
    pub seed: u32,
}

impl Request {
    pub fn parse(args: impl IntoIterator<Item = String>) -> Result<Option<Self>, String> {
        let mut arguments = args.into_iter();
        let _program = arguments.next();
        let Some(first) = arguments.next() else {
            return Ok(None);
        };
        if first != "--simulate" {
            return Err(format!("unknown argument {first:?}\n\n{}", usage()));
        }
        let scenario = arguments
            .next()
            .ok_or_else(|| format!("--simulate requires a scenario\n\n{}", usage()))?
            .parse()?;
        let mut request = Self {
            scenario,
            target_step: u64::from(SIMULATION_HZ),
            screenshot: None,
            trace: None,
            seed: 0x5444_4b50,
        };

        while let Some(option) = arguments.next() {
            let value = arguments
                .next()
                .ok_or_else(|| format!("{option} requires a value"))?;
            match option.as_str() {
                "--at" => request.target_step = seconds_to_step(&value)?,
                "--step" => {
                    request.target_step = value
                        .parse()
                        .map_err(|_| format!("invalid step {value:?}"))?;
                }
                "--screenshot" => request.screenshot = Some(PathBuf::from(value)),
                "--trace" => request.trace = Some(PathBuf::from(value)),
                "--seed" => {
                    request.seed = value
                        .parse()
                        .map_err(|_| format!("invalid seed {value:?}"))?;
                }
                _ => return Err(format!("unknown simulation option {option:?}")),
            }
        }
        if request.target_step > MAX_SIMULATION_STEPS {
            return Err(format!(
                "requested step {} exceeds the ten-minute validation limit {MAX_SIMULATION_STEPS}",
                request.target_step
            ));
        }
        Ok(Some(request))
    }
}

#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
fn seconds_to_step(value: &str) -> Result<u64, String> {
    let seconds: f64 = value
        .parse()
        .map_err(|_| format!("invalid simulation time {value:?}"))?;
    if !seconds.is_finite() || seconds < 0.0 {
        return Err("simulation time must be a finite non-negative number".to_owned());
    }
    if seconds > 600.0 {
        return Err("simulation time exceeds the ten-minute validation limit".to_owned());
    }
    let step = (seconds * SIMULATION_HZ_F64).round();
    Ok(step as u64)
}

pub const fn usage() -> &'static str {
    "Usage:\n  tdkpin-rs\n  tdkpin-rs --simulate SCENARIO [--at SECONDS | --step N] [--screenshot FILE.png] [--trace FILE.json] [--seed N]\n\nScenarios: loading, attract, highscores, name-entry, effect-7, autoplay, launcher, flippers, panel, targets, claw-1, claw-6, claw-7, claw-18"
}

#[derive(Debug, Serialize, PartialEq)]
pub struct Snapshot {
    pub scenario: &'static str,
    pub step: u64,
    pub seconds: f64,
    pub ball: BallSnapshot,
    pub secondary_ball: Option<BallSnapshot>,
    pub claw: ClawSnapshot,
    pub flippers: FlipperSnapshot,
    pub launcher_charge: f32,
    pub collision_id: Option<u8>,
    pub panel_frame: Option<u16>,
    pub events: Vec<String>,
}

#[derive(Debug, Serialize, PartialEq)]
pub struct BallSnapshot {
    pub x: f32,
    pub y: f32,
    pub velocity_x: f32,
    pub velocity_y: f32,
    pub in_launcher: bool,
}

#[derive(Debug, Serialize, PartialEq, Eq)]
pub struct ClawSnapshot {
    pub active: bool,
    pub frame: u8,
    pub bank: &'static str,
    pub ball_suspended: bool,
}

#[derive(Debug, Serialize, PartialEq, Eq)]
pub struct FlipperSnapshot {
    pub left_raised: bool,
    pub right_raised: bool,
}

pub struct Simulation {
    scenario: Scenario,
    step: u64,
    game: Game,
    trace: Vec<Snapshot>,
    autoplay_launcher_frames: u32,
}

impl Simulation {
    pub fn new(scenario: Scenario, seed: u32) -> Self {
        let mut game = Game::new_with_seed(1, seed);
        if scenario == Scenario::Panel {
            game.panel_frame = Some(0);
            game.wheel_holes.fill(true);
        }
        if scenario == Scenario::Targets {
            game.begin_target_scenario();
        }
        if scenario == Scenario::Effect7 {
            game.begin_effect_scenario(7);
        }
        let initial_events = scenario
            .claw_terminal()
            .map_or_else(Vec::new, |terminal| game.begin_claw_scenario(terminal));
        let mut simulation = Self {
            scenario,
            step: 0,
            game,
            trace: Vec::new(),
            autoplay_launcher_frames: 0,
        };
        simulation.record(initial_events);
        simulation
    }

    pub fn advance_to(&mut self, target_step: u64) {
        while self.step < target_step {
            let controls = if self.scenario == Scenario::Autoplay {
                self.autoplay_controls()
            } else {
                controls_for(self.scenario, self.step)
            };
            let events = if matches!(
                self.scenario,
                Scenario::Loading | Scenario::Attract | Scenario::HighScores | Scenario::NameEntry
            ) {
                Vec::new()
            } else {
                self.game.update(SIMULATION_DT, 3, controls)
            };
            self.step += 1;
            self.record(events);
        }
    }

    pub fn game(&self) -> &Game {
        &self.game
    }

    pub fn final_snapshot(&self) -> &Snapshot {
        self.trace
            .last()
            .expect("a simulation always records step zero")
    }

    pub fn write_trace(&self, path: &Path) -> Result<(), String> {
        let json = serde_json::to_vec_pretty(&self.trace)
            .map_err(|error| format!("could not encode simulation trace: {error}"))?;
        fs::write(path, json)
            .map_err(|error| format!("could not write trace {}: {error}", path.display()))
    }

    fn record(&mut self, events: Vec<Event>) {
        let step = u32::try_from(self.step).expect("simulation step is limited to 72000");
        let claw_bank = match self.game.claw.bank {
            ClawSpriteBank::Closing => "closing",
            ClawSpriteBank::Opening => "opening",
        };
        self.trace.push(Snapshot {
            scenario: self.scenario.name(),
            step: self.step,
            seconds: f64::from(step) / SIMULATION_HZ_F64,
            ball: BallSnapshot {
                x: self.game.ball.position.x,
                y: self.game.ball.position.y,
                velocity_x: self.game.ball.velocity.x,
                velocity_y: self.game.ball.velocity.y,
                in_launcher: self.game.ball.in_launcher,
            },
            secondary_ball: self.game.secondary_ball.map(|ball| BallSnapshot {
                x: ball.position.x,
                y: ball.position.y,
                velocity_x: ball.velocity.x,
                velocity_y: ball.velocity.y,
                in_launcher: ball.in_launcher,
            }),
            claw: ClawSnapshot {
                active: self.game.claw.active,
                frame: self.game.claw.frame,
                bank: claw_bank,
                ball_suspended: self.game.claw.ball_suspended,
            },
            flippers: FlipperSnapshot {
                left_raised: self.game.flippers.left_raised,
                right_raised: self.game.flippers.right_raised,
            },
            launcher_charge: self.game.launcher_charge,
            collision_id: self.game.last_collision_id,
            panel_frame: self.game.panel_frame,
            events: events
                .into_iter()
                .map(|event| format!("{event:?}"))
                .collect(),
        });
    }

    fn autoplay_controls(&mut self) -> Controls {
        if self.game.ball.in_launcher {
            let launch_down = self.autoplay_launcher_frames < 132;
            self.autoplay_launcher_frames += 1;
            Controls {
                launch_down,
                ..Controls::default()
            }
        } else {
            self.autoplay_launcher_frames = 0;
            Controls {
                left_flipper: self.game.ball.position.y > 340.0
                    && self.game.ball.position.x < 175.0,
                right_flipper: self.game.ball.position.y > 340.0
                    && self.game.ball.position.x > 140.0,
                ..Controls::default()
            }
        }
    }
}

fn controls_for(scenario: Scenario, step: u64) -> Controls {
    match scenario {
        Scenario::Autoplay => unreachable!("autoplay controls need mutable simulation state"),
        Scenario::Launcher => Controls {
            launch_down: step < u64::from(SIMULATION_HZ),
            ..Controls::default()
        },
        Scenario::Flippers => {
            let phase = step % u64::from(SIMULATION_HZ);
            Controls {
                left_flipper: (15..45).contains(&phase),
                right_flipper: (60..90).contains(&phase),
                ..Controls::default()
            }
        }
        Scenario::Loading
        | Scenario::Attract
        | Scenario::HighScores
        | Scenario::NameEntry
        | Scenario::Effect7
        | Scenario::Panel
        | Scenario::Targets
        | Scenario::Claw1
        | Scenario::Claw6
        | Scenario::Claw7
        | Scenario::Claw18 => Controls::default(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn args(values: &[&str]) -> Vec<String> {
        values.iter().map(|value| (*value).to_owned()).collect()
    }

    #[test]
    fn interactive_run_has_no_simulation_request() {
        assert_eq!(Request::parse(args(&["tdkpin-rs"])), Ok(None));
    }

    #[test]
    fn simulation_request_accepts_time_outputs_and_seed() {
        let request = Request::parse(args(&[
            "tdkpin-rs",
            "--simulate",
            "claw-7",
            "--at",
            "0.15",
            "--screenshot",
            "frame.png",
            "--trace",
            "trace.json",
            "--seed",
            "42",
        ]))
        .expect("request should parse")
        .expect("simulation should be selected");

        assert_eq!(request.scenario, Scenario::Claw7);
        assert_eq!(request.target_step, 18);
        assert_eq!(request.screenshot, Some(PathBuf::from("frame.png")));
        assert_eq!(request.trace, Some(PathBuf::from("trace.json")));
        assert_eq!(request.seed, 42);
    }

    #[test]
    fn seeded_simulations_have_identical_step_traces() {
        let mut first = Simulation::new(Scenario::Claw6, 7);
        let mut second = Simulation::new(Scenario::Claw6, 7);
        first.advance_to(90);
        second.advance_to(90);

        assert_eq!(first.trace, second.trace);
    }

    #[test]
    fn autoplay_stays_finite_for_two_minutes_of_exact_physics() {
        let mut simulation = Simulation::new(Scenario::Autoplay, 1);
        simulation.advance_to(u64::from(SIMULATION_HZ) * 120);
        let event_count = |name: &str| {
            simulation
                .trace
                .iter()
                .flat_map(|snapshot| &snapshot.events)
                .filter(|event| event.as_str() == name)
                .count()
        };

        assert!(event_count("Launch") >= 1);
        assert!(event_count("FlipperMove") >= 2);
        assert!(event_count("Target") >= 1);
        assert!(event_count("Bumper") >= 1);
        assert!(simulation.trace.iter().all(|snapshot| {
            snapshot.ball.x.is_finite()
                && snapshot.ball.y.is_finite()
                && snapshot.ball.velocity_x.is_finite()
                && snapshot.ball.velocity_y.is_finite()
                && snapshot.secondary_ball.as_ref().is_none_or(|ball| {
                    ball.x.is_finite()
                        && ball.y.is_finite()
                        && ball.velocity_x.is_finite()
                        && ball.velocity_y.is_finite()
                })
        }));
    }

    #[test]
    fn launcher_scenario_holds_for_one_second_then_releases() {
        let mut simulation = Simulation::new(Scenario::Launcher, 1);
        simulation.advance_to(60);
        assert!(simulation.game.ball.in_launcher);
        assert!((0.19..=0.21).contains(&simulation.game.launcher_charge));

        simulation.advance_to(121);
        assert!(!simulation.game.ball.in_launcher);
        assert!(simulation.trace[121].events.contains(&"Launch".to_owned()));
    }

    #[test]
    fn flipper_scenario_records_both_movement_edges() {
        let mut simulation = Simulation::new(Scenario::Flippers, 1);
        simulation.advance_to(50);

        let edge_steps = simulation
            .trace
            .iter()
            .filter(|snapshot| snapshot.events.iter().any(|event| event == "FlipperMove"))
            .collect::<Vec<_>>();
        assert_eq!(edge_steps.len(), 2);
        assert!(
            edge_steps
                .iter()
                .all(|snapshot| snapshot.events == ["FlipperMove", "Sound(2021)"])
        );
        assert!(!simulation.game.flippers.left_raised);
    }

    #[test]
    fn panel_scenario_runs_all_frames_and_cleans_up() {
        let mut simulation = Simulation::new(Scenario::Panel, 1);
        simulation.advance_to(4);
        assert_eq!(simulation.game.panel_frame, Some(1));
        assert_eq!(simulation.trace[4].events, ["Sound(2013)"]);

        simulation.advance_to(1_020);
        assert_eq!(simulation.game.panel_frame, None);
        assert_eq!(simulation.game.wheel_holes, [false; 5]);
        let events = simulation
            .trace
            .iter()
            .flat_map(|snapshot| snapshot.events.iter().map(String::as_str))
            .collect::<Vec<_>>();
        assert_eq!(
            events.iter().filter(|event| **event == "StopSound").count(),
            3
        );
        assert_eq!(
            events
                .iter()
                .filter(|event| **event == "Sound(2013)")
                .count(),
            4
        );
        assert_eq!(
            events
                .iter()
                .filter(|event| **event == "Sound(2012)")
                .count(),
            2
        );
    }

    #[test]
    fn target_scenario_runs_all_six_rotation_states() {
        let mut simulation = Simulation::new(Scenario::Targets, 1);
        simulation.advance_to(4);
        assert_eq!(simulation.game.target_rotation_state, Some(1));
        assert_eq!(simulation.trace[4].events, ["Sound(2011)"]);

        simulation.advance_to(22);
        assert_eq!(simulation.game.target_rotation_state, Some(6));
        assert_eq!(
            simulation.game.wheel_holes,
            [false, true, false, false, true]
        );
        simulation.advance_to(26);
        assert_eq!(simulation.game.target_rotation_state, None);
    }

    #[test]
    fn every_claw_scenario_completes_its_seeded_capture() {
        for scenario in [
            Scenario::Claw1,
            Scenario::Claw6,
            Scenario::Claw7,
            Scenario::Claw18,
        ] {
            let mut simulation = Simulation::new(scenario, 1);
            for target in 1..=120 {
                simulation.advance_to(target);
                let index = usize::try_from(target).unwrap_or(120);
                if simulation.trace[index]
                    .events
                    .iter()
                    .any(|event| event == "ClawRelease")
                {
                    break;
                }
            }
            let events: Vec<&str> = simulation
                .trace
                .iter()
                .flat_map(|snapshot| snapshot.events.iter().map(String::as_str))
                .collect();

            let captures = events
                .iter()
                .filter(|event| **event == "ClawCapture")
                .count();
            let releases = events
                .iter()
                .filter(|event| **event == "ClawRelease")
                .count();
            assert!(captures >= 1, "{} must enter the claw", scenario.name());
            assert_eq!(
                releases,
                1,
                "{} must release its seeded capture",
                scenario.name()
            );
            assert!(captures >= releases);
            assert!(simulation.game.ball.position.is_finite());
            assert!(simulation.game.ball.velocity.is_finite());
        }
    }
}