from functools import cache
from operator import add, sub, mul, ifloordiv
from subprocess import run, PIPE
from utils import open_day

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 = { '+': add, '-': sub, '*': mul, '/': ifloordiv }

@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'))

@cache
def to_sexp(m):
    if m == 'humn': return m
    expr = monkeys[m]
    if isinstance(expr, int):
        return expr
    a, b = to_sexp(expr[0]), to_sexp(expr[2])
    if isinstance(a, int) and isinstance(b, int):
        return ops[expr[1]](a, b)
    return f'({expr[1]} {to_sexp(expr[0])} {to_sexp(expr[2])})'

filename = '21.z3'
with open(filename, 'w') as f:
    print(f'(declare-const humn Int)', file=f)
    root = monkeys['root']
    print(f'(assert (= {to_sexp(root[0])} {to_sexp(root[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])