import pygame
import random
import math

SCREEN_WIDTH = 1280
SCREEN_HEIGHT = 720
FPS = 60

# =========================
# DIFFICULTY SETTINGS
# =========================
difficulty = 2  # 1 = Easy, 2 = Medium, 3 = Hard
difficulty_flash_timer = 0
difficulty_cooldown_timer = 0  # 10-second cooldown

# =========================
# POLARITY SWITCH COOLDOWN
# =========================
polarity_cooldown = 0
POLARITY_COOLDOWN_TIME = int(FPS * 0.7)
polarity_flash_timer = 0

pygame.init()
pygame.mixer.init()

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

pygame.mixer.music.load("sprites/music.ogg")
pygame.mixer.music.set_volume(0.5)
pygame.mixer.music.play(-1)

player_surf = pygame.image.load("sprites/orb_purple.png").convert_alpha()
player_surf = pygame.transform.scale(player_surf, (50, 50))
player_rect = player_surf.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2))
player_polarity = 1
player_lives = 3

# =========================
# INITIAL ENEMIES BASED ON DIFFICULTY (easier)
# =========================
if difficulty == 1:
    NUM_ENEMIES = 12
    difficulty_text = "Easy"
    difficulty_color = (0, 200, 0)
elif difficulty == 2:
    NUM_ENEMIES = 22
    difficulty_text = "Medium"
    difficulty_color = (255, 200, 0)
elif difficulty == 3:
    NUM_ENEMIES = 35
    difficulty_text = "Hard"
    difficulty_color = (255, 60, 60)

# =========================
# SPAWN ENEMIES LIST
# =========================
enemies = []
for _ in range(NUM_ENEMIES):
    polarity = random.choice([1, -1])
    surf = pygame.image.load("sprites/orb_red.png" if polarity == 1 else "sprites/orb_blue.png").convert_alpha()
    surf = pygame.transform.scale(surf, (40, 40))
    side = random.choice(["top", "bottom", "left", "right"])
    if side == "top":
        x = random.randint(0, SCREEN_WIDTH)
        y = -100
    elif side == "bottom":
        x = random.randint(0, SCREEN_WIDTH)
        y = SCREEN_HEIGHT + 100
    elif side == "left":
        x = -100
        y = random.randint(0, SCREEN_HEIGHT)
    else:
        x = SCREEN_WIDTH + 100
        y = random.randint(0, SCREEN_HEIGHT)
    rect = surf.get_rect(center=(x, y))
    enemies.append({
        "surf": surf,
        "rect": rect,
        "polarity": polarity,
        "speed_x": random.uniform(-3.5, 3.5),
        "speed_y": random.uniform(-3.5, 3.5)
    })

score = 0
hit_cooldown = 0
game_started = False
difficulty_timer = 0

