Allow filtering Mega, Shadow, and Legendary Pokémon

src/pogo_scaled_estimators/calculator.py

@@ -4,18 +4,29 @@
 # license that can be found in the LICENSE file or at
 # https://opensource.org/licenses/MIT.
 
+import math
+from enum import Flag, auto
 from typing import final
 
 from rich.progress import Progress
 
-from pogo_scaled_estimators.pokebattler_proxy import Moveset, PokebattlerProxy, Raid
+from pogo_scaled_estimators.pokebattler_proxy import MovesetResult, PokebattlerProxy, Raid
 from pogo_scaled_estimators.utilities import format_move_name, format_pokemon_name
 
 
+class Filter(Flag):
+    NO_FILTER = 0
+    DISALLOW_LEGACY_MOVES = auto()
+    DISALLOW_MEGA_POKEMON = auto()
+    DISALLOW_SHADOW_POKEMON = auto()
+    DISALLOW_LEGENDARY_POKEMON = auto()
+
+
 @final
 class Calculator:
-    def __init__(self, attacker_types: list[str]) -> None:
+    def __init__(self, attacker_types: list[str], filters: Filter = Filter.NO_FILTER) -> None:
         self.attacker_types = attacker_types
+        self.filters = filters
         self._pokebattler_proxy = PokebattlerProxy()
         self._progress = Progress()
 
@@ -24,6 +35,7 @@ class Calculator:
         attackers = {
             attacker: {"RAID_LEVEL_3": [], "RAID_LEVEL_5": [], "RAID_LEVEL_MEGA": []}
             for attacker in self._pokebattler_proxy.with_charged_moves(self.attacker_types)
+            if self._allowed_attacker(attacker)
         }
 
         with self._progress:
@@ -44,9 +56,7 @@ class Calculator:
                     raid = Raid(raid_tier, defender, level, party)
                     results = {
                         attacker: [
-                            moveset
-                            for moveset in movesets
-                            if self._pokebattler_proxy.move_type(moveset.charged_move) in self.attacker_types
+                            moveset for moveset in movesets if self._allowed_move(attacker, moveset.charged_move)
                         ]
                         for attacker, movesets in self._pokebattler_proxy.simulate(raid).items()
                         if attacker in attackers
@@ -64,28 +74,64 @@ class Calculator:
 
         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"]
-                or not raid_tier_results["RAID_LEVEL_5"]
-                or not raid_tier_results["RAID_LEVEL_MEGA"]
-            ):
+            fast_move, charged_move = self._best_moves(raid_tier_results)
+            ase = self._ase(raid_tier_results)
+            if not math.isfinite(ase):
                 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, fast_move, charged_move, ase))
+            ase_by_attacker.append((attacker, fast_move, charged_move, self._ase(raid_tier_results)))
         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)
-            fast_move_type = self._pokebattler_proxy.move_type(fast_move)
-            charged_move_type = self._pokebattler_proxy.move_type(charged_move)
+            fast_move_type = self._pokebattler_proxy.find_move(fast_move).typing
+            charged_move_type = self._pokebattler_proxy.find_move(charged_move).typing
             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)
+    def _allowed_attacker(self, pokemon_id: str) -> bool:
+        if Filter.DISALLOW_MEGA_POKEMON in self.filters and ("_MEGA" in pokemon_id or "_PRIMAL" in pokemon_id):
+            return False
+        if Filter.DISALLOW_SHADOW_POKEMON in self.filters and "_SHADOW_FORM" in pokemon_id:
+            return False
+        pokemon = self._pokebattler_proxy.find_pokemon(pokemon_id)
+        if Filter.DISALLOW_LEGENDARY_POKEMON in self.filters and pokemon.get("rarity") in [
+            "POKEMON_RARITY_LEGENDARY",
+            "POKEMON_RARITY_MYTHIC",
+            "POKEMON_RARITY_ULTRA_BEAST",
+        ]:
+            return False
+        return True
+
+    def _allowed_move(self, pokemon_id: str, move_id: str) -> bool:
+        move = self._pokebattler_proxy.find_move(move_id)
+        if move.typing not in self.attacker_types:
+            return False
+        if Filter.DISALLOW_LEGACY_MOVES in self.filters:
+            pokemon = self._pokebattler_proxy.find_pokemon(pokemon_id)
+            move_id = move_id.removesuffix("_PLUS_PLUS")
+            if move_id in pokemon.get("eliteQuickMove", []) or move_id in pokemon.get("eliteCinematicMove", []):
+                return False
+        return True
+
+    def _best_moves(self, moveset_results_by_tier: dict[str, list[MovesetResult]]) -> tuple[str, str]:
+        movesets: set[tuple[str, str]] = set()
+        for moveset_results in moveset_results_by_tier.values():
+            movesets.update((moveset.fast_move, moveset.charged_move) for moveset in moveset_results)
+        if not movesets:
+            return ("MOVE_NONE", "MOVE_NONE")
+        return min(movesets, key=lambda moveset: self._ase(moveset_results_by_tier, only=moveset))
+
+    def _ase(self, moveset_results_by_tier: dict[str, list[MovesetResult]], only=None) -> float:
+        try:
+            return (
+                0.15 * self._average_estimator(moveset_results_by_tier["RAID_LEVEL_3"], only)
+                + 0.50 * self._average_estimator(moveset_results_by_tier["RAID_LEVEL_5"], only)
+                + 0.35 * self._average_estimator(moveset_results_by_tier["RAID_LEVEL_MEGA"], only)
+            )
+        except ZeroDivisionError:
+            return float("inf")
+
+    def _average_estimator(self, moveset_results: list[MovesetResult], only: tuple[str, str] | None = None) -> float:
+        if only:
+            moveset_results = [m for m in moveset_results if m.fast_move == only[0] and m.charged_move == only[1]]
+        return sum(moveset.estimator for moveset in moveset_results) / len(moveset_results)

