#!/usr/bin/env python3

# import numpy as np
from functools import reduce
from re import findall
from copy import deepcopy
import sys

# filename = "in/day11.ref"
filename = "in/day11.pzl"
data = open(filename).read()

num_mon = len([None for _ in data.split('\n\n')])

mons = list()
itmss = list()
for m in data.split('\n\n'):
    lines = [line for line in m.rstrip('\n').split('\n')]
    mon = dict()
    # mon['n'] = int(lines[0][-2])
    # mon['items'] = [int(i.strip()) for i in lines[1].split(':')[-1].split(',')]
    itmss.append([int(i.strip()) for i in lines[1].split(':')[-1].split(',')])
    mon['op'] = lines[2].split('=')[-1].strip()
    mon['test'] = int(lines[3].split(' ')[-1])
    mon['true'] = int(lines[4].split(' ')[-1])
    mon['false'] = int(lines[5].split(' ')[-1])
    mons.append(mon)

# for mon,itms in zip(mons,itmss):
#     print(mon)
#     print(itms)

def get_next_one_ape(n, old_timss, sums, div=0, mod=0):
    new_itmss = [list() for _ in old_timss]

    for ape,itms in enumerate(old_timss):
        for wl in itms:
            if n == ape:
                old = wl
                wl = eval(mons[ape]['op'])
                if div != 0:
                    wl = int(wl/div)
                if mod != 0:
                    wl = wl % mod
                sums[ape] += 1
                if wl % mons[ape]['test'] == 0:
                    new_itmss[mons[ape]['true']].append(wl)
                else:
                    new_itmss[mons[ape]['false']].append(wl)
            else:
                new_itmss[ape].append(wl)
    return new_itmss, sums


def get_next(itmss, sums, div=0, mod=0):
    for ape,itms in enumerate(itmss):
        itmss, sums = get_next_one_ape(ape, itmss, sums, div=div, mod=mod)
    return itmss, sums


res1 = 0
div1 = 3
sums1 = [0 for _ in itmss]
itmss1 = itmss
for i in range(20):
    itmss1, sums1 = get_next(itmss1, sums1, div=div1)
res1 = reduce(lambda x,y:x*y, sorted(sums1, reverse=True)[:2])


res2 = 0
mod2 = reduce(lambda x,y:x*y, map(lambda a:a['test'],mons))
sums2 = [0 for _ in itmss]
itmss2 = itmss
for i in range(1, 10001):
    itmss2, sums2 = get_next(itmss2, sums2, mod=mod2)
    # if i in [1, 20, 100, 200, 300, 400, 500, 1000, 2000, 10000]:
    #     print(i, sums2)
res2 = reduce(lambda x,y:x*y, sorted(sums2, reverse=True)[:2])

print('res1:', res1)
print('res2:', res2)