#!/usr/bin/env python
# coding: utf-8

# In[1]:


class Jump(Exception):
    def __init__(self, offset):
        self.offset = offset
        super().__init__(offset)


def opcode(operands):
    def decorator(f):
        class Opcode:
            def __set_name__(self, owner, name):
                self.opcode = name[3:]
                owner.opcodes[self.opcode] = self

            def __repr__(self):
                return f"<opcode {self.opcode} {operands!r}>"

            def value(self, operand, type_):
                if type_ == "r":
                    return operand
                try:
                    return int(operand)
                except ValueError:
                    return self.registers[operand]

            def __call__(self, cpu, *ops):
                self.registers = cpu.registers
                try:
                    f(cpu, *map(self.value, ops, operands))
                    cpu.pos += 1
                except Jump as j:
                    cpu.pos += j.offset
                return self.opcode

        return Opcode()

    return decorator


class Proc:
    opcodes = {}

    def __init__(self, debug=True):
        self.reset(debug)

    def reset(self, debug=True):
        self.registers = dict.fromkeys("abcdefgh", 0)
        self.debug = debug
        if not debug:
            self.registers["a"] = 1
        self.pos = 0

    def run(self, instructions):
        if not self.debug:
            instructions = self.optimise(instructions)
        while 0 <= self.pos < len(instructions):
            opcode, *ops = instructions[self.pos]
            yield self.opcodes[opcode](self, *ops)

    @opcode("")
    def op_nop(self):
        pass

    @opcode("rv")
    def op_set(self, x, y):
        self.registers[x] = y

    @opcode("rv")
    def op_sub(self, x, y):
        self.registers[x] -= y

    @opcode("rv")
    def op_mul(self, x, y):
        self.registers[x] *= y

    @opcode("rv")
    def op_mod(self, x, y):
        self.registers[x] %= y

    @opcode("vv")
    def op_jnz(self, x, y):
        if x:
            raise Jump(y)

    def optimise(self, instructions):
        # modulus operation over two registers, setting a third flag register
        # using two working registers. If the flag register
        # is set, jump out of the outer loop
        operand1, operand2 = instructions[13][2], instructions[11][2]
        workreg = instructions[11][1]
        flagreg = instructions[15][1]
        return (
            instructions[:10]
            + [
                ("set", workreg, operand1),
                ("mod", workreg, operand2),
                ("jnz", workreg, "8"),
                ("set", flagreg, "0"),
                ("jnz", "1", "11"),
                ("jnz", "1", "5"),
            ]
            + [("nop",)] * 4
            + instructions[20:]
        )


# In[2]:


import aocd

data = aocd.get_data(day=23, year=2017)
instructions = [line.split() for line in data.splitlines()]


# In[3]:


proc = Proc()
print("Part 1:", sum(1 for opcode in proc.run(instructions) if opcode == "mul"))


# In[4]:


from collections import deque

proc = Proc(debug=False)
deque(proc.run(instructions), 0)
print("Part 2:", proc.registers["h"])


# ### Optimisations
# 
# The following set of instructions is basically setting `f` to `0` (true) if `b % d == 0` is true, using `g` as the stack for operands, using multiplication with e, incrementing in single steps:
# 
#     0.  set e 2  > for e in range(2, b):
#     1.  set g d  v
#     2.  mul g e  v
#     3.  sub g b  v
#     4.  jnz g 2  >     if d * e == b:
#     5.  set f 0  >         f = 0
#     6.  sub e -1 0
#     7.  set g e  0
#     8.  sub g b  0
#     9.  jnz g -8 0
# 
# We can replace those 10 with a simple `mod` operand instead, filling out the rest with `nop` codes:
# 
#     0.  set g b  v
#     1.  mod g d  v
#     2.  jnz g 8  > if b % d == 0:
#     3.  set f 0  >      f = 0
#     4.  jnz 1 6  > skip the remaining nops
#     5-9 nop
# 
# Next, the outer loop is inefficient; it keeps on looping while only the first `f == 0` setting is important:
# 
#     -2.  set f 1   > f = 1
#     -1.  set d 2   > for d in range(2, b):
#       ... the inner loop above, so repeated b - 1 times
#     10. sub d -1  1
#     11. set g d   1
#     12. sub g b   1
#     13. jnz g -13 1
#     14. jnz f 2   > if not f:
#     15. sub h -1  >    h += 1
# 
# That outer loop can be jumped out of when we set f = 0, so in the inner loop, after `set f 0`, we can add a jump to the `h += 1` instruction, stepping out of the outer loop:
# 
#     0.  set g b  v
#     1.  mod g d  v
#     2.  jnz g 8  > if b % d == 0:
#     3.  set f 0  >      f = 0
#     4.  jnz 1 11 >      break (and skip the if not f test)
#     5.  jnz 1 5  > skip the remaining nops
#     6-9 nop
# 

# In[5]:


# Cheating option, just run Python code
lower = (99 * 100) + 100000
upper = lower + 17000
h = 0
for b in range(lower, upper + 1, 17):
    f = 1
    for d in range(2, b):
        if b % d == 0:
            f = 0
            break
    if not f:
        h += 1
print(h)