src/pogo_scaled_estimators/cli.py

@@ -5,9 +5,11 @@
 # https://opensource.org/licenses/MIT.
 
 import argparse
+import operator
 import sys
+from functools import reduce
 
-import pogo_scaled_estimators.calculator
+from pogo_scaled_estimators.calculator import Calculator, Filter
 
 
 def main_cli():
@@ -15,9 +17,16 @@ def main_cli():
     _ = parser.add_argument("type", nargs="+", help="an attacker type")
     _ = parser.add_argument("--level", type=int, default=40)
     _ = parser.add_argument("--party", type=int, default=1)
+    _ = parser.add_argument("--no-legacy", dest="filters", action="append_const", const=Filter.DISALLOW_LEGACY_MOVES)
+    _ = parser.add_argument("--no-mega", dest="filters", action="append_const", const=Filter.DISALLOW_MEGA_POKEMON)
+    _ = parser.add_argument("--no-shadow", dest="filters", action="append_const", const=Filter.DISALLOW_SHADOW_POKEMON)
+    _ = parser.add_argument(
+        "--no-legendary", dest="filters", action="append_const", const=Filter.DISALLOW_LEGENDARY_POKEMON
+    )
 
     args = parser.parse_args()
-    calculator = pogo_scaled_estimators.calculator.Calculator(args.type)
+    filters = reduce(operator.or_, args.filters, Filter.NO_FILTER)
+    calculator = Calculator(args.type, filters)
     calculator.calculate(level=args.level, party=args.party)
 
 

src/pogo_scaled_estimators/pokebattler_proxy.py

@@ -25,13 +25,19 @@ class Raid:
 
 
 @dataclass
-class Moveset:
+class Move:
+    move_id: str
+    typing: str
+
+
+@dataclass
+class MovesetResult:
     fast_move: str
     charged_move: str
     estimator: float
 
     def scale(self, factor: float):
-        return Moveset(self.fast_move, self.charged_move, self.estimator / factor)
+        return MovesetResult(self.fast_move, self.charged_move, self.estimator / factor)
 
 
 @final
@@ -62,7 +68,7 @@ class PokebattlerProxy:
     def resists(self) -> dict:
         return self.cached_session.get(f"{BASE_URL}/resists").json()
 
-    def simulate(self, raid: Raid) -> dict[str, list[Moveset]]:
+    def simulate(self, raid: Raid) -> dict[str, list[MovesetResult]]:
         query_string = {
             "sort": "ESTIMATOR",
             "weatherCondition": "NO_WEATHER",
@@ -78,10 +84,10 @@ class PokebattlerProxy:
             query_string["friendshipLevel"] = "FRIENDSHIP_LEVEL_4"
         url = f"{BASE_URL}/raids/defenders/{raid.defender}/levels/{raid.tier}/attackers/levels/{raid.level}/strategies/CINEMATIC_ATTACK_WHEN_POSSIBLE/DEFENSE_RANDOM_MC?{urllib.parse.urlencode(query_string, doseq=True)}"
         response = self._cached_session.get(url)
-        results: dict[str, list[Moveset]] = {}
+        results: dict[str, list[MovesetResult]] = {}
         for attacker in response.json()["attackers"][0]["randomMove"]["defenders"]:
             results[attacker["pokemonId"]] = [
-                Moveset(
+                MovesetResult(
                     attacker_moves["move1"], attacker_moves["move2"], cast(float, attacker_moves["result"]["estimator"])
                 )
                 for attacker_moves in attacker["byMove"]
@@ -162,8 +168,12 @@ class PokebattlerProxy:
             if any(moveset["cinematicMove"] in charged_moves for moveset in mon["movesets"])
         ]
 
-    def pokemon_type(self, name: str) -> str:
-        return next(filter(lambda mon: mon["pokemonId"] == name, self.pokemon))["type"]
+    def find_pokemon(self, name: str) -> dict:
+        return next(filter(lambda mon: mon["pokemonId"] == name, self.pokemon))
 
-    def move_type(self, name: str) -> str:
-        return next(filter(lambda move: "moveId" in move and move["moveId"] == name, self.moves))["type"]
+    def pokemon_type(self, name: str) -> str:
+        return self.find_pokemon(name)["type"]
+
+    def find_move(self, move_id: str) -> Move:
+        move = next(filter(lambda move: "moveId" in move and move["moveId"] == move_id, self.moves))
+        return Move(move_id, move["type"])