Regenerate the processor code after fixing the beq and bne instructions. Make sure to address the instructions in a way that makes sense and is correct. Right now beq and bne don't operate properly and lead to 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
        self.registers[0] = 0  # $zero is always 0

    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": 0, "$v0": 2, "$v1": 3,
            "$a0": 4, "$a1": 5, "$a2": 6, "$a3": 7,
            "$t0": 8, "$t1": 9, "$t2": 10, "$t3": 11,
            "$t4": 12, "$t5": 13, "$t6": 14, "$t7": 15,
            "$s0": 16, "$s1": 17, "$s2": 18, "$s3": 19,
            "$s4": 20, "$s5": 21, "$s6": 22, "$s7": 23,
            "$ra": 31, "$1": 1, "$at": 1
        }
        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, current_instruction_address):
        instruction_line = instruction_line.strip().replace(',', '')
        parts = instruction_line.split()
        opcode = self.mapping(parts[0])

        # Handle pseudo-instruction 'li' by converting it to 'addi'
        if parts[0] == 'li':
            parts = ['addi', parts[1], '$zero', parts[2]]
            opcode = self.opcodes['addi']

        # R-type instructions
        if parts[0] in self.R_type:
            rd = format(self.registers[parts[1]], '05b')
            rs = format(self.registers[parts[2]], '05b')
            rt = format(self.registers[parts[3]], '05b')
            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 = format(self.registers[parts[1]], '05b')
                var_name = parts[2]
                if var_name in self.var_addresses:
                    var_address = self.var_addresses[var_name]
                    base_reg = format(0, '05b')  # $zero
                    offset = var_address
                    offset_bin = format(offset & 0xFFFF, '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 = format(self.registers[parts[1]], '05b')
                rs = format(self.registers[parts[2]], '05b')
                immediate = format(int(parts[3]) & 0xFFFF, '016b')
                return f"{opcode}{rs}{rt}{immediate}"
            elif parts[0] == 'beq':
                rs = format(self.registers[parts[1]], '05b')
                rt = format(self.registers[parts[2]], '05b')
                offset = self.calculate_offset(parts[3], current_instruction_address)
                return f"{opcode}{rs}{rt}{offset}"
            elif parts[0] == 'bne':
                rs = format(self.registers[parts[1]], '05b')
                rt = format(self.registers[parts[2]], '05b')
                offset = self.calculate_offset(parts[3], current_instruction_address)
                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, current_instruction_address):
        """Calculate the branch offset for labels."""
        target_address = self.labels[label]
        offset = (target_address - (current_instruction_address + 4)) // 4  # PC-relative offset
        # Convert offset to two's complement for negative values
        offset = offset & 0xFFFF
        return format(offset, '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 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
            current_instruction_address = self.current_address
            machine_code = self.to_machine_code(line, current_instruction_address)
            if machine_code:
                binary_instructions.append(machine_code)
                self.instruction_map[current_instruction_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 self.control_unit.Jump == 1:  # Jump instruction
            self.program_counter.count = immediate * 4
        elif self.control_unit.Branch == 1:
            branch_offset = immediate * 4
            if (opcode == '000100' and alu_result == 0) or (opcode == '000101' and alu_result != 0):
                self.program_counter.count = pc_value + 4 + branch_offset
            else:
                self.program_counter.count = pc_value + 4
        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

            # Get the assembly instruction
            assembly_instruction = self.instructions_obj.instruction_map.get(pc_value, 'Unknown Instruction')
            print(f"Executing instruction at PC {pc_value}: {assembly_instruction}")

            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) if rs is not None else 0
            rt_val = self.registers.read_register(rt) if rt is not None else 0

            alu_control = self.determine_alu_control(opcode, funct)

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

            # Print the ALU operation
            print(f"ALU operation: {alu_control}, operands: {rs_val}, {rt_val}, result: {alu_result}")

            mem_data = self.memory_access(alu_result, rt_val)

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

            print(f"OPCODE: {opcode}, ALU Result: {alu_result}, Immediate: {immediate_value}, Branch: {self.control_unit.Branch}")

            self.update_pc(opcode, alu_result, immediate_value)

            self.debug()

    def debug(self):
        print(f"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
    '    lw $t4, 1, ',
    '    addi $t1, $zero, 0',   # Initialize sum to 0 (replacing li)
    '    addi $t2, $zero, 1',   # Initialize counter to 1 (replacing li)
    'loop:',
    '    slt $t3, $t0, $t2', # Set $t3 = 1 if $t0 < $t2 (i.e., if n < counter)
    '    beq $t3, $t4, 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()