undercooked/tests/sim/_cook_policy.py

"""The scripted cook: the deterministic policy behind the task 04 long runs.

`test_freshness.py`'s 1 000-tick conservation sequence and
`test_no_deadlock.py`'s 2 000-tick worst case drive the same kitchen the
same way: one cook, two fixed positions one tile apart (the reach zone
gives four "first target" rights per position — the front tile, then the
two diagonal-front tiles, then underfoot, and a press resolves to the
*first station cell* of its facing's walk), every station in reach by a
facing change, and an action is always *face, then press, then (hold),
then release*. The whole run is a pure function of the game state and a
small harness memory (`Harness`) — no RNG, no wall clock: a re-run of the
run is bit-identical, which is what lets the tests assert per tick.

The layout (a 5 x 5 floor, the cook's two positions (2.5, 3.5) and
(2.5, 4.5)):

```text
position 1 (2.5, 3.5):          position 2 (2.5, 4.5):
  UP front     shelf (2, 2)       LEFT front   —  (out of map)
  UP diag      —                  LEFT 2nd     prep (1, 4)
  LEFT front   board (1, 3)       LEFT 3rd     board (1, 3)
  RIGHT front  stove (3, 3)       RIGHT 2nd    window (3, 4)
  UP 3rd       sink (3, 2)        RIGHT 3rd    stove (3, 3)
  (the sink is reachable from position 1; the cook never washes there —
  the in-place wash is the policy's one fix)
```

The policy is the worst case the task names: **serve everything possible,
never wash until forced.** With the small stock (two plates, two pots)
that means a dirty prep or a dirty stove (a stuck prop) or an empty clean
stock stops the next dish, and the one fix is the in-place wash — the
prop's stock paid back on completion. A full hand at a stuck prop sets
its item down on the board (a table — the board works raw only, so a
chopped or cooked set-down stays a marked object), washes, and picks it
back up.

```text
shelf ─raw─▶ board ─chopped─▶ stove ─cooked─▶ (walk) ─prep ─plated─▶ (walk) ─window
            (18 ticks)       (120 ticks)                  (12 ticks)
                                        │  the dirty prep / stove (a stuck
                                        │  prop) is washed in place (24 ticks)
                                        ▼
                                    the stock pair is paid back
```
"""

from __future__ import annotations

from dataclasses import dataclass

from undercooked.sim.grid import Grid
from undercooked.sim.model import (
    Facing,
    Item,
    ItemForm,
    ItemKind,
    KitchenState,
    Player,
    Station,
    StationKind,
    StationState,
    TileKind,
    Vec2,
)
from undercooked.sim.step import STEP_DT, PlayerInput, step

__all__ = [
    "POS1_Y",
    "POS2_Y",
    "Harness",
    "build_kitchen",
    "choose_plan",
    "next_frame",
    "run_policy",
]

#: The cook's two positions (the y coordinate; x is 2.5 at both).
POS1_Y = 3.5  # the shelf / board / stove half
POS2_Y = 4.5  # the prep / window half

#: The six stations, by id, at their cells.
SHELF = "shelf-2-2"
BOARD = "board-1-3"
STOVE = "stove-3-3"
PREP = "prep-1-4"
SERVE = "serve-3-4"
SINK = "sink-3-2"

FACE_SH = Facing.UP  # shelf (2, 2): position 1's UP front, 1.0 centre-to-centre
BOARD_FACE = Facing.LEFT  # board (1, 3): position 1's LEFT front, 1.0
STOVE_FACE = Facing.RIGHT  # stove (3, 3): position 1's RIGHT front, 1.0
PREP_FACE = Facing.LEFT  # prep (1, 4): position 2's LEFT 2nd, 1.0
SERVE_FACE = Facing.RIGHT  # window (3, 4): position 2's RIGHT 2nd, 1.0
SINK_FACE = Facing.RIGHT  # sink (3, 2): position 1's RIGHT 2nd, 1.414


def _station(
    station_id: str,
    kind: StationKind,
    x: int,
    y: int,
    contents: tuple[Item, ...] = (),
) -> Station:
    return Station(id=station_id, kind=kind, x=x, y=y, contents=contents)


