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

# Copyright 2024 Zoé Cassiopée Gauthier.
#
# Use of this source code is governed by an MIT-style
# license that can be found in the LICENSE file or at
# https://opensource.org/licenses/MIT.

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

import requests_cache

from pogo_scaled_estimators.typing import PokemonType

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 Move:
    move_id: str
    typing: PokemonType


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

    def scale(self, factor: float) -> "MovesetResult":
        return MovesetResult(self.fast_move, self.charged_move, self.estimator / factor)


class PokebattlerMove(TypedDict):
    moveId: str
    type: PokemonType


class PokebattlerPokemon(TypedDict):
    pokemonId: str
    type: PokemonType
    type2: NotRequired[PokemonType]
    quickMoves: list[str]
    cinematicMoves: list[str]
    movesets: list[dict[str, str]]


class PokebattlerRaid(TypedDict):
    pokemonId: str


class PokebattlerRaidTierInfo(TypedDict):
    guessTier: str


class PokebattlerRaidTier(TypedDict):
    tier: str
    info: PokebattlerRaidTierInfo
    raids: list[PokebattlerRaid]


@final
class PokebattlerProxy:
    def __init__(self, log_level="INFO"):
        self._cached_session = requests_cache.CachedSession("pokebatter_cache", cache_control=True, use_cache_dir=True)

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

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

    @cached_property
    def raids(self) -> list[PokebattlerRaidTier]:
        return self._cached_session.get(f"{BASE_URL}/raids").json()["tiers"]

    @cached_property
    def resists(self) -> dict[str, list[float]]:
        return self._cached_session.get(f"{BASE_URL}/resists").json()

    def simulate(self, raid: Raid) -> dict[str, list[MovesetResult]]:
        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[MovesetResult]] = {}
        response_json = response.json()
        for attacker in response_json["attackers"][0]["randomMove"]["defenders"]:
            results[attacker["pokemonId"]] = [
                MovesetResult(
                    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[PokemonType]) -> dict[str, list[str]]:
        raid_tiers: list[str] = []
        raid_bosses: dict[str, list[str]] = {}

        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):
            for boss in (raid["pokemonId"] for raid in tier["raids"]):
                if boss.endswith("_FORM"):
                    continue
                boss_pokemon: PokebattlerPokemon = next(filter(lambda mon: mon["pokemonId"] == boss, self.pokemon))
                boss_types: tuple[PokemonType, PokemonType] = (
                    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: PokemonType, defender_types: tuple[PokemonType, PokemonType]) -> 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[PokemonType]) -> list[str]:
        charged_moves: list[str] = [
            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 find_pokemon(self, name: str) -> PokebattlerPokemon:
        return next(filter(lambda mon: mon["pokemonId"] == name, self.pokemon))

    def pokemon_type(self, name: str) -> PokemonType:
        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"])