harness.pyΒΆ

Part of Mr. Rescue.

  1"""Scripted headless playthrough that captures one screenshot per stage.
  2
  3:func:`capture_playthrough` drives a real :class:`MrRescueRoot` with
  4:class:`~simvx.core.InputSimulator` and writes one PNG per stage into
  5``screenshots/``.
  6
  7The pilot below plays from live world state rather than from a fixed key
  8timeline: it presses "up" on the frame it is actually standing on a ladder,
  9steps off once it has climbed a full storey, and presses "grab" on the frame a
 10civilian comes into reach. That is what makes each screenshot show the outcome
 11its name promises instead of the key that was sent.
 12
 13Module-level ``random`` is seeded first: civilians, fire bugs and screenshake all
 14draw from it, so without a seed the run would not repeat frame for frame and the
 15capture points would drift.
 16
 17Run it from the port directory with ``python harness.py``.
 18"""
 19
 20from __future__ import annotations
 21
 22import random
 23from collections.abc import Callable
 24from pathlib import Path
 25
 26from nodes.building import FLOOR_HEIGHT
 27from nodes.tile_grid import TILE_SIZE
 28
 29from simvx.core import DeviceButton, InputSimulator, Key, Node
 30from simvx.graphics import App, save_png
 31
 32# Height of one storey in pixels: the exact distance a ladder spans.
 33STOREY_PX = FLOOR_HEIGHT * TILE_SIZE
 34
 35RUN_SEED = 42  # scene generator + module-level ``random``
 36START_FRAME = 30  # confirm on the title screen
 37RUN_FRAME = 50  # start running right
 38SPRAY_FRAME = 70  # open the hose
 39SPRAY_STOP_X = 400.0  # close it again short of the ladder, so climbing can start
 40CARRY_WALK = 40  # frames spent carrying the civilian clear of the group
 41CARRY_HOLD = 60  # frames spent carrying before the suit is pushed over
 42RESTART_DELAY = 40  # frames on the end screen before confirming
 43
 44# Frame at which each screenshot is taken. The pilot is state-driven but the run
 45# is deterministic, so these land on the stages listed in ``capture_playthrough``.
 46CAPTURES = {
 47    20: "01_title.png",
 48    48: "02_world_overview.png",
 49    157: "03_water_spray.png",
 50    330: "04_ladder_climb.png",
 51    500: "05_upper_floor.png",
 52    580: "06_civilian_grab.png",
 53    620: "07_end_screen.png",
 54    664: "08_post_restart.png",
 55}
 56TOTAL_FRAMES = 700
 57
 58
 59class _Pilot:
 60    """Plays one run, deciding each frame's input from the live scene."""
 61
 62    def __init__(self, sim: InputSimulator):
 63        self._sim = sim
 64        self._held: set[Key] = set()
 65        self._taps: dict[DeviceButton, int] = {}
 66        self.stage = "title"
 67        self._mount_y = 0.0
 68        self._carry_frame = 0
 69        self._end_frame = -1
 70
 71    # ------------------------------------------------------------- keys
 72
 73    def _hold(self, key: DeviceButton):
 74        if key not in self._held:
 75            self._sim.press_key(key)
 76            self._held.add(key)
 77
 78    def _let_go(self, key: DeviceButton):
 79        if key in self._held:
 80            self._sim.release_key(key)
 81            self._held.discard(key)
 82
 83    def _tap(self, idx: int, key: DeviceButton):
 84        """Press *key* now and release it two frames later (a one-shot action)."""
 85        self._hold(key)
 86        self._taps[key] = idx + 2
 87
 88    # ------------------------------------------------------------- driver
 89
 90    def step(self, idx: int, root: Node):
 91        for key, due in list(self._taps.items()):
 92            if idx >= due:
 93                self._let_go(key)
 94                del self._taps[key]
 95
 96        if root.phase == "end":
 97            self._end(idx)
 98            return
 99        if root.phase == "menu":