def build_kitchen(
    shelf_stock: int = 20,
    plates: int = 2,
    pots: int = 2,
) -> KitchenState:
    """The scripted kitchen: a 5 x 5 floor, the six stations, one cook.

    The stock counters are seeded from the (test's stand-in for)
    `level.station_limits` — the round a phase-05 room would build from a
    level with those limits — so the conservation the tests assert per tick
    is against a finite total: `…_clean + …_dirty == plates / pots`.
    """
    grid = Grid(width=5, height=5, tiles=(TileKind.FLOOR,) * 25)
    stations = (
        _station(
            SHELF,
            StationKind.SPAWN_SHELF,
            2,
            2,
            (Item(ItemKind.TOMATO, ItemForm.RAW),) * shelf_stock,
        ),
        _station(BOARD, StationKind.CHOP_BOARD, 1, 3),
        _station(STOVE, StationKind.STOVE, 3, 3),
        _station(PREP, StationKind.PREP, 1, 4),
        _station(SERVE, StationKind.SERVE, 3, 4),
        _station(SINK, StationKind.SINK, 3, 2),
    )
    cook = Player(id="p0", pos=Vec2(2.5, POS1_Y), facing=Facing.UP, seat=0)
    return KitchenState(
        level_id="04_policy",
        seed=2024,
        grid=grid,
        stations=stations,
        players=(cook,),
        order_pool=("soup",),
        order_concurrent=2,
        order_spawn_min=100,  # 5 s
        order_spawn_max=140,  # 7 s
        order_deadline=1200,  # 60 s
        order_penalty=10,
        next_order_at=30,  # the 1.5 s first-order delay
        next_order_id=1,
        target_score=1000,
        time_left=300.0,
        counters={
            "plates_clean": plates,
            "plates_dirty": 0,
            "pots_clean": pots,
            "pots_dirty": 0,
        },
        rng_state=2024,
    )


@dataclass
class Harness:
    """The policy's memory: what it is about to do, and the press bookkeeping.

    The press is a rising edge, so between two presses there is always a
    release tick; a facing change costs a face tick (the action gate targets
    with the *pre-movement* facing, so the face lands one tick before the
    press can use it); a position change costs a walk (a few `move` frames,
    the one-tile step between the positions).
    """

    face: Facing = Facing.UP
    plan: tuple[Facing, str, str, str] | None = None  # (face, kind, station, pos)
    holding: bool = False
    hold_station: str = ""  # the station the hold advances (work or wash)
    release_pending: bool = True  # the first tick is a release: the gate is clean
    tick: int = 0


def _stations_by(state: KitchenState) -> dict[str, Station]:
    return {station.id: station for station in state.stations}


def _dirty(station: Station) -> bool:
    return station.dirty or station.state is StationState.DIRTY


def _position(state: KitchenState) -> str:
    """The cook's position: `"1"`, `"2"`, or `"mid"` (walking)."""
    y = state.players[0].pos.y
    if y > 4.0:
        return "2"
    if y < 3.5:
        return "1"
    return "mid"


def choose_plan(state: KitchenState) -> tuple[Facing, str, str, str] | None:
    """The policy's decision: `(face, kind, station, pos)`, or `None` (idle).

    `kind` is `"press"` (one press, one effect), `"hold"` (the press starts
    the work or the wash, and the hold advances it while the station is
    `WORKING`) or `"walk"` (the one-tile step to the other position);
    `pos` is the position the action needs.
    """
    p = state.players[0]
    s = _stations_by(state)
    board, stove, prep = s[BOARD], s[STOVE], s[PREP]
    shelf = s[SHELF]
    plates_clean = state.counters.get("plates_clean", 0)

    # The forced washes: a stuck prop (a dirty prep or stove) with a free
    # hand is washed in place. Never washed earlier — this is the only wash
    # the policy ever does.
    if _dirty(prep) and p.carrying is None:
        return (PREP_FACE, "hold", PREP, "2")
    if _dirty(stove) and p.carrying is None:
        return (STOVE_FACE, "hold", STOVE, "1")
    # A stuck prop with a full hand: the item goes to the board (a table —
    # the board works raw only, so a chopped or cooked set-down stays a
    # marked object), the hand frees, and the wash above runs next. The
    # board is reachable from both positions.
    if (_dirty(prep) or _dirty(stove)) and p.carrying is not None:
        return (
            BOARD_FACE if _position(state) != "2" else Facing.UP,
            "press",
            BOARD,
            _position(state) if _position(state) != "mid" else "1",
        )

    carried = p.carrying
    if carried is None:
        # A work in flight is held (the cook started it — the hold is what
        # advances it), a ready output is picked, and a tabled object (a
        # set-down on the board, the wash's temporary) is picked back.
        if board.state is StationState.WORKING:
            return (BOARD_FACE, "hold", BOARD, "1")
        if stove.state is StationState.WORKING:
            return (STOVE_FACE, "hold", STOVE, "1")
        if prep.state is StationState.WORKING:
            return (PREP_FACE, "hold", PREP, "2")
        if board.state is StationState.READY and board.contents:
            return (BOARD_FACE, "press", BOARD, "1")
        if stove.state is StationState.READY and stove.contents:
            return (STOVE_FACE, "press", STOVE, "1")
        if prep.state is StationState.READY and prep.contents:
            return (PREP_FACE, "press", PREP, "2")
        if board.state is StationState.IDLE and board.contents:
            return (BOARD_FACE, "press", BOARD, "1")  # the tabled object
        if shelf.contents:
            return (FACE_SH, "press", SHELF, "1")
        return None
    if carried.form is ItemForm.RAW:
        if board.state is StationState.WORKING:
            return (BOARD_FACE, "hold", BOARD, "1")
        return (BOARD_FACE, "press", BOARD, "1")  # the set-down starts the chop
    if carried.form is ItemForm.CHOPPED:
        if stove.state is StationState.WORKING:
            return (STOVE_FACE, "hold", STOVE, "1")
        return (STOVE_FACE, "press", STOVE, "1")  # the set-down starts the cook
    if carried.form is ItemForm.COOKED:
        if prep.state is StationState.WORKING:
            return (PREP_FACE, "hold", PREP, "2")
        if _dirty(prep) or plates_clean <= 0:
            return (Facing.UP, "press", BOARD, "2")  # the table (position 2)
        return (PREP_FACE, "press", PREP, "2")  # the press starts the plate
    # A `PLATED` dish: serve while the rail has an order, wait while it
    # does not — a plated dish does not decay in the hand, so the wait is
    # free (it is the "serve everything possible" half of the policy).
    if state.orders:
        return (SERVE_FACE, "press", SERVE, "2")
    return None


