from lib import str_to_ints, get_data


def addr(regs, a, b, c):
    regs[c] = regs[a] + regs[b]


def addi(regs, a, b, c):
    regs[c] = regs[a] + b


def mulr(regs, a, b, c):
    regs[c] = regs[a] * regs[b]


def muli(regs, a, b, c):
    regs[c] = regs[a] * b


def banr(regs, a, b, c):
    regs[c] = regs[a] & regs[b]


def bani(regs, a, b, c):
    regs[c] = regs[a] & b


def borr(regs, a, b, c):
    regs[c] = regs[a] | regs[b]


def bori(regs, a, b, c):
    regs[c] = regs[a] | b


def setr(regs, a, _, c):
    regs[c] = regs[a]


def seti(regs, a, _, c):
    regs[c] = a


def gtir(regs, a, b, c):
    regs[c] = 1 if a > regs[b] else 0


def gtri(regs, a, b, c):
    regs[c] = 1 if regs[a] > b else 0


def gtrr(regs, a, b, c):
    regs[c] = 1 if regs[a] > regs[b] else 0


def eqir(regs, a, b, c):
    regs[c] = 1 if a == regs[b] else 0


def eqri(regs, a, b, c):
    regs[c] = 1 if regs[a] == b else 0


def eqrr(regs, a, b, c):
    regs[c] = 1 if regs[a] == regs[b] else 0


OPS = [
    addr,
    addi,
    mulr,
    muli,
    banr,
    bani,
    borr,
    bori,
    setr,
    seti,
    gtir,
    gtri,
    gtrr,
    eqir,
    eqri,
    eqrr,
]


def part_1(data):
    examples = []

    lines = data.splitlines()
    test_prog = []
    last_i = None
    for i in range(len(lines)):
        line = lines[i]
        if line.startswith("Before:"):
            regs = str_to_ints(line)
            inst = str_to_ints(lines[i + 1])
            regs_after = str_to_ints(lines[i + 2])
            examples.append((regs, inst, regs_after))
            last_i = i + 2

    assert last_i is not None
    for line in lines[last_i:]:
        if line.strip() != "":
            test_prog.append(str_to_ints(line))

    r = 0
    for before_orig, inst, after in examples:
        ops_correct = []
        for op in OPS:
            before = list(before_orig)
            op(before, *inst[1:])
            if before == after:
                ops_correct.append(op)
        if len(ops_correct) >= 3:
            r += 1
    print(r)

    code_to_op = {}
    while len(OPS) > 0:
        for before_orig, inst, after in examples:
            ops_correct = []
            for op in OPS:
                before = list(before_orig)
                op(before, *inst[1:])
                if before == after:
                    ops_correct.append(op)
            if len(ops_correct) == 1:
                code_to_op[inst[0]] = ops_correct[0]
                OPS.remove(ops_correct[0])

    regs = [0, 0, 0, 0]
    for line in test_prog:
        op_code = line[0]
        vals = line[1:]
        code_to_op[op_code](regs, *vals)
    print(regs[0])


def main():
    data = get_data(__file__)
    part_1(data)


if __name__ == "__main__":
    main()