100            if self.stage == "title" and idx >= START_FRAME:
101                self._tap(idx, Key.ENTER)
102                self.stage = "advance"
103            return
104
105        scene = root.scene
106        if self.stage == "advance":
107            self._advance(idx, scene)
108        elif self.stage == "climb":
109            self._climb(scene.player)
110        elif self.stage == "cross":
111            self._cross(idx, scene.player)
112        elif self.stage == "carry":
113            self._carry(idx, scene.player)
114
115    # ------------------------------------------------------------- stages
116
117    def _advance(self, idx: int, scene):
118        """Run right along the ground floor with the hose open, up to the ladder."""
119        player, grid = scene.player, scene.grid
120        if idx >= RUN_FRAME:
121            self._hold(Key.RIGHT)
122        if SPRAY_FRAME <= idx and player.position.x < SPRAY_STOP_X:
123            self._hold(Key.LEFT_SHIFT)
124        else:
125            self._let_go(Key.LEFT_SHIFT)
126        # Mount the moment the gun-height cell is a rung: holding "up" over a
127        # ladder is exactly how a player climbs.
128        if grid.is_ladder(*grid.cell_for(player.position.x, player.position.y - 11)):
129            self._let_go(Key.RIGHT)
130            self._hold(Key.UP)
131            self._mount_y = float(player.position.y)
132            self.stage = "climb"
133
134    def _climb(self, player):
135        """Hold "up" for one full storey, then step off sideways onto the floor."""
136        if player.position.y <= self._mount_y - STOREY_PX:
137            self._let_go(Key.UP)
138            self._hold(Key.LEFT)
139            self.stage = "cross"
140
141    def _cross(self, idx: int, player):
142        """Walk the upper floor hosing the fires down, and grab the first civilian."""
143        self._hold(Key.LEFT)
144        self._hold(Key.LEFT_SHIFT)
145        if player.can_grab:
146            self._let_go(Key.LEFT)
147            self._let_go(Key.LEFT_SHIFT)
148            self._tap(idx, Key.E)
149            self._carry_frame = idx
150            self.stage = "carry"
151
152    def _carry(self, idx: int, player):
153        """Carry her back the way we came, then let the suit cook."""
154        held = idx - self._carry_frame
155        if held < CARRY_WALK:
156            self._hold(Key.RIGHT)
157        else:
158            self._let_go(Key.RIGHT)
159        # Push the suit past its heat limit so the end screen is reachable inside
160        # the frame budget; real play gets there through fire exposure.
161        if held >= CARRY_HOLD:
162            player.temperature = player.max_temperature + 1.0
163
164    def _end(self, idx: int):
165        if self._end_frame < 0:
166            self._end_frame = idx
167            for key in list(self._held):
168                self._let_go(key)
169        elif idx == self._end_frame + RESTART_DELAY:
170            self._tap(idx, Key.ENTER)
171
172
173def capture_playthrough(
174    root_factory: Callable[..., Node],
175    *,
176    window_size: tuple[int, int],
177    bg_colour,
178    out_dir: Path,
179) -> None:
180    """Run the scripted playthrough headlessly and save the stage screenshots.
181
182    The stages, in the order they are captured:
183
184    ``01_title``
185        the title screen.
186    ``02_world_overview``
187        the generated building, camera on the player at his ground-floor spawn.
188    ``03_water_spray``
189        running right with the hose open, the jet landing on a burning
190        ground-floor fire (a second one is already out: score 20, fires 5).
191    ``04_ladder_climb``
192        on the ladder, part-way between the ground floor and the one above.
193    ``05_upper_floor``
194        stepped off the ladder and crossing the floor above, hosing down the two
195        fires that were burning on it.
196    ``06_civilian_grab``
197        carrying a civilian, picked up on the frame she came into reach.
198    ``07_end_screen``
199        the overheat end screen with the run summary.
200    ``08_post_restart``
201        back on the title screen after confirming a restart.
202    """
203    out_dir.mkdir(exist_ok=True)
204    random.seed(RUN_SEED)
205
206    width, height = window_size
207    app = App(title="Mr. Rescue (test)", width=width, height=height, visible=False, bg_colour=bg_colour)
208    root = root_factory(seed=RUN_SEED)
209    pilot: _Pilot | None = None  # built on the first frame, once the root is in its tree
210
211    def drive(idx: int, _t: float):
212        nonlocal pilot
213        if pilot is None:
214            pilot = _Pilot(InputSimulator(root.tree))
215        pilot.step(idx, root)
216        return None
217
218    capture_frames = sorted(CAPTURES)
219    frames = app.run_headless(
220        root,
221        frames=TOTAL_FRAMES,
222        on_frame=drive,
223        capture_frames=capture_frames,
224    )
225    for idx, frame in zip(capture_frames, frames, strict=False):
226        if frame is None:
227            continue
228        # Headless captures come back with the scene's own alpha; force opaque
229        # so the PNGs composite predictably in docs and image viewers.
230        frame[..., 3] = 255
231        path = out_dir / CAPTURES[idx]
232        save_png(frame, path)
233        print(f"saved {path}")
234
235
236if __name__ == "__main__":
237    from main import WINDOW_H, WINDOW_W, MrRescueRoot
238    from nodes import colours
239
240    capture_playthrough(
241        MrRescueRoot,
242        window_size=(WINDOW_W, WINDOW_H),
243        bg_colour=colours.BG,
244        out_dir=Path(__file__).resolve().parent / "screenshots",
245    )