def run_policy(
    state: KitchenState,
    book: object,
    ticks: int,
) -> tuple[list[KitchenState], list]:
    """Drive `ticks` substeps of the policy through the whole `step()`.

    Returns the per-tick states (the state after each substep, tick 0
    excluded) and the per-tick event lists, so a test can assert per tick
    without re-running the run (the run is deterministic — a re-run is
    bit-identical, so the metrics a test computes are the round's).
    """
    harness = Harness()
    states: list[KitchenState] = []
    event_lists: list[list] = []
    for tick in range(ticks):
        harness.tick = tick
        frame = next_frame(state, harness)
        state, events = step(state, STEP_DT, (frame,), book)  # type: ignore[arg-type]
        states.append(state)
        event_lists.append(list(events))
    return states, event_lists


def next_frame(state: KitchenState, harness: Harness) -> PlayerInput:
    """The one `PlayerInput` this tick: a face, a press, the hold, or a walk.

    The frame's `face` is explicit, so the facing is exactly the policy's
    plan, never a movement side effect; `move` is the walk's one-tile step
    (zero otherwise).
    """
    p = state.players[0]
    if harness.release_pending:
        harness.release_pending = False
        return PlayerInput(
            player_id="p0",
            move=Vec2(0.0, 0.0),
            face=harness.face,
            received_at_tick=harness.tick,
        )
    if harness.holding:
        station = _stations_by(state).get(harness.hold_station)
        if station is not None and station.state is StationState.WORKING:
            return PlayerInput(
                player_id="p0",
                move=Vec2(0.0, 0.0),
                action=True,
                face=harness.face,
                received_at_tick=harness.tick,
            )
        # The work is done: release, then plan again.
        harness.holding = False
        harness.plan = None
        harness.release_pending = True
        return PlayerInput(
            player_id="p0",
            move=Vec2(0.0, 0.0),
            face=harness.face,
            received_at_tick=harness.tick,
        )
    plan = harness.plan
    if plan is None:
        plan = choose_plan(state)
        harness.plan = plan
        harness.face = plan[0] if plan is not None else p.facing
    if plan is None:
        return PlayerInput(
            player_id="p0",
            move=Vec2(0.0, 0.0),
            face=p.facing,
            received_at_tick=harness.tick,
        )
    face, _kind, _station_id, pos = plan
    # The walk: the action needs the other position; the one-tile step runs
    # to within 0.1 of the target centre (the facing tick / press follow).
    if pos is not None:
        here = _position(state)
        if here not in ("mid", pos):
            target_y = POS2_Y if pos == "2" else POS1_Y
            return PlayerInput(
                player_id="p0",
                move=Vec2(0.0, 1.0 if target_y > p.pos.y else -1.0),
                face=face,
                received_at_tick=harness.tick,
            )
        if here == "mid":
            # Still walking: keep stepping until the arrival band.
            target_y = POS2_Y if pos == "2" else POS1_Y
            if abs(p.pos.y - target_y) > 0.1:
                return PlayerInput(
                    player_id="p0",
                    move=Vec2(0.0, 1.0 if target_y > p.pos.y else -1.0),
                    face=face,
                    received_at_tick=harness.tick,
                )
    if p.facing is not face:
        # The face tick: the action gate targets with the pre-movement
        # facing, so the face must land one tick before the press.
        return PlayerInput(
            player_id="p0",
            move=Vec2(0.0, 0.0),
            face=face,
            received_at_tick=harness.tick,
        )
    # The press: a rising edge (the previous tick was a release / face /
    # idle tick, all `action=False` by construction).
    if _kind == "hold":
        harness.holding = True
        harness.hold_station = _station_id
    harness.plan = None
    harness.release_pending = True
    return PlayerInput(
        player_id="p0",
        move=Vec2(0.0, 0.0),
        action=True,
        face=face,
        received_at_tick=harness.tick,
    )