Advent of Code - day 22 in Python 3

day22.py

#!/usr/bin/env python3

class Spell():
    def __init__(self):
        self.timer = self.duration
    def copy(self):
        a_copy = type(self)()
        a_copy.timer = self.timer
        return a_copy
    def apply_dispel(self, state):
        pass

class MagicMissile(Spell):
    cost = 53
    duration = 0
    @staticmethod
    def apply_effect(state):
        state.boss_hit_points -= 4
    
class Drain(Spell):
    cost = 73
    duration = 0
    @staticmethod
    def apply_effect(state):
        state.boss_hit_points -= 2
        state.player_hit_points += 2

class Shield(Spell):
    cost = 113
    duration = 6
    def apply_effect(self, state):
        if self.timer == self.duration:
            state.player_armor += 7
    def apply_dispel(self, state):
        state.player_armor -= 7

class Poison(Spell):
    cost = 173
    duration = 6
    @staticmethod
    def apply_effect(state):
        state.boss_hit_points -= 3

class Recharge(Spell):
    cost = 229
    duration = 5
    @staticmethod
    def apply_effect(state):
        state.player_mana += 101

SPELLS = [
    MagicMissile,
    Drain,
    Shield,
    Poison,
    Recharge,
]

IMMEDIATE_SPELLS = [
    spell
    for spell in SPELLS
    if spell.duration == 0
]

class State():
    def __init__(self, other_state=None):
        if other_state:
            self.player_turn = other_state.player_turn
            self.player_hit_points = other_state.player_hit_points
            self.player_mana = other_state.player_mana
            self.player_mana_spent = other_state.player_mana_spent
            self.player_armor = other_state.player_armor
            self.active_spells = []
            self.new_spell = None
            self.boss_hit_points = other_state.boss_hit_points
            self.boss_damage = other_state.boss_damage

    @staticmethod
    def start_state(boss_hit_points, boss_damage):
        state = State()
        state.player_turn = False # This is *before* the first turn
        state.player_hit_points = 50
        state.player_mana = 500
        state.player_mana_spent = 0
        state.player_armor = 0
        state.active_spells = []
        state.new_spell = None
        state.boss_hit_points = boss_hit_points
        state.boss_damage = boss_damage
        return state


def get_next_state(prev_state, chosen_spell=None):
    new_state = State(prev_state)
    new_state.player_turn = not prev_state.player_turn

    assert not new_state.player_turn or chosen_spell is not None
    assert new_state.player_turn or chosen_spell is None
    assert new_state.boss_hit_points > 0 and new_state.player_hit_points > 0

    # Rule for part 2: Lose 1 hit-point every turn
    if new_state.player_turn:
        new_state.player_hit_points -= 1
        if new_state.player_hit_points <= 0:
            # The new state has malformed fields but we don't care
            return new_state

    # Apply effects of active spells
    for spell in prev_state.active_spells:
        if spell.timer == 0:
            spell.apply_dispel(new_state)
            continue
        spell_copy = spell.copy()
        assert spell_copy.duration != spell_copy.timer
        spell_copy.apply_effect(new_state)
        spell_copy.timer -= 1
        if spell_copy.timer != 0 and chosen_spell == type(spell_copy):
            return None
        new_state.active_spells.append(spell_copy)

    # Apply effect of spell cast prev turn
    if prev_state.new_spell and prev_state.new_spell not in IMMEDIATE_SPELLS:
        prev_spell = prev_state.new_spell()
        prev_spell.apply_effect(new_state)
        prev_spell.timer -= 1
        if prev_spell.timer != 0:
            new_state.active_spells.append(prev_spell)

    # Check if anyone has died as a result of the spells
    if new_state.boss_hit_points < 0 or new_state.player_hit_points < 0:
        return new_state

    if new_state.player_turn:
        # Cast new spell
        spell_cost = chosen_spell.cost
        if prev_state.player_mana < spell_cost:
            return None
        new_state.player_mana -= spell_cost
        new_state.player_mana_spent += spell_cost
        new_state.new_spell = chosen_spell
        if chosen_spell in IMMEDIATE_SPELLS:
            chosen_spell.apply_effect(new_state)
    else:
        # Attack
        new_state.new_spell = None
        assert new_state.player_armor in (0, 7)
        hit_damage = max(new_state.boss_damage - new_state.player_armor, 1)
        new_state.player_hit_points -= hit_damage

    return new_state

def main():
    input_fname = 'day22_input.txt'
    with open(input_fname) as f:
        input_boss_hit_points = int(f.readline().split(': ')[1])
        input_boss_damage = int(f.readline().split(': ')[1])

    min_mana = None

    start_state = State.start_state(input_boss_hit_points, input_boss_damage)
    games_in_progress = [start_state]
    while games_in_progress:
        next_iteration = []
        for state in games_in_progress:
            assert state is not None
            if state.player_hit_points <= 0 or state.boss_hit_points <= 0:
                if state.player_hit_points > 0:
                    if min_mana is None or state.player_mana_spent < min_mana:
                        min_mana = state.player_mana_spent
                    print('Player won with {}. Min: {}'.format(state.player_mana_spent, min_mana))
            else:
                # Drop games that have exceeded our current minimum
                if min_mana is not None and state.player_mana_spent > min_mana:
                    continue

                if state.player_turn:
                    next_iteration.append(get_next_state(state))
                else:
                    for spell in SPELLS:
                        next_state = get_next_state(state, spell)
                        if next_state:
                            next_iteration.append(next_state)
        games_in_progress = next_iteration

    print('Min mana:', min_mana)

if __name__ == '__main__':
    main()