import pygame
import random
import math

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

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

# 粒子类（烟花炸开后的碎片）
class Particle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.color = color
        # 随机速度
        speed = random.uniform(2, 6)
        angle = random.uniform(0, 2 * math.pi)
        self.vx = math.cos(angle) * speed
        self.vy = math.sin(angle) * speed
        self.life = 100  # 存活帧数
        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-2), max(0, g-2), max(0, b-2))

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

# 烟花火箭类（升空的引线）
class Firework:
    def __init__(self, target_x):
        self.x = random.randint(100, WIDTH - 100)
        self.y = HEIGHT
        self.target_x = target_x
        self.target_y = random.randint(80, HEIGHT // 2)
        self.speed = 6
        self.exploded = False
        self.particles = []
        # 随机烟花颜色
        self.color = (
            random.randint(100, 255),
            random.randint(100, 255),
            random.randint(100, 255)
        )

    def update(self):
        if not self.exploded:
            # 飞向目标点
            dx = self.target_x - self.x
            dy = self.target_y - self.y
            dist = math.hypot(dx, dy)
            if dist < self.speed:
                # 到达目标，爆炸生成粒子
                self.explode()
            else:
                self.x += dx / dist * self.speed
                self.y += dy / dist * self.speed
        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
        # 生成80个爆炸粒子
        for _ in range(80):
            self.particles.append(Particle(self.x, self.y, self.color))

    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 = []

# 主循环
running = True
while running:
    # 半透明拖影，实现残影效果
    screen.fill((0, 0, 0, 20), special_flags=pygame.BLEND_RGBA_MIN)

    # 事件监听
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        # 鼠标点击发射烟花
        if event.type == pygame.MOUSEBUTTONDOWN:
            mx, my = pygame.mouse.get_pos()
            fireworks.append(Firework(mx))

    # 更新并绘制所有烟花
    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()