pogo-scaled-estimators/src/pogo_scaled_estimators/pokebattler_proxy.py

import urllib.parse
from dataclasses import dataclass
from functools import cached_property
from typing import cast, final

import requests_cache

BASE_URL = "https://fight.pokebattler.com"
WEAKNESS = 1.6
DOUBLE_WEAKNESS = WEAKNESS * WEAKNESS


@dataclass(frozen=True)
class Raid:
    tier: str
    defender: str
    level: int = 40
    party: int = 1


@dataclass
class Moveset:
    fast_move: str
    charged_move: str
    estimator: float

    def scale(self, factor: float):
        return Moveset(self.fast_move, self.charged_move, self.estimator / factor)


@final
class PokebattlerProxy:
    def __init__(self):
        self._cached_session = requests_cache.CachedSession("pokebatter_cache", cache_control=True)
        self._pokemon: dict | None = None
        self._raids: dict | None = None
        self._resists: dict | None = None

    @property
    def cached_session(self):
        return self._cached_session

    @cached_property
    def moves(self) -> dict:
        return self.cached_session.get(f"{BASE_URL}/moves").json()["move"]

    @cached_property
    def pokemon(self) -> dict:
        return self.cached_session.get(f"{BASE_URL}/pokemon").json()["pokemon"]

    @cached_property
    def raids(self) -> dict:
        return self.cached_session.get(f"{BASE_URL}/raids").json()

    @cached_property
    def resists(self) -> dict:
        return self.cached_session.get(f"{BASE_URL}/resists").json()

    def simulate(self, raid: Raid) -> dict[str, list[Moveset]]:
        query_string = {
            "sort": "ESTIMATOR",
            "weatherCondition": "NO_WEATHER",
            "dodgeStrategy": "DODGE_REACTION_TIME",
            "aggregation": "AVERAGE",
            "includeLegendary": "true",
            "includeShadow": "true",
            "includeMegas": "true",
            "primalAssistants": "",
            "numParty": str(raid.party),
        }
        if raid.tier != "RAID_LEVEL_3":
            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]] = {}
        for attacker in response.json()["attackers"][0]["randomMove"]["defenders"]:
            results[attacker["pokemonId"]] = [
                Moveset(
                    attacker_moves["move1"], attacker_moves["move2"], cast(float, attacker_moves["result"]["estimator"])
                )
                for attacker_moves in attacker["byMove"]
            ]
        return results

    def raid_bosses(self, attacker_types: list[str]) -> dict:
        raid_tiers = []
        raid_bosses = {}

        for raid_level in ["3", "5", "MEGA", "MEGA_5", "ULTRA_BEAST"]:
            tier = f"RAID_LEVEL_{raid_level}"
            raid_tiers.extend(
                [
                    tier,
                    f"{tier}_LEGACY",
                    f"{tier}_FUTURE",
                ]
            )
            raid_bosses[tier] = []

        for tier in filter(lambda tier: tier["tier"] in raid_tiers, self.raids["tiers"]):
            for boss in (raid["pokemon"] for raid in tier["raids"]):
                if boss.endswith("_FORM"):
                    continue
                boss_pokemon = next(filter(lambda mon: mon["pokemonId"] == boss, self.pokemon))
                if ("candyToEvolve" in boss_pokemon or boss in ["SEADRA", "SEALEO"]) and boss not in [
                    "KELDEO",
                    "LUMINEON",
                    "MANAPHY",
                    "PHIONE",
                    "STUNFISK",
                    "TERRAKION",
                ]:
                    continue
                boss_types = (
                    boss_pokemon["type"],
                    boss_pokemon.get("type2", "POKEMON_TYPE_NONE"),
                )
                if any(self._is_weak(attacker_type, boss_types) for attacker_type in attacker_types):
                    raid_bosses[tier["info"]["guessTier"]].append(boss)

        return raid_bosses

    def _is_weak(self, attacker_type: str, defender_types: tuple[str, str]) -> bool:
        pokemon_types = list(self.resists.keys())
        defender_type_indices = (
            pokemon_types.index(defender_types[0]),
            pokemon_types.index(defender_types[1]),
        )
        attack_resist = (
            self.resists[attacker_type][defender_type_indices[0]]
            * self.resists[attacker_type][defender_type_indices[1]]
        )

        # Check for double weakness.
        if attack_resist >= DOUBLE_WEAKNESS:
            return True

        # Checkout for single weakness if not double weak to anything else.
        any_double_weaknesses = any(
            r[defender_type_indices[0]] * r[defender_type_indices[1]] >= DOUBLE_WEAKNESS for r in self.resists.values()
        )
        if not any_double_weaknesses and attack_resist >= WEAKNESS:
            return True

        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]
        return [
            mon["pokemonId"]
            for mon in self.pokemon
            if any(moveset["cinematicMove"] in charged_moves for moveset in mon["movesets"])
        ]