src/pogo_scaled_estimators/calculator.py

@@ -4,29 +4,18 @@
 # 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 MovesetResult, PokebattlerProxy, Raid
+from pogo_scaled_estimators.pokebattler_proxy import Moveset, 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], filters: Filter = Filter.NO_FILTER) -> None:
+    def __init__(self, attacker_types: list[str]) -> None:
         self.attacker_types = attacker_types
-        self.filters = filters
         self._pokebattler_proxy = PokebattlerProxy()
         self._progress = Progress()
 
@@ -35,7 +24,6 @@ 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:
@@ -56,7 +44,9 @@ class Calculator:
                     raid = Raid(raid_tier, defender, level, party)
                     results = {
                         attacker: [
-                            moveset for moveset in movesets if self._allowed_move(attacker, moveset.charged_move)
+                            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
@@ -74,64 +64,28 @@ class Calculator:
 
         ase_by_attacker: list[tuple[str, str, str, float]] = []
         for attacker, raid_tier_results in attackers.items():
-            fast_move, charged_move = self._best_moves(raid_tier_results)
+            if (
-            ase = self._ase(raid_tier_results)
+                not raid_tier_results["RAID_LEVEL_3"]
-            if not math.isfinite(ase):
+                or not raid_tier_results["RAID_LEVEL_5"]
+                or not raid_tier_results["RAID_LEVEL_MEGA"]
+            ):
                 continue
-            ase_by_attacker.append((attacker, fast_move, charged_move, self._ase(raid_tier_results)))
+            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.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.find_move(fast_move).typing
+            fast_move_type = self._pokebattler_proxy.move_type(fast_move)
-            charged_move_type = self._pokebattler_proxy.find_move(charged_move).typing
+            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 _allowed_attacker(self, pokemon_id: str) -> bool:
+    def _average_estimator(self, movesets: list[Moveset]) -> float:
-        if Filter.DISALLOW_MEGA_POKEMON in self.filters and ("_MEGA" in pokemon_id or "_PRIMAL" in pokemon_id):
+        return sum(moveset.estimator for moveset in movesets) / len(movesets)
-            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,11 +5,9 @@
 # https://opensource.org/licenses/MIT.
 
 import argparse
-import operator
 import sys
-from functools import reduce
 
-from pogo_scaled_estimators.calculator import Calculator, Filter
+import pogo_scaled_estimators.calculator
 
 
 def main_cli():
@@ -17,16 +15,9 @@ 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()
-    filters = reduce(operator.or_, args.filters, Filter.NO_FILTER)
+    calculator = pogo_scaled_estimators.calculator.Calculator(args.type)
-    calculator = Calculator(args.type, filters)
     calculator.calculate(level=args.level, party=args.party)
 
 

src/pogo_scaled_estimators/pokebattler_proxy.py

@@ -25,19 +25,13 @@ class Raid:
 
 
 @dataclass
-class Move:
+class Moveset:
-    move_id: str
-    typing: str
-
-
-@dataclass
-class MovesetResult:
     fast_move: str
     charged_move: str
     estimator: float
 
     def scale(self, factor: float):
-        return MovesetResult(self.fast_move, self.charged_move, self.estimator / factor)
+        return Moveset(self.fast_move, self.charged_move, self.estimator / factor)
 
 
 @final
@@ -68,7 +62,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[MovesetResult]]:
+    def simulate(self, raid: Raid) -> dict[str, list[Moveset]]:
         query_string = {
             "sort": "ESTIMATOR",
             "weatherCondition": "NO_WEATHER",
@@ -84,10 +78,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[MovesetResult]] = {}
+        results: dict[str, list[Moveset]] = {}
         for attacker in response.json()["attackers"][0]["randomMove"]["defenders"]:
             results[attacker["pokemonId"]] = [
-                MovesetResult(
+                Moveset(
                     attacker_moves["move1"], attacker_moves["move2"], cast(float, attacker_moves["result"]["estimator"])
                 )
                 for attacker_moves in attacker["byMove"]
@@ -168,12 +162,8 @@ class PokebattlerProxy:
             if any(moveset["cinematicMove"] in charged_moves for moveset in mon["movesets"])
         ]
 
-    def find_pokemon(self, name: str) -> dict:
-        return next(filter(lambda mon: mon["pokemonId"] == name, self.pokemon))
-
     def pokemon_type(self, name: str) -> str:
-        return self.find_pokemon(name)["type"]
+        return next(filter(lambda mon: mon["pokemonId"] == name, self.pokemon))["type"]
 
-    def find_move(self, move_id: str) -> Move:
+    def move_type(self, name: str) -> str:
-        move = next(filter(lambda move: "moveId" in move and move["moveId"] == move_id, self.moves))
+        return next(filter(lambda move: "moveId" in move and move["moveId"] == name, self.moves))["type"]
-        return Move(move_id, move["type"])