data/mock_env.py

"""Built-in mock ARC-AGI-3 game: multi-level push-box ("sokoban-lite").

Purpose: the FULL QwenJev pipeline (collect -> build-data -> distill ->
train-s1 -> eval) must run on any machine, including CPU-only laptops with
no `arc-agi` package and no game downloads. `MockGame` reproduces the parts
of the official environment contract that the pipeline touches:

    reset() / step(action) / observe()
    FrameData-like fields: grids (<=64x64, values 0..15),
    levels_completed, win_levels, available_actions, state
    actions: RESET, ACTION1..ACTION5, ACTION6(x,y), ACTION7 (undo)

Mechanics (sparse-reward, locality-sensitive - good training signal):
  * grid values: 0 floor, 1 wall, 2 player, 3 box, 4 goal;
  * ACTION1..4 move the player (up/down/left/right); walking into a box
    pushes it one cell if the far side is free;
  * a level is complete when every box sits on a goal -> levels_completed+1
    and a NEW (larger, more boxes) level is generated;
  * ACTION5 waits (no-op), ACTION6(x,y) steps onto an adjacent walkable cell,
    ACTION7 undoes up to 8 recent frames;
  * running out of a per-level action budget -> GAME_OVER;
  * completing `win_levels` levels -> WIN.

Deterministic under `seed`; solvability approximated with BFS reachability
from player to each box and from each box to a goal (walls only).
"""

from __future__ import annotations

import random
from dataclasses import dataclass, field
from typing import Sequence

GAME_ID = "mock-pushbox"

# cell values (ARC palette indices; 0..15 legal, we use 0,1,2,3,4)
FLOOR, WALL, PLAYER, BOX, GOAL = 0, 1, 2, 3, 4

DIRS = {
    "ACTION1": (-1, 0),   # up
    "ACTION2": (1, 0),    # down
    "ACTION3": (0, -1),   # left
    "ACTION4": (0, 1),    # right
}
ALL_ACTIONS = (
    "RESET", "ACTION1", "ACTION2", "ACTION3", "ACTION4",
    "ACTION5", "ACTION6", "ACTION7",
)
UNDO_DEPTH = 8


@dataclass
class MockFrame:
    """FrameData-compatible observation."""

    grids: list[list[list[int]]]
    levels_completed: int = 0
    win_levels: int = 3
    available_actions: list[str] = field(default_factory=lambda: list(ALL_ACTIONS))
    state: str = "NOT_PLAYED"      # NOT_PLAYED | NOT_FINISHED | WIN | GAME_OVER


@dataclass
class _Level:
    grid: list[list[int]]
    boxes: list[tuple[int, int]]
    goals: list[tuple[int, int]]
    player: tuple[int, int]
    budget: int
    steps_used: int = 0


