Change the way the instructions are processed in this processor. Create a instructions map, which contains the adresses of each instruction, thereby fix the jump instruction from running into an infinite loop.

class Control_Signals():
    def __init__(self):
        self.RegDst = 0
        self.RegWrite = 0
        self.ALUsrc = 0
        self.MemWrite = 0
        self.MemRead = 0
        self.MemtoReg = 0
        self.Branch = 0
        self.Jump = 0
        self.ALUOp = "00"  # ALUOp for ALU control

class Program_Counter():
    def __init__(self):
        self.count = 0

class Data_Memory():
    def __init__(self):
        self.data = {}  # Use a dictionary for memory
        self.next_free_address = 0  # Start addresses from 0

    def add_variable(self, var_name, value):
        address = self.next_free_address
        self.data[address] = value
        self.next_free_address += 4  # Word-aligned
        return address

    def data_read(self, address):
        return self.data.get(address, 0)

    def data_write(self, address, data_val):
        self.data[address] = data_val

class Instruction_Memory():
    def __init__(self, machine_code):
        self.instructions = list(machine_code)

    def read_instruction(self, idx):
        return self.instructions[idx]

class Registers():
    def __init__(self):
        self.registers = [0] * 32

    def read_register(self, reg_num):
        return self.registers[reg_num]

    def write_register(self, reg_num, data):
        if reg_num != 0 and reg_num is not None:  # Prevent writing to $zero
            self.registers[reg_num] = data

class ALU():
    def execute(self, op1, op2, alu_control):
        if alu_control == "add":
            return op1 + op2
        elif alu_control == "sub":
            return op1 - op2
        elif alu_control == "and":
            return op1 & op2
        elif alu_control == "or":
            return op1 | op2
        elif alu_control == "slt":
            return 1 if op1 < op2 else 0

class Instructions:
    def __init__(self):
        self.R_type = ['add', 'sub', 'and', 'or', 'slt']
        self.I_type = ['lw', 'sw', 'beq', 'addi', 'bne']
        self.J_type = ['j']
        self.opcodes = {
            'add': '000000',
            'sub': '000000',
            'and': '000000',
            'or': '000000',
            'slt': '000000',
            'lw': '100011',
            'sw': '101011',
            'beq': '000100',
            'bne': '000101',
            'addi': '001000',
            'j': '000010'
        }
        self.funct= {
            'add': '100000', 'sub': '100010', 'and': '100100', 'or': '100101', 'slt': '101010'
        }
        self.registers={
            "$zero": '00000', "$v0": '00010', "$v1": '00011',
            "$a0": '00100', "$a1": '00101', "$a2": '00110', "$a3": '00111',
            "$t0": '01000', "$t1": '01001', "$t2": '01010', "$t3": '01011',
            "$t4": '01100', "$t5": '01101', "$t6": '01110', "$t7": '01111',
            "$s0": '10000', "$s1": '10001', "$s2": '10010', "$s3": '10011',
            "$s4": '10100', "$s5": '10101', "$s6": '10110', "$s7": '10111',
            "$ra": '11111', "$1": '00001', "$at": '00001'
        }
        self.labels= {}
        self.current_address= 0
        self.memory= Data_Memory()
        self.instruction_map= {}
        self.var_addresses = {}  # Mapping from variable names to addresses

    def mapping(self, instruction):
        if instruction in self.R_type:
            return '000000'
        elif instruction in self.I_type or instruction in self.J_type:
            return self.opcodes[instruction]
        else:
            return None

    def to_machine_code(self, instruction_line):
        instruction_line = instruction_line.strip().replace(',', '')
        parts = instruction_line.split()
        opcode = self.mapping(parts[0])

        # R-type instructions
        if parts[0] in self.R_type:
            rd = self.registers[parts[1]]
            rs = self.registers[parts[2]]
            rt = self.registers[parts[3]]
            funct = self.funct[parts[0]]
            shamt = '00000'
            return f"{opcode}{rs}{rt}{rd}{shamt}{funct}"

        # I-type instructions
        elif parts[0] in self.I_type:
            if parts[0] == 'lw' or parts[0] == 'sw':
                rt = self.registers[parts[1]]
                var_name = parts[2]
                if var_name in self.var_addresses:
                    var_address = self.var_addresses[var_name]
                    base_reg = '00000'  # $zero
                    offset = var_address
                    offset_bin = format(offset, '016b')
                    return f"{opcode}{base_reg}{rt}{offset_bin}"
                else:
                    print(f"Error: Variable {var_name} not found in memory.")
                    return None
            elif parts[0] == 'addi':
                rt = self.registers[parts[1]]
                rs = self.registers[parts[2]]
                immediate = format(int(parts[3]), '016b')
                return f"{opcode}{rs}{rt}{immediate}"
            elif parts[0] == 'beq':
                rs = self.registers[parts[1]]
                rt = self.registers[parts[2]]
                offset = self.calculate_offset(parts[3])
                return f"{opcode}{rs}{rt}{offset}"
            elif parts[0] == 'bne':
                rs = self.registers[parts[1]]
                rt = self.registers[parts[2]]
                offset = self.calculate_offset(parts[3])
                return f"{opcode}{rs}{rt}{offset}"

        # J-type instructions
        elif parts[0] in self.J_type:
            address = self.calculate_jump_address(parts[1])
            return f"{opcode}{address}"

    def calculate_offset(self, label):
        """Calculate the branch offset for labels."""
        target_address = self.labels[label]
        offset = (target_address - self.current_address - 4) // 4  # PC-relative offset
        return format(offset & 0xFFFF, '016b')

    def calculate_jump_address(self, label):
        """Calculate the jump address for labels."""
        target_address = self.labels[label] // 4
        return format(target_address, '026b')

    def first_pass(self, program):
        address = 0
        data_section = False
        for line in program:
            line= line.strip()
            if not line:
                continue
            if line.startswith('.data'):
                data_section= True
                continue
            elif line.startswith('.text'):
                data_section= False
                continue

            if line.endswith(':'):
                label= line[:-1].strip()
                self.labels[label] = address
                print(f"Address for label '{label}': {address}")
                continue

            elif data_section:
                parts= line.strip().split()
                var_name= parts[0].split(':')[0].strip()
                value = int(parts[-1])  # Get value after .word
                assigned_address = self.memory.add_variable(var_name, value)
                self.var_addresses[var_name] = assigned_address
                print(f"Variable '{var_name}' assigned address {assigned_address}")
                continue
            else:
                print(f"Address: {address} for instruction: {line}")
                address += 4

        for label in self.labels:
            print(f"Label '{label}' at address {self.labels[label]}")

    def load_program_from_file(self, filename):
        """Load MIPS program from a text file."""
        with open(filename, 'r') as file:
            program = file.readlines()
        return program

    def compile_program(self, program):
        """Compile the full program into machine code."""
        binary_instructions = []
        self.first_pass(program)

        self.current_address = 0
        for line in program:
            line = line.strip()
            if not line or line.startswith('.'):
                continue
            if line.endswith(':'):
                continue
            machine_code = self.to_machine_code(line)
            if machine_code:
                binary_instructions.append(machine_code)
                self.instruction_map[self.current_address] = line
                self.current_address += 4

        return binary_instructions

