from functools import cache
from utils import open_day
from subprocess import run, PIPE

monkeys = {}
with open_day(21) as f:
    for line in f:
        target, expr = line.rstrip().split(': ')
        expr = expr.split(' ')
        if len(expr) == 1:
            expr = int(expr[0])
        monkeys[target] = expr

ops = {
    '+': lambda a, b: a + b,
    '-': lambda a, b: a - b,
    '*': lambda a, b: a * b,
    '/': lambda a, b: a // b,
}

@cache
def eval_monkey(m):
    expr = monkeys[m]
    if isinstance(expr, int):
        return expr
    return ops[expr[1]](eval_monkey(expr[0]), eval_monkey(expr[2]))

print(eval_monkey('root'))

filename = '21.z3'
with open(filename, 'w') as f:
    f.write(''.join(f'(declare-const {monkey} Int)\n' for monkey in monkeys.keys()))
    for k, v in monkeys.items():
        if k == 'humn': continue
        if k == 'root':
            print(f'(assert (= {v[0]} {v[2]}))', file=f)
            continue
        if isinstance(v, int):
            print(f'(assert (= {k} {v}))', file=f)
        else:
            print(f'(assert (= {k} ({v[1]} {v[0]} {v[2]})))', file=f)
    print('(check-sat)\n(get-value (humn))', file=f)
r = run(['z3', filename], stdout=PIPE)
print(r.stdout.splitlines()[1].decode().strip('()').split()[1])