import pygame
import random
import math
import sys
from pygame.locals import *

pygame.init()

# 屏幕设置
WIDTH, HEIGHT = 1000, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("绚丽烟花模拟 - 按空格键发射烟花 | 按R重置 | 按1-5选择颜色")

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 50, 50)
GREEN = (50, 255, 50)
BLUE = (50, 100, 255)
YELLOW = (255, 255, 50)
PURPLE = (200, 50, 255)
CYAN = (50, 255, 255)
ORANGE = (255, 150, 50)
PINK = (255, 100, 180)
COLORS = [RED, GREEN, BLUE, YELLOW, PURPLE, CYAN, ORANGE, PINK]

# 尝试加载字体，如果失败则创建简单字体
def create_font(size):
    try:
        # 尝试使用系统默认字体
        font = pygame.font.Font(None, size)
        return font
    except:
        # 如果失败，创建简单的位图字体
        return None

# 创建简单文本渲染函数
def draw_text(surface, text, x, y, color=WHITE, size=24, centered=False):
    # 使用默认字体
    font = pygame.font.Font(None, size)
    text_surface = font.render(text, True, color)
    if centered:
        surface.blit(text_surface, (x - text_surface.get_width() // 2, y - text_surface.get_height() // 2))
    else:
        surface.blit(text_surface, (x, y))
    return text_surface.get_width()

# 烟花粒子类
class Particle:
    def __init__(self, x, y, color, is_trail=False):
        self.x = x
        self.y = y
        self.color = color
        self.size = random.randint(2, 4) if not is_trail else 1
        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.original_color = color
        
    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += self.gravity
        self.life -= self.decay
        
        # 添加空气阻力
        self.vx *= 0.99
        self.vy *= 0.99
        
        # 如果粒子是拖尾，减小大小
        if self.trail and self.life < 50:
            self.size = max(0.5, self.size * 0.97)
            
        # 颜色随生命变化
        if self.life < 30:
            r, g, b = self.color
            r = max(0, r - 5)
            g = max(0, g - 5)
            b = max(0, b - 5)
            self.color = (r, g, b)
            
    def draw(self, surface):
        if self.life > 0:
            # 根据生命值计算透明度
            alpha = min(255, int(self.life * 2.5))
            
            # 创建临时surface绘制带透明度的粒子
            if self.trail:
                # 拖尾粒子更小
                pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), max(1, int(self.size)))
            else:
                # 主爆炸粒子
                radius = max(1, int(self.size))
                pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), radius)
                
                # 添加光晕效果
                if radius > 1:
                    # 创建光晕表面
                    glow_radius = radius * 2
                    glow_surface = pygame.Surface((glow_radius*2, glow_radius*2), pygame.SRCALPHA)
                    # 绘制半透明圆形作为光晕
                    pygame.draw.circle(glow_surface, (*self.color[:3], 50), (glow_radius, glow_radius), glow_radius)
                    # 绘制到主表面
                    surface.blit(glow_surface, (int(self.x - glow_radius), int(self.y - glow_radius)))
    
    def is_dead(self):
        return self.life <= 0

# 烟花类
class Firework:
    def __init__(self, x, y, color=None):
        self.x = x
        self.y = y
        self.color = color if color else random.choice(COLORS)
        self.speed = random.uniform(3, 6)
        self.particles = []
        self.exploded = False
        self.trail = []
        self.max_height = random.randint(50, HEIGHT//2)
        
    def update(self):
        if not self.exploded:
            # 更新位置
            self.y -= self.speed
            self.speed *= 0.98  # 逐渐减速
            
            # 添加拖尾粒子
            if random.random() < 0.7:
                trail_color = (
                    min(255, self.color[0] + 50),
                    min(255, self.color[1] + 50),
                    min(255, self.color[2] + 50)
                )
                self.trail.append(Particle(self.x, self.y, trail_color, is_trail=True))
            
            # 检查是否到达爆炸高度
            if self.speed < 0.5 or self.y < self.max_height:
                self.explode()
        else:
            # 更新爆炸粒子
            for particle in self.particles[:]:
                particle.update()
                if particle.is_dead():
                    self.particles.remove(particle)
        
        # 更新拖尾粒子
        for trail in self.trail[:]:
            trail.update()
            if trail.is_dead():
                self.trail.remove(trail)
    
    def explode(self):
        self.exploded = True
        # 创建爆炸粒子
        num_particles = random.randint(50, 150)
        for _ in range(num_particles):
            # 创建稍微变化的颜色
            r = max(0, min(255, self.color[0] + random.randint(-30, 30)))
            g = max(0, min(255, self.color[1] + random.randint(-30, 30)))
            b = max(0, min(255, self.color[2] + random.randint(-30, 30)))
            particle_color = (r, g, b)
            
            self.particles.append(Particle(self.x, self.y, particle_color))
            
        # 添加第二次爆炸效果
        if random.random() < 0.3:  # 30%几率有第二次爆炸
            for _ in range(random.randint(20, 40)):
                r = max(0, min(255, self.color[0] + random.randint(-50, 50)))
                g = max(0, min(255, self.color[1] + random.randint(-50, 50)))
                b = max(0, min(255, self.color[2] + random.randint(-50, 50)))
                particle_color = (r, g, b)
                
                # 第二次爆炸粒子速度较慢
                particle = Particle(self.x, self.y, particle_color)
                particle.speed = random.uniform(1, 3)
                particle.vx = math.cos(particle.angle) * particle.speed
                particle.vy = math.sin(particle.angle) * particle.speed
                particle.life = random.randint(60, 100)
                self.particles.append(particle)
    
    def draw(self, surface):
        # 绘制拖尾
        for trail in self.trail:
            trail.draw(surface)
            
        if not self.exploded:
            # 绘制上升的烟花
            pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), 3)
            # 添加上升烟花的光晕
            glow_radius = 5
            pygame.draw.circle(surface, (*self.color[:3], 100), (int(self.x), int(self.y)), glow_radius)
        else:
            # 绘制爆炸粒子
            for particle in self.particles:
                particle.draw(surface)
    
    def is_dead(self):
        return self.exploded and len(self.particles) == 0

