import pygame
import random

# Initialize pygame
pygame.init()

# Screen dimensions
WIDTH, HEIGHT = 600, 400
CELL_SIZE = 20

# Colors
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLACK = (0, 0, 225)

# Set up the screen
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Snake Game")

# Font for score
tfont = pygame.font.Font(None, 36)

# Snake setup
snake = [(WIDTH // 2, HEIGHT // 2)]
direction = (CELL_SIZE, 0)

# Apple setup
num_apples = 5
apples = [(random.randint(0, WIDTH // CELL_SIZE - 1) * CELL_SIZE, 
           random.randint(0, HEIGHT // CELL_SIZE - 1) * CELL_SIZE) for _ in range(num_apples)]

# Score
score = 1

clock = pygame.time.Clock()

running = True
while running:
    screen.fill(BLACK)
    
    # Event handling
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP and direction != (0, CELL_SIZE):
                direction = (0, -CELL_SIZE)
            elif event.key == pygame.K_DOWN and direction != (0, -CELL_SIZE):
                direction = (0, CELL_SIZE)
            elif event.key == pygame.K_LEFT and direction != (CELL_SIZE, 0):
                direction = (-CELL_SIZE, 0)
            elif event.key == pygame.K_RIGHT and direction != (-CELL_SIZE, 0):
                direction = (CELL_SIZE, 0)
    
    # Move the snake
    new_head = (snake[0][0] + direction[0], snake[0][1] + direction[1])
    
    # Check for collisions with walls or itself
    if (new_head[0] < 0 or new_head[0] >= WIDTH or
        new_head[1] < 0 or new_head[1] >= HEIGHT or
        new_head in snake):
        running = False
    
    snake.insert(0, new_head)
    
    # Check if the snake eats an apple
    if new_head in apples:
        apples.remove(new_head)
        apples.append((random.randint(0, WIDTH // CELL_SIZE - 1) * CELL_SIZE, 
                       random.randint(0, HEIGHT // CELL_SIZE - 1) * CELL_SIZE))
        score += 1
    else:
        snake.pop()
    
    # Draw the snake
    for segment in snake:
        pygame.draw.rect(screen, GREEN, (*segment, CELL_SIZE, CELL_SIZE))
    
    # Draw the apples
    for apple in apples:
        pygame.draw.rect(screen, RED, (*apple, CELL_SIZE, CELL_SIZE))
    
    # Display score
    score_text = tfont.render(f"Score: {score}", True, WHITE)
    screen.blit(score_text, (10, 10))
    
    pygame.display.flip()
    clock.tick(10)

pygame.quit()