import pygame
import sys
import random

# Initialize Pygame
pygame.init()

# Constants
WIDTH, HEIGHT = 600, 400
FPS = 60
BALL_RADIUS = 20
ENEMY_RADIUS = 15
PLAYER_SPEED = 5
ENEMY_SPEED = 3

# Colors
WHITE = (255, 255, 255)
RED = (255, 0, 0)
BLUE = (0, 0, 255)

# Create the screen
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Avoid the Balls")

# Clock to control the frame rate
clock = pygame.time.Clock()

# Player ball
player_pos = [WIDTH // 2, HEIGHT // 2]
player_speed = [0, 0]

# Enemy balls
enemies = [{'pos': [random.randint(0, WIDTH), random.randint(0, HEIGHT)],
            'speed': [random.choice([-1, 1]) * ENEMY_SPEED, random.choice([-1, 1]) * ENEMY_SPEED]}
           for _ in range(5)]

def draw_player(pos):
    pygame.draw.circle(screen, BLUE, (int(pos[0]), int(pos[1])), BALL_RADIUS)

def draw_enemy(pos):
    pygame.draw.circle(screen, RED, (int(pos[0]), int(pos[1])), ENEMY_RADIUS)

def main():
    global player_pos, player_speed

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

        keys = pygame.key.get_pressed()
        player_speed = [0, 0]

        if keys[pygame.K_LEFT] and player_pos[0] > BALL_RADIUS:
            player_speed[0] = -PLAYER_SPEED
        if keys[pygame.K_RIGHT] and player_pos[0] < WIDTH - BALL_RADIUS:
            player_speed[0] = PLAYER_SPEED
        if keys[pygame.K_UP] and player_pos[1] > BALL_RADIUS:
            player_speed[1] = -PLAYER_SPEED
        if keys[pygame.K_DOWN] and player_pos[1] < HEIGHT - BALL_RADIUS:
            player_speed[1] = PLAYER_SPEED

        # Update player position
        player_pos[0] += player_speed[0]
        player_pos[1] += player_speed[1]

        # Update enemy positions
        for enemy in enemies:
            enemy['pos'][0] += enemy['speed'][0]
            enemy['pos'][1] += enemy['speed'][1]

            # Check for collisions with player
            distance = pygame.math.Vector2(player_pos[0] - enemy['pos'][0], player_pos[1] - enemy['pos'][1]).length()
            if distance < BALL_RADIUS + ENEMY_RADIUS:
                print("Game Over!")
                pygame.quit()
                sys.exit()

            # Wrap enemies around the screen
            if enemy['pos'][0] < 0 or enemy['pos'][0] > WIDTH:
                enemy['pos'][0] = WIDTH - enemy['pos'][0]
            if enemy['pos'][1] < 0 or enemy['pos'][1] > HEIGHT:
                enemy['pos'][1] = HEIGHT - enemy['pos'][1]

        # Draw everything
        screen.fill(WHITE)
        draw_player(player_pos)
        for enemy in enemies:
            draw_enemy(enemy['pos'])

        pygame.display.flip()
        clock.tick(FPS)

if __name__ == "__main__":
    main()