# 创建背景星星
def create_stars(count=100):
    stars = []
    for _ in range(count):
        x = random.randint(0, WIDTH)
        y = random.randint(0, HEIGHT)
        brightness = random.randint(150, 255)
        size = random.uniform(0.1, 1.5)
        stars.append((x, y, brightness, size))
    return stars

# 绘制UI
def draw_ui(selected_color, fireworks_count, particles_count):
    # 绘制半透明UI背景
    ui_surface = pygame.Surface((WIDTH, 100), pygame.SRCALPHA)
    pygame.draw.rect(ui_surface, (20, 20, 30, 200), (0, 0, WIDTH, 100))
    screen.blit(ui_surface, (0, 0))
    
    # 绘制标题
    draw_text(screen, "绚丽烟花模拟", WIDTH//2, 15, YELLOW, 36, centered=True)
    
    # 绘制控制说明
    draw_text(screen, "空格: 发射烟花 | R: 重置 | 1-8: 选择颜色 | ESC: 退出", WIDTH//2, 55, WHITE, 24, centered=True)
    
    # 绘制颜色选择器
    color_box_size = 30
    for i, color in enumerate(COLORS[:8]):
        x = 20 + i * (color_box_size + 10)
        y = HEIGHT - 50
        pygame.draw.rect(screen, color, (x, y, color_box_size, color_box_size))
        
        # 绘制选中边框
        if i == selected_color:
            pygame.draw.rect(screen, WHITE, (x-2, y-2, color_box_size+4, color_box_size+4), 2)
            
        # 绘制颜色编号
        draw_text(screen, str(i+1), x + color_box_size//2, y + color_box_size//2, WHITE, 20, centered=True)
    
    # 绘制统计数据
    stats_text = f"烟花数量: {fireworks_count} | 粒子总数: {particles_count}"
    stats_width = draw_text(screen, stats_text, WIDTH - 20, HEIGHT - 40, CYAN, 20)
    draw_text(screen, stats_text, WIDTH - stats_width - 20, HEIGHT - 40, CYAN, 20)
    
    # 绘制颜色选择标签
    draw_text(screen, "选择颜色 (1-8):", 20, HEIGHT - 80, WHITE, 20)

# 主函数
def main():
    clock = pygame.time.Clock()
    fireworks = []
    stars = create_stars(150)
    selected_color = 0  # 默认选择第一个颜色
    
    # 创建渐变背景
    background = pygame.Surface((WIDTH, HEIGHT))
    for y in range(HEIGHT):
        # 从深蓝色渐变到黑色
        color_value = int(10 + (y / HEIGHT) * 20)
        pygame.draw.line(background, (color_value, color_value, 30 + color_value), (0, y), (WIDTH, y))
    
    # 自动发射烟花计时器
    auto_launch_timer = 0
    
    running = True
    while running:
        # 事件处理
        for event in pygame.event.get():
            if event.type == QUIT:
                running = False
            elif event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    running = False
                elif event.key == K_SPACE:
                    # 空格键发射烟花
                    x = random.randint(100, WIDTH-100)
                    fireworks.append(Firework(x, HEIGHT, COLORS[selected_color]))
                elif event.key == K_r:
                    # R键重置
                    fireworks.clear()
                elif event.key in (K_1, K_2, K_3, K_4, K_5, K_6, K_7, K_8):
                    # 数字键选择颜色
                    selected_color = event.key - K_1
            elif event.type == MOUSEBUTTONDOWN:
                # 鼠标点击发射烟花
                if event.button == 1:  # 左键
                    fireworks.append(Firework(event.pos[0], HEIGHT, COLORS[selected_color]))
                elif event.button == 3:  # 右键随机颜色
                    fireworks.append(Firework(event.pos[0], HEIGHT))
        
        # 自动发射烟花
        auto_launch_timer += 1
        if auto_launch_timer > 60:  # 每秒自动发射一个
            if random.random() < 0.3:  # 30%几率自动发射
                x = random.randint(100, WIDTH-100)
                fireworks.append(Firework(x, HEIGHT, random.choice(COLORS)))
                auto_launch_timer = 0
        
        # 绘制背景
        screen.blit(background, (0, 0))
        
        # 绘制星星
        for x, y, brightness, size in stars:
            # 星星闪烁效果
            twinkle = brightness + random.randint(-20, 20)
            twinkle = max(150, min(255, twinkle))
            color = (twinkle, twinkle, twinkle)
            pygame.draw.circle(screen, color, (x, y), size)
        
        # 更新和绘制烟花
        particles_count = 0
        for firework in fireworks[:]:
            firework.update()
            firework.draw(screen)
            particles_count += len(firework.particles) + len(firework.trail)
            
            if firework.is_dead():
                fireworks.remove(firework)
        
        # 绘制UI
        draw_ui(selected_color, len(fireworks), particles_count)
        
        # 更新显示
        pygame.display.flip()
        clock.tick(60)
    
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()