import pygame
import numpy as np

# Initialize pygame
pygame.init()

# Define constants
WIDTH, HEIGHT = 800, 600
CELL_SIZE = 10
GRID_WIDTH, GRID_HEIGHT = WIDTH // CELL_SIZE, HEIGHT // CELL_SIZE

# Set up the display
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Conway's Game of Life")

# Initialize the grid
def empty_grid(width, height):
    return np.zeros((width, height), dtype=int)

def random_grid(width, height):
    return np.random.randint(2, size=(width, height), dtype=int)

grid = random_grid(GRID_WIDTH, GRID_HEIGHT)

# Game of Life logic
def update_grid(grid):
    new_grid = grid.copy()
    for i in range(grid.shape[0]):
        for j in range(grid.shape[1]):
            # Count living neighbors
            live_neighbors = np.sum(grid[i-1:i+2, j-1:j+2]) - grid[i, j]

            if grid[i, j] == 1 and (live_neighbors < 2 or live_neighbors > 3):
                new_grid[i, j] = 0
            elif grid[i, j] == 0 and live_neighbors == 3:
                new_grid[i, j] = 1
    return new_grid

# Game Loop
running = True
clock = pygame.time.Clock()

while running:
    screen.fill((0, 0, 0))

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Update the grid and draw the new state
    grid = update_grid(grid)

    for i in range(GRID_WIDTH):
        for j in range(GRID_HEIGHT):
            color = (255, 255, 255) if grid[i, j] == 1 else (0, 0, 0)
            rect = (i * CELL_SIZE, j * CELL_SIZE, CELL_SIZE, CELL_SIZE)
            pygame.draw.rect(screen, color, rect)

    pygame.display.flip()
    clock.tick(10) # control the simulation speed

# Quit Pygame
pygame.quit()