import pygame
import random
import math

# 初始化
pygame.init()
pygame.mixer.init()

# 全屏获取屏幕分辨率
screen_info = pygame.display.Info()
W, H = screen_info.current_w, screen_info.current_h
screen = pygame.display.set_mode((W, H), pygame.FULLSCREEN)
pygame.display.set_caption("顶级动态烟花")
clock = pygame.time.Clock()
FPS = 60

# 超高颜值霓虹渐变色系
FIRE_COLORS = [
    (255, 30, 30), (255, 120, 20), (255, 220, 20),
    (80, 255, 80), (20, 180, 255), (130, 80, 255),
    (255, 80, 220), (255, 255, 255), (0, 255, 255)
]

# 烟花粒子类
class Spark:
    def __init__(self, x, y, color, speed_angle, speed_len):
        self.x = x
        self.y = y
        self.vx = math.cos(speed_angle) * speed_len
        self.vy = math.sin(speed_angle) * speed_len
        self.color = color
        self.life = 120
        self.max_life = 120
        self.gravity = 0.06
        self.fade = 1.8

    def update(self):
        self.vy += self.gravity
        self.x += self.vx
        self.y += self.vy
        self.life -= 1

    def draw(self, surf):
        alpha = int((self.life / self.max_life) * 255)
        r, g, b = self.color
        color = (r, g, b, alpha)
        size = max(1, self.life // 25)
        temp = pygame.Surface((size*2, size*2), pygame.SRCALPHA)
        pygame.draw.circle(temp, color, (size, size), size)
        surf.blit(temp, (self.x-size, self.y-size))

# 升空烟花弹
class Rocket:
    def __init__(self):
        self.x = random.randint(80, W-80)
        self.y = H
        self.speed = random.uniform(-7.5, -5.5)
        self.target_h = random.randint(80, H//2)
        self.color = random.choice(FIRE_COLORS)
        self.explode = False
        self.sparks = []
        self.trail = []

    def update(self):
        if not self.explode:
            self.trail.append((self.x, self.y))
            if len(self.trail) > 15:
                self.trail.pop(0)
            self.y += self.speed
            self.speed += 0.12
            if self.y <= self.target_h or self.speed >= -1:
                self.burst()
        else:
            alive = []
            for s in self.sparks:
                s.update()
                if s.life > 0:
                    alive.append(s)
            self.sparks = alive

    def burst(self):
        self.explode = True
        count = random.randint(120, 220)
        for i in range(count):
            angle = math.pi * 2 * i / count + random.uniform(-0.2, 0.2)
            speed = random.uniform(2, 6)
            self.sparks.append(Spark(self.x, self.y, self.color, angle, speed))

    def draw(self, surf):
        if not self.explode:
            # 绘制升空拖尾
            for idx, (tx, ty) in enumerate(self.trail):
                a = int(60 + idx*12)
                pygame.draw.circle(surf, (*self.color, a), (int(tx), int(ty)), 2)
            pygame.draw.circle(surf, self.color, (int(self.x), int(self.y)), 4)
        else:
            for s in self.sparks:
                s.draw(surf)

def main():
    rockets = []
    run = True
    while run:
        clock.tick(FPS)
        # 纯黑夜空背景
        screen.fill((0, 0, 8))

        # 随机自动发射烟花
        if random.randint(0, 18) == 0:
            rockets.append(Rocket())

        # 事件监听 ESC退出
        for e in pygame.event.get():
            if e.type == pygame.QUIT:
                run = False
            if e.type == pygame.KEYDOWN and e.key == pygame.K_ESCAPE:
                run = False

        # 更新渲染所有烟花
        live_rockets = []
        for r in rockets:
            r.update()
            r.draw(screen)
            if not r.explode or len(r.sparks) > 0:
                live_rockets.append(r)
        rockets = live_rockets

        pygame.display.update()
    pygame.quit()

if __name__ == "__main__":
    main()