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 动态烟花盛宴")

# 定义颜色 (RGB格式)
COLORS = [
    (255, 0, 0), (255, 165, 0), (255, 255, 0), 
    (0, 255, 0), (0, 255, 255), (0, 0, 255), 
    (255, 0, 255), (255, 255, 255)
]

# 烟花粒子类 (爆炸后的小火花)
class Particle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.color = color
        self.radius = random.randint(1, 3)
        # 随机向各个方向爆炸
        angle = random.uniform(0, 2 * math.pi)
        speed = random.uniform(1, 6)
        self.vx = math.cos(angle) * speed
        self.vy = math.sin(angle) * speed
        self.lifetime = random.randint(40, 80)  # 粒子存活时间
        self.alpha = 255  # 透明度

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += 0.08  # 模拟重力
        self.lifetime -= 1
        # 随着存活时间减少，逐渐变透明
        self.alpha = max(0, self.lifetime * 4)

    def draw(self, surface):
        if self.lifetime > 0:
            # 创建一个带透明度的临时 surface 来绘制粒子
            particle_surface = pygame.Surface((self.radius * 2, self.radius * 2), pygame.SRCALPHA)
            particle_color_with_alpha = (*self.color, self.alpha)

            surface.blit(particle_surface, (int(self.x - self.radius), int(self.y - self.radius)))

    def is_alive(self):
        return self.lifetime > 0

# 烟花类 (升空阶段)
class Firework:
    def __init__(self):
        self.x = random.randint(100, WIDTH - 100)
        self.y = HEIGHT
        self.target_y = random.randint(100, HEIGHT // 2) # 随机爆炸高度
        self.color = random.choice(COLORS)
        self.speed = random.randint(4, 7)
        self.exploded = False
        self.particles = []

    def update(self):
        if not self.exploded:
            self.y -= self.speed
            # 到达目标高度后爆炸
            if self.y <= self.target_y:
                self.exploded = True
                # 生成爆炸粒子
                for _ in range(100):
                    self.particles.append(Particle(self.x, self.y, self.color))
        else:
            # 更新所有粒子
            for particle in self.particles:
                particle.update()
            # 移除已经消失的粒子
            self.particles = [p for p in self.particles if p.is_alive()]

    def draw(self, surface):
        if not self.exploded:
            # 绘制升空的烟花
            pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), 4)
        else:
            # 绘制爆炸的粒子
            for particle in self.particles:
                particle.draw(surface)

    def is_finished(self):
        # 烟花爆炸且所有粒子都消失后，判定为结束
        return self.exploded and len(self.particles) == 0

# 主程序循环
def main():
    clock = pygame.time.Clock()
    fireworks = []
    running = True

    while running:
        # 处理退出事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # 用半透明黑色填充背景，制造出烟花的拖尾残影效果
        screen.fill((0, 0, 0))
        
        # 随机生成新的烟花
        if random.random() < 0.03:  # 约每30帧生成一个
            fireworks.append(Firework())

        # 更新和绘制所有烟花
        for firework in fireworks[:]:
            firework.update()
            firework.draw(screen)
            if firework.is_finished():
                fireworks.remove(firework)

        pygame.display.flip()  # 刷新屏幕
        clock.tick(60)  # 限制帧率为60 FPS

    pygame.quit()

if __name__ == "__main__":
    main()