import pygame
import random
import math
import colorsys

# 初始化pygame
pygame.init()

# 设置窗口尺寸
WIDTH, HEIGHT = 1000, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("动态烟花模拟")

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
BACKGROUND = (5, 5, 20)  # 深蓝色背景，类似夜空

# 烟花粒子类
class Particle:
    def __init__(self, x, y, color=None, is_trail=False):
        self.x = x
        self.y = y
        self.color = color if color else self.random_color()
        self.size = random.uniform(2, 4) if not is_trail else random.uniform(1, 2)
        self.speed = random.uniform(2, 8)
        self.angle = random.uniform(0, 2 * math.pi)
        self.vx = math.cos(self.angle) * self.speed
        self.vy = math.sin(self.angle) * self.speed
        self.gravity = 0.1
        self.life = 100  # 粒子寿命
        self.decay = random.uniform(0.5, 1.5)  # 生命衰减速度
        self.trail = is_trail
        self.alpha = 255  # 透明度

    def random_color(self):
        # 生成鲜艳的颜色
        h = random.random()
        s = 0.8 + random.random() * 0.2
        v = 0.8 + random.random() * 0.2
        r, g, b = colorsys.hsv_to_rgb(h, s, v)
        return (int(r * 255), int(g * 255), int(b * 255))

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += self.gravity
        self.life -= self.decay
        self.alpha = max(0, int(self.life * 2.55))
        
        # 减慢粒子的水平速度
        self.vx *= 0.99
        
        return self.life > 0

    def draw(self, surface):
        if self.alpha > 0:
            # 创建一个带有透明度的表面
            particle_surface = pygame.Surface((int(self.size * 2), int(self.size * 2)), pygame.SRCALPHA)
            if self.trail:
                # 轨迹粒子
                pygame.draw.circle(particle_surface, 
                                  (*self.color, self.alpha), 
                                  (int(self.size), int(self.size)), 
                                  int(self.size))
            else:
                # 爆炸粒子
                pygame.draw.circle(particle_surface, 
                                  (*self.color, self.alpha), 
                                  (int(self.size), int(self.size)), 
                                  int(self.size))
                # 添加光晕效果
                pygame.draw.circle(particle_surface, 
                                  (*self.color, self.alpha // 3), 
                                  (int(self.size), int(self.size)), 
                                  int(self.size * 1.5))
            
            surface.blit(particle_surface, (int(self.x - self.size), int(self.y - self.size)))

# 烟花类
class Firework:
    def __init__(self, x=None, y=None, target_y=None):
        self.x = x if x else random.randint(50, WIDTH - 50)
        self.y = y if y else HEIGHT
        self.target_y = target_y if target_y else random.randint(50, HEIGHT // 2)
        self.speed = random.uniform(3, 6)
        self.color = self.random_color()
        self.particles = []
        self.exploded = False
        self.trail_particles = []
        self.size = 3
        
    def random_color(self):
        h = random.random()
        s = 0.9 + random.random() * 0.1
        v = 0.9 + random.random() * 0.1
        r, g, b = colorsys.hsv_to_rgb(h, s, v)
        return (int(r * 255), int(g * 255), int(b * 255))
    
    def update(self):
        if not self.exploded:
            # 更新烟花上升轨迹
            self.y -= self.speed
            
            # 添加轨迹粒子
            if random.random() < 0.7:
                self.trail_particles.append(Particle(
                    self.x + random.uniform(-2, 2), 
                    self.y + random.uniform(-2, 2), 
                    self.color, 
                    is_trail=True
                ))
            
            # 检查是否到达爆炸点
            if self.y <= self.target_y:
                self.explode()
                
        # 更新轨迹粒子
        for i in range(len(self.trail_particles)-1, -1, -1):
            if not self.trail_particles[i].update():
                del self.trail_particles[i]
        
        # 更新爆炸粒子
        for i in range(len(self.particles)-1, -1, -1):
            if not self.particles[i].update():
                del self.particles[i]
                
        return len(self.particles) > 0 or not self.exploded
    
    def explode(self):
        self.exploded = True
        # 创建爆炸粒子
        num_particles = random.randint(50, 150)
        for _ in range(num_particles):
            self.particles.append(Particle(self.x, self.y, self.color))
    
    def draw(self, surface):
        # 绘制轨迹粒子
        for particle in self.trail_particles:
            particle.draw(surface)
        
        # 绘制爆炸粒子
        for particle in self.particles:
            particle.draw(surface)
        
        # 如果还没有爆炸，绘制上升的烟花
        if not self.exploded:
            # 绘制烟花主体
            pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), self.size)
            # 绘制烟花光晕
            pygame.draw.circle(surface, (*self.color, 100), (int(self.x), int(self.y)), self.size * 2)

