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

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

class Particle:
    """单个烟花粒子"""
    def __init__(self, x, y, color=None):
        self.x = x
        self.y = y
        # 随机生成速度向量
        angle = random.uniform(0, 2 * math.pi)
        speed = random.uniform(2, 8)
        self.vx = math.cos(angle) * speed
        self.vy = math.sin(angle) * speed
        # 生命周期 (帧数)
        self.life = random.randint(40, 80)
        self.max_life = self.life
        # 颜色
        if color:
            self.color = color
        else:
            self.color = (
                random.randint(150, 255),
                random.randint(100, 200),
                random.randint(50, 150)
            )
        # 大小
        self.size = random.uniform(2, 4)
        # 重力影响
        self.gravity = 0.05

    def update(self):
        """更新粒子状态"""
        self.x += self.vx
        self.y += self.vy
        self.vy += self.gravity  # 模拟重力
        self.life -= 1
        # 逐渐变小
        self.size *= 0.98

    def draw(self, screen):
        """绘制粒子"""
        if self.life > 0:
            # 根据剩余生命计算透明度
            alpha = int(255 * (self.life / self.max_life))
            # 创建带透明度的颜色
            color_with_alpha = (*self.color, alpha)
            # 由于 pygame.draw.circle 不支持 alpha，我们创建一个临时 surface
            surf = pygame.Surface((int(self.size*3), int(self.size*3)), pygame.SRCALPHA)
            pygame.draw.circle(surf, color_with_alpha, (int(self.size*1.5), int(self.size*1.5)), int(self.size))
            screen.blit(surf, (int(self.x - self.size*1.5), int(self.y - self.size*1.5)))

    def is_dead(self):
        return self.life <= 0 or self.y > HEIGHT + 20


class Firework:
    """烟花主体 (上升阶段)"""
    def __init__(self):
        self.x = random.randint(100, WIDTH - 100)
        self.y = HEIGHT
        self.target_y = random.randint(50, HEIGHT // 2)
        self.speed = random.uniform(4, 7)
        self.exploded = False
        self.particles = []
        # 烟花上升时的颜色 (尾部光效)
        self.trail_color = (
            random.randint(180, 255),
            random.randint(120, 220),
            random.randint(60, 160)
        )

    def update(self):
        """更新上升状态"""
        if not self.exploded:
            # 上升运动
            self.y -= self.speed
            # 轻微水平飘动
            self.x += random.uniform(-0.5, 0.5)
            # 到达目标高度则爆炸
            if self.y <= self.target_y:
                self.explode()
        else:
            # 更新所有粒子
            for p in self.particles[:]:
                p.update()
                if p.is_dead():
                    self.particles.remove(p)

    def explode(self):
        """爆炸产生粒子"""
        self.exploded = True
        # 生成 80~150 个粒子
        num_particles = random.randint(80, 150)
        # 主色调
        main_color = (
            random.randint(100, 255),
            random.randint(80, 230),
            random.randint(50, 200)
        )
        for _ in range(num_particles):
            # 部分粒子颜色微调
            offset = random.randint(-30, 30)
            color = (
                max(0, min(255, main_color[0] + offset)),
                max(0, min(255, main_color[1] + offset)),
                max(0, min(255, main_color[2] + offset))
            )
            self.particles.append(Particle(self.x, self.y, color))

    def draw(self, screen):
        """绘制烟花"""
        if not self.exploded:
            # 绘制上升轨迹 (小圆点)
            pygame.draw.circle(screen, self.trail_color, (int(self.x), int(self.y)), 3)
            # 拖尾效果: 画几个更暗的点
            for i in range(1, 4):
                pygame.draw.circle(
                    screen,
                    (max(0, self.trail_color[0]-40*i),
                     max(0, self.trail_color[1]-40*i),
                     max(0, self.trail_color[2]-40*i)),
                    (int(self.x - i*1.5), int(self.y + i*3)),
                    max(1, 3-i)
                )
        else:
            for p in self.particles:
                p.draw(screen)

    def is_dead(self):
        """判断烟花是否完全消失"""
        if not self.exploded:
            return False
        return len(self.particles) == 0


def main():
    running = True
    fireworks = []
    # 控制烟花生成间隔
    firework_timer = 0

    while running:
        screen.fill(BLACK)

        # 事件处理
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                # 点击鼠标左键发射烟花
                if event.button == 1:
                    fw = Firework()
                    fw.x = event.pos[0]
                    fw.y = HEIGHT
                    fw.target_y = random.randint(50, HEIGHT//2)
                    fireworks.append(fw)

        # 自动生成烟花 (每 20~40 帧)
        firework_timer += 1
        if firework_timer >= random.randint(20, 45):
            fireworks.append(Firework())
            firework_timer = 0

        # 更新所有烟花
        for fw in fireworks[:]:
            fw.update()
            if fw.is_dead():
                fireworks.remove(fw)

        # 绘制所有烟花
        for fw in fireworks:
            fw.draw(screen)

        # 显示当前烟花数量 (调试用，可注释掉)
        font = pygame.font.Font(None, 24)
        text = font.render(f"烟花数: {len(fireworks)}", True, WHITE)
        screen.blit(text, (10, 10))

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

    pygame.quit()


if __name__ == "__main__":
    main()