advent-of-code/src/holt59/aoc/2015/day22.py

from __future__ import annotations

import heapq
from typing import Any, Iterator, Literal, TypeAlias, cast

from ..base import BaseSolver

PlayerType: TypeAlias = Literal["player", "boss"]
SpellType: TypeAlias = Literal["magic missile", "drain", "shield", "poison", "recharge"]
BuffType: TypeAlias = Literal["shield", "poison", "recharge"]
Node: TypeAlias = tuple[
    PlayerType,
    int,
    int,
    int,
    int,
    int,
    tuple[tuple[BuffType, int], ...],
    tuple[tuple[SpellType, int], ...],
]

ATTACK_SPELLS: list[tuple[SpellType, int, int, int]] = [
    ("magic missile", 53, 4, 0),
    ("drain", 73, 2, 2),
]
BUFF_SPELLS: list[tuple[BuffType, int, int]] = [
    ("shield", 113, 6),
    ("poison", 173, 6),
    ("recharge", 229, 5),
]


def play(
    player_hp: int,
    player_mana: int,
    player_armor: int,
    boss_hp: int,
    boss_attack: int,
    hard_mode: bool,
) -> tuple[tuple[SpellType, int], ...]:
    winning_node: tuple[tuple[SpellType, int], ...] | None = None

    visited: set[
        tuple[PlayerType, int, int, int, int, tuple[tuple[BuffType, int], ...]]
    ] = set()
    nodes: list[Node] = [
        ("player", 0, player_hp, player_mana, player_armor, boss_hp, (), ())
    ]

    while winning_node is None:
        (
            player,
            mana,
            player_hp,
            player_mana,
            player_armor,
            boss_hp,
            buffs,
            spells,
        ) = heapq.heappop(nodes)

        if (player, player_hp, player_mana, player_armor, boss_hp, buffs) in visited:
            continue

        visited.add((player, player_hp, player_mana, player_armor, boss_hp, buffs))
        new_buffs: list[tuple[BuffType, int]] = []
        for buff, length in buffs:
            length = length - 1
            match buff:
                case "poison":
                    boss_hp = max(boss_hp - 3, 0)
                case "shield":
                    if length == 0:
                        player_armor -= 7
                case "recharge":
                    player_mana += 101

            if length > 0:
                new_buffs.append((buff, length))

        if hard_mode and player == "player":
            player_hp = player_hp - 1

        if player_hp <= 0:
            continue

        if boss_hp <= 0:
            winning_node = spells
            continue

        buffs = tuple(new_buffs)

        if player == "boss":
            heapq.heappush(
                nodes,
                (
                    "player",
                    mana,
                    max(0, player_hp - max(boss_attack - player_armor, 1)),
                    player_mana,
                    player_armor,
                    boss_hp,
                    buffs,
                    spells,
                ),
            )
        else:
            buff_types = {b for b, _ in buffs}

            for spell, cost, damage, regeneration in ATTACK_SPELLS:
                if player_mana < cost:
                    continue

                heapq.heappush(
                    nodes,
                    (
                        "boss",
                        mana + cost,
                        player_hp + regeneration,
                        player_mana - cost,
                        player_armor,
                        max(0, boss_hp - damage),
                        buffs,
                        spells + cast("tuple[tuple[SpellType, int]]", ((spell, cost),)),
                    ),
                )

            for buff_type, buff_cost, buff_length in BUFF_SPELLS:
                if buff_type in buff_types:
                    continue

                if player_mana < buff_cost:
                    continue

                heapq.heappush(
                    nodes,
                    (
                        "boss",
                        mana + buff_cost,
                        player_hp,
                        player_mana - buff_cost,
                        player_armor + 7 * (buff_type == "shield"),
                        boss_hp,
                        buffs
                        + cast(
                            "tuple[tuple[BuffType, int]]", ((buff_type, buff_length),)
                        ),
                        spells
                        + cast(
                            "tuple[tuple[SpellType, int]]", ((buff_type, buff_cost),)
                        ),
                    ),
                )

    return winning_node


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any]:
        lines = input.splitlines()

        player_hp = 50
        player_mana = 500
        player_armor = 0

        boss_hp = int(lines[0].split(":")[1].strip())
        boss_attack = int(lines[1].split(":")[1].strip())

        yield sum(
            c
            for _, c in play(
                player_hp, player_mana, player_armor, boss_hp, boss_attack, False
            )
        )

        # 1242 (not working)
        yield sum(
            c
            for _, c in play(
                player_hp, player_mana, player_armor, boss_hp, boss_attack, True
            )
        )