gameboy_worlds.emulation.hamtaro.parsers

  1import os
  2from enum import Enum
  3
  4import numpy as np
  5
  6from gameboy_worlds.utils import verify_parameters, log_error, log_warn
  7from gameboy_worlds.emulation.parser import NamedScreenRegion, StateParser
  8
  9
 10class _BaseHamtaroParser(StateParser):
 11    """
 12    Minimal parser scaffold for Hamtaro variants.
 13
 14    This only configures rom_data_path so the base StateParser can run.
 15    """
 16
 17    VARIANT = ""
 18
 19    def __init__(self, pyboy, parameters):
 20        verify_parameters(parameters)
 21        variant = self.VARIANT
 22        if f"{variant}_rom_data_path" not in parameters:
 23            log_error(
 24                f"ROM data path not found for variant: {variant}. Add {variant}_rom_data_path to config files.",
 25                parameters,
 26            )
 27        self.rom_data_path = parameters[f"{variant}_rom_data_path"]
 28        super().__init__(pyboy, parameters)
 29
 30    def __repr__(self) -> str:
 31        return f"{self.__class__.__name__}(variant={self.VARIANT})"
 32
 33
 34class AgentState(Enum):
 35    FREE_ROAM = 0
 36    IN_DIALOGUE = 1
 37    IN_MENU = 2
 38
 39
 40class HamtaroStateParser(_BaseHamtaroParser):
 41    """
 42    Hamtaro parser with optional reference-region matching.
 43
 44    If captured reference regions are available under rom_data/captures, they are
 45    used for menu and dialogue detection. Otherwise the parser falls back to
 46    broad visual heuristics so dev_play remains usable before bootstrapping.
 47    """
 48
 49    REGIONS = [
 50        ("dialogue_top_border", 8, 96, 144, 4),
 51        ("dialogue_bottom_border", 8, 140, 144, 4),
 52        ("menu_top_band", 45, 6, 25, 94),
 53        ("menu_bottom_band", 10, 100, 135, 10),
 54    ]
 55
 56    DIALOGUE_TOP_Y = 96
 57    DIALOGUE_BOTTOM_Y = 140
 58
 59    def __init__(self, pyboy, parameters):
 60        named_regions = []
 61        super().__init__(pyboy, parameters)
 62        captures_dir = os.path.join(self.rom_data_path, "captures")
 63        for region_name, start_x, start_y, width, height in self.REGIONS:
 64            target_path = os.path.join(captures_dir, region_name)
 65            if not os.path.exists(f"{target_path}.npy"):
 66                target_path = None
 67            region = NamedScreenRegion(
 68                region_name,
 69                start_x,
 70                start_y,
 71                width,
 72                height,
 73                parameters=parameters,
 74                target_path=target_path,
 75            )
 76            expected_shape = (height, width, 1)
 77            if region.target is not None and region.target.shape != expected_shape:
 78                log_warn(
 79                    f"Ignoring stale capture for {region_name}: expected shape {expected_shape}, found {region.target.shape}. Re-capture this region in dev_play.",
 80                    parameters,
 81                )
 82                region.target = None
 83                region.target_path = None
 84            named_regions.append(region)
 85        self.named_screen_regions.update(
 86            {region.name: region for region in named_regions}
 87        )
 88
 89    def _has_target(self, region_name: str) -> bool:
 90        region = self.named_screen_regions.get(region_name)
 91        return region is not None and region.target is not None
 92
 93    def _matches_all_available(self, current_screen: np.ndarray, region_names) -> bool:
 94        available_regions = [name for name in region_names if self._has_target(name)]
 95        if len(available_regions) == 0:
 96            return False
 97        try:
 98            return all(
 99                self.named_region_matches_target(current_screen, region_name)
100                for region_name in available_regions
101            )
102        except (RuntimeError, ValueError):
103            return False
104
105    def _row_dark_fraction(self, frame: np.ndarray, row_idx: int) -> float:
106        row = frame[row_idx, :, 0]
107        return float(np.mean(row < 40))
108
109    def _region_dark_fraction(
110        self, frame: np.ndarray, x0: int, y0: int, x1: int, y1: int
111    ) -> float:
112        region = frame[y0:y1, x0:x1, 0]
113        return float(np.mean(region < 90))
114
115    def _region_bright_fraction(
116        self, frame: np.ndarray, x0: int, y0: int, x1: int, y1: int
117    ) -> float:
118        region = frame[y0:y1, x0:x1, 0]
119        return float(np.mean(region > 180))
120
121    def is_in_dialogue(self, current_screen: np.ndarray) -> bool:
122        if self._matches_all_available(
123            current_screen, ["dialogue_top_border", "dialogue_bottom_border"]
124        ):
125            return True
126        top_border = self._row_dark_fraction(current_screen, self.DIALOGUE_TOP_Y)
127        bottom_border = self._row_dark_fraction(
128            current_screen, self.DIALOGUE_BOTTOM_Y
129        )
130        text_density = self._region_dark_fraction(current_screen, 12, 108, 148, 132)
131        background_brightness = self._region_bright_fraction(
132            current_screen, 12, 108, 148, 132
133        )
134        return (
135            top_border > 0.45
136            and bottom_border > 0.45
137            and text_density > 0.06
138            and background_brightness > 0.30
139        )
140
141    def is_in_menu(self, current_screen: np.ndarray) -> bool:
142        if self._matches_all_available(
143            current_screen, ["menu_top_band", "menu_bottom_band"]
144        ):
145            return True
146        if self.is_in_dialogue(current_screen):
147            return False
148        top_band_dark = self._region_dark_fraction(current_screen, 0, 0, 160, 14)
149        bottom_band_dark = self._region_dark_fraction(current_screen, 0, 104, 160, 144)
150        center_bright = self._region_bright_fraction(current_screen, 20, 20, 140, 124)
151        return top_band_dark > 0.35 and bottom_band_dark > 0.35 and center_bright > 0.2
152
153    def get_agent_state(self, current_screen: np.ndarray) -> AgentState:
154        if self.is_in_menu(current_screen):
155            return AgentState.IN_MENU
156        if self.is_in_dialogue(current_screen):
157            return AgentState.IN_DIALOGUE
158        return AgentState.FREE_ROAM
159
160
161class HamtaroHamHamHeartbreakParser(HamtaroStateParser):
162    VARIANT = "hamtaro_ham_ham_heartbreak"
163
164
165class HamtaroHamHamsUniteParser(HamtaroStateParser):
166    VARIANT = "hamtaro_ham_hams_unite"