from machine import Pin, I2C
from time import sleep, ticks_ms,ticks_diff
from sh1106 import SH1106_I2C

# Initialize I2C interface
i2c = I2C(0, scl=Pin(22), sda=Pin(21))  # Adjust pins as per your hardware setup

OLED_WIDTH = 64
OLED_HEIGHT = 128
OLED_EDGE = 5

SQUARE_SIZE = 5
SQUARE_SPEED = 2
TIME_INTERVAL = 50


oled_display = SH1106_I2C(OLED_WIDTH, OLED_HEIGHT, i2c)

def draw_square(x, y):
    # oled_display.fill(0)  # Clear the screen
    oled_display.rect(x, y, SQUARE_SIZE, SQUARE_SIZE, 1)
    oled_display.show()

def move_square(square, ticks_ms_prev):
    if ticks_diff(ticks_ms(), ticks_ms_prev) > TIME_INTERVAL:
        x, y, dx, dy = square
        ticks_ms_prev = ticks_ms()
        draw_square(x, y)
        # Update square position
        x += dx
        y += dy

        # Check for collisions with the edges
        if x <= OLED_EDGE or x >= OLED_WIDTH - OLED_EDGE:
            dx = -dx  # Reverse X direction
        if y <= OLED_EDGE or y >= OLED_HEIGHT - OLED_EDGE:
            dy = -dy  # Reverse Y direction
        square = [x, y, dx, dy]

    return square, ticks_ms_prev

def main():
    sq1_prev_time = sq2_prev_time = ticks_ms()
    sq1 = [10, 10, 2, 4]      # Initial position of square 1 [x, y, dx, dy]
    sq2 = [20, 30, -3, -2]    # Initial position of square 2 [x, y, dx, dy]

    while True:
        oled_display.fill(0)  # Clear the screen
        sq1, sq1_prev_time = move_square(sq1, sq1_prev_time)
        sq2, sq2_prev_time = move_square(sq2, sq2_prev_time)

if __name__ == "__main__":

  main()