running = True
while running:
    clock.tick(FPS)
    screen.fill((20, 20, 20))

    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

            # POLARITY SWITCH
            if event.key == pygame.K_SPACE and game_started and polarity_cooldown == 0:
                player_polarity *= -1
                polarity_cooldown = POLARITY_COOLDOWN_TIME
                polarity_flash_timer = int(FPS * 0.2)

            # DIFFICULTY CHANGE
            if game_started and difficulty_cooldown_timer == 0:
                if event.key in [pygame.K_1, pygame.K_2, pygame.K_3]:
                    if event.key == pygame.K_1:
                        difficulty = 1
                    elif event.key == pygame.K_2:
                        difficulty = 2
                    elif event.key == pygame.K_3:
                        difficulty = 3

                    # Easier versions
                    if difficulty == 1:
                        NUM_ENEMIES = 12
                        difficulty_text = "Easy"
                        difficulty_color = (0, 200, 0)
                    elif difficulty == 2:
                        NUM_ENEMIES = 22
                        difficulty_text = "Medium"
                        difficulty_color = (255, 200, 0)
                    elif difficulty == 3:
                        NUM_ENEMIES = 35
                        difficulty_text = "Hard"
                        difficulty_color = (255, 60, 60)

                    # RESET ENEMIES
                    enemies = []
                    for _ in range(NUM_ENEMIES):
                        polarity = random.choice([1, -1])
                        surf = pygame.image.load("sprites/orb_red.png" if polarity == 1 else "sprites/orb_blue.png").convert_alpha()
                        surf = pygame.transform.scale(surf, (40, 40))
                        side = random.choice(["top", "bottom", "left", "right"])
                        if side == "top":
                            x = random.randint(0, SCREEN_WIDTH)
                            y = -100
                        elif side == "bottom":
                            x = random.randint(0, SCREEN_WIDTH)
                            y = SCREEN_HEIGHT + 100
                        elif side == "left":
                            x = -100
                            y = random.randint(0, SCREEN_HEIGHT)
                        else:
                            x = SCREEN_WIDTH + 100
                            y = random.randint(0, SCREEN_HEIGHT)
                        rect = surf.get_rect(center=(x, y))
                        enemies.append({
                            "surf": surf,
                            "rect": rect,
                            "polarity": polarity,
                            "speed_x": random.uniform(-3.5, 3.5),
                            "speed_y": random.uniform(-3.5, 3.5)
                        })

                    difficulty_flash_timer = FPS // 2
                    difficulty_cooldown_timer = FPS * 10

    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

    if not game_started and (dx or dy):
        game_started = True

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

    if game_started:
        for enemy in enemies:
            # Polarity-based movement
            if enemy["polarity"] == player_polarity:
                enemy["rect"].x += -4 if enemy["rect"].centerx < player_rect.centerx else 3
                enemy["rect"].y += -4 if enemy["rect"].centery < player_rect.centery else 3
            else:
                enemy["rect"].x += 4 if enemy["rect"].centerx < player_rect.centerx else -3
                enemy["rect"].y += 4 if enemy["rect"].centery < player_rect.centery else -3

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

            # RESPAWN IF OFF SCREEN (+/-100 pixels)
            if (enemy["rect"].right < -100 or enemy["rect"].left > SCREEN_WIDTH + 100 or
                    enemy["rect"].bottom < -100 or enemy["rect"].top > SCREEN_HEIGHT + 100):
                side = random.choice(["top", "bottom", "left", "right"])
                if side == "top":
                    x = random.randint(0, SCREEN_WIDTH)
                    y = -100
                elif side == "bottom":
                    x = random.randint(0, SCREEN_WIDTH)
                    y = SCREEN_HEIGHT + 100
                elif side == "left":
                    x = -100
                    y = random.randint(0, SCREEN_HEIGHT)
                else:
                    x = SCREEN_WIDTH + 100
                    y = random.randint(0, SCREEN_HEIGHT)
                enemy["rect"].center = (x, y)
                enemy["speed_x"] = random.uniform(-3.5, 3.5)
                enemy["speed_y"] = random.uniform(-3.5, 3.5)

            # Collision with player
            if hit_cooldown == 0 and player_rect.colliderect(enemy["rect"]):
                player_lives -= 1
                hit_cooldown = FPS

                # RESET ALL ENEMIES WHEN HIT
                enemies = []
                for _ in range(NUM_ENEMIES):
                    polarity = random.choice([1, -1])
                    surf = pygame.image.load("sprites/orb_red.png" if polarity == 1 else "sprites/orb_blue.png").convert_alpha()
                    surf = pygame.transform.scale(surf, (40, 40))
                    side = random.choice(["top", "bottom", "left", "right"])
                    if side == "top":
                        x = random.randint(0, SCREEN_WIDTH)
                        y = -100
                    elif side == "bottom":
                        x = random.randint(0, SCREEN_WIDTH)
                        y = SCREEN_HEIGHT + 100
                    elif side == "left":
                        x = -100
                        y = random.randint(0, SCREEN_HEIGHT)
                    else:
                        x = SCREEN_WIDTH + 100
                        y = random.randint(0, SCREEN_HEIGHT)
                    rect = surf.get_rect(center=(x, y))
                    enemies.append({
                        "surf": surf,
                        "rect": rect,
                        "polarity": polarity,
                        "speed_x": random.uniform(-3.5, 3.5),
                        "speed_y": random.uniform(-3.5, 3.5)
                    })
                break

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

        # Score multiplier based on difficulty
        score_multiplier = 1 + (difficulty - 1) * 0.5
        if moved:
            score += score_multiplier

        if hit_cooldown > 0:
            hit_cooldown -= 1

    # TIMERS
    if difficulty_flash_timer > 0:
        difficulty_flash_timer -= 1
    if polarity_cooldown > 0:
        polarity_cooldown -= 1
    if polarity_flash_timer > 0:
        polarity_flash_timer -= 1
    if difficulty_cooldown_timer > 0:
        difficulty_cooldown_timer -= 1

    # DRAW PLAYER WITH FLASH EFFECT
    if polarity_flash_timer > 0:
        flash_color = (255, 255, 255)
        flash_surf = player_surf.copy()
        flash_surf.fill(flash_color, special_flags=pygame.BLEND_ADD)
        screen.blit(flash_surf, player_rect)
    else:
        screen.blit(player_surf, player_rect)

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

    # Difficulty flash effect
    if difficulty_flash_timer > 0 and difficulty_flash_timer % 10 < 5:
        diff_color = (255, 255, 255)
    else:
        diff_color = difficulty_color

    difficulty_hud = font.render(f"Difficulty: {difficulty_text}", True, diff_color)
    screen.blit(difficulty_hud, (20, 60))

    # Show difficulty cooldown timer
    if difficulty_cooldown_timer > 0:
        seconds_left = difficulty_cooldown_timer // FPS
        cooldown_hud = font.render(f"Change in: {seconds_left}s", True, (200, 200, 255))
        screen.blit(cooldown_hud, (20, 100))

    # Show player polarity at top right (- for Red, + for Blue)
    polarity_text = "-" if player_polarity == 1 else "+"
    polarity_hud = font.render(f"Polarity: {polarity_text}", True, (255, 255, 255))
    screen.blit(polarity_hud, (SCREEN_WIDTH - polarity_hud.get_width() - 20, 20))

    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: {int(score)}", True, (255, 255, 255))

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()

game_over = True
while game_over:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            game_over = False

pygame.quit()