import re
from string import ascii_lowercase, ascii_uppercase

def extract_single_digit(line: str) -> int:
    r = re.compile(r"\d+")
    for m in r.findall(line):
        return int(m)
    raise Exception("No single digit sequence in '{line}'")

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

EXAMPLE = """
Monkey 0:
  Starting items: 79, 98
  Operation: new = old * 19
  Test: divisible by 23
    If true: throw to monkey 2
    If false: throw to monkey 3

Monkey 1:
  Starting items: 54, 65, 75, 74
  Operation: new = old + 6
  Test: divisible by 19
    If true: throw to monkey 2
    If false: throw to monkey 0

Monkey 2:
  Starting items: 79, 60, 97
  Operation: new = old * old
  Test: divisible by 13
    If true: throw to monkey 1
    If false: throw to monkey 3

Monkey 3:
  Starting items: 74
  Operation: new = old + 3
  Test: divisible by 17
    If true: throw to monkey 0
    If false: throw to monkey 1
"""

def clean(text: str) -> list[str]:
    return list(filter(lambda l: l.strip() != "", text.splitlines()))

def get_monkesy(lines):
    monkeys = []
    for i, line in enumerate(lines):
        if line.startswith("Monkey"):
            d = extract_single_digit(line)
            monkeys.append([d])
        elif line.strip().startswith("Starting"):
            d = extract_digits_list(line)
            monkeys[-1].append(d)
        elif line.strip().startswith("Operation"):
            if "old * old" in line:
                f = lambda y: y * y
                monkeys[-1].append(f)
            elif "*" in line:
                d = extract_single_digit(line)
                f = lambda x: (lambda y: x * y)
                f = f(d)
                monkeys[-1].append(f)
            elif "+" in line:
                d = extract_single_digit(line)
                f = lambda x: (lambda y: x + y)
                f = f(d)
                monkeys[-1].append(f)
            else:
                raise Exception(line)
        elif line.strip().startswith("Test"):
            d = extract_single_digit(line)
            monkeys[-1].append(d)
        elif line.strip().startswith("If true"):
            d = extract_single_digit(line)
            monkeys[-1].append(d)
        elif line.strip().startswith("If false"):
            d = extract_single_digit(line)
            monkeys[-1].append(d)
        else:
            print(line)
            raise Exception("Unexpected line.")
    for m in monkeys:
        m.append(0)
    return monkeys

def solve(lines: list[str]):
    monkeys = get_monkesy(lines)
    for _ in range(20):
        for monkey in monkeys:
            id, items, op, test, if_true, if_false, count = monkey
            while items:
                item = items.pop()
                monkey[6] += 1
                worry = item
                worry = op(worry)
                worry //= 3
                if worry % test == 0:
                    monkeys[if_true][1].append(worry)
                else:
                    monkeys[if_false][1].append(worry)
                monkeys[id][1] = monkeys[id][1][1:]
    counts = sorted([m[6] for m in monkeys])[-2:]
    return counts[0] * counts[1]

def solve2(lines: list[str]):
    monkeys = get_monkesy(lines)
    mod = 1
    for m in monkeys:
        mod *= m[3]

    for _ in range(10000):
        for monkey in monkeys:
            id, items, op, test, if_true, if_false, count = monkey
            for item in list(items):
                monkey[6] += 1
                worry = item
                worry = op(worry)
                worry %= mod
                if worry % test == 0:
                    monkeys[if_true][1].append(worry)
                else:
                    monkeys[if_false][1].append(worry)
                monkeys[id][1] = monkeys[id][1][1:]
    counts = sorted([m[6] for m in monkeys])[-2:]
    return counts[0] * counts[1]

def main():
    example = clean(EXAMPLE)
    print("Example  1:", solve(example))

    data = clean(open("i11.txt").read())
    print("Solution 1:", solve(data))

    example = clean(EXAMPLE)
    print("Example  2:", solve2(example))

    data = clean(open("i11.txt").read())
    print("Solution 2:", solve2(data))

if __name__ == "__main__":
    main()