import machine
import ssd1306
import keypad
# Initialize the I2C bus, OLED display and Keypad
i2c = machine.I2C(scl=machine.Pin(22), sda=machine.Pin(21))
display = ssd1306.SSD1306_I2C(64, 32, i2c)
kp_rows = [machine.Pin(0), machine.Pin(5), machine.Pin(18), machine.Pin(19)]
kp_cols = [machine.Pin(23), machine.Pin(22), machine.Pin(21), machine.Pin(17)]
keys = keypad.Keys(kp_rows, kp_cols, ['1','2','3','A','4','5','6','B','7','8','9','C','*','0','#','D'])
# Define code to show on the display and initialize user input placeholder variables
display_code = ''
user_code = ''
# Loop until user enters correct code
while True:
    # Check if user has entered any key from the keypad
    key = keys.wait()
    if key != '':
        display_code += key  # Add key to display code
        display.fill(0)      # Clear the display
        display.text(display_code, 0, 0)  # Display the code on display
        display.show()       # Update the display
        # Check if user has completed entering the code
        if len(display_code) == 4:
            user_code = display_code
            display_code = ''
            display.fill(0)
            display.text('Code entered:', 0, 0)  # Display message that code is entered
            display.text(user_code, 0, 10)      # Display the final code on display
            display.show()
            # Check if the user entered the correct code
            if user_code == '1234':
                # Display success message and exit
                display.fill(0)
                display.text("Success!", 0, 0)
                display.show()
                break
            else:
                # Display error message and return to start
                display.fill(0)
                display.text("Incorrect code!", 0, 0)
                display.show()
                user_code = ''
                display_code = ''