import pygame
import random
import math

pygame.init()

WIDTH, HEIGHT = 1600, 900
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("星空烟花 - 多造型随机版")
clock = pygame.time.Clock()

# 爆炸音效
try:
    explosion_sound = pygame.mixer.Sound(buffer=
        b'RIFF$\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00\x80>\x00\x00\x80>\x00\x00\x01\x00\x08\x00data\x00\x00\x00\x00'
        b'\x80\x80\x80\x80\x85\x87\x89\x8a\x87\x80y&\xd0\x90\xf0\x08\xe0\x02\xcc\x00\x80\x00\x00\x80\x00\x00\x80\x00\x00\x80'
    )
    explosion_sound.set_volume(0.8)
except:
    explosion_sound = None

# 真实星空
stars = []
for _ in range(800):
    x = random.randint(0, WIDTH)
    y = random.randint(0, HEIGHT)
    size = random.uniform(0.3, 1.8)
    base_bright = random.randint(60, 140)
    flicker_speed = random.uniform(0.01, 0.03)
    stars.append([x, y, size, base_bright, flicker_speed, 0])

# 长条粒子
class Particle:
    def __init__(self, x, y, vx, vy, color):
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.angle = math.atan2(vy, vx)
        self.color = color
        self.life = random.randint(70, 140)
        self.max_life = self.life
        self.length = random.uniform(8, 18)
        self.width = random.uniform(1, 2.5)
        self.gravity = 0.12

    def update(self):
        self.vy += self.gravity
        self.x += self.vx
        self.y += self.vy
        self.vx *= 0.98
        self.vy *= 0.98
        self.life -= 1
        self.length = max(2, self.length - 0.1)

    def draw(self, surface):
        if self.life > 0:
            alpha = int(255 * (self.life / self.max_life))
            end_x = self.x + math.cos(self.angle) * self.length
            end_y = self.y + math.sin(self.angle) * self.length
            pygame.draw.line(surface, (*self.color, alpha), (self.x, self.y), (end_x, end_y), int(self.width))

# 烟花 + 多造型爆炸
class Firework:
    def __init__(self, x):
        self.x = x
        self.y = HEIGHT
        self.target_y = random.randint(120, 450)
        self.speed = random.uniform(6, 9)
        color_list = [
            (255,0,0),(255,140,0),(255,255,0),
            (0,255,0),(0,220,255),(0,100,255),
            (200,0,255),(255,0,220),(255,255,255)
        ]
        self.color = random.choice(color_list)
        self.exploded = False
        self.particles = []
        self.trail_points = []
        self.shape = random.choice(['circle','heart','ring','cross','star','spiral'])

    def update(self):
        if not self.exploded:
            self.y -= self.speed
            self.trail_points.append((self.x, self.y))
            if len(self.trail_points) > 60:
                self.trail_points.pop(0)
            if self.y <= self.target_y:
                self.explode()

    def explode(self):
        self.exploded = True
        if explosion_sound: explosion_sound.play()
        s = self.shape
        x, y, c = self.x, self.y, self.color

        # 圆形
        if s == 'circle':
            for _ in range(200):
                a = random.uniform(0, math.pi*2)
                r = random.uniform(4,9)
                self.particles.append(Particle(x,y, math.cos(a)*r, math.sin(a)*r, c))

        # 圆环
        elif s == 'ring':
            for i in range(200):
                a = i/200 * math.pi*2
                r = 7
                self.particles.append(Particle(x,y, math.cos(a)*r, math.sin(a)*r, c))

        # 十字
        elif s == 'cross':
            for a in [0, math.pi/2, math.pi, math.pi*3/2]:
                for d in range(8):
                    r = 2 + d*0.8
                    self.particles.append(Particle(x,y, math.cos(a)*r, math.sin(a)*r, c))

        # 爱心
        elif s == 'heart':
            for t in [i*0.1 for i in range(100)]:
                a = 2*math.pi*t
                rx = 16 * math.sin(a)**3
                ry = -13 * math.cos(a) +5*math.cos(2*a)+2*math.cos(3*a)+math.cos(4*a)
                self.particles.append(Particle(x,y, rx*0.3, ry*0.3, c))

        # 五角星
        elif s == 'star':
            for i in range(5):
                a1 = i*math.pi*2/5
                a2 = (i+0.5)*math.pi*2/5
                for d in range(12):
                    r = d*0.7
                    self.particles.append(Particle(x,y, math.cos(a1)*r, math.sin(a1)*r, c))
                    self.particles.append(Particle(x,y, math.cos(a2)*r*0.4, math.sin(a2)*r*0.4, c))

        # 螺旋
        elif s == 'spiral':
            for i in range(250):
                a = i*0.2
                r = 1 + i*0.04
                self.particles.append(Particle(x,y, math.cos(a)*r, math.sin(a)*r, c))

    def draw(self, surface):
        if not self.exploded:
            for i, (tx, ty) in enumerate(self.trail_points):
                alpha = int(200 * i/len(self.trail_points))
                w = max(1, int(3*i/len(self.trail_points)))
                s = pygame.Surface((w,w), pygame.SRCALPHA)
                pygame.draw.circle(s, (*self.color, alpha), (w//2,w//2), w//2)
                surface.blit(s, (int(tx), int(ty)))
            pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), 4)

fireworks = []
running = True

while running:
    screen.fill((0,0,10))
    for star in stars:
        x,y,size,b,spd,t = star
        t += spd
        br = b + math.sin(t)*25
        pygame.draw.circle(screen, (br,br,br), (int(x),int(y)), int(size))
        star[5] = t

    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            running = False
        if e.type == pygame.MOUSEBUTTONDOWN:
            mx, my = pygame.mouse.get_pos()
            fireworks.append(Firework(mx))

    for fw in fireworks[:]:
        fw.update()
        fw.draw(screen)
        for p in fw.particles[:]:
            p.update()
            p.draw(screen)
            if p.life <= 0:
                fw.particles.remove(p)
        if fw.exploded and not fw.particles:
            fireworks.remove(fw)

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

pygame.quit()