TypeSafe Mario — Jev picks NES controller inputs from telemetry, not pixels

"""TypeSafe-powered direct controller for Super Mario Bros."""

from .actions import Action
from .state import MarioSnapshot, MarioStateParser

__all__ = ["Action", "MarioSnapshot", "MarioStateParser"]

from __future__ import annotations

import argparse
import json
from collections.abc import Sequence
from pathlib import Path

from .policy import HeuristicPolicy, TypeSafePolicy
from .runner import run_episode
from .state import MarioStateParser


def _demo_ram() -> bytearray:
    ram = bytearray(0x0800)
    # Mario position and one goomba ahead.
    ram[0x006D] = 0
    ram[0x0086] = 172
    ram[0x000F] = 1
    ram[0x0016] = 0x06
    ram[0x006E] = 0
    ram[0x0087] = 214
    ram[0x00CF] = 79
    # Fill the tile-map row immediately below Mario with solid ground.
    ground_row = (79 + 16 - 32) // 16
    for column in range(16):
        ram[0x0500 + ground_row * 16 + column] = 1
    return ram


def state_demo() -> int:
    info = {
        "world": 1,
        "stage": 1,
        "area": 1,
        "x_pos": 172,
        "y_pos": 79,
        "y_pixel": 79,
        "left_x_pos": 60,
        "progress": 172,
        "progress_max": 172,
        "status": "small",
        "player_state": 8,
        "life": 2,
        "coins": 0,
        "score": 200,
        "time": 387,
        "death": False,
        "clear": False,
    }
    snapshot = MarioStateParser().parse(info, _demo_ram(), previous_action="right")
    print(json.dumps(snapshot.to_state(), indent=2))
    print("\n--- TEXT VIEW ---\n")
    print(snapshot.to_text())
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="typesafe-mario")
    subparsers = parser.add_subparsers(dest="command", required=True)
    subparsers.add_parser("state-demo", help="Show structured and text state without API calls")

    play = subparsers.add_parser("play", help="Run a Mario episode")
    play.add_argument("--env", default="SuperMarioBros-1-1-v0")
    play.add_argument(
        "--frames-per-decision",
        type=int,
        default=8,
        help="Minimum macro duration in emulator frames",
    )
    play.add_argument("--max-decisions", type=int, default=2000)
    play.add_argument("--seed", type=int, default=123)
    play.add_argument("--policy", choices=("typesafe", "heuristic"), default="typesafe")
    play.add_argument(
        "--display",
        choices=("dashboard", "game", "none"),
        default="dashboard",
        help="Combined telemetry dashboard, plain game window, or headless mode",
    )
    play.add_argument("--artifacts-dir", type=Path, default=Path("artifacts"))
    play.add_argument(
        "--screenshot",
        type=Path,
        help="Save the first populated dashboard frame as a PNG",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    if args.command == "state-demo":
        return state_demo()
    if args.command == "play":
        policy = TypeSafePolicy() if args.policy == "typesafe" else HeuristicPolicy()
        log_path = run_episode(
            env_id=args.env,
            policy=policy,
            frames_per_decision=args.frames_per_decision,
            max_decisions=args.max_decisions,
            seed=args.seed,
            artifacts_dir=args.artifacts_dir,
            display=args.display,
            screenshot_path=args.screenshot,
        )
        print(f"Run log: {log_path.resolve()}")
        return 0
    raise AssertionError(f"Unhandled command: {args.command}")


if __name__ == "__main__":
    raise SystemExit(main())

from __future__ import annotations

from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from typing import Any

from .actions import ACTION_DESCRIPTIONS, Action
from .policy import Decision
from .state import MarioSnapshot

Color = tuple[int, int, int]


class DashboardCommand(StrEnum):
    CONTINUE = "continue"
    RESTART = "restart"
    QUIT = "quit"


@dataclass(frozen=True)
class DashboardTheme:
    canvas: Color = (15, 17, 21)
    surface: Color = (22, 25, 31)
    raised: Color = (29, 33, 40)
    line: Color = (48, 54, 64)
    text: Color = (236, 239, 243)
    muted: Color = (141, 150, 163)
    accent: Color = (137, 166, 193)
    accent_soft: Color = (72, 89, 105)
    warning: Color = (210, 167, 100)
    danger: Color = (199, 112, 101)


@dataclass(frozen=True)
class DashboardConfig:
    width: int = 1440
    height: int = 810
    panel_width: int = 472
    fps: int = 60


def clamp01(value: float | None) -> float:
    if value is None:
        return 0.0
    return max(0.0, min(1.0, float(value)))


class LiveDashboard:
    """One-window game feed and TypeSafe telemetry display."""

    def __init__(
        self,
        config: DashboardConfig | None = None,
        theme: DashboardTheme | None = None,
    ) -> None:
        try:
            import pygame
        except ImportError as exc:
            raise RuntimeError(
                "The dashboard needs pygame. Install the mario extras first."
            ) from exc

        self.pg = pygame
        self.config = config or DashboardConfig()
        self.theme = theme or DashboardTheme()
        pygame.init()
        pygame.display.set_caption("TypeSafe plays Super Mario Bros.")
        self.screen = pygame.display.set_mode((self.config.width, self.config.height))
        self.clock = pygame.time.Clock()
        self.font_title = pygame.font.SysFont("Segoe UI Semibold", 25)
        self.font_action = pygame.font.SysFont("Segoe UI Semibold", 31)
        self.font_body = pygame.font.SysFont("Segoe UI", 17)
        self.font_small = pygame.font.SysFont("Segoe UI", 14)
        self.font_label = pygame.font.SysFont("Segoe UI Semibold", 14)
        self.font_mono = pygame.font.SysFont("Cascadia Mono", 14)

    def close(self) -> None:
        self.pg.quit()

    def save(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        self.pg.image.save(self.screen, str(path))

    def _text(
        self,
        value: str,
        font: Any,
        color: Color,
        x: int,
        y: int,
    ) -> None:
        self.screen.blit(font.render(value, True, color), (x, y))

    def _rule(self, x: int, y: int, width: int) -> None:
        self.pg.draw.line(self.screen, self.theme.line, (x, y), (x + width, y), 1)

    def _bar(
        self,
        *,
        label: str,
        value: float,
        x: int,
        y: int,
        width: int,
        selected: bool = False,
        color: Color | None = None,
    ) -> None:
        value = clamp01(value)
        bar_color = color or (self.theme.accent if selected else self.theme.accent_soft)
        self._text(label, self.font_small, self.theme.text if selected else self.theme.muted, x, y)
        value_text = f"{value * 100:4.1f}%"
        text_width = self.font_small.size(value_text)[0]
        self._text(value_text, self.font_small, self.theme.muted, x + width - text_width, y)
        bar_y = y + 21
        self.pg.draw.rect(self.screen, self.theme.raised, (x, bar_y, width, 7), border_radius=3)
        fill = max(2, round(width * value)) if value > 0 else 0
        if fill:
            self.pg.draw.rect(self.screen, bar_color, (x, bar_y, fill, 7), border_radius=3)

    def _draw_game(self, frame: Any, x: int, y: int, width: int, height: int) -> None:
        self.pg.draw.rect(self.screen, self.theme.surface, (x, y, width, height), border_radius=4)
        if frame is None:
            self._text(
                "Waiting for emulator frame", self.font_body, self.theme.muted, x + 28, y + 28
            )
            return
        try:
            import numpy as np
        except ImportError as exc:
            raise RuntimeError("The dashboard needs numpy.") from exc

        pixels = np.asarray(frame)
        if pixels.ndim != 3 or pixels.shape[2] < 3:
            raise ValueError(f"Expected an RGB frame, received shape {pixels.shape!r}")
        pixels = pixels[:, :, :3]
        source = self.pg.surfarray.make_surface(pixels.swapaxes(0, 1))
        source_ratio = source.get_width() / source.get_height()
        target_ratio = width / height
        if source_ratio > target_ratio:
            target_width = width
            target_height = round(width / source_ratio)
        else:
            target_height = height
            target_width = round(height * source_ratio)
        scaled = self.pg.transform.scale(source, (target_width, target_height))
        target_x = x + (width - target_width) // 2
        target_y = y + (height - target_height) // 2
        self.screen.blit(scaled, (target_x, target_y))
        self.pg.draw.rect(self.screen, self.theme.line, (x, y, width, height), 1, border_radius=4)

    def _restart_rect(self) -> Any:
        return self.pg.Rect(self.config.width - 144, 13, 120, 36)

    def draw(
        self,
        frame: Any,
        snapshot: MarioSnapshot,
        decision: Decision | None,
        *,
        decision_index: int,
        episode_reward: float,
        waiting: bool = False,
        run_ended: bool = False,
    ) -> DashboardCommand:
        restart_rect = self._restart_rect()
        command = DashboardCommand.CONTINUE
        for event in self.pg.event.get():
            if event.type == self.pg.QUIT:
                return DashboardCommand.QUIT
            if event.type == self.pg.KEYDOWN and event.key in (self.pg.K_ESCAPE, self.pg.K_q):
                return DashboardCommand.QUIT
            if event.type == self.pg.KEYDOWN and event.key == self.pg.K_r:
                command = DashboardCommand.RESTART
            if (
                event.type == self.pg.MOUSEBUTTONDOWN
                and event.button == 1
                and restart_rect.collidepoint(event.pos)
            ):
                command = DashboardCommand.RESTART

        c = self.config
        t = self.theme
        self.screen.fill(t.canvas)
        margin = 24
        header_h = 58
        panel_x = c.width - c.panel_width
        game_x = margin
        game_y = header_h + 14
        game_w = panel_x - margin * 2
        game_h = c.height - game_y - margin

        self._text("TypeSafe plays Mario", self.font_title, t.text, margin, 18)
        if run_ended:
            status = "Run ended"
        else:
            status = "Waiting for Jev" if waiting else "Live decision loop"
        status_width = self.font_small.size(status)[0]
        dot_x = panel_x - status_width - 24
        status_color = t.danger if run_ended else (t.warning if waiting else t.accent)
        self.pg.draw.circle(self.screen, status_color, (dot_x, 31), 4)
        self._text(status, self.font_small, t.muted, dot_x + 12, 21)
        hovered = restart_rect.collidepoint(self.pg.mouse.get_pos())
        button_color = t.accent_soft if hovered else t.raised
        self.pg.draw.rect(self.screen, button_color, restart_rect, border_radius=5)
        self.pg.draw.rect(self.screen, t.accent, restart_rect, 1, border_radius=5)
        label = "Restart"
        label_width, label_height = self.font_label.size(label)
        self._text(
            label,
            self.font_label,
            t.text,
            restart_rect.centerx - label_width // 2,
            restart_rect.centery - label_height // 2,
        )
        self._rule(margin, header_h, c.width - margin * 2)
        self._draw_game(frame, game_x, game_y, game_w, game_h)

        x = panel_x + 18
        width = c.panel_width - 42
        y = game_y
        self._text("World 1–1", self.font_label, t.muted, x, y)
        self._text(f"Decision {decision_index:04d}", self.font_small, t.muted, x + width - 100, y)
        y += 27

        action_name = (
            "Reading state…" if decision is None else decision.action.value.replace("_", " ")
        )
        self._text(action_name, self.font_action, t.text, x, y)
        y += 42
        if decision is None:
            description = "Jev is evaluating the next bounded controller action."
        else:
            description = ACTION_DESCRIPTIONS[decision.action]
        self._text(description, self.font_small, t.muted, x, y)
        y += 32
        self._rule(x, y, width)
        y += 18

        confidence = decision.confidence if decision else 0.0
        latency = decision.latency_ms if decision else 0.0
        metric_width = width // 3
        metrics = (
            ("Confidence", f"{confidence * 100:.0f}%"),
            ("Latency", f"{latency:.0f} ms"),
            ("Reward", f"{episode_reward:+.1f}"),
        )
        for index, (label, value) in enumerate(metrics):
            metric_x = x + index * metric_width
            self._text(label, self.font_small, t.muted, metric_x, y)
            self._text(value, self.font_body, t.text, metric_x, y + 19)
        y += 58
        self._text("Controller probabilities", self.font_label, t.text, x, y)
        y += 24

        probabilities = decision.probabilities if decision else {}
        selected_action = decision.action.value if decision else None
        for action in Action:
            probability = float(probabilities.get(action.value, 0.0))
            self._bar(
                label=action.value.replace("_", " "),
                value=probability,
                x=x,
                y=y,
                width=width,
                selected=action.value == selected_action,
            )
            y += 36

        self._rule(x, y + 2, width)
        y += 17
        self._text("Situation", self.font_label, t.text, x, y)
        y += 24
        jump = decision.jump_needed_probability if decision else None
        danger = clamp01((decision.danger_score or 0.0) / 2.0) if decision else 0.0
        self._bar(label="Jump useful now", value=clamp01(jump), x=x, y=y, width=width)
        y += 40
        danger_color = t.danger if danger >= 0.66 else t.warning
        self._bar(label="Immediate danger", value=danger, x=x, y=y, width=width, color=danger_color)
        y += 43

        enemies = ", ".join(enemy.kind for enemy in snapshot.enemies) or "none"
        state_lines = (
            f"Position  {snapshot.x}, {snapshot.y}",
            f"Motion    {snapshot.direction.replace('_', ' ')}",
            f"Nearby    {enemies}",
            f"Time      {snapshot.time_left}",
        )
        for line in state_lines:
            self._text(line, self.font_mono, t.muted, x, y)
            y += 20

        footer = (
            "Click Restart or press R to play again" if run_ended else "R restarts · Esc or Q quits"
        )
        self._text(footer, self.font_small, t.muted, x, c.height - 28)
        self.pg.display.flip()
        self.clock.tick(c.fps)
        return command

from __future__ import annotations

import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Protocol

from .actions import ACTION_DESCRIPTIONS, Action
from .state import MarioSnapshot


@dataclass(frozen=True)
class Decision:
    action: Action
    confidence: float
    probabilities: Mapping[str, float]
    latency_ms: float
    jump_needed_probability: float | None = None
    danger_score: float | None = None


class Policy(Protocol):
    def choose(self, snapshot: MarioSnapshot, actions: Sequence[Action]) -> Decision: ...


class TypeSafePolicy:
    def __init__(self) -> None:
        try:
            from typesafe_sdk import Choice, Noul, Score, TypeSafeClient
        except ImportError as exc:
            raise RuntimeError(
                "typesafe-sdk is not installed. Install the project before using Jev."
            ) from exc
        self._Choice = Choice
        self._Noul = Noul
        self._Score = Score
        self._client = TypeSafeClient()

    def close(self) -> None:
        close = getattr(self._client, "close", None)
        if callable(close):
            close()

    @staticmethod
    def _answer(response: Any, question_id: str, typed_collection: str) -> Any:
        collection = getattr(response, typed_collection, None)
        if collection is not None and question_id in collection:
            return collection[question_id]
        answers = getattr(response, "answers", None)
        if answers is not None and question_id in answers:
            return answers[question_id]
        raise KeyError(f"TypeSafe response omitted {question_id!r}")

    def choose(self, snapshot: MarioSnapshot, actions: Sequence[Action]) -> Decision:
        criteria = {action.value: ACTION_DESCRIPTIONS[action] for action in actions}
        questions = {
            "next_action": self._Choice(
                instructions={
                    "question": "Which controller macro should Mario commit to next?",
                    "goal": "Advance toward the stage flag while avoiding death.",
                    "timing": "The selected action is held for at least 8 emulator frames.",
                    "geometry": (
                        "Use `terrain.observation_reliability`. While airborne, prefer "
                        "`terrain.last_grounded_preview` over low-reliability current geometry. "
                        "A trusted obstacle or gap within three tiles requires a forward jump."
                    ),
                    "trajectory": (
                        "Use `trajectory`. If `crossing_known_gap` is true, preserve forward "
                        "speed and keep a forward jump held while rising. Do not switch to noop "
                        "or left over a gap."
                    ),
                    "stall": (
                        "If `episode.stalled_frames` is increasing and "
                        "`recent_control.outcome` is blocked, the current non-jump action failed."
                    ),
                    "enemy_timing": (
                        "Use `hazard` projections, not distance alone. Code has already accounted "
                        "for inference delay, action cadence, and the frames needed to clear an "
                        "enemy. If `jump_must_start_this_decision` is true, choose a forward jump "
                        "now; another right-run decision will miss the takeoff deadline. If "
                        "`contact_within_reaction_horizon` is true, also jump immediately. If "
                        "`will_land_before_contact` is true, the current jump will not clear the "
                        "enemy and another takeoff will be needed after landing. Use "
                        "`upcoming_enemies` and their spacing to avoid landing on a second or "
                        "third enemy hidden behind the nearest one."
                    ),
                    "delay": (
                        "`reaction_timing` describes how far the world moves before this choice "
                        "takes effect. Judge urgency from projected rather than current distance."
                    ),
                },
                criteria=criteria,
            ),
            "jump_needed": self._Noul(
                instructions=(
                    "Do trusted `terrain`, projected `hazard`, `trajectory`, and "
                    "`player.jump_phase` indicate that a forward jump should begin or remain "
                    "held now? `hazard.jump_must_start_this_decision=true` is unambiguously yes. "
                    "Also count a trusted obstacle/gap within three tiles, immediate projected "
                    "contact, or a rising jump over a known gap as yes."
                )
            ),
            "danger": self._Score(
                instructions="How dangerous is Mario's immediate situation?",
                criteria=[
                    "Safe open movement",
                    "Potential obstacle or enemy soon",
                    "Immediate collision, fall, or enemy threat",
                ],
            ),
        }
        started = time.perf_counter()
        response = self._client.system_one(state=snapshot.to_state(), questions=questions)
        latency_ms = (time.perf_counter() - started) * 1000

        action_answer = self._answer(response, "next_action", "choices")
        jump_answer = self._answer(response, "jump_needed", "nouls")
        danger_answer = self._answer(response, "danger", "scores")
        action = Action(str(action_answer.choice))
        probabilities = {
            str(key): float(value) for key, value in dict(action_answer.probabilities).items()
        }
        return Decision(
            action=action,
            confidence=float(action_answer.confidence),
            probabilities=probabilities,
            latency_ms=latency_ms,
            jump_needed_probability=float(jump_answer.noul),
            danger_score=float(danger_answer.score),
        )


class HeuristicPolicy:
    """Offline smoke-test policy; not intended as the Mario benchmark baseline."""

    def choose(self, snapshot: MarioSnapshot, actions: Sequence[Action]) -> Decision:
        allowed = set(actions)
        action = Action.RIGHT_RUN if Action.RIGHT_RUN in allowed else actions[0]
        if snapshot.stalled_steps >= 2 and Action.RIGHT_RUN_JUMP in allowed:
            action = Action.RIGHT_RUN_JUMP
        return Decision(
            action=action,
            confidence=1.0,
            probabilities={candidate.value: float(candidate == action) for candidate in actions},
            latency_ms=0.0,
        )

from __future__ import annotations

import json
from concurrent.futures import Future, ThreadPoolExecutor
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from .actions import ACTION_TO_INDEX, JUMP_ACTIONS, JUMP_RELEASE_ACTION, Action
from .dashboard import DashboardCommand, LiveDashboard
from .policy import Decision, Policy
from .state import MarioStateParser


def _unwrap_ram(env: Any) -> Any:
    current = env
    visited: set[int] = set()
    while id(current) not in visited:
        visited.add(id(current))
        ram = getattr(current, "ram", None)
        if ram is not None:
            return ram
        next_env = getattr(current, "env", None)
        if next_env is None:
            break
        current = next_env
    return None


def create_mario_env(env_id: str, render_mode: str = "human") -> Any:
    try:
        import gym_super_mario_bros  # noqa: F401
        import gymnasium as gym
        from gym_super_mario_bros.actions import SIMPLE_MOVEMENT
        from nes_py.wrappers import JoypadSpace
    except ImportError as exc:
        raise RuntimeError(
            'Mario dependencies are missing. Install with: pip install -e ".[mario]"'
        ) from exc

    env = gym.make(env_id, render_mode=render_mode)
    return JoypadSpace(env, SIMPLE_MOVEMENT)


def _record_decision(
    log: Any,
    *,
    decision_index: int,
    snapshot: Any,
    decision: Decision,
    reward: float,
    terminated: bool,
    truncated: bool,
) -> None:
    record = {
        "decision": decision_index,
        "state": snapshot.to_state(),
        "debug_state": snapshot.to_debug_state(),
        "state_text": snapshot.to_text(),
        "action": decision.action.value,
        "confidence": decision.confidence,
        "probabilities": dict(decision.probabilities),
        "jump_needed_probability": decision.jump_needed_probability,
        "danger_score": decision.danger_score,
        "latency_ms": decision.latency_ms,
        "reward": reward,
        "terminated": bool(terminated),
        "truncated": bool(truncated),
    }
    log.write(json.dumps(record, separators=(",", ":")) + "\n")
    log.flush()


def _run_realtime_dashboard(
    *,
    env: Any,
    dashboard: LiveDashboard,
    policy: Policy,
    parser: MarioStateParser,
    frame: Any,
    info: dict[str, Any],
    log: Any,
    frames_per_decision: int,
    max_decisions: int,
    screenshot_path: Path | None,
) -> None:
    actions = tuple(Action)
    active_decision: Decision | None = None
    pending: Future[Decision] | None = None
    pending_snapshot: Any = None
    pending_index = 0
    pending_request_frame = 0
    next_index = 0
    frame_index = 0
    last_request_frame = -frames_per_decision
    episode_reward = 0.0
    reward_since_decision = 0.0
    previous_action: Action | None = None
    previous_reward = 0.0
    previous_latency_ms = 0.0
    previous_response_delay_frames = 0
    screenshot_saved = False
    terminated = truncated = False

    with ThreadPoolExecutor(max_workers=1, thread_name_prefix="typesafe-jev") as executor:
        while True:
            decision_updated = False
            snapshot = parser.parse(
                info,
                _unwrap_ram(env),
                previous_action=previous_action.value if previous_action else None,
                previous_reward=previous_reward,
                previous_latency_ms=previous_latency_ms,
                previous_response_delay_frames=previous_response_delay_frames,
            )

            if pending is not None and pending.done():
                active_decision = pending.result()
                _record_decision(
                    log,
                    decision_index=pending_index,
                    snapshot=pending_snapshot,
                    decision=active_decision,
                    reward=reward_since_decision,
                    terminated=terminated,
                    truncated=truncated,
                )
                print(
                    f"#{pending_index:04d} x={pending_snapshot.x:04d} "
                    f"action={active_decision.action.value:<15} "
                    f"confidence={active_decision.confidence:.2f} "
                    f"latency={active_decision.latency_ms:.0f}ms"
                )
                pending = None
                reward_since_decision = 0.0
                decision_updated = True
                previous_latency_ms = active_decision.latency_ms
                previous_response_delay_frames = frame_index - pending_request_frame

            run_ended = bool(
                snapshot.dead
                or snapshot.clear
                or terminated
                or truncated
                or (next_index >= max_decisions and pending is None)
            )

            if (
                not run_ended
                and pending is None
                and next_index < max_decisions
                and frame_index - last_request_frame >= frames_per_decision
            ):
                pending_snapshot = snapshot
                pending_index = next_index
                pending = executor.submit(policy.choose, snapshot, actions)
                pending_request_frame = frame_index
                next_index += 1
                last_request_frame = frame_index

            if not run_ended:
                action = active_decision.action if active_decision else Action.NOOP
                if decision_updated and action in JUMP_ACTIONS and snapshot.grounded:
                    # A new jump macro needs a button-up edge before A is pressed again.
                    action = JUMP_RELEASE_ACTION[action]
                frame, reward, terminated, truncated, info = env.step(ACTION_TO_INDEX[action])
                previous_action = action
                previous_reward = float(reward)
                reward_since_decision += float(reward)
                episode_reward += float(reward)
                frame_index += 1

            command = dashboard.draw(
                frame,
                snapshot,
                active_decision,
                decision_index=max(0, next_index - 1),
                episode_reward=episode_reward,
                waiting=pending is not None,
                run_ended=run_ended,
            )
            if screenshot_path is not None and active_decision is not None and not screenshot_saved:
                dashboard.save(screenshot_path)
                screenshot_saved = True
            if command == DashboardCommand.QUIT:
                break
            if command == DashboardCommand.RESTART:
                if pending is not None:
                    pending.cancel()
                frame, info = env.reset()
                parser.reset()
                active_decision = None
                pending = None
                pending_snapshot = None
                pending_index = 0
                pending_request_frame = 0
                next_index = 0
                frame_index = 0
                last_request_frame = -frames_per_decision
                episode_reward = 0.0
                reward_since_decision = 0.0
                previous_action = None
                previous_reward = 0.0
                previous_latency_ms = 0.0
                previous_response_delay_frames = 0
                terminated = truncated = False
                print("--- Restarted ---")


def run_episode(
    *,
    env_id: str,
    policy: Policy,
    frames_per_decision: int,
    max_decisions: int,
    seed: int,
    artifacts_dir: Path,
    display: str = "dashboard",
    screenshot_path: Path | None = None,
) -> Path:
    if frames_per_decision < 1:
        raise ValueError("frames_per_decision must be at least 1")

    if display not in {"dashboard", "game", "none"}:
        raise ValueError("display must be dashboard, game, or none")
    render_mode = "human" if display == "game" else "rgb_array"
    env = create_mario_env(env_id, render_mode=render_mode)
    dashboard = LiveDashboard() if display == "dashboard" else None
    parser = MarioStateParser(decision_horizon_frames=frames_per_decision)
    artifacts_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    log_path = artifacts_dir / f"run-{timestamp}.jsonl"

    frame, info = env.reset(seed=seed)

    try:
        with log_path.open("w", encoding="utf-8") as log:
            if dashboard:
                _run_realtime_dashboard(
                    env=env,
                    dashboard=dashboard,
                    policy=policy,
                    parser=parser,
                    frame=frame,
                    info=info,
                    log=log,
                    frames_per_decision=frames_per_decision,
                    max_decisions=max_decisions,
                    screenshot_path=screenshot_path,
                )
                return log_path

            previous_action: Action | None = None
            previous_reward = 0.0
            previous_latency_ms = 0.0
            actions = tuple(Action)
            for decision_index in range(max_decisions):
                snapshot = parser.parse(
                    info,
                    _unwrap_ram(env),
                    previous_action=previous_action.value if previous_action else None,
                    previous_reward=previous_reward,
                    previous_latency_ms=previous_latency_ms,
                    previous_response_delay_frames=0,
                )
                if snapshot.dead or snapshot.clear:
                    break

                decision = policy.choose(snapshot, actions)
                action_index = ACTION_TO_INDEX[decision.action]
                total_reward = 0.0
                terminated = truncated = False
                for _ in range(frames_per_decision):
                    frame, reward, terminated, truncated, info = env.step(action_index)
                    total_reward += float(reward)
                    if terminated or truncated:
                        break

                _record_decision(
                    log,
                    decision_index=decision_index,
                    snapshot=snapshot,
                    decision=decision,
                    reward=total_reward,
                    terminated=terminated,
                    truncated=truncated,
                )
                print(
                    f"#{decision_index:04d} x={snapshot.x:04d} "
                    f"action={decision.action.value:<15} "
                    f"confidence={decision.confidence:.2f} latency={decision.latency_ms:.0f}ms"
                )

                previous_action = decision.action
                previous_reward = total_reward
                previous_latency_ms = decision.latency_ms
                if terminated or truncated:
                    break
    finally:
        if dashboard:
            dashboard.close()
        env.close()
        close = getattr(policy, "close", None)
        if callable(close):
            close()

    return log_path

from __future__ import annotations

import unittest

from typesafe_mario.actions import JUMP_RELEASE_ACTION, Action
from typesafe_mario.dashboard import DashboardCommand, clamp01


class DashboardTests(unittest.TestCase):
    def test_clamp01_handles_missing_and_out_of_range_values(self) -> None:
        self.assertEqual(clamp01(None), 0.0)
        self.assertEqual(clamp01(-0.2), 0.0)
        self.assertEqual(clamp01(0.42), 0.42)
        self.assertEqual(clamp01(1.8), 1.0)

    def test_jump_macros_have_non_jump_release_frames(self) -> None:
        self.assertEqual(JUMP_RELEASE_ACTION[Action.RIGHT_JUMP], Action.RIGHT)
        self.assertEqual(JUMP_RELEASE_ACTION[Action.RIGHT_RUN_JUMP], Action.RIGHT_RUN)
        self.assertEqual(JUMP_RELEASE_ACTION[Action.JUMP], Action.NOOP)

    def test_dashboard_commands_are_stable_strings(self) -> None:
        self.assertEqual(DashboardCommand.RESTART, "restart")
        self.assertEqual(DashboardCommand.QUIT, "quit")


if __name__ == "__main__":
    unittest.main()

from __future__ import annotations

import unittest

from typesafe_mario.state import MarioStateParser


class MarioStateParserTests(unittest.TestCase):
    def base_info(self, **overrides: object) -> dict[str, object]:
        info: dict[str, object] = {
            "world": 1,
            "stage": 1,
            "area": 1,
            "x_pos": 100,
            "y_pos": 80,
            "y_pixel": 80,
            "progress": 100,
            "time": 390,
            "life": 2,
            "status": "small",
        }
        info.update(overrides)
        return info

    def test_first_snapshot_is_not_stalled(self) -> None:
        snapshot = MarioStateParser().parse(self.base_info())

        self.assertEqual(snapshot.stalled_steps, 0)
        self.assertEqual(snapshot.direction, "nearly_stationary")

    def test_motion_and_stall_are_derived_across_frames(self) -> None:
        parser = MarioStateParser()
        parser.parse(self.base_info())

        moving = parser.parse(self.base_info(x_pos=108, progress=108, y_pos=75))
        stalled = parser.parse(self.base_info(x_pos=108, progress=108, y_pos=75))

        self.assertEqual(moving.direction, "moving_right")
        self.assertEqual(moving.vertical_motion, "falling")
        self.assertEqual(moving.best_progress, 108)
        self.assertEqual(stalled.stalled_steps, 1)

    def test_ram_enemy_and_grid_become_structured_state(self) -> None:
        ram = bytearray(0x0800)
        ram[0x000F] = 1
        ram[0x0016] = 0x06
        ram[0x006E] = 0
        ram[0x0087] = 142
        ram[0x00CF] = 80
        ram[0x0010] = 1
        ram[0x0017] = 0x06
        ram[0x006F] = 0
        ram[0x0088] = 180
        ram[0x00D0] = 80

        snapshot = MarioStateParser().parse(self.base_info(), ram)
        state = snapshot.to_state()
        debug = snapshot.to_debug_state()

        self.assertEqual(state["hazard"]["nearest_enemy_kind"], "goomba")
        self.assertEqual(state["hazard"]["nearest_enemy_distance_pixels"], 42)
        self.assertEqual(len(state["hazard"]["upcoming_enemies"]), 2)
        self.assertEqual(state["hazard"]["spacing_to_second_enemy_pixels"], 38)
        self.assertEqual(len(debug["local_grid"]["rows"]), 9)
        self.assertIn("M", "".join(debug["local_grid"]["rows"]))
        self.assertIn("goomba 42px ahead", snapshot.to_text())

        moving = MarioStateParser()
        moving.parse(self.base_info(), ram)
        next_state = moving.parse(
            self.base_info(x_pos=108, progress=108),
            ram,
            previous_latency_ms=100,
        ).to_state()
        self.assertEqual(next_state["hazard"]["relative_velocity_x"], -8)
        self.assertEqual(next_state["hazard"]["estimated_contact_frames"], 4)
        self.assertTrue(next_state["hazard"]["contact_within_reaction_horizon"])
        self.assertTrue(next_state["hazard"]["takeoff_window_already_missed"])
        self.assertFalse(next_state["hazard"]["jump_must_start_this_decision"])
        self.assertEqual(next_state["reaction_timing"]["last_inference_delay_frames"], 6)
        self.assertEqual(next_state["hazard"]["projected_distance_after_reaction_pixels"], 0)

        deadline_ram = bytearray(0x0800)
        deadline_ram[0x000F] = 1
        deadline_ram[0x0016] = 0x06
        deadline_ram[0x006E] = 0
        deadline_ram[0x0087] = 236
        deadline_ram[0x00CF] = 80
        for column in range(16):
            deadline_ram[0x0500 + 5 * 16 + column] = 1
        deadline_parser = MarioStateParser()
        deadline_parser.parse(self.base_info(), deadline_ram)
        deadline_state = deadline_parser.parse(
            self.base_info(x_pos=108, progress=108),
            deadline_ram,
            previous_latency_ms=100,
            previous_response_delay_frames=8,
        ).to_state()
        self.assertEqual(deadline_state["hazard"]["estimated_contact_frames"], 16)
        self.assertEqual(deadline_state["hazard"]["takeoff_deadline_frames"], 0)
        self.assertEqual(deadline_state["reaction_timing"]["last_inference_delay_frames"], 8)
        self.assertTrue(deadline_state["hazard"]["jump_must_start_this_decision"])

    def test_pipe_geometry_is_exposed_as_navigation_state(self) -> None:
        ram = bytearray(0x0800)
        for column in range(16):
            ram[0x0500 + 5 * 16 + column] = 1
        pipe_column = 7
        ram[0x0500 + 3 * 16 + pipe_column] = 1
        ram[0x0500 + 4 * 16 + pipe_column] = 1

        snapshot = MarioStateParser().parse(self.base_info(), ram)
        terrain = snapshot.to_state()["terrain"]

        self.assertTrue(terrain["obstacle_ahead"])
        self.assertEqual(terrain["obstacle_distance_tiles"], 1)
        self.assertEqual(terrain["obstacle_height_tiles"], 2)
        self.assertNotIn("summary", terrain)

    def test_recent_control_tracks_macro_outcome(self) -> None:
        parser = MarioStateParser()
        parser.parse(self.base_info(), previous_action="right_run")
        state = parser.parse(
            self.base_info(x_pos=106, progress=106),
            previous_action="right_run",
        ).to_state()

        self.assertEqual(state["recent_control"]["action"], "right_run")
        self.assertEqual(state["recent_control"]["frames_observed"], 2)
        self.assertEqual(state["recent_control"]["progress_gained_pixels"], 6)

    def test_airborne_state_exposes_trajectory_and_low_reliability_geometry(self) -> None:
        ram = bytearray(0x0800)
        for column in range(16):
            ram[0x0500 + 5 * 16 + column] = 1
        parser = MarioStateParser()
        parser.parse(self.base_info(), ram)

        state = parser.parse(self.base_info(x_pos=106, y_pos=90, progress=106), ram).to_state()

        self.assertEqual(state["trajectory"]["airborne_frames"], 1)
        self.assertEqual(state["trajectory"]["horizontal_distance_since_takeoff_pixels"], 6)
        self.assertEqual(state["terrain"]["observation_reliability"], "low_airborne")

    def test_reset_clears_episode_history(self) -> None:
        parser = MarioStateParser(goal="Finish safely")
        parser.parse(self.base_info(x_pos=300, progress=300), previous_action="right_run")

        parser.reset()
        snapshot = parser.parse(self.base_info(x_pos=40, progress=40))

        self.assertEqual(snapshot.goal, "Finish safely")
        self.assertEqual(snapshot.best_progress, 40)
        self.assertEqual(snapshot.stalled_steps, 0)
        self.assertIsNone(snapshot.previous_action)


if __name__ == "__main__":
    unittest.main()