from tqdm import tqdm
import re
import math
from queue import Queue


with open("11.txt") as f:
  data = f.read()

ex_data = """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"""

re1 = re.compile(r"^Monkey (\d+):$")
re2 = re.compile(r"^\s*Starting items: (\d+(, \d+)*)$")
re3 = re.compile(r"^\s*Operation: new = (.*)$")
re4 = re.compile(r"^\s*Test: divisible by (\d+)$")
re5 = re.compile(r"^\s*If true: throw to monkey (\d+)$")
re6 = re.compile(r"^\s*If false: throw to monkey (\d+)$")
def parse_data(data):
  monkeys = dict()
  monkey_data = data.split("\n\n")
  for segment in monkey_data:
    lines = segment.splitlines()
    m = re1.match(lines[0])
    monkey_num = int(m.group(1))
    m = re2.match(lines[1])
    starting_items = list(map(int, m.group(1).split(", ")))
    m = re3.match(lines[2])
    expr = m.group(1)
    operation = eval(f"lambda old: {expr}")
    m = re4.match(lines[3])
    div_by = int(m.group(1))
    m = re5.match(lines[4])
    throw_true = int(m.group(1))
    m = re6.match(lines[5])
    throw_false = int(m.group(1))
    monkey = dict(starting_items = starting_items,
                  operation = operation,
                  div_by = div_by,
                  throw_true = throw_true,
                  throw_false = throw_false)
    monkeys[monkey_num] = monkey

  return monkeys

def solve1(data):
  data = parse_data(data)

  cur_items = []
  for monkey_num in range(len(data)):
    monkey_info = data[monkey_num]
    q = []
    for item in monkey_info["starting_items"]:
      q.append(item)
    cur_items.append(q)
  # print(cur_items)

  cts = [0] * len(data)
  
  for i in range(20):
    for monkey_num, monkey_items in enumerate(cur_items):
      for item in monkey_items:
        new_worry = data[monkey_num]["operation"](item)
        new_worry //= 3
        is_divisible = new_worry % data[monkey_num]["div_by"] == 0
        next_monkey = data[monkey_num]["throw_true" if is_divisible else "throw_false"]
        cur_items[next_monkey].append(new_worry)
      cts[monkey_num] += len(monkey_items)
      cur_items[monkey_num] = []

    # print(cur_items)

  print(math.prod(sorted(cts, reverse=True)[:2]))

def solve2(data):
  data = parse_data(data)

  cur_items = []
  modulus = 1
  for monkey_num in range(len(data)):
    monkey_info = data[monkey_num]
    modulus *= monkey_info["div_by"]
    q = []
    for item in monkey_info["starting_items"]:
      q.append(item)
    cur_items.append(q)
  # print(cur_items)

  # print(modulus)
  cts = [0] * len(data)
  
  for i in range(10000):
    for monkey_num, monkey_items in enumerate(cur_items):
      for item in monkey_items:
        new_worry = data[monkey_num]["operation"](item)
        new_worry %= modulus
        is_divisible = new_worry % data[monkey_num]["div_by"] == 0
        next_monkey = data[monkey_num]["throw_true" if is_divisible else "throw_false"]
        cur_items[next_monkey].append(new_worry)
      cts[monkey_num] += len(monkey_items)
      cur_items[monkey_num] = []

    # if i % 100 == 0: print(cur_items)
    # print(cur_items)

  print(math.prod(sorted(cts, reverse=True)[:2]))

solve1(ex_data)
solve1(data)

solve2(ex_data)
solve2(data)