import pygame
import random
import math

# 初始化pygame
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Python动态烟花")
clock = pygame.time.Clock()
FPS = 60

# 颜色常量
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

# 随机生成烟花彩色
def random_color():
    return (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))

# 烟花粒子类
class Particle:
    def __init__(self, x, y, color, speed, angle):
        self.x = x
        self.y = y
        self.color = color
        # 速度分解
        self.vx = math.cos(angle) * speed
        self.vy = math.sin(angle) * speed
        self.life = random.randint(40, 80)  # 粒子存活帧数
        self.gravity = 0.08  # 重力下坠

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += self.gravity
        self.life -= 1
        # 颜色随生命变淡
        r, g, b = self.color
        self.color = (max(0, r-3), max(0, g-3), max(0, b-3))

    def draw(self):
        pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), 2)

# 升空炮弹类
class Firework:
    def __init__(self):
        self.x = random.randint(100, WIDTH - 100)
        self.y = HEIGHT
        self.speed = random.randint(7, 10)
        self.target_y = random.randint(80, 250)  # 爆炸高度
        self.color = random_color()
        self.exploded = False
        self.particles = []

    def update(self):
        if not self.exploded:
            # 向上飞行
            self.y -= self.speed
            if self.y <= self.target_y:
                self.explode()
        else:
            # 更新爆炸粒子
            for p in self.particles:
                p.update()
            # 移除死亡粒子
            self.particles = [p for p in self.particles if p.life > 0]

    def explode(self):
        self.exploded = True
        # 生成一圈扩散粒子
        particle_num = random.randint(80, 120)
        for i in range(particle_num):
            angle = math.radians(i * (360 / particle_num))
            speed = random.uniform(2, 6)
            self.particles.append(Particle(self.x, self.y, self.color, speed, angle))

    def draw(self):
        if not self.exploded:
            # 绘制升空光点
            pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), 3)
        else:
            for p in self.particles:
                p.draw()

# 主程序列表存储所有烟花
fireworks = []
launch_timer = 0  # 发射计时器

running = True
while running:
    # 事件监听
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # 黑色夜空背景（半透明拖尾，营造光晕）
    screen.fill((0, 0, 0, 30), special_flags=pygame.BLEND_RGBA_MIN)

    # 定时自动发射烟花
    launch_timer += 1
    if launch_timer > 25:
        fireworks.append(Firework())
        launch_timer = 0

    # 更新并绘制所有烟花
    for fw in fireworks:
        fw.update()
        fw.draw()

    # 清除已经爆炸完毕的烟花
    fireworks = [fw for fw in fireworks if not (fw.exploded and len(fw.particles) == 0)]

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

pygame.quit()