# 星星背景类
class Star:
    def __init__(self):
        self.x = random.randint(0, WIDTH)
        self.y = random.randint(0, HEIGHT)
        self.size = random.uniform(0.1, 1.5)
        self.brightness = random.uniform(0.3, 1.0)
        self.twinkle_speed = random.uniform(0.01, 0.05)
        self.twinkle_offset = random.uniform(0, math.pi * 2)
        
    def update(self):
        # 星星闪烁效果
        self.brightness = 0.5 + 0.5 * math.sin(pygame.time.get_ticks() * self.twinkle_speed + self.twinkle_offset)
        
    def draw(self, surface):
        color_value = int(200 * self.brightness)
        color = (color_value, color_value, color_value)
        pygame.draw.circle(surface, color, (int(self.x), int(self.y)), self.size)

# 创建烟花列表
fireworks = []
# 创建星星背景
stars = [Star() for _ in range(100)]

# 字体设置
font = pygame.font.SysFont(None, 28)
big_font = pygame.font.SysFont(None, 48)

# 控制变量
clock = pygame.time.Clock()
running = True
last_firework_time = 0
firework_interval = 800  # 毫秒
auto_launch = True
show_instructions = True
instruction_time = pygame.time.get_ticks()

# 主游戏循环
while running:
    current_time = pygame.time.get_ticks()
    
    # 处理事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                # 空格键手动发射烟花
                fireworks.append(Firework())
            elif event.key == pygame.K_a:
                # A键切换自动发射
                auto_launch = not auto_launch
            elif event.key == pygame.K_c:
                # C键清空所有烟花
                fireworks.clear()
            elif event.key == pygame.K_ESCAPE:
                running = False
        elif event.type == pygame.MOUSEBUTTONDOWN:
            # 鼠标点击发射烟花
            if event.button == 1:  # 左键
                x, y = event.pos
                fireworks.append(Firework(x=x, y=HEIGHT, target_y=y))
            elif event.button == 3:  # 右键
                # 在鼠标位置直接爆炸
                x, y = event.pos
                fw = Firework(x=x, y=y, target_y=y)
                fw.exploded = True
                fw.explode()
                fireworks.append(fw)
    
    # 自动发射烟花
    if auto_launch and current_time - last_firework_time > firework_interval:
        fireworks.append(Firework())
        last_firework_time = current_time
        # 随机化下一次发射间隔
        firework_interval = random.randint(500, 1500)
    
    # 更新星星
    for star in stars:
        star.update()
    
    # 更新烟花
    for i in range(len(fireworks)-1, -1, -1):
        if not fireworks[i].update():
            del fireworks[i]
    
    # 绘制背景
    screen.fill(BACKGROUND)
    
    # 绘制星星
    for star in stars:
        star.draw(screen)
    
    # 绘制烟花
    for firework in fireworks:
        firework.draw(screen)
    
    # 显示标题
    title_text = big_font.render("动态烟花模拟", True, (255, 255, 200))
    screen.blit(title_text, (WIDTH // 2 - title_text.get_width() // 2, 20))
    
    # 显示控制说明
    if show_instructions:
        instructions = [
            "控制说明:",
            "空格键: 发射随机烟花",
            "鼠标左键: 在指定高度发射烟花",
            "鼠标右键: 在点击位置直接爆炸",
            "A键: 切换自动发射 (" + ("开" if auto_launch else "关") + ")",
            "C键: 清除所有烟花",
            "ESC键: 退出程序"
        ]
        
        for i, line in enumerate(instructions):
            text = font.render(line, True, (200, 220, 255))
            screen.blit(text, (20, 80 + i * 30))
        
        # 5秒后隐藏说明
        if current_time - instruction_time > 10000:
            show_instructions = False
    else:
        # 显示简化的控制提示
        hint_text = font.render("按A键显示控制说明", True, (150, 150, 200))
        screen.blit(hint_text, (20, HEIGHT - 40))
    
    # 显示烟花计数
    count_text = font.render(f"当前烟花数量: {len(fireworks)}", True, (200, 255, 200))
    screen.blit(count_text, (WIDTH - count_text.get_width() - 20, 20))
    
    # 显示自动发射状态
    auto_text = font.render(f"自动发射: {'开' if auto_launch else '关'}", True, (255, 200, 200) if auto_launch else (200, 200, 200))
    screen.blit(auto_text, (WIDTH - auto_text.get_width() - 20, 50))
    
    # 显示提示
    hint_text = font.render("点击或按空格键发射更多烟花！", True, (255, 255, 200))
    screen.blit(hint_text, (WIDTH // 2 - hint_text.get_width() // 2, HEIGHT - 30))
    
    # 更新显示
    pygame.display.flip()
    
    # 控制帧率
    clock.tick(60)

# 退出pygame
pygame.quit()