Remove attacker movesets that don't match the requested type

src/pogo_scaled_estimators/calculator.py

@@ -9,7 +9,7 @@ from typing import final
 from rich.progress import Progress
 
 from pogo_scaled_estimators.pokebattler_proxy import Moveset, PokebattlerProxy, Raid
-from pogo_scaled_estimators.utilities import format_pokemon_name
+from pogo_scaled_estimators.utilities import format_move_name, format_pokemon_name
 
 
 @final
@@ -43,19 +43,22 @@ class Calculator:
                     self._progress.update(task, advance=1)
                     raid = Raid(raid_tier, defender, level, party)
                     results = {
-                        attacker: movesets
+                        attacker: [moveset for moveset in movesets if self._pokebattler_proxy.move_type(moveset.charged_move) in self.attacker_types]
                         for attacker, movesets in self._pokebattler_proxy.simulate(raid).items()
                         if attacker in attackers
                     }
+                    if not results:
+                        continue
                     best_movesets = {
                         attacker: min(movesets, key=lambda moveset: moveset.estimator)
                         for attacker, movesets in results.items()
+                        if movesets
                     }
                     best_estimator = min(best_movesets.values(), key=lambda moveset: moveset.estimator).estimator
                     for attacker, moveset in best_movesets.items():
                         attackers[attacker][simplified_raid_tier].append(moveset.scale(best_estimator))
 
-        ase_by_attacker: list[tuple[str, float]] = []
+        ase_by_attacker: list[tuple[str, str, str, float]] = []
         for attacker, raid_tier_results in attackers.items():
             if (
                 not raid_tier_results["RAID_LEVEL_3"]
@@ -63,16 +66,20 @@ class Calculator:
                 or not raid_tier_results["RAID_LEVEL_MEGA"]
             ):
                 continue
+            fast_move = raid_tier_results["RAID_LEVEL_5"][0].fast_move
+            charged_move = raid_tier_results["RAID_LEVEL_5"][0].charged_move
             ase = (
                 0.15 * self._average_estimator(raid_tier_results["RAID_LEVEL_3"])
                 + 0.50 * self._average_estimator(raid_tier_results["RAID_LEVEL_5"])
                 + 0.35 * self._average_estimator(raid_tier_results["RAID_LEVEL_MEGA"])
             )
-            ase_by_attacker.append((attacker, ase))
-        ase_by_attacker.sort(key=lambda item: item[1])
-        for attacker, ase in ase_by_attacker:
+            ase_by_attacker.append((attacker, fast_move, charged_move, ase))
+        ase_by_attacker.sort(key=lambda item: item[3])
+        for attacker, fast_move, charged_move, ase in ase_by_attacker:
             attacker_type = self._pokebattler_proxy.pokemon_type(attacker)
-            self._progress.console.print(f"{format_pokemon_name(attacker, attacker_type)}: {ase:.2f}")
+            fast_move_type = self._pokebattler_proxy.move_type(fast_move)
+            charged_move_type = self._pokebattler_proxy.move_type(charged_move)
+            self._progress.console.print(f"[bold]{format_pokemon_name(attacker, attacker_type)}[/bold] ({format_move_name(fast_move, fast_move_type)}/{format_move_name(charged_move, charged_move_type)}): {ase:.2f}")
 
     def _average_estimator(self, movesets: list[Moveset]) -> float:
         return sum(moveset.estimator for moveset in movesets) / len(movesets)

src/pogo_scaled_estimators/pokebattler_proxy.py

@@ -145,12 +145,15 @@ class PokebattlerProxy:
         return False
 
     def with_charged_moves(self, attacker_types: list[str]) -> list[str]:
-        charged_moves = [move["moveId"] for move in self.moves if "type" in move and move["type"] in attacker_types]
+        charged_moves = [move["moveId"] for move in self.moves if "moveId" in move and "type" in move and move["type"] in attacker_types]
         return [
             mon["pokemonId"]
             for mon in self.pokemon
             if any(moveset["cinematicMove"] in charged_moves for moveset in mon["movesets"])
         ]
 
-    def pokemon_type(self, name: str) -> str | None:
+    def pokemon_type(self, name: str) -> str:
         return next(filter(lambda mon: mon["pokemonId"] == name, self.pokemon))["type"]
+
+    def move_type(self, name: str) -> str:
+        return next(filter(lambda move: move["moveId"] == name, self.moves))["type"]

src/pogo_scaled_estimators/utilities.py

@@ -36,8 +36,11 @@ def format_pokemon_name(name: str, pokemon_type: str | None = None):
     return formatted_name
 
 
-def format_move_name(name):
+def format_move_name(name, move_type: str | None = None):
     parts = [part.capitalize() for part in name.split("_")]
     if parts[-1] == "Fast":
         parts = parts[:-1]
-    return " ".join(parts)
+    formatted_name = " ".join(parts)
+    if move_type:
+        return f"[{POKEMON_TYPE_COLORS[move_type]}]{formatted_name}[/]"
+    return formatted_name