class Ram(object):
    def __init__(self, instr, input_registers):
        self.instr = instr # list of instructions
        self.current = 0 # current instruction
        self.input_registers = input_registers # I registers
        self.output_registers = [] # O registers
        self.work_registers = [] # R registers

    def read_register(self, type_register, index, ref_origin=None, ref_target=None):
        '''
        Read from a register, return the value
        '''
        if type_register == "i": # input
            return self.input_registers[index]
        elif type_register == "r": # work
            value = self.work_registers[index]
            if value is None:
                raise ValueError("register empty")
            else:
                return self.work_registers[index]
        elif type_register == "o": # output
            raise TypeError("output registers are write only")
        elif type_register == "value":
            return index
        elif type_register == "@": # reference
            target_index = self.read_register(ref_origin, index)
            return self.read_register(ref_target, target_index)

    def write_register(self, value, type_register, index, ref_origin=None, ref_target=None):
        '''
        Write to a register
        '''
        if type_register == "i": # input
            raise TypeError("input registers are read only")
        elif type_register == "r": # work
            if len(self.work_registers) <= index:
                self.work_registers.extend([None] * (index + 1 - len(self.work_registers))) # extend with uninitialized values
            self.work_registers[index] = value
        elif type_register == "o": #output
            if len(self.output_registers) <= index:
                self.output_registers.extend([None] * (index + 1 - len(self.output_registers))) # extend with uninitialized values
            self.output_registers[index] = value
        elif type_register == "value":
            raise TypeError("cannot write on value")
        elif type_register == "@": # reference
            target_index = self.read_register(ref_origin, index)
            self.write_register(value, ref_target, target_index)

    def op(self, type_op, r1, r2, r3):
        '''
        Arithmetic operation (addition, subtraction, division, multiplication)
        r3 = r1 op r2
        '''
        self.current += 1
        if type_op == 'ADD':
            value_1 = self.read_register(*r1)
            value_2 = self.read_register(*r2)
            value_3 = value_1 + value_2
            self.write_register(value_3, *r3)
        elif type_op == 'SUB':
            value_1 = self.read_register(*r1)
            value_2 = self.read_register(*r2)
            value_3 = value_1 - value_2
            self.write_register(value_3, *r3)
        elif type_op == 'DIV':
            value_1 = self.read_register(*r1)
            value_2 = self.read_register(*r2)
            value_3 = value_1 / value_2
            self.write_register(value_3, *r3)
        elif type_op == 'MULT':
            value_1 = self.read_register(*r1)
            value_2 = self.read_register(*r2)
            value_3 = value_1 * value_2
            self.write_register(value_3, *r3)

    def op_ctrl(self, type_op, z, r1=None, r2=None):
        '''
        Control operation (Jump, Jump if equal, Jump if less)
        '''
        if type_op == 'JUMP':
            if isinstance(z, int):
                self.current += z
            else:
               raise ValueError("wrong operand type (should be integer)")
        if type_op == 'JE':
            value_1 = self.read_register(*r1)
            value_2 = self.read_register(*r2)
            if value_1 == value_2:
                self.op_ctrl('JUMP', z)
            else:
                self.current += 1
        if type_op == 'JL':
            value_1 = self.read_register(*r1)
            value_2 = self.read_register(*r2)
            if value_1 > value_2:
                self.op_ctrl('JUMP', z)
            else:
                self.current += 1

    def __register_to_str(register):
        if register[0] == '@':
            return f'{register[3]}@{register[2]}{register[1]}'
        elif register[0] == 'value':
            return f'{register[1]}'
        else:
            return f'{register[0]}{register[1]}'

    def __instr_to_str(self, instr_number):
        instr = self.instr[instr_number]
        op = instr["args"][0]
        number = f'{instr_number} :'
        instr_str = f' {number: <6}{op}('
        if op == 'JUMP':
            instr_str += f'{instr["args"][1]})'
        elif op == 'JL' or op == 'JE':
            instr_str += f'{instr["args"][1]}, '
            instr_str += f'{Ram.__register_to_str(instr["args"][2])}, '
            instr_str += f'{Ram.__register_to_str(instr["args"][3])})'
        else:
            instr_str += f'{Ram.__register_to_str(instr["args"][1])}, '
            instr_str += f'{Ram.__register_to_str(instr["args"][2])}, '
            instr_str += f'{Ram.__register_to_str(instr["args"][3])})'
        return instr_str


    def execute(self):
        '''
        Run machine to completion
        '''
        while self.current < len(self.instr):
            instr_str = self.__instr_to_str(self.current)
            self.instr[self.current]['op'](self, *self.instr[self.current]['args'])
            print(f'{instr_str: <28}{self.work_registers}')

    def generate_graph(self):
        '''
        Return precedance graph edges between instructions
        (instr1, instr2) exists if instr2 is reachable from instr1 in 1 step
        '''
        graph_edges = []
        n_instr = len(self.instr)
        for (n,instr) in enumerate(self.instr):
            if ((instr['op'] == Ram.op) and (n<n_instr-1)):
                graph_edges.append((n, n+1))
            elif (instr['args'][0] == 'JUMP'):
                dest = n+instr['args'][1]
                if dest<n_instr:
                    graph_edges.append((n, dest))
            elif ((instr['args'][0] == 'JL') or (instr['args'][0] == 'JE')):
                if (n<n_instr-1):
                    graph_edges.append((n, n+1))
                graph_edges.append((n, n+instr['args'][1]))
        return graph_edges

    def accessible_instr(self):
        '''
        Return a list of accessible instructions in a machine
        '''
        graph_edges = self.generate_graph()
        accessible_from_instr = dict()
        for edge in graph_edges:
            accessible_from_instr.setdefault(edge[0], []).append(edge[1])
        waitlist = [0]
        accessible = []
        while len(waitlist)>0:
            current = waitlist.pop()
            accessible.append(current)
            try:
                temp = accessible_from_instr[current]
            except KeyError: # no node reachable from this node
                pass
            for elem in temp:
                if elem not in accessible and elem not in waitlist:
                    waitlist.append(elem)
        return sorted(accessible)

    def remove_unreachable_instr(self):
        '''
        Remove dead code in a machine (unreachable instructions)
        Print all the instructions still in the machine after
        '''
        accessible = self.accessible_instr()
        new_instr = [instr for (i,instr) in enumerate(self.instr) if i in accessible]
        for instr in new_instr:
            print(instr)
        self.instr = new_instr

    def predecessors(self, node):
        graph_edges = self.generate_graph()
        return [edge[0] for edge in graph_edges if edge[1] == node]

    def merge_instructions(self):
        for (n,instr) in enumerate(self.instr):
            if instr['op'] == Ram.op and (instr['args'][0] == 'ADD' or instr['args'][0] == 'MULT'):
                registry = instr['args'][3]
                op = instr['args'][0]
                if (instr['args'][1] == registry and instr['args'][2][0] == 'value') or (instr['args'][2] == registry and instr['args'][1][0] == 'value'):
                    predecessors = self.predecessors(n)
                    if len(predecessors) == 1:
                        predecessor_instr = self.instr[predecessors[0]]
                        if predecessor_instr['op'] == Ram.op:
                            if op == predecessor_instr['args'][0] and predecessor_instr['args'][1][0] == 'value' and predecessor_instr['args'][2][0] == 'value' and predecessor_instr['args'][3] == registry:
                                if op == 'ADD':
                                    value_int = predecessor_instr['args'][1][1] + predecessor_instr['args'][2][1]
                                    if instr['args'][1][0] == 'value':
                                        value_int += instr['args'][1][1]
                                    else:
                                        value_int += instr['args'][2][1]
                                    self.instr[n] = {"op": Ram.op, "args": (op, ('value', value_int), ('value', 0), registry)}
                                elif op == 'MULT':
                                    value_int = predecessor_instr['args'][1][1] * predecessor_instr['args'][2][1]
                                    if instr['args'][1][0] == 'value':
                                        value_int *= instr['args'][1][1]
                                    else:
                                        value_int *= instr['args'][2][1]
                                    self.instr[n] = {"op": Ram.op, "args": (op, ('value', value_int), ('value', 1), registry)}
                                self.instr.pop(predecessors[0])
        for instr in self.instr:
            print(instr)