class MockGame:
    """One game instance. `step` accepts an action NAME (str) - the mock has
    no dependency on the official `arcengine` enums."""

    def __init__(self, seed: int = 0, win_levels: int = 3, base_size: int = 10):
        self.rng = random.Random(seed)
        self.win_levels = win_levels
        self.base_size = base_size
        self.levels_completed = 0
        self.state = "NOT_PLAYED"
        self._level: _Level | None = None
        self._undo: list[list[list[int]]] = []
        self._last_frame: MockFrame | None = None

    # ------------------------------------------------------------------ #
    # level generation                                                     #
    # ------------------------------------------------------------------ #
    def _gen_level(self, index: int) -> _Level:
        rng = self.rng
        size = min(64, self.base_size + 2 * index)          # grow with levels
        n_boxes = min(3, 1 + index // 2)
        wall_ratio = 0.10

        grid = [[FLOOR] * size for _ in range(size)]
        for i in range(size):                                # border walls
            grid[0][i] = grid[size - 1][i] = WALL
            grid[i][0] = grid[i][size - 1] = WALL
        inner = [(y, x) for y in range(1, size - 1)
                 for x in range(1, size - 1)]
        for (y, x) in rng.sample(inner, int(len(inner) * wall_ratio)):
            grid[y][x] = WALL

        free = [(y, x) for (y, x) in inner if grid[y][x] == FLOOR]
        rng.shuffle(free)

        def far(cell, others, min_dist):
            for (y, x) in others:
                if abs(y - cell[0]) + abs(x - cell[1]) < min_dist:
                    return False
            return True

        player = free.pop()
        boxes, goals = [], []
        tries = 0
        while len(boxes) < n_boxes and free and tries < 200:
            tries += 1
            b = free.pop()
            if not far(b, [player], 3):
                continue
            if not far(b, boxes + goals, 2):
                continue
            g = free.pop()
            if not far(g, boxes + goals, 2) or g == b:
                continue
            if not (self._reachable(grid, player, b)
                    and self._reachable(grid, b, g)):
                continue
            boxes.append(b)
            goals.append(g)

        if not boxes:                                        # degenerate: retry
            return self._gen_level(index)

        for (y, x) in boxes:
            grid[y][x] = BOX
        for (y, x) in goals:
            if grid[y][x] == FLOOR:
                grid[y][x] = GOAL
        grid[player[0]][player[1]] = PLAYER
        return _Level(
            grid=grid, boxes=boxes, goals=goals, player=player,
            budget=6 * size * size // max(1, n_boxes),
        )

    @staticmethod
    def _reachable(grid, src, dst) -> bool:
        """BFS on floor cells (walls only) - walls-only solvability proxy."""
        size = len(grid)
        seen = {src}
        queue = [src]
        while queue:
            y, x = queue.pop()
            if (y, x) == dst:
                return True
            for dy, dx in ((-1, 0), (1, 0), (0, -1), (0, 1)):
                ny, nx = y + dy, x + dx
                if (0 <= ny < size and 0 <= nx < size
                        and grid[ny][nx] != WALL and (ny, nx) not in seen):
                    seen.add((ny, nx))
                    queue.append((ny, nx))
        return False

    # ------------------------------------------------------------------ #
    # environment contract                                                 #
    # ------------------------------------------------------------------ #
    def observe(self) -> MockFrame:
        if self._level is None:
            return MockFrame(grids=[[[FLOOR]]],
                             levels_completed=self.levels_completed,
                             win_levels=self.win_levels,
                             state=self.state)
        return MockFrame(
            grids=[[list(row) for row in self._level.grid]],
            levels_completed=self.levels_completed,
            win_levels=self.win_levels,
            available_actions=list(ALL_ACTIONS),
            state=self.state,
        )

    def reset(self) -> MockFrame:
        self.levels_completed = 0
        self.state = "NOT_FINISHED"
        self._level = self._gen_level(0)
        self._undo = []
        self._last_frame = self.observe()
        return self._last_frame

    def step(self, action: str, x: int | None = None, y: int | None = None) -> MockFrame:
        if self.state in ("WIN", "GAME_OVER"):
            return self.observe()
        if action == "RESET":
            return self.reset()

        lvl = self._level
        self._undo.append([list(row) for row in lvl.grid])
        if len(self._undo) > UNDO_DEPTH:
            self._undo.pop(0)

        moved = False
        if action in DIRS:
            moved = self._move(lvl, *DIRS[action])
        elif action == "ACTION5":
            moved = False                                     # explicit no-op
        elif action == "ACTION6":
            moved = self._click(lvl, x, y)
        elif action == "ACTION7":
            if self._undo:
                lvl.grid = self._undo.pop()                   # restore state
                self._sync_entities(lvl)
                moved = True
        # unknown action names: treated as no-op (frame unchanged)

        lvl.steps_used += 1
        if self._all_boxes_on_goals(lvl):
            self.levels_completed += 1
            if self.levels_completed >= self.win_levels:
                self.state = "WIN"
            else:
                self._level = self._gen_level(self.levels_completed)
                self._undo = []
        elif lvl.steps_used >= lvl.budget:
            self.state = "GAME_OVER"

        self._last_frame = self.observe()
        return self._last_frame

    # ------------------------------------------------------------------ #
    # mechanics                                                            #
    # ------------------------------------------------------------------ #
    def _move(self, lvl: _Level, dy: int, dx: int) -> bool:
        size = len(lvl.grid)
        py, px = lvl.player
        ny, nx = py + dy, px + dx
        if not (0 <= ny < size and 0 <= nx < size) or lvl.grid[ny][nx] == WALL:
            return False
        if lvl.grid[ny][nx] == BOX:
            by, bx = ny + dy, nx + dx
            if not (0 <= by < size and 0 <= bx < size):
                return False
            if lvl.grid[by][bx] in (WALL, BOX):
                return False
            lvl.grid[by][bx] = BOX
            lvl.boxes[lvl.boxes.index((ny, nx))] = (by, bx)
        lvl.grid[py][px] = FLOOR
        lvl.grid[ny][nx] = PLAYER
        lvl.player = (ny, nx)
        return True

    def _click(self, lvl: _Level, x: int | None, y: int | None) -> bool:
        """ACTION6(x, y): one step toward an adjacent walkable cell."""
        if x is None or y is None:
            return False
        size = len(lvl.grid)
        if not (0 <= y < size and 0 <= x < size):
            return False
        py, px = lvl.player
        if abs(y - py) + abs(x - px) != 1:
            return False
        if lvl.grid[y][x] in (WALL, BOX):
            return False
        lvl.grid[py][px] = FLOOR
        lvl.grid[y][x] = PLAYER
        lvl.player = (y, x)
        return True

    @staticmethod
    def _sync_entities(lvl: _Level) -> None:
        """Rebuild player/box coordinates after an undo (grid is truth)."""
        lvl.boxes = []
        for y, row in enumerate(lvl.grid):
            for x, v in enumerate(row):
                if v == PLAYER:
                    lvl.player = (y, x)
                    lvl.grid[y][x] = FLOOR                 # re-stamp below
                elif v == BOX:
                    lvl.boxes.append((y, x))
        y, x = lvl.player
        lvl.grid[y][x] = PLAYER

    @staticmethod
    def _all_boxes_on_goals(lvl: _Level) -> bool:
        goal_set = set(lvl.goals)
        return all(b in goal_set for b in lvl.boxes)


# --------------------------------------------------------------------- #
# adapter used by collect / eval                                        #
# --------------------------------------------------------------------- #


def make_mock_env(seed: int = 0, win_levels: int = 3) -> MockGame:
    return MockGame(seed=seed, win_levels=win_levels)


def step_env(env, action_name: str, x: int | None = None, y: int | None = None):
    """Dispatch one action to either a MockGame (name-based) or the official
    offline engine (GameAction enum). Keeps callers engine-agnostic."""
    if isinstance(env, MockGame):
        return env.step(action_name, x=x, y=y)
    from arcengine import GameAction                          # official toolkit

    action = GameAction.from_name(action_name)
    if x is not None and y is not None and action.is_complex():
        # official signature: step(action, data={"x": .., "y": ..})
        return env.step(action, data={"x": int(x), "y": int(y)})
    return env.step(action)