class Registers:
    """
    ARM7TDMI register set. All registers are 32 bits wide.

    Only 16 registers are visible at any given time.
    There are 20 "banked" registers that get swapped in whenever the CPU
    switches to various privileged modes.
    """

    def __init__(self):
        self.SP = 13
        self.LR = 14
        self.PC = 15
        self.BANKED_GPR_RANGE_START = 8
        self.BANKED_GPR_RANGE_END = 12
        self.BANKED_GPR_RANGE_LEN = self.BANKED_GPR_RANGE_END - self.BANKED_GPR_RANGE_START + 1
        self.regs = array('I', [0] * 16)
        self.cpsr = ProgramStatusRegister()
        self.spsr = ProgramStatusRegister()
        self.banked_old_gpr = array('I', [0] * self.BANKED_GPR_RANGE_LEN)
        self.banked_fiq_gpr = array('I', [0] * self.BANKED_GPR_RANGE_LEN)
        self.banked_sp = array('I', [0] * 6)
        self.banked_lr = array('I', [0] * 6)
        self.banked_spsr_user = ProgramStatusRegister()
        self.banked_spsr_fiq = ProgramStatusRegister()
        self.banked_spsr_irq = ProgramStatusRegister()
        self.banked_spsr_swi = ProgramStatusRegister()
        self.banked_spsr_abort = ProgramStatusRegister()
        self.banked_spsr_undefined = ProgramStatusRegister()

    def get(self, reg: int):
        return self.regs[reg]

    def set(self, reg: int, value: int):
        self.regs[reg] = value

    def get_sp(self):
        return self.regs[self.SP]

    def set_sp(self, value: int):
        self.regs[self.SP] = value

    def get_lr(self):
        return self.regs[self.LR]

    def set_lr(self, value: int):
        self.regs[self.LR] = value

    def get_pc(self):
        return self.regs[self.PC]

    def set_pc(self, value: int):
        self.regs[self.PC] = value

    def switch_mode(self, new_mode: CPUMode):
        """
        Switch from the current mode to the specified mode.

        Saves the current register values to the current mode's bank and
        replaces them with values from the new mode's bank.
        """
        if new_mode == self.cpsr.get_mode():
            return
        old_mode = self.cpsr.get_mode()
        self.cpsr.set_mode(new_mode)
        old_bank_index = self.get_bank_index(old_mode)
        new_bank_index = self.get_bank_index(new_mode)
        self.banked_sp[old_bank_index] = self.get_sp()
        self.banked_lr[old_bank_index] = self.get_lr()
        self.get_banked_spsr(old_bank_index).reg = self.spsr.reg
        self.set_sp(self.banked_sp[new_bank_index])
        self.set_lr(self.banked_lr[new_bank_index])
        self.spsr.reg = self.get_banked_spsr(new_bank_index).reg
        if new_mode == CPUMode.FIQ:
            for i in range(self.BANKED_GPR_RANGE_LEN):
                self.banked_old_gpr[i] = self.regs[i + self.BANKED_GPR_RANGE_START]
                self.regs[i + self.BANKED_GPR_RANGE_START] = self.banked_fiq_gpr[i]
        elif old_mode == CPUMode.FIQ:
            for i in range(self.BANKED_GPR_RANGE_LEN):
                self.banked_fiq_gpr[i] = self.regs[i + self.BANKED_GPR_RANGE_START]
                self.regs[i + self.BANKED_GPR_RANGE_START] = self.banked_old_gpr[i]

    def get_bank_index(self, mode: CPUMode) -> BankIndex:
        if mode == CPUMode.SYSTEM or mode == CPUMode.USER:
            return BankIndex.BANK_SYSTEM_USER
        elif mode == CPUMode.FIQ:
            return BankIndex.BANK_FIQ
        elif mode == CPUMode.IRQ:
            return BankIndex.BANK_IRQ
        elif mode == CPUMode.SWI:
            return BankIndex.BANK_SWI
        elif mode == CPUMode.ABORT:
            return BankIndex.BANK_ABORT
        elif mode == CPUMode.UNDEFINED:
            return BankIndex.BANK_UNDEFINED

    def get_banked_spsr(self, bank_index: BankIndex) -> ProgramStatusRegister:
        if bank_index == BankIndex.BANK_SYSTEM_USER:
            return self.banked_spsr_user
        if bank_index == BankIndex.BANK_FIQ:
            return self.banked_spsr_fiq
        elif bank_index == BankIndex.BANK_IRQ:
            return self.banked_spsr_irq
        elif bank_index == BankIndex.BANK_SWI:
            return self.banked_spsr_swi
        elif bank_index == BankIndex.BANK_ABORT:
            return self.banked_spsr_abort
        elif bank_index == BankIndex.BANK_UNDEFINED:
            return self.banked_spsr_undefined

class ProgramStatusRegister:

    def __init__(self, value: int=0):
        self.reg = value

    def get_mode(self) -> CPUMode | int:
        return self.reg & 31

    def set_mode(self, mode: CPUMode):
        self.reg = self.reg & ~31 | mode

    def get_state(self) -> CPUState | int:
        return get_bit(self.reg, 5)

    def set_state(self, state: CPUState):
        self.reg = set_bit(self.reg, 5, state)

    def get_sign_flag(self) -> bint:
        return get_bit(self.reg, 31)

    def set_sign_flag(self, value: bint):
        self.reg = set_bit(self.reg, 31, value)

    def get_zero_flag(self) -> bint:
        return get_bit(self.reg, 30)

    def set_zero_flag(self, value: bint):
        self.reg = set_bit(self.reg, 30, value)

    def get_carry_flag(self) -> bint:
        return get_bit(self.reg, 29)

    def set_carry_flag(self, value: bint):
        self.reg = set_bit(self.reg, 29, value)

    def get_overflow_flag(self) -> bint:
        return get_bit(self.reg, 28)

    def set_overflow_flag(self, value: bint):
        self.reg = set_bit(self.reg, 28, value)

    def get_sticky_overflow_flag(self) -> bint:
        return get_bit(self.reg, 27)

    def set_sticky_overflow_flag(self, value: bint):
        self.reg = set_bit(self.reg, 27, value)

    def get_irq_disable(self) -> bint:
        return get_bit(self.reg, 7)

    def set_irq_disable(self, value: bint):
        self.reg = set_bit(self.reg, 7, value)

    def get_fiq_disable(self) -> bint:
        return get_bit(self.reg, 6)

    def set_fiq_disable(self, value: bint):
        self.reg = set_bit(self.reg, 6, value)