import pygame
import sys
import random
import math
import numpy as np

# Initialize Pygame
pygame.init()
pygame.mixer.init(frequency=22050, size=-16, channels=2, buffer=512)

# Set up the display
width, height = 800, 600
screen = pygame.display.set_mode((width, height), pygame.RESIZABLE)
pygame.display.set_caption("Asteroids")

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Ship
ship_size = 20
ship_pos = [float(width // 2), float(height // 2)]
ship_angle = 0
ship_speed = 0

# Bullets
bullets = []
bullet_speed = 10

# Asteroids
asteroids = []
asteroid_count = 10

# Particles
particles = []

# Sound generation functions
def generate_sine_wave(frequency, duration, sample_rate=22050, volume=0.3):
    frames = int(duration * sample_rate)
    arr = np.zeros((frames, 2), dtype=np.int16)
    for i in range(frames):
        t = float(i) / sample_rate
        wave = volume * 32767.0 * np.sin(2 * np.pi * frequency * t)
        arr[i] = [wave, wave]
    return pygame.sndarray.make_sound(arr)

def generate_laser_sound():
    sample_rate = 22050
    duration = 0.1
    frames = int(duration * sample_rate)
    arr = np.zeros((frames, 2), dtype=np.int16)
    
    for i in range(frames):
        t = float(i) / sample_rate
        # Descending frequency sweep
        frequency = 800 * (1 - t / duration * 0.5)
        envelope = np.exp(-t * 20)  # Quick decay
        wave = envelope * 0.2 * 32767.0 * np.sin(2 * np.pi * frequency * t)
        arr[i] = [wave, wave]
    
    return pygame.sndarray.make_sound(arr)

def generate_explosion_sound():
    sample_rate = 22050
    duration = 0.3
    frames = int(duration * sample_rate)
    arr = np.zeros((frames, 2), dtype=np.int16)
    
    for i in range(frames):
        t = float(i) / sample_rate
        envelope = np.exp(-t * 5)  # Exponential decay
        # White noise with low-pass filter effect
        noise = np.random.normal(0, 1) * envelope * 0.15 * 32767.0
        arr[i] = [noise, noise]
    
    return pygame.sndarray.make_sound(arr)

def generate_thrust_sound():
    sample_rate = 22050
    duration = 0.4
    frames = int(duration * sample_rate)
    arr = np.zeros((frames, 2), dtype=np.int16)
    
    for i in range(frames):
        t = float(i) / sample_rate
        # Smooth low rumble without noise
        wave = (0.15 * 32767.0 * np.sin(2 * np.pi * 30 * t) +
                0.08 * 32767.0 * np.sin(4 * np.pi * 30 * t) +
                0.04 * 32767.0 * np.sin(6 * np.pi * 30 * t))
        arr[i] = [wave, wave]
    
    return pygame.sndarray.make_sound(arr)

def generate_background_rumble():
    sample_rate = 22050
    duration = 2.0
    frames = int(duration * sample_rate)
    arr = np.zeros((frames, 2), dtype=np.int16)
    
    for i in range(frames):
        t = float(i) / sample_rate
        # Very low frequency with random variations
        frequency = 20 + np.random.normal(0, 2)
        wave = 0.05 * 32767.0 * np.sin(2 * np.pi * frequency * t)
        arr[i] = [wave, wave]
    
    return pygame.sndarray.make_sound(arr)

# Create sound effects
laser_sound = generate_laser_sound()
explosion_sound = generate_explosion_sound()
thrust_sound = generate_thrust_sound()
background_rumble = generate_background_rumble()

# Sound channels
thrust_channel = pygame.mixer.Channel(0)
effects_channel = pygame.mixer.Channel(1)
background_channel = pygame.mixer.Channel(2)

# Background rumble disabled - uncomment to enable
# background_channel.play(background_rumble, loops=-1)
# background_channel.set_volume(0.2)

# Score
score = 0
font = pygame.font.Font(None, 36)

# Game over flag
game_over = False

def create_asteroid():
    size = random.randint(20, 50)
    x = random.choice([random.randint(0, width // 4), random.randint(3 * width // 4, width)])
    y = random.choice([random.randint(0, height // 4), random.randint(3 * height // 4, height)])
    speed = random.uniform(0.5, 2)
    angle = random.uniform(0, 2 * math.pi)
    sides = random.choice([5, 6])
    return {"pos": [float(x), float(y)], "size": size, "speed": speed, "angle": angle, "sides": sides}

def create_explosion(x, y, color=WHITE):
    for _ in range(20):
        angle = random.uniform(0, 2 * math.pi)
        speed = random.uniform(1, 5)
        size = random.randint(1, 4)
        lifetime = random.randint(20, 40)
        particles.append({
            "pos": [float(x), float(y)],
            "vel": [math.cos(angle) * speed, math.sin(angle) * speed],
            "size": size,
            "lifetime": lifetime,
            "max_lifetime": lifetime,
            "color": color
        })

def create_thrust_particles():
    if ship_speed > 0:
        for _ in range(3):
            angle = ship_angle + math.pi + random.uniform(-0.5, 0.5)
            speed = random.uniform(2, 4)
            particles.append({
                "pos": [ship_pos[0] - math.cos(ship_angle) * ship_size,
                       ship_pos[1] + math.sin(ship_angle) * ship_size],
                "vel": [math.cos(angle) * speed, math.sin(angle) * speed],
                "size": random.randint(2, 3),
                "lifetime": random.randint(10, 20),
                "max_lifetime": random.randint(10, 20),
                "color": (255, 165, 0)
            })

def draw_polygon(surface, color, pos, size, sides, angle=0):
    points = []
    for i in range(sides):
        angle_point = angle + (2 * math.pi * i / sides)
        point = (pos[0] + size * math.cos(angle_point),
                 pos[1] + size * math.sin(angle_point))
        points.append(point)
    pygame.draw.polygon(surface, color, points, 2)

# Create initial asteroids
for _ in range(asteroid_count):
    asteroids.append(create_asteroid())

# Game loop
clock = pygame.time.Clock()

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        elif event.type == pygame.VIDEORESIZE:
            width, height = event.w, event.h
            screen = pygame.display.set_mode((width, height), pygame.RESIZABLE)
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_q or event.key == pygame.K_ESCAPE:
                pygame.quit()
                sys.exit()
            elif event.key == pygame.K_SPACE and not game_over:
                bullets.append([ship_pos[0], ship_pos[1], ship_angle])
                effects_channel.play(laser_sound)

    if not game_over:
        # Ship movement
        keys = pygame.key.get_pressed()
        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            ship_angle += 0.1
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            ship_angle -= 0.1
        if keys[pygame.K_UP] or keys[pygame.K_w]:
            ship_speed = min(ship_speed + 0.1, 5)
        elif keys[pygame.K_s]:
            ship_speed = max(ship_speed - 0.15, 0)
        else:
            ship_speed = max(ship_speed - 0.05, 0)

        ship_pos[0] += math.cos(ship_angle) * ship_speed
        ship_pos[1] -= math.sin(ship_angle) * ship_speed

        # Wrap ship around screen
        if ship_pos[0] < 0:
            ship_pos[0] += width
        elif ship_pos[0] > width:
            ship_pos[0] -= width
        
        if ship_pos[1] < 0:
            ship_pos[1] += height
        elif ship_pos[1] > height:
            ship_pos[1] -= height

        # Update bullets
        for bullet in bullets[:]:
            bullet[0] += math.cos(bullet[2]) * bullet_speed
            bullet[1] -= math.sin(bullet[2]) * bullet_speed
            if bullet[0] < 0 or bullet[0] > width or bullet[1] < 0 or bullet[1] > height:
                bullets.remove(bullet)

        # Update particles
        for particle in particles[:]:
            particle["pos"][0] += particle["vel"][0]
            particle["pos"][1] += particle["vel"][1]
            particle["lifetime"] -= 1
            if particle["lifetime"] <= 0:
                particles.remove(particle)

        # Create thrust particles
        create_thrust_particles()
        
        # Handle thrust sound
        if ship_speed > 0:
            if not thrust_channel.get_busy():
                thrust_channel.play(thrust_sound, loops=-1)
                thrust_channel.set_volume(0.2)
        else:
            thrust_channel.stop()

        # Update asteroids
        for asteroid in asteroids[:]:
            asteroid["pos"][0] += math.cos(asteroid["angle"]) * asteroid["speed"]
            asteroid["pos"][1] += math.sin(asteroid["angle"]) * asteroid["speed"]

            # Wrap asteroids around screen
            if asteroid["pos"][0] < 0:
                asteroid["pos"][0] += width
            elif asteroid["pos"][0] > width:
                asteroid["pos"][0] -= width
            
            if asteroid["pos"][1] < 0:
                asteroid["pos"][1] += height
            elif asteroid["pos"][1] > height:
                asteroid["pos"][1] -= height

            # Check collision with ship
            if (abs(asteroid["pos"][0] - ship_pos[0]) < asteroid["size"] and
                abs(asteroid["pos"][1] - ship_pos[1]) < asteroid["size"]):
                create_explosion(ship_pos[0], ship_pos[1], (255, 0, 0))
                effects_channel.play(explosion_sound)
                game_over = True

            # Check collision with bullets
            for bullet in bullets[:]:
                if (abs(asteroid["pos"][0] - bullet[0]) < asteroid["size"] and
                    abs(asteroid["pos"][1] - bullet[1]) < asteroid["size"]):
                    create_explosion(asteroid["pos"][0], asteroid["pos"][1])
                    effects_channel.play(explosion_sound)
                    asteroids.remove(asteroid)
                    bullets.remove(bullet)
                    score += 1
                    asteroids.append(create_asteroid())
                    break

    # Draw everything
    screen.fill(BLACK)

    if not game_over:
        # Draw ship with better design
        # Main body
        points = [
            (ship_pos[0] + math.cos(ship_angle) * ship_size,
             ship_pos[1] - math.sin(ship_angle) * ship_size),
            (ship_pos[0] + math.cos(ship_angle + 2.5) * ship_size * 0.7,
             ship_pos[1] - math.sin(ship_angle + 2.5) * ship_size * 0.7),
            (ship_pos[0] + math.cos(ship_angle + 3.14) * ship_size * 0.3,
             ship_pos[1] - math.sin(ship_angle + 3.14) * ship_size * 0.3),
            (ship_pos[0] + math.cos(ship_angle - 2.5) * ship_size * 0.7,
             ship_pos[1] - math.sin(ship_angle - 2.5) * ship_size * 0.7),
        ]
        pygame.draw.polygon(screen, WHITE, points, 2)
        
        # Cockpit
        cockpit_pos = (
            ship_pos[0] + math.cos(ship_angle) * ship_size * 0.3,
            ship_pos[1] - math.sin(ship_angle) * ship_size * 0.3
        )
        pygame.draw.circle(screen, WHITE, (int(cockpit_pos[0]), int(cockpit_pos[1])), 3)
        
        # Engine glow when thrusting
        if ship_speed > 0:
            engine_pos = (
                ship_pos[0] - math.cos(ship_angle) * ship_size * 0.8,
                ship_pos[1] + math.sin(ship_angle) * ship_size * 0.8
            )
            pygame.draw.circle(screen, (255, 165, 0), (int(engine_pos[0]), int(engine_pos[1])), 4)

        # Draw bullets
        for bullet in bullets:
            pygame.draw.circle(screen, WHITE, (int(bullet[0]), int(bullet[1])), 2)
        
        # Draw particles
        for particle in particles:
            alpha = particle["lifetime"] / particle["max_lifetime"]
            color = tuple(max(0, min(255, int(c * alpha))) for c in particle["color"])
            pygame.draw.circle(screen, color, 
                             (int(particle["pos"][0]), int(particle["pos"][1])), 
                             particle["size"])

    # Draw asteroids
    for asteroid in asteroids:
        draw_polygon(screen, WHITE, asteroid["pos"], asteroid["size"], asteroid["sides"])

    # Draw score
    score_text = font.render(f"Score: {score}", True, WHITE)
    screen.blit(score_text, (10, 10))

    if game_over:
        game_over_text = font.render("GAME OVER", True, WHITE)
        screen.blit(game_over_text, (width // 2 - 70, height // 2 - 18))

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