import pygame
import random
import math

# -----------------------------
# Settings
# -----------------------------
SCREEN_WIDTH = 1280
SCREEN_HEIGHT = 720
FPS = 80
NUM_ENEMIES = 14

# -----------------------------
# Initialize Pygame
# -----------------------------
pygame.init()
pygame.mixer.init()  # Initialize mixer for music

screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Polarity Field")
clock = pygame.time.Clock()
font = pygame.font.SysFont(None, 36)

# -----------------------------
# Load background music
# -----------------------------
pygame.mixer.music.load("sprites/music.ogg")
pygame.mixer.music.set_volume(0.5)
pygame.mixer.music.play(-1)  # Loop indefinitely

# -----------------------------
# Load sprites AFTER display is ready
# -----------------------------
player_surf = pygame.image.load("sprites/orb_purple.png").convert_alpha()
player_surf = pygame.transform.scale(player_surf, (60, 60))
player_rect = player_surf.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2))
player_polarity = 1
player_lives = 3

enemies = []
for _ in range(NUM_ENEMIES):
    polarity = random.choice([1, -1])
    if polarity == 1:
        surf = pygame.image.load("sprites/orb_red.png").convert_alpha()
    else:
        surf = pygame.image.load("sprites/orb_blue.png").convert_alpha()
    surf = pygame.transform.scale(surf, (50, 50))

    # Spawn 500 pixels away from player
    while True:
        x = random.randint(0, SCREEN_WIDTH)
        y = random.randint(0, SCREEN_HEIGHT)
        dx = x - player_rect.centerx
        dy = y - player_rect.centery
        dist = math.hypot(dx, dy)
        if dist >= 500:
            break

    rect = surf.get_rect(center=(x, y))
    speed_x = 0  # stationary until game starts
    speed_y = 0
    enemies.append({
        "surf": surf,
        "rect": rect,
        "polarity": polarity,
        "speed_x": speed_x,
        "speed_y": speed_y
    })

# -----------------------------
# Score & cooldown
# -----------------------------
score = 0
hit_cooldown = 0
game_started = False  # Game begins only after first movement

# -----------------------------
# Main loop
# -----------------------------
running = True
while running:
    clock.tick(FPS)
    screen.fill((20, 20, 20))

    # -----------------------------
    # Events
    # -----------------------------
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                running = False
            if event.key == pygame.K_SPACE and game_started:
                player_polarity *= -1

    # -----------------------------
    # Player movement
    # -----------------------------
    keys = pygame.key.get_pressed()
    dx = dy = 0
    if keys[pygame.K_w] or keys[pygame.K_UP]:
        dy -= 6
    if keys[pygame.K_s] or keys[pygame.K_DOWN]:
        dy += 6
    if keys[pygame.K_a] or keys[pygame.K_LEFT]:
        dx -= 6
    if keys[pygame.K_d] or keys[pygame.K_RIGHT]:
        dx += 6

    # Start game on first movement
    if not game_started and (dx != 0 or dy != 0):
        game_started = True
        # Activate enemy movement now
        for enemy in enemies:
            enemy["speed_x"] = random.uniform(-3, 3)
            enemy["speed_y"] = random.uniform(-3, 3)

    player_rect.move_ip(dx, dy)
    player_rect.clamp_ip(pygame.Rect(0, 0, SCREEN_WIDTH, SCREEN_HEIGHT))
    moved = dx != 0 or dy != 0

    # -----------------------------
    # Enemy updates & scoring (only if game started)
    # -----------------------------
    if game_started:
        for enemy in enemies:
            # Polarity effect
            if enemy["polarity"] == player_polarity:  # repel
                if enemy["rect"].centerx < player_rect.centerx:
                    enemy["rect"].x -= 4
                else:
                    enemy["rect"].x += 4
                if enemy["rect"].centery < player_rect.centery:
                    enemy["rect"].y -= 4
                else:
                    enemy["rect"].y += 4
            else:  # attract
                if enemy["rect"].centerx < player_rect.centerx:
                    enemy["rect"].x += 4
                else:
                    enemy["rect"].x -= 4
                if enemy["rect"].centery < player_rect.centery:
                    enemy["rect"].y += 4
                else:
                    enemy["rect"].y -= 4

            # Random drift
            enemy["rect"].x += enemy["speed_x"] * 0.3
            enemy["rect"].y += enemy["speed_y"] * 0.3

            # Bounce on screen edges
            if enemy["rect"].left < 0 or enemy["rect"].right > SCREEN_WIDTH:
                enemy["speed_x"] *= -1
            if enemy["rect"].top < 0 or enemy["rect"].bottom > SCREEN_HEIGHT:
                enemy["speed_y"] *= -1

            # Prevent stacking
            for other in enemies:
                if other != enemy:
                    dx_sep = enemy["rect"].centerx - other["rect"].centerx
                    dy_sep = enemy["rect"].centery - other["rect"].centery
                    dist_sep = math.hypot(dx_sep, dy_sep)
                    min_dist = 30
                    if dist_sep < min_dist and dist_sep != 0:
                        push_x = (dx_sep / dist_sep) * (min_dist - dist_sep) / 2
                        push_y = (dy_sep / dist_sep) * (min_dist - dist_sep) / 2
                        enemy["rect"].move_ip(push_x, push_y)
                        other["rect"].move_ip(-push_x, -push_y)

            # Collision
            if hit_cooldown == 0 and player_rect.colliderect(enemy["rect"]):
                player_lives -= 1
                hit_cooldown = FPS
                for e in enemies:
                    dx_push = e["rect"].centerx - player_rect.centerx
                    dy_push = e["rect"].centery - player_rect.centery
                    e["rect"].move_ip(dx_push//2, dy_push//2)

            # Draw enemy
            screen.blit(enemy["surf"], enemy["rect"])

        # Update score
        if moved:
            score += 1

        # Update cooldown
        if hit_cooldown > 0:
            hit_cooldown -= 1

    # -----------------------------
    # Draw player
    # -----------------------------
    screen.blit(player_surf, player_rect)

    # -----------------------------
    # Draw HUD
    # -----------------------------
    hud = font.render(f"Lives: {player_lives}  Score: {score}", True, (255, 255, 255))
    screen.blit(hud, (20, 20))

    # -----------------------------
    # Check for game over
    # -----------------------------
    if player_lives <= 0:
        running = False

    pygame.display.flip()

# -----------------------------
# Game Over screen
# -----------------------------
screen.fill((20, 20, 20))
game_over_text = font.render("GAME OVER", True, (255, 0, 0))
score_text = font.render(f"Final Score: {score}", True, (255, 255, 255))

# Center text
screen.blit(game_over_text, (SCREEN_WIDTH//2 - game_over_text.get_width()//2, SCREEN_HEIGHT//2 - 40))
screen.blit(score_text, (SCREEN_WIDTH//2 - score_text.get_width()//2, SCREEN_HEIGHT//2 + 10))

pygame.display.flip()

# Wait until the player closes the window
game_over = True
while game_over:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            game_over = False

pygame.quit()
