Revamp this code by using adresses in the form of a map for the instructions and the labels as well. Correct the implementations of beq, bne and jump properly. Make sure that the jumps and branching are handled properly and make sure infinite loop isn't caused.

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
    '    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)
    '    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()