class Processor():
    def __init__(self, instructions_obj, program):
        self.control_unit = Control_Signals()
        self.program_counter = Program_Counter()
        self.instructions_obj = instructions_obj
        machine_code = instructions_obj.compile_program(program)
        self.instruction_memory = Instruction_Memory(machine_code)
        self.data_memory = instructions_obj.memory  # Use the same data memory
        self.registers = Registers()
        self.alu = ALU()

    def instruction_fetch(self):
        pc_value = self.program_counter.count
        instruction_index = pc_value // 4  # Assuming word-aligned instructions
        if instruction_index >= len(self.instruction_memory.instructions):
            return None  # No more instructions
        instruction = self.instruction_memory.read_instruction(instruction_index)
        return instruction

    def instruction_decode(self, instruction):
        # Reset control signals
        self.control_unit = Control_Signals()

        opcode = instruction[0:6]
        if opcode == '000000':  # R-type
            rs = int(instruction[6:11], 2)
            rt = int(instruction[11:16], 2)
            rd = int(instruction[16:21], 2)
            shamt = int(instruction[21:26], 2)
            funct = instruction[26:32]
            immediate_value = None
        elif opcode == '000010':  # Jump instruction
            address = instruction[6:32]
            immediate_value = int(address, 2)
            rs = rt = rd = shamt = funct = None
            # Set control signals for Jump
            self.control_unit.Jump = 1
        else:  # I-type instructions
            rs = int(instruction[6:11], 2)
            rt = int(instruction[11:16], 2)
            immediate = instruction[16:32]
            # Sign-extend immediate value
            if immediate[0] == '1':  # Negative immediate
                immediate_value = int(immediate, 2) - (1 << 16)
            else:
                immediate_value = int(immediate, 2)
            rd = None
            shamt = None
            funct = None

        # Set control signals based on opcode
        if opcode == "000000":  # R-type
            self.control_unit.RegDst = 1
            self.control_unit.RegWrite = 1
            self.control_unit.ALUOp = "10"
        elif opcode == "100011":  # lw
            self.control_unit.ALUsrc = 1
            self.control_unit.MemRead = 1
            self.control_unit.RegWrite = 1
            self.control_unit.MemtoReg = 1
        elif opcode == "101011":  # sw
            self.control_unit.ALUsrc = 1
            self.control_unit.MemWrite = 1
        elif opcode == "000100":  # beq
            self.control_unit.Branch = 1
            self.control_unit.ALUOp = "01"
        elif opcode == "000101":  # bne
            self.control_unit.Branch = 1
            self.control_unit.ALUOp = "01"
        elif opcode == "001000":  # addi
            self.control_unit.ALUsrc = 1
            self.control_unit.RegWrite = 1

        return opcode, rs, rt, rd, immediate_value, shamt, funct

    def determine_alu_control(self, opcode, funct):
        if self.control_unit.ALUOp == "10":  # R-type
            if funct == '100000':  # add
                return 'add'
            elif funct == '100010':  # sub
                return 'sub'
            elif funct == '100100':  # and
                return 'and'
            elif funct == '100101':  # or
                return 'or'
            elif funct == '101010':  # slt
                return 'slt'
            else:
                return 'add'  # Default action for unknown funct
        elif self.control_unit.ALUOp == "00":  # lw, sw, addi
            return 'add'
        elif self.control_unit.ALUOp == "01":  # beq, bne
            return 'sub'
        else:
            return 'add'  # Default action

    def execution(self, rs_val, rt_val, immediate, alu_control):
        if self.control_unit.ALUsrc == 1 and immediate is not None:
            alu_result = self.alu.execute(rs_val, immediate, alu_control)
        else:
            alu_result = self.alu.execute(rs_val, rt_val, alu_control)
        return alu_result

    def memory_access(self, alu_result, rt_val):
        if self.control_unit.MemWrite:
            self.data_memory.data_write(alu_result, rt_val)
            return None
        elif self.control_unit.MemRead:
            return self.data_memory.data_read(alu_result)
        return None

    def write_back(self, alu_result, mem_data, rd, rt):
        if self.control_unit.RegWrite:
            write_data = mem_data if self.control_unit.MemtoReg else alu_result
            write_reg = rd if self.control_unit.RegDst else rt
            self.registers.write_register(write_reg, write_data)

    def update_pc(self, opcode, alu_result, immediate):
        pc_value = self.program_counter.count
        if opcode == '000010':  # Jump instruction
            self.program_counter.count = immediate * 4
        elif opcode == '000100' and alu_result == 0:  # beq
            branch_offset = immediate * 4
            self.program_counter.count = pc_value + 4 + branch_offset
        elif opcode == '000101' and alu_result != 0:  # bne
            branch_offset = immediate * 4
            self.program_counter.count = pc_value + 4 + branch_offset
        else:
            self.program_counter.count = pc_value + 4

    def run(self):
        while True:
            pc_value = self.program_counter.count
            instruction = self.instruction_fetch()
            if instruction is None:
                break  # No more instructions

            opcode, rs, rt, rd, immediate_value, shamt, funct = self.instruction_decode(instruction)

            # For jump instruction, skip ALU and other stages
            if self.control_unit.Jump == 1:
                self.update_pc(opcode, None, immediate_value)
                self.debug()
                continue

            rs_val = self.registers.read_register(rs)
            rt_val = self.registers.read_register(rt)

            alu_control = self.determine_alu_control(opcode, funct)

            alu_result = self.execution(rs_val, rt_val, immediate_value, alu_control)

            mem_data = self.memory_access(alu_result, rt_val)

            self.write_back(alu_result, mem_data, rd, rt)

            self.update_pc(opcode, alu_result, immediate_value)

            self.debug()

    def debug(self):
        print(f"PC: {self.program_counter.count}, Registers: {self.registers.registers}")
        print(f"Data Memory: {self.data_memory.data}")
        print('-----------------------')

# Sample MIPS code
program = [
    '.data',
    '    n: .word 10',       # Change this value to compute sum of different n
    '    sum: .word 0',
    '.text',
    '    lw $t0, n',         # Load n into $t0
    '    li $t1, 0',         # Initialize sum to 0
    '    li $t2, 1',         # Initialize counter to 1
    'loop:',
    '    slt $t3, $t0, $t2', # Set $t3 = 1 if $t0 < $t2 (i.e., if n < counter)
    '    bne $t3, $zero, end', # If $t3 != 0, jump to end
    '    add $t1, $t1, $t2', # sum += counter
    '    addi $t2, $t2, 1',  # Increment counter
    '    j loop',            # Repeat the loop
    'end:',
    '    sw $t1, sum'       # Store result in sum
]

# Create an Instructions object
instructions_obj = Instructions()
# Create a Processor and run the program
processor = Processor(instructions_obj, program)
processor.run()