from copy import deepcopy
from dataclasses import dataclass
from functools import partial, reduce
from operator import mul as mul_op
from typing import Callable

def identity(o):
    return o
def value(v, _):
    return v
def add(a, b, o):
    return a(o) + b(o)
def mul(a, b, o):
    return a(o) * b(o)

@dataclass
class Monkey:
    items: list[int]
    op: Callable[[int], int]
    test_div: int
    target: tuple[int, int]
    inspections: int = 0
    def from_str(s):
        items, op, test_div, true_tgt, false_tgt = [l.split(': ')[1] for l in s.split('\n')[1:]]
        items = [int(i) for i in items.split(', ')]
        op = op.split(' ')[2:]
        lhs = identity if op[0] == 'old' else partial(value, int(op[0]))
        rhs = identity if op[2] == 'old' else partial(value, int(op[2]))
        op = partial(add, lhs, rhs) if op[1] == '+' else partial(mul, lhs, rhs)
        test_div = int(test_div.split(' ')[-1])
        true_tgt = int(true_tgt.split(' ')[-1])
        false_tgt = int(false_tgt.split(' ')[-1])
        return Monkey(items, op, test_div, (true_tgt, false_tgt))

with open('11.in') as f:
    monkeys = [Monkey.from_str(m) for m in f.read().rstrip().split('\n\n')]


def solve(monkeys, rounds, do_div):
    monkeys = deepcopy(monkeys)
    modulo = reduce(mul_op, (m.test_div for m in monkeys), 1)
    for _ in range(rounds):
        for m in monkeys:
            while m.items:
                m.inspections += 1
                i = m.items.pop(0)
                i = m.op(i)
                if do_div: i //= 3
                target = m.target[min(i % m.test_div, 1)]
                monkeys[target].items.append(i % modulo)
    return reduce(mul_op, sorted(m.inspections for m in monkeys)[-2:], 1)

print(solve(monkeys, 20, True))
print(solve(monkeys, 10000, False))