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()

# 颜色随机函数
def random_color():
    return (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))

# 烟花弹（上升阶段）
class Rocket:
    def __init__(self, x):
        self.x = x
        self.y = HEIGHT
        self.speed_y = random.uniform(-7, -11)
        self.speed_x = random.uniform(-1, 1)
        self.color = random_color()
        self.explode_y = random.randint(120, 300)
        self.exploded = False

    def update(self):
        self.x += self.speed_x
        self.y += self.speed_y
        # 到达指定高度爆炸
        if self.y <= self.explode_y:
            self.exploded = True

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

# 爆炸粒子
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, 6)
        self.vx = math.cos(angle) * speed
        self.vy = math.sin(angle) * speed
        self.life = 255  # 透明度

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += 0.12  # 重力
        self.life -= 4

    def draw(self):
        if self.life > 0:
            c = (self.color[0], self.color[1], self.color[2], self.life)
            surf = pygame.Surface((6, 6), pygame.SRCALPHA)
            pygame.draw.circle(surf, c, (3, 3), 3)
            screen.blit(surf, (int(self.x)-3, int(self.y)-3))

rockets = []
particles = []

running = True
while running:
    # 黑色背景
    screen.fill((0, 0, 0))
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        # 鼠标点击发射烟花
        if event.type == pygame.MOUSEBUTTONDOWN:
            mx, my = pygame.mouse.get_pos()
            rockets.append(Rocket(mx))

    # 更新烟花弹
    new_particles = []
    for rocket in rockets[:]:
        rocket.update()
        rocket.draw()
        if rocket.exploded:
            rockets.remove(rocket)
            # 生成爆炸粒子
            color = rocket.color
            for _ in range(60):
                particles.append(Particle(rocket.x, rocket.y, color))

    # 更新粒子
    for p in particles[:]:
        p.update()
        p.draw()
        if p.life <= 0:
            particles.remove(p)

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

pygame.quit()