from machine import Pin, I2C
from ssd1306 import SSD1306_I2C
import framebuf, sys, utime

# OLED display resolution
pix_res_x = 128
pix_res_y = 64

# Grid size (32x16) with each cell being 4x4 pixels
grid_width = 32
grid_height = 16
cell_size = 4

# Initialize grid for Game of Life
grid = [[0] * grid_width for _ in range(grid_height)]
simulation_running = False

def init_i2c(scl_pin, sda_pin):
    # Initialize I2C device
    i2c_dev = I2C(1, scl=Pin(scl_pin), sda=Pin(sda_pin), freq=200000)
    i2c_addr = [hex(ii) for ii in i2c_dev.scan()]
    
    if not i2c_addr:
        print('No I2C Display Found')
        sys.exit()
    else:
        print("I2C Address      : {}".format(i2c_addr[0]))
        print("I2C Configuration: {}".format(i2c_dev))
    
    return i2c_dev

def initialize_rpentomino():
    global grid
    # Clear the grid
    for y in range(grid_height):
        for x in range(grid_width):
            grid[y][x] = 0
    # Initialize the R-pentomino pattern in the center
    cx, cy = grid_width // 2, grid_height // 2
    grid[cy][cx] = 1
    grid[cy][cx + 1] = 1
    grid[cy - 1][cx - 1] = 1
    grid[cy - 1][cx] = 1
    grid[cy + 1][cx] = 1

def draw_grid(oled):
    oled.fill(0)  # Clear the display
    for y in range(grid_height):
        for x in range(grid_width):
            if grid[y][x] == 1:
                oled.fill_rect(x * cell_size, y * cell_size, cell_size, cell_size, 1)
    oled.show()

def apply_game_of_life_rules():
    global grid
    next_grid = [[0] * grid_width for _ in range(grid_height)]
    
    for y in range(grid_height):
        for x in range(grid_width):
            # Count alive neighbors
            alive_neighbors = 0
            for i in range(-1, 2):
                for j in range(-1, 2):
                    if i == 0 and j == 0:
                        continue
                    ny, nx = y + i, x + j
                    if 0 <= nx < grid_width and 0 <= ny < grid_height:
                        alive_neighbors += grid[ny][nx]
            
            # Apply rules
            if grid[y][x] == 1 and (alive_neighbors < 2 or alive_neighbors > 3):
                next_grid[y][x] = 0  # Cell dies
            elif grid[y][x] == 0 and alive_neighbors == 3:
                next_grid[y][x] = 1  # Cell becomes alive
            else:
                next_grid[y][x] = grid[y][x]  # Cell stays the same
    
    grid = next_grid

def main():
    # Initialize I2C and OLED
    i2c_dev = init_i2c(scl_pin=27, sda_pin=26)
    oled = SSD1306_I2C(pix_res_x, pix_res_y, i2c_dev)
    
    # Set up the button
    button = Pin(15, Pin.IN, Pin.PULL_UP)
    last_button_state = button.value()
    
    # Initialize the R-pentomino pattern
    initialize_rpentomino()
    draw_grid(oled)

    while True:
        # Check button press to toggle simulation
        current_button_state = button.value()
        if last_button_state == 1 and current_button_state == 0:
            global simulation_running
            simulation_running = not simulation_running
            if simulation_running:
                print("Simulation started.")
            else:
                print("Simulation paused. Resetting to R-pentomino...")
                initialize_rpentomino()
                draw_grid(oled)
        
        last_button_state = current_button_state

        # Run Game of Life if simulation is active
        if simulation_running:
            apply_game_of_life_rules()
            draw_grid(oled)
            utime.sleep_ms(0)  # Control the simulation speed

if __name__ == '__main__':
    main()