from lib import get_data, str_to_ints
import d16


def print_regs(regs):
    print(" ".join([f"regs[{i}]={regs[i]:10}" for i in range(len(regs))]) + " (ip)")


def print_inst(i, insts, ip):
    op, arg1, arg2, dest = insts[i]
    dest = f"regs[{dest}]" if dest != ip else "ip"
    match op:
        case "addi":
            print(f"{i:2} {dest} = regs[{arg1}] + {arg2}")
        case "addr":
            print(f"{i:2} {dest} = regs[{arg1}] + regs[{arg2}]")
        case "muli":
            print(f"{i:2} {dest} = regs[{arg1}] * {arg2}")
        case "mulr":
            print(f"{i:2} {dest} = regs[{arg1}] * regs[{arg2}]")
        case "seti":
            print(f"{i:2} {dest} = {arg1}")
        case "setr":
            print(f"{i:2} {dest} = regs[{arg1}]")
        case "eqrr":
            print(f"{i:2} {dest} = (regs[{arg1}] == regs[{arg2}]) ? 1 : 0")
        case "gtrr":
            print(f"{i:2} {dest} = (regs[{arg1}] > regs[{arg2}]) ? 1 : 0")
        case _:
            print(f"{i:2} {dest} = {op}({arg1}, {arg2})")


def sum_of_divisors(n):
    total = 1
    for i in range(2, int(n**0.5) + 1):
        if n % i == 0:
            total += i
            if i != n // i:  # do not count square root twice
                total += n // i
    if n != 1:
        total += n
    return total


def run(data, break_after=None, reg_zero_init=0):
    ip = None
    regs = [0 for _ in range(6)]
    regs[0] = reg_zero_init

    insts = []
    for line in data.splitlines():
        if line.startswith("#"):
            (ip,) = str_to_ints(line)
        else:
            fs = line.split()
            vals = str_to_ints(line)
            insts.append([fs[0]] + vals)

    count = 0
    assert ip is not None
    while regs[ip] < len(insts):
        if break_after is not None and count > break_after:
            break
        inst = insts[regs[ip]]
        f = getattr(d16, inst[0])
        f(regs, *inst[1:])
        regs[ip] += 1
        count += 1
    return regs


def part_1(data):
    regs = run(data)
    print(regs[0])


def part_2(data):
    regs = run(data, 10_000, 1)
    # by analysing the code we can see that the int computer counts the sum of
    # the number of divisors
    print(sum_of_divisors(regs[1]))


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


if __name__ == "__main__":
    main()


listing = """
ip=5
 0 ip = regs[5] + 16
 1 regs[2] = 1
 2 regs[4] = 1
 3 regs[3] = regs[2] * regs[4]              |
 4 regs[3] = (regs[3] == regs[1]) ? 1 : 0   |
 5 ip = regs[3] + regs[5]                   | jump to 7 if regs[3] == regs[1]
 6 ip = regs[5] + 1                         | skip 7
 7 regs[0] = regs[2] + regs[0]
 8 regs[4] = regs[4] + 1                    | regs[4] += 1
 9 regs[3] = (regs[4] > regs[1]) ? 1 : 0    |
10 ip = regs[5] + regs[3]                   | jump to 12 if regs[4] > regs[1]
11 ip = 2                                   | goto 3
12 regs[2] = regs[2] + 1                    | regs[2] += 1
13 regs[3] = (regs[2] > regs[1]) ? 1 : 0    | 
14 ip = regs[3] + regs[5]                   | jump to 16 if regs[2] > regs[1]
15 ip = 1                                   | goto 2
16 ip = regs[5] * regs[5]
17 regs[1] = regs[1] + 2
18 regs[1] = regs[1] * regs[1]
19 regs[1] = regs[5] * regs[1]
20 regs[1] = regs[1] * 11
21 regs[3] = regs[3] + 6
22 regs[3] = regs[3] * regs[5]
23 regs[3] = regs[3] + 15
24 regs[1] = regs[1] + regs[3]
25 ip = regs[5] + regs[0]
26 ip = 0
27 regs[3] = regs[5]
28 regs[3] = regs[3] * regs[5]
29 regs[3] = regs[5] + regs[3]
30 regs[3] = regs[5] * regs[3]
31 regs[3] = regs[3] * 14
32 regs[3] = regs[3] * regs[5]
33 regs[1] = regs[1] + regs[3]
34 regs[0] = 0
35 ip = 0

hypothesis: usm of numbers that divide the number in regs[0]
"""