import pygame
import numpy as np
import time

# Define constants
WIDTH, HEIGHT = 800, 600  # Screen dimensions
CELL_SIZE = 10  # Size of each cell
GRID_WIDTH = WIDTH // CELL_SIZE
GRID_HEIGHT = HEIGHT // CELL_SIZE
FPS = 10  # Frames per second (speed of the simulation)

# Define colors
ALIVE_COLOR = (0, 255, 0)
DEAD_COLOR = (0, 0, 0)

# Initialize pygame
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Game of Life")

# Create the grid (2D array)
grid = np.zeros((GRID_HEIGHT, GRID_WIDTH))

# Function to draw the grid
def draw_grid():
    for y in range(GRID_HEIGHT):
        for x in range(GRID_WIDTH):
            color = ALIVE_COLOR if grid[y, x] == 1 else DEAD_COLOR
            pygame.draw.rect(screen, color, (x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE, CELL_SIZE))
            pygame.draw.rect(screen, (50, 50, 50), (x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE, CELL_SIZE), 1)

# Function to count live neighbors
def count_neighbors(x, y):
    neighbors = 0
    for i in range(-1, 2):
        for j in range(-1, 2):
            if i == 0 and j == 0:
                continue
            nx, ny = (x + i) % GRID_WIDTH, (y + j) % GRID_HEIGHT
            neighbors += grid[ny, nx]
    return neighbors

# Function to update the grid based on the rules
def update_grid():
    global grid
    new_grid = np.copy(grid)
    
    for y in range(GRID_HEIGHT):
        for x in range(GRID_WIDTH):
            neighbors = count_neighbors(x, y)
            if grid[y, x] == 1:  # Live cell
                if neighbors < 2 or neighbors > 3:
                    new_grid[y, x] = 0  # Cell dies
            else:  # Dead cell
                if neighbors == 3:
                    new_grid[y, x] = 1  # Cell becomes alive

    grid = new_grid

# Function to handle mouse click events to toggle cells
def toggle_cell(x, y):
    grid[y, x] = 1 - grid[y, x]

# Main loop
def main():
    running = True
    clock = pygame.time.Clock()

    while running:
        screen.fill((0, 0, 0))  # Clear the screen
        draw_grid()  # Draw the grid

        # Event handling
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                x, y = event.pos
                grid_x, grid_y = x // CELL_SIZE, y // CELL_SIZE
                toggle_cell(grid_x, grid_y)

        # Update the grid
        update_grid()

        # Update the display
        pygame.display.flip()

        # Control the frame rate
        clock.tick(FPS)
        time.sleep(0.1)  # Slow down the simulation slightly for better viewing

    pygame.quit()

if __name__ == "__main__":
    main()