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)

# 随机生成烟花彩色
def random_color():
    colors = [
        (255, 50, 50), (50, 255, 50), (50, 50, 255),
        (255, 255, 50), (255, 50, 255), (50, 255, 255),
        (255, 165, 0), (220, 100, 220), (100, 220, 220)
    ]
    return random.choice(colors)

# 烟花粒子类（爆炸后的碎片）
class Particle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.color = color
        # 随机角度与速度
        angle = random.uniform(0, math.pi * 2)
        speed = random.uniform(2, 7)
        self.vx = math.cos(angle) * speed
        self.vy = math.sin(angle) * speed
        self.life = 255  # 透明度
        self.gravity = 0.08  # 重力

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += self.gravity
        self.life -= 4  # 逐渐变暗

    def draw(self):
        if self.life > 0:
            # 带透明度绘制粒子
            surface = pygame.Surface((4, 4), pygame.SRCALPHA)
            pygame.draw.circle(surface, (*self.color, self.life), (2, 2), 2)
            screen.blit(surface, (self.x, self.y))

# 升空火箭类
class Rocket:
    def __init__(self, target_x, target_y):
        self.x = random.randint(100, WIDTH - 100)
        self.y = HEIGHT
        self.target_x = target_x
        self.target_y = target_y
        self.speed = 6
        self.color = random_color()
        self.exploded = False
        self.particles = []

    def update(self):
        # 向目标点飞行
        dx = self.target_x - self.x
        dy = self.target_y - self.y
        dist = math.hypot(dx, dy)

        if dist < self.speed:
            # 到达目标，炸开
            self.exploded = True
            # 生成一堆爆炸粒子
            for _ in range(80):
                self.particles.append(Particle(self.target_x, self.target_y, self.color))
        else:
            self.x += (dx / dist) * self.speed
            self.y += (dy / dist) * self.speed

        # 更新爆炸粒子
        for p in self.particles:
            p.update()
        # 移除消亡粒子
        self.particles = [p for p in self.particles if p.life > 0]

    def draw(self):
        if not self.exploded:
            # 绘制升空的火箭光点
            pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), 3)
        # 绘制爆炸粒子
        for p in self.particles:
            p.draw()

# 存储所有火箭
rockets = []
running = True

# 主循环
while running:
    # 黑色半透明填充，制造拖尾余晖效果
    bg_surface = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    bg_surface.fill((0, 0, 0, 25))
    screen.blit(bg_surface, (0, 0))

    # 事件监听
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        # 鼠标左键点击发射烟花
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            mx, my = pygame.mouse.get_pos()
            rockets.append(Rocket(mx, my))

    # 更新绘制所有火箭
    for rocket in rockets:
        rocket.update()
        rocket.draw()

    # 【这里是修复的关键一行】过滤清理完全消失的烟花
    rockets = [r for r in rockets if (not r.exploded) or len(r.particles) > 0]

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

pygame.quit()