import pygame
import random
import math

# 初始化
pygame.init()
WIDTH, HEIGHT = 900, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🎆 Python 动态烟花")

clock = pygame.time.Clock()

# 颜色
COLORS = [
    (255, 50, 50),
    (50, 255, 50),
    (50, 50, 255),
    (255, 255, 50),
    (255, 50, 255),
    (50, 255, 255)
]

# 粒子类
class Particle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        angle = random.uniform(0, 2 * math.pi)
        speed = random.uniform(2, 7)
        self.vx = math.cos(angle) * speed
        self.vy = math.sin(angle) * speed
        self.life = random.randint(40, 80)
        self.color = color

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

    def draw(self):
        alpha = max(self.life / 80, 0)
        r, g, b = self.color
        color = (int(r * alpha), int(g * alpha), int(b * alpha))
        pygame.draw.circle(screen, color, (int(self.x), int(self.y)), 2)

# 烟花类
class Firework:
    def __init__(self, x, y):
        self.particles = []
        color = random.choice(COLORS)
        for _ in range(120):
            self.particles.append(Particle(x, y, color))

    def update(self):
        for p in self.particles:
            p.update()
        self.particles = [p for p in self.particles if p.life > 0]

    def draw(self):
        for p in self.particles:
            p.draw()

# 主循环
fireworks = []
running = True

while running:
    clock.tick(60)
    screen.fill((10, 10, 20))

    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, my))

    # 自动烟花
    if random.random() < 0.03:
        fireworks.append(Firework(random.randint(100, WIDTH - 100), HEIGHT))

    for fw in fireworks:
        fw.update()
        fw.draw()

    fireworks = [fw for fw in fireworks if fw.particles]

    pygame.display.flip()

pygame.quit()