import shared
import operator
import math

from dataclasses import dataclass
from typing import Callable


@dataclass
class Monkey:
    number: int
    items: list
    div: int
    t: int
    f: int
    by: int
    op: operator
    inspect_count: int = 0


def load_monkeys(lines):
    monkeys = []
    count = 0
    for line in lines:
        if line.lstrip().startswith("Starting"):
            items = list(map(int, line.split(": ")[-1].split(",")))
        elif line.lstrip().startswith("Operation"):
            op = line.split(": ")[-1]
            try:
                by = int(line.split(" ")[-1])
            except ValueError:
                by = None  # know None means by self
            if "*" in op:
                operand = operator.mul
            elif "+" in op:
                operand = operator.add
        elif line.lstrip().startswith("Test"):
            div = int(line.split(" ")[-1])
        elif line.lstrip().startswith("If true"):
            true = int(line.split(" ")[-1])
        elif line.lstrip().startswith("If false"):
            false = int(line.split(" ")[-1])
            monkey = Monkey(
                number=count, items=items, op=operand, by=by, div=div, t=true, f=false
            )
            monkeys.append(monkey)
            count += 1
    lcm = math.lcm(*[m.div for m in monkeys])
    return monkeys, lcm


def part1(lines):
    monkeys, _ = load_monkeys(lines)

    for current_round in range(20):
        for monkey in monkeys:
            for item in monkey.items:
                monkey.inspect_count += 1
                if monkey.by is None:
                    # Self
                    item = monkey.op(item, item)
                else:
                    item = monkey.op(item, monkey.by)
                item = item // 3
                if item % monkey.div == 0:
                    to = monkey.t
                else:
                    to = monkey.f
                monkeys[to].items.append(item)
            monkey.items = []

    counts = list(sorted([monkey.inspect_count for monkey in monkeys]))
    total = counts[-1] * counts[-2]
    print(total)


def part2(lines):
    monkeys, lcm = load_monkeys(lines)

    for current_round in range(10000):
        for monkey in monkeys:
            for item in monkey.items:
                monkey.inspect_count += 1
                if monkey.by is None:
                    # Self
                    item = monkey.op(item, item)
                else:
                    item = monkey.op(item, monkey.by)
                item = item % lcm
                if item % monkey.div == 0:
                    to = monkey.t
                else:
                    to = monkey.f
                monkeys[to].items.append(item)
            monkey.items = []

    counts = list(sorted([monkey.inspect_count for monkey in monkeys]))
    total = counts[-1] * counts[-2]
    print(total)


def main():
    rows = [row for row in shared.load(11)]
    part1(rows)
    part2(rows)


if __name__ == "__main__":
    main()