import pygame
import random
import math

# --- 初始化设置 ---
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Python 形状烟花 (心形/五角星)")

# --- 常量定义 ---
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
SHAPE_CIRCLE = 0
SHAPE_STAR = 1
SHAPE_HEART = 2

# --- 烟花类 ---
class Firework:
    def __init__(self, shape_type=SHAPE_CIRCLE):
        self.x = random.randint(100, WIDTH - 100)
        self.y = HEIGHT
        self.color = (random.randint(100, 255), random.randint(100, 255), random.randint(100, 255))
        self.vx = 0
        self.vy = random.uniform(-14, -10)
        self.exploded = False
        self.particles = []
        self.shape_type = shape_type

    def update(self):
        if not self.exploded:
            self.x += self.vx
            self.y += self.vy
            self.vy += 0.18 # 重力
            
            # 到达最高点爆炸
            if self.vy >= 0:
                self.explode()
                return True
        else:
            for p in self.particles[:]:
                p.update()
                if p.life <= 0:
                    self.particles.remove(p)
            if len(self.particles) == 0:
                return True
        return False

    def explode(self):
        self.exploded = True
        
        # --- 核心逻辑：根据形状生成粒子 ---
        if self.shape_type == SHAPE_CIRCLE:
            self._create_circle_particles()
        elif self.shape_type == SHAPE_STAR:
            self._create_star_particles()
        elif self.shape_type == SHAPE_HEART:
            self._create_heart_particles()

    def _create_circle_particles(self):
        # 圆形：随机角度
        count = 100
        for _ in range(count):
            angle = random.uniform(0, 2 * math.pi)
            speed = random.uniform(2, 6)
            vx = math.cos(angle) * speed
            vy = math.sin(angle) * speed
            self.particles.append(Particle(self.x, self.y, self.color, vx, vy))

    def _create_star_particles(self):
        # 五角星：利用三角函数计算五个角的顶点，并在角之间填充粒子
        points = 5
        count = 150
        # 外层半径（角尖）和内层半径（凹陷处）
        outer_r = 7
        inner_r = 3.5 
        
        for i in range(count):
            # 这里的逻辑是：将360度分成5份，每一份里随机一个角度
            # 并且让半径在 inner_r 和 outer_r 之间跳动，形成星形
            sector = i % points
            angle_step = (2 * math.pi) / points
            angle = sector * angle_step + random.uniform(-0.2, 0.2) # 增加一点随机扩散
            
            # 简单的星形算法：半径随角度变化
            # 这是一个简化的星形极坐标公式
            r = random.uniform(inner_r, outer_r)
            # 为了让它更像爆炸，我们混合一些随机性
            if random.random() > 0.5:
                 r = outer_r # 强制一部分粒子飞到最远端形成角尖
            
            vx = math.cos(angle) * r
            vy = math.sin(angle) * r
            self.particles.append(Particle(self.x, self.y, self.color, vx, vy))

    def _create_heart_particles(self):
        # 心形：使用心形曲线参数方程
        # x = 16sin^3(t)
        # y = 13cos(t) - 5cos(2t) - 2cos(3t) - cos(4t)
        count = 150
        for i in range(count):
            t = (i / count) * 2 * math.pi # t 从 0 到 2π
            
            # 计算心形坐标
            heart_x = 16 * (math.sin(t) ** 3)
            heart_y = 13 * math.cos(t) - 5 * math.cos(2*t) - 2 * math.cos(3*t) - math.cos(4*t)
            
            # 缩放系数 (控制心形的大小)
            scale = 0.45 
            
            # 将计算出的坐标作为速度向量
            # 注意：pygame的y轴是向下的，所以 heart_y 要取反，否则心形是倒着的
            vx = heart_x * scale
            vy = -heart_y * scale 
            
            # 增加一点随机扰动，让烟花看起来不那么僵硬
            vx += random.uniform(-0.5, 0.5)
            vy += random.uniform(-0.5, 0.5)

            self.particles.append(Particle(self.x, self.y, self.color, vx, vy))

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

# --- 粒子类 ---
class Particle:
    def __init__(self, x, y, color, vx, vy):
        self.x = x
        self.y = y
        self.color = color
        self.vx = vx
        self.vy = vy
        self.life = 100
        self.decay = random.randint(2, 5)
        self.size = random.randint(2, 4)

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += 0.12  # 重力
        self.vx *= 0.95  # 空气阻力
        self.vy *= 0.95
        self.life -= self.decay

    def draw(self):
        if self.life > 0:
            size = max(1, int(self.size * (self.life / 100)))
            pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), size)

# --- 主程序 ---
def main():
    clock = pygame.time.Clock()
    running = True
    fireworks = []
    current_shape = SHAPE_CIRCLE # 默认形状
    auto_fire_timer = 0

    while running:
        screen.fill(BLACK) # 这里用fill清除，如果想要拖尾可以改成半透明覆盖
        
        # 1. 事件处理
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_1:
                    current_shape = SHAPE_CIRCLE
                    print("切换为：圆形")
                elif event.key == pygame.K_2:
                    current_shape = SHAPE_STAR
                    print("切换为：五角星")
                elif event.key == pygame.K_3:
                    current_shape = SHAPE_HEART
                    print("切换为：心形")
            if event.type == pygame.MOUSEBUTTONDOWN:
                # 点击发射当前选定形状的烟花
                fw = Firework(current_shape)
                fw.x, fw.y = pygame.mouse.get_pos()
                fw.vy = 0
                fw.explode()
                fireworks.append(fw)

        # 2. 自动发射逻辑
        auto_fire_timer += 1
        # 每隔一段时间自动发射一个随机形状或当前形状的烟花
        if auto_fire_timer > 40: 
            fw = Firework(current_shape)
            fireworks.append(fw)
            auto_fire_timer = 0

        # 3. 更新与绘制
        for fw in fireworks[:]:
            if fw.update():
                fireworks.remove(fw)
            fw.draw()

        # 4. 显示UI提示
        font = pygame.font.SysFont("Arial", 24)
        shape_text = ["CIRCLE", "STAR", "HEART"]
        text = font.render(f"Shape: {shape_text[current_shape]} (Press 1, 2, 3)", True, WHITE)
        screen.blit(text, (10, 10))

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

    pygame.quit()

if __name__ == "__main__":
    main()