from collections import defaultdict
from math import lcm
from tools.aoc import AOCDay
from typing import Any


class Monkey:
    def __init__(self, monkey_id: int):
        self.monkey_id = monkey_id
        self.items = []
        self.operation = None
        self.test = None
        self.action_true = None
        self.action_false = None

    def turn(self, part2: bool = False):
        for _ in range(len(self.items)):
            item = self.items.pop(0)
            op, right = self.operation

            if op == "+":
                item += int(right)
            elif right == "old":
                item *= item
            else:
                item *= int(right)

            if not part2:
                item //= 3

            if item % self.test == 0:
                yield item, self.action_true
            else:
                yield item, self.action_false


class Day(AOCDay):
    inputs = [
        [
            (10605, "input11_test"),
            (56595, "input11_dennis"),
            (98280, "input11"),
        ],
        [
            (2713310158, "input11_test"),
            (15693274740, "input11_dennis"),
            (17673687232, "input11"),
        ]
    ]

    def get_monkeys(self) -> list:
        monkeys = {}
        current_monkey = None
        for line in self.getInput():
            if not line:
                continue

            if line.startswith("Monkey"):
                current_monkey = Monkey(int(line[:-1].split(" ")[1]))
                monkeys[current_monkey.monkey_id] = current_monkey
            elif line.startswith("  Starting items:"):
                current_monkey.items = list(map(int, line.split(": ")[1].split(", ")))
            elif line.startswith("  Operation"):
                current_monkey.operation = line.split(": ")[1].split(" = ")[1].split(" ")[1:]
            elif line.startswith("  Test:"):
                current_monkey.test = int(line.split(" ")[-1])
            elif line.startswith("    If true:"):
                current_monkey.action_true = int(line.split(" ")[-1])
            elif line.startswith("    If false:"):
                current_monkey.action_false = int(line.split(" ")[-1])
            else:
                print("Don't know how to parse '", line, "'")

        return [monkeys[x] for x in sorted(monkeys)]

    def get_monkey_business(self, amount: int = 20, part2: bool = False) -> int:
        monkeys = self.get_monkeys()
        monkey_inspects = defaultdict(int)
        modder = lcm(*[m.test for m in monkeys])

        for _ in range(amount):
            for monkey in monkeys:
                for item, target in monkey.turn(part2):
                    monkey_inspects[monkey.monkey_id] += 1
                    if part2:
                        monkeys[target].items.append(item % modder)
                    else:
                        monkeys[target].items.append(item)

        f = sorted(monkey_inspects.values())[-2:]
        return f[0] * f[1]

    def part1(self) -> Any:
        return self.get_monkey_business(20, False)

    def part2(self) -> Any:
        return self.get_monkey_business(10_000, True)


if __name__ == '__main__':
    day = Day(2022, 11)
    day.run(verbose=True)