aoc-2019-05-2.py

data = map(int, open("input.txt").readlines()[0].strip().split(","))


class Computer(object):
    def __init__(self, data, stdio=None):
        self.ram = list(data)
        self.pc = 0
        self.halt = False
        self.stdio = stdio

        self.ops = {
            1: self.add,
            2: self.mul,
            3: self.inp,
            4: self.out,
            5: self.ift,
            6: self.iff,
            7: self.lt,
            8: self.eq,
            99: self.hlt,
        }

    def __repr__(self):
        return ",".join(str(val) if idx != self.pc else f"[{val}]" for idx, val in enumerate(self.ram)) + f"\nHALT: {self.halt}"

    def run(self, pc=None, fn=None):
        if pc is not None:
            self.pc = pc
        while True:
            if fn is not None:
                fn(self)
            if not self.step():
                break

    def step(self):
        if self.halt or self.pc >= len(self.ram):
            return False

        opcode = self.ram[self.pc] % 100
        if opcode not in self.ops:
            print(f"ERROR! Unexpected opcode {opcode} at index {self.pc}")
            return False

        return self.ops[opcode]()

    def load_args(self, n_args):
        mode = self.ram[self.pc]
        mode = ("0" * (n_args+1) + str(mode))[::-1][2:]
        args = list()
        for i in range(n_args):
            args.append(self.ram[self.ram[self.pc+i+1]]
                        if mode[i] == "0" else self.ram[self.pc + i+1])
        return args

    def add(self):
        a, b = self.load_args(2)
        addr = self.ram[self.pc + 3]
        self.ram[addr] = a + b
        self.pc += 4
        return True

    def mul(self):
        a, b = self.load_args(2)
        addr = self.ram[self.pc + 3]
        self.ram[addr] = a * b
        self.pc += 4
        return True

    def inp(self):
        addr = self.ram[self.pc+1]
        if not self.stdio:
            print(f"ERROR! No input for input operation at {self.pc}")
            return False
        self.ram[addr] = self.stdio.pop(0)
        self.pc += 2
        return True

    def out(self):
        addr = self.ram[self.pc+1]
        if addr < 0 or addr >= len(self.ram):
            print(
                f"ERROR! Requested address {addr} is not available, requested at position {self.pc}")
            return False
        print(self.ram[addr])
        self.pc += 2
        return True

    def ift(self):
        val, addr = self.load_args(2)
        if val:
            self.pc = addr
        else:
            self.pc += 3
        return True

    def iff(self):
        val, addr = self.load_args(2)
        if val == 0:
            self.pc = addr
        else:
            self.pc += 3
        return True

    def lt(self):
        a, b = self.load_args(2)
        addr = self.ram[self.pc + 3]
        self.ram[addr] = int(a < b)
        self.pc += 4
        return True

    def eq(self):
        a, b = self.load_args(2)
        addr = self.ram[self.pc + 3]
        self.ram[addr] = int(a == b)
        self.pc += 4
        return True

    def hlt(self):
        self.halt = True
        return False


# Computer(data, [5]).run(fn=print)
Computer(data, [5]).run()