Add per-condition encounter rates to seed data (#26)

Co-authored-by: Julian Tabel <juliantabel.jt@gmail.com>
Co-committed-by: Julian Tabel <juliantabel.jt@gmail.com>

tools/import-pokedb/import_pokedb/mappings.py

@@ -84,8 +84,8 @@ ENCOUNTER_METHOD_MAP: dict[str, str] = {
     "cave-spot": "walk",
     "bubble-spot": "surf",
     "sand-spot": "walk",
-    "horde": "walk",
-    "sos-encounter": "walk",
+    "horde": "horde",
+    "sos-encounter": "sos",
     "ambush": "walk",
     # Seaweed / diving
     "diving": "surf",
@@ -105,7 +105,7 @@ ENCOUNTER_METHOD_MAP: dict[str, str] = {
     "dust-cloud": "walk",
     "hidden-grotto": "static",
     "hidden-encounter": "walk",
-    "horde-encounter": "walk",
+    "horde-encounter": "horde",
     "shaking-trees": "walk",
     "shaking-ore-deposits": "walk",
     "island-scan": "static",

tools/import-pokedb/import_pokedb/models.py

@@ -13,16 +13,22 @@ class Encounter:
     encounter_rate: int
     min_level: int
     max_level: int
+    conditions: dict[str, int] | None = None
 
     def to_dict(self) -> dict:
-        return {
+        d: dict = {
             "pokeapi_id": self.pokeapi_id,
             "pokemon_name": self.pokemon_name,
             "method": self.method,
-            "encounter_rate": self.encounter_rate,
             "min_level": self.min_level,
             "max_level": self.max_level,
         }
+        if self.conditions:
+            d["encounter_rate"] = None
+            d["conditions"] = self.conditions
+        else:
+            d["encounter_rate"] = self.encounter_rate
+        return d
 
 
 @dataclass

tools/import-pokedb/import_pokedb/processing.py

@@ -65,61 +65,75 @@ def parse_rate(value: str | None) -> int | None:
     return None
 
 
-def extract_encounter_rate(record: dict[str, Any], generation: int) -> int:
-    """Extract a single encounter_rate from a PokeDB encounter record.
+def extract_encounter_data(
+    record: dict[str, Any],
+    generation: int,
+) -> tuple[int, dict[str, int] | None]:
+    """Extract encounter rate and per-condition rates from a PokeDB record.
 
-    Flattens generation-specific rate variants into a single value.
+    Returns (rate, conditions) where:
+    - rate is the max/overall rate (used for sorting and backward compat)
+    - conditions is a dict of {condition_name: rate} or None for flat rates
     """
-    # Gen 1/3/6: rate_overall
+    # Gen 1/3/6: rate_overall — flat rate, no conditions
     rate_overall = parse_rate(record.get("rate_overall"))
     if rate_overall is not None:
-        return rate_overall
+        return rate_overall, None
 
-    # Gen 2/4: time-of-day rates — take the max
-    time_rates = [
-        parse_rate(record.get("rate_morning")),
-        parse_rate(record.get("rate_day")),
-        parse_rate(record.get("rate_night")),
-    ]
-    time_rates = [r for r in time_rates if r is not None]
-    if time_rates:
-        return max(time_rates)
+    # Gen 2/4/7: time-of-day rates
+    time_fields = {
+        "morning": parse_rate(record.get("rate_morning")),
+        "day": parse_rate(record.get("rate_day")),
+        "night": parse_rate(record.get("rate_night")),
+    }
+    time_conditions = {k: v for k, v in time_fields.items() if v is not None}
+    if time_conditions:
+        rate = max(time_conditions.values())
+        return rate, time_conditions
 
-    # Gen 5: seasonal rates — take the max
-    season_rates = [
-        parse_rate(record.get("rate_spring")),
-        parse_rate(record.get("rate_summer")),
-        parse_rate(record.get("rate_autumn")),
-        parse_rate(record.get("rate_winter")),
-    ]
-    season_rates = [r for r in season_rates if r is not None]
-    if season_rates:
-        return max(season_rates)
+    # Gen 5: seasonal rates
+    season_fields = {
+        "spring": parse_rate(record.get("rate_spring")),
+        "summer": parse_rate(record.get("rate_summer")),
+        "autumn": parse_rate(record.get("rate_autumn")),
+        "winter": parse_rate(record.get("rate_winter")),
+    }
+    season_conditions = {
+        k: v for k, v in season_fields.items() if v is not None
+    }
+    if season_conditions:
+        rate = max(season_conditions.values())
+        return rate, season_conditions
 
-    # Gen 8 Sw/Sh: weather rates — take the max
-    weather_rates = []
+    # Gen 8 Sw/Sh: weather rates
+    weather_conditions: dict[str, int] = {}
     for key, val in record.items():
         if key.startswith("weather_") and key.endswith("_rate") and val:
             parsed = parse_rate(val)
             if parsed is not None:
-                weather_rates.append(parsed)
-    if weather_rates:
-        return max(weather_rates)
+                # "weather_clear_rate" -> "clear"
+                condition_name = key[len("weather_"):-len("_rate")]
+                weather_conditions[condition_name] = parsed
+    if weather_conditions:
+        rate = max(weather_conditions.values())
+        return rate, weather_conditions
 
-    # Gen 8 Legends Arceus: boolean conditions → presence-based
-    if record.get("during_any_time") or record.get("during_morning") or \
-       record.get("during_day") or record.get("during_evening") or record.get("during_night"):
-        return 100  # Present under conditions
+    # Gen 8 Legends Arceus: boolean conditions — presence-based
+    if (
+        record.get("during_any_time")
+        or record.get("during_morning")
+        or record.get("during_day")
+        or record.get("during_evening")
+        or record.get("during_night")
+    ):
+        return 100, None
 
-    # Gen 9 Sc/Vi: probability weights → normalize
+    # Gen 9 Sc/Vi: probability weights — normalize
     prob_overall = record.get("probability_overall")
     if prob_overall:
         parsed = parse_rate(prob_overall)
         if parsed is not None:
-            # These are spawn weights (e.g. "20", "300"), not percentages.
-            # We'll normalize them later during aggregation when we have
-            # all encounters for a location. For now, store the raw weight.
-            return parsed
+            return parsed, None
 
     # Check time-based probability variants
     prob_rates = [
@@ -130,10 +144,10 @@ def extract_encounter_rate(record: dict[str, Any], generation: int) -> int:
     ]
     prob_rates = [r for r in prob_rates if r is not None]
     if prob_rates:
-        return max(prob_rates)
+        return max(prob_rates), None
 
     # Fallback: gift/trade/static encounters with no rate
-    return 100
+    return 100, None
 
 
 # ---------------------------------------------------------------------------
@@ -212,8 +226,8 @@ def process_encounters(
         # Parse levels
         min_level, max_level = parse_levels(record.get("levels"))
 
-        # Extract rate
-        encounter_rate = extract_encounter_rate(record, generation)
+        # Extract rate and conditions
+        encounter_rate, conditions = extract_encounter_data(record, generation)
 
         # Location area
         area_id = record.get("location_area_identifier", "")
@@ -227,6 +241,7 @@ def process_encounters(
             encounter_rate=encounter_rate,
             min_level=min_level,
             max_level=max_level,
+            conditions=conditions,
         )
 
         by_area.setdefault(area_id, []).append(enc)
@@ -234,10 +249,28 @@ def process_encounters(
     return by_area
 
 
+def _merge_conditions(
+    a: dict[str, int] | None,
+    b: dict[str, int] | None,
+) -> dict[str, int] | None:
+    """Merge two condition dicts by summing rates per key."""
+    if a is None and b is None:
+        return None
+    merged = dict(a or {})
+    for k, v in (b or {}).items():
+        merged[k] = merged.get(k, 0) + v
+    return merged
+
+
+def _cap_conditions(conditions: dict[str, int]) -> dict[str, int]:
+    """Cap each condition rate at 100."""
+    return {k: min(v, 100) for k, v in conditions.items()}
+
+
 def aggregate_encounters(encounters: list[Encounter]) -> list[Encounter]:
     """Aggregate encounters by (pokeapi_id, method), merging level ranges and summing rates.
 
-    Replicates the Go tool's aggregation logic.
+    Preserves per-condition rates through aggregation.
     """
     key_type = tuple[int, str]
     agg: dict[key_type, Encounter] = {}
@@ -250,8 +283,10 @@ def aggregate_encounters(encounters: list[Encounter]) -> list[Encounter]:
             existing.encounter_rate += enc.encounter_rate
             existing.min_level = min(existing.min_level, enc.min_level)
             existing.max_level = max(existing.max_level, enc.max_level)
+            existing.conditions = _merge_conditions(
+                existing.conditions, enc.conditions
+            )
         else:
-            # Copy so we don't mutate the original
             agg[k] = Encounter(
                 pokeapi_id=enc.pokeapi_id,
                 pokemon_name=enc.pokemon_name,
@@ -259,6 +294,7 @@ def aggregate_encounters(encounters: list[Encounter]) -> list[Encounter]:
                 encounter_rate=enc.encounter_rate,
                 min_level=enc.min_level,
                 max_level=enc.max_level,
+                conditions=dict(enc.conditions) if enc.conditions else None,
             )
             order.append(k)
 
@@ -266,6 +302,9 @@ def aggregate_encounters(encounters: list[Encounter]) -> list[Encounter]:
     for k in order:
         e = agg[k]
         e.encounter_rate = min(e.encounter_rate, 100)
+        if e.conditions:
+            e.conditions = _cap_conditions(e.conditions)
+            e.encounter_rate = max(e.conditions.values())
         result.append(e)
 
     # Sort by rate descending, then name ascending