Advent of Code 2022 - Day 11

from __future__ import annotations
from collections import deque
from dataclasses import dataclass
from functools import reduce
import operator
import re
from typing import Optional, Union


@dataclass
class WorryTest:
    div_by: int
    if_true: int
    if_false: int


@dataclass
class Operation:
    op: Union[operator.add, operator.mul]
    rhs: str | int


@dataclass
class Monkey:
    number: int
    items: deque[int]
    operation: Operation
    test: WorryTest
    inspected: int = 0

    def catch(self, item):
        self.items.append(item)


def inspect_item(item: int, o: Operation) -> int:
    if type(o.rhs) is str:
        return o.op(item, item)
    else:
        return o.op(item, o.rhs)


def throw_to_target(item: int, t: WorryTest) -> int:
    return t.if_true if (item % t.div_by) == 0 else t.if_false


def reduce_worry(
    item: int, op: Union[operator.floordiv, operator.mod], amount: int = 3
) -> int:
    return op(item, amount)


def parse_monkey(line: str) -> int:
    pattern = r"Monkey (\d):"
    result = re.search(pattern, line)
    return int(result.groups()[0])


def parse_starting_items(line: str) -> deque[int]:
    pattern = r"(\d+)"
    result = re.findall(pattern, line)
    return deque(list(map(int, result)))


def parse_operation(line: str) -> Operation:
    op: Optional[Union[operator.add, operator.mul]] = None
    instruction = line.split()[4:]
    match instruction[0]:
        case "+":
            op = operator.add
        case "*":
            op = operator.mul
    try:
        rhs = int(instruction[1])
    except ValueError:
        rhs = instruction[1]
    return Operation(op, rhs)


def parse_test(lines: list[str]) -> WorryTest:
    div_by = int(lines[0].strip().split()[-1])
    if_true = int(lines[1].strip()[-1])
    if_false = int(lines[2].strip()[-1])
    return WorryTest(div_by, if_true, if_false)


def parse_file(filepath) -> list[Monkey]:
    monkeys = []
    with open(filepath, "r") as f:
        for chunk in f.read().split("\n\n"):
            text = chunk.splitlines()
            monkey_n = parse_monkey(text[0])
            starting_items = parse_starting_items(text[1])
            operation = parse_operation(text[2])
            test = parse_test(text[3:])
            monkeys.append(Monkey(monkey_n, starting_items, operation, test))

    return monkeys


def play_round_one(monkeys: list[Monkey]):
    for i, monkey in enumerate(monkeys):
        while monkey.items:
            item = monkey.items.popleft()
            item = inspect_item(item, monkey.operation)
            monkey.inspected += 1
            item = reduce_worry(item, operator.floordiv, 3)
            target = throw_to_target(item, monkey.test)
            monkeys[target].catch(item)


def play_round_two(monkeys: list[Monkey], modulo: int):
    for i, monkey in enumerate(monkeys):
        while monkey.items:
            item = monkey.items.popleft()
            item = inspect_item(item, monkey.operation)
            monkey.inspected += 1
            item = reduce_worry(item, operator.mod, modulo)
            target = throw_to_target(item, monkey.test)
            monkeys[target].catch(item)


def part_one(filepath) -> int:
    monkeys = parse_file(filepath)
    for _ in range(20):
        play_round_one(monkeys)
    n_inspected = sorted([m.inspected for m in monkeys])
    return n_inspected[-1] * n_inspected[-2]


def part_two(filepath) -> int:
    monkeys = parse_file(filepath)
    modulo = reduce(operator.mul, [m.test.div_by for m in monkeys])
    for _ in range(10_000):
        play_round_two(monkeys, modulo)
    n_inspected = sorted([m.inspected for m in monkeys])
    return n_inspected[-1] * n_inspected[-2]