import pygame
import random
import math
from pygame import gfxdraw

# 初始化
pygame.init()
pygame.display.set_caption("🎆 Python 动态烟花")

# 屏幕设置
WIDTH, HEIGHT = 900, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()
FPS = 60

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

NEON_COLORS = [
    (255, 50, 50),    # 红
    (50, 255, 50),    # 绿
    (50, 150, 255),   # 蓝
    (255, 255, 50),   # 黄
    (255, 50, 255),   # 紫
    (255, 150, 50),   # 橙
]

# ------------------ 粒子系统 ------------------

class Particle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.color = color
        self.alpha = 255
        self.size = random.randint(2, 4)

        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.gravity = 0.05
        self.trail = []

    def update(self):
        self.trail.append((self.x, self.y))
        if len(self.trail) > 10:
            self.trail.pop(0)

        self.vy += self.gravity
        self.x += self.vx
        self.y += self.vy
        self.alpha -= 3

    def draw(self, surface):
        # 拖尾
        for i, (tx, ty) in enumerate(self.trail):
            alpha = int(self.alpha * (i / len(self.trail)))
            pygame.draw.circle(
                surface,
                (*self.color, alpha),
                (int(tx), int(ty)),
                max(1, self.size - 1)
            )

        # 粒子本体
        if self.alpha > 0:
            pygame.draw.circle(
                surface,
                (*self.color, self.alpha),
                (int(self.x), int(self.y)),
                self.size
            )

    def is_dead(self):
        return self.alpha <= 0


class Firework:
    def __init__(self, x, target_y, color):
        self.x = x
        self.y = HEIGHT
        self.target_y = target_y
        self.color = color
        self.speed = 6
        self.exploded = False
        self.particles = []

    def update(self):
        if not self.exploded:
            self.y -= self.speed
            if self.y <= self.target_y:
                self.explode()
        else:
            self.particles = [p for p in self.particles if not p.is_dead()]
            for p in self.particles:
                p.update()

    def explode(self):
        self.exploded = True
        for _ in range(80):
            self.particles.append(Particle(self.x, self.y, self.color))

    def draw(self, surface):
        if not self.exploded:
            pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), 4)
        else:
            for p in self.particles:
                p.draw(surface)

    def is_dead(self):
        return self.exploded and len(self.particles) == 0


# ------------------ 主程序 ------------------

def draw_background(surface):
    surface.fill((10, 10, 30))


def main():
    fireworks = []
    running = True

    while running:
        clock.tick(FPS)
        draw_background(screen)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            elif event.type == pygame.MOUSEBUTTONDOWN:
                x, y = pygame.mouse.get_pos()
                color = random.choice(NEON_COLORS)
                fireworks.append(Firework(x, y, color))

            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    x = random.randint(100, WIDTH - 100)
                    y = random.randint(50, HEIGHT // 2)
                    color = random.choice(NEON_COLORS)
                    fireworks.append(Firework(x, y, color))

        # 更新 & 绘制
        fireworks = [fw for fw in fireworks if not fw.is_dead()]
        for fw in fireworks:
            fw.update()
            fw.draw(screen)

        pygame.display.flip()

    pygame.quit()


if __name__ == "__main__":
    main()