import pygame
import sys
import math
import random
import numpy as np
from pygame.locals import *

# 初始化pygame
pygame.init()

# 设置窗口尺寸
WIDTH, HEIGHT = 1000, 800
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("复杂交互式分形宇宙 - 腾讯元宝")

# 创建离屏表面用于高效绘制
fractal_surface = pygame.Surface((WIDTH, HEIGHT))
particle_surface = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
ui_surface = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)

# 字体处理 - 尝试加载不同字体


def load_font(size, bold=False):
    """尝试多种字体，回退到系统字体"""
    font_names = [
        "Microsoft YaHei",  # 微软雅黑
        "SimHei",           # 黑体
        "SimSun",           # 宋体
        "Arial",
        None  # 回退到默认字体
    ]

    for font_name in font_names:
        try:
            if bold:
                return pygame.font.SysFont(font_name, size, bold=True)
            else:
                return pygame.font.SysFont(font_name, size)
        except:
            continue

    # 如果都失败，使用默认字体
    return pygame.font.Font(None, size)


# 加载字体
font_small = load_font(20)
font_medium = load_font(24)
font_large = load_font(36, bold=True)
font_title = load_font(48, bold=True)

# 颜色定义
COLORS = [
    (0, 7, 100),       # 深蓝
    (32, 107, 203),    # 天蓝
    (237, 255, 255),   # 浅蓝白
    (255, 170, 0),     # 橙色
    (255, 0, 128)      # 粉色
]

UI_COLORS = {
    'background': (0, 0, 0, 180),     # 半透明白色背景
    'text': (255, 255, 255),         # 白色文字
    'text_shadow': (0, 0, 0, 150),   # 文字阴影
    'highlight': (255, 200, 50),     # 高亮色
    'button': (50, 50, 50, 200),     # 按钮背景
    'button_hover': (80, 80, 80, 200)  # 按钮悬停
}

# 复杂分形参数


class FractalParams:
    def __init__(self):
        self.center_x = 0.0
        self.center_y = 0.0
        self.zoom = 0.5
        self.max_iter = 100
        self.color_speed = 0.1
        self.time = 0.0
        self.mode = 0  # 0: Mandelbrot, 1: Julia, 2: Burning Ship, 3: Tricorn

        # Julia集参数
        self.julia_cx = -0.7
        self.julia_cy = 0.27015

        # 动画参数
        self.animate = True
        self.pulse = 0.0
        self.rotation = 0.0


fractal_params = FractalParams()

# 绘制带阴影的文本


def draw_text_with_shadow(surface, text, font, x, y, color=UI_COLORS['text'],
                          shadow_color=UI_COLORS['text_shadow'], shadow_offset=2):
    """绘制带阴影的文本，提高可读性"""
    # 绘制阴影
    shadow = font.render(text, True, shadow_color)
    surface.blit(shadow, (x + shadow_offset, y + shadow_offset))

    # 绘制主文本
    text_surface = font.render(text, True, color)
    surface.blit(text_surface, (x, y))

    return text_surface

# 绘制按钮


def draw_button(surface, text, font, x, y, width, height, is_hovered=False):
    """绘制一个按钮"""
    # 按钮背景
    button_color = UI_COLORS['button_hover'] if is_hovered else UI_COLORS['button']
    button_rect = pygame.Rect(x, y, width, height)

    # 绘制按钮背景
    pygame.draw.rect(surface, button_color, button_rect, border_radius=8)
    pygame.draw.rect(
        surface, UI_COLORS['highlight'], button_rect, 2, border_radius=8)

    # 绘制按钮文字
    text_surf = font.render(text, True, UI_COLORS['text'])
    text_rect = text_surf.get_rect(center=button_rect.center)
    surface.blit(text_surf, text_rect)

    return button_rect

# 粒子系统


class Particle:
    def __init__(self):
        self.reset()
        self.color = random.choice(COLORS[1:])

    def reset(self):
        side = random.randint(0, 3)
        if side == 0:  # 上
            self.x = random.uniform(0, WIDTH)
            self.y = 0
            self.vx = random.uniform(-0.5, 0.5)
            self.vy = random.uniform(0.5, 2)
        elif side == 1:  # 右
            self.x = WIDTH
            self.y = random.uniform(0, HEIGHT)
            self.vx = random.uniform(-2, -0.5)
            self.vy = random.uniform(-0.5, 0.5)
        elif side == 2:  # 下
            self.x = random.uniform(0, WIDTH)
            self.y = HEIGHT
            self.vx = random.uniform(-0.5, 0.5)
            self.vy = random.uniform(-2, -0.5)
        else:  # 左
            self.x = 0
            self.y = random.uniform(0, HEIGHT)
            self.vx = random.uniform(0.5, 2)
            self.vy = random.uniform(-0.5, 0.5)

        self.size = random.uniform(1, 3)
        self.life = random.uniform(50, 200)
        self.max_life = self.life
        self.alpha = 255

    def update(self):
        self.x += self.vx
        self.y += self.vy

        # 添加一些随机运动
        self.vx += random.uniform(-0.05, 0.05)
        self.vy += random.uniform(-0.05, 0.05)

        # 限制速度
        speed = math.sqrt(self.vx**2 + self.vy**2)
        if speed > 3:
            self.vx = self.vx / speed * 3
            self.vy = self.vy / speed * 3

        self.life -= 1
        self.alpha = int(255 * (self.life / self.max_life))

        if self.life <= 0 or self.x < -100 or self.x > WIDTH + 100 or self.y < -100 or self.y > HEIGHT + 100:
            self.reset()

    def draw(self, surface):
        alpha_color = (*self.color, self.alpha)
        pygame.draw.circle(surface, alpha_color,
                           (int(self.x), int(self.y)), int(self.size))


# 创建粒子
particles = [Particle() for _ in range(500)]

# 轨道粒子


class OrbitalParticle:
    def __init__(self):
        self.angle = random.uniform(0, 2 * math.pi)
        self.distance = random.uniform(200, 400)
        self.speed = random.uniform(0.002, 0.01)
        self.size = random.uniform(2, 4)
        self.color = random.choice(COLORS[1:])
        self.alpha = random.randint(100, 200)

    def update(self, time):
        self.angle += self.speed
        self.x = WIDTH // 2 + math.cos(self.angle + time) * self.distance
        self.y = HEIGHT // 2 + math.sin(self.angle + time) * self.distance

        # 脉动效果
        pulse_size = 0.5 + 0.5 * math.sin(time * 2 + self.angle)
        self.current_size = self.size * pulse_size

    def draw(self, surface):
        alpha_color = (*self.color, self.alpha)
        pygame.draw.circle(surface, alpha_color, (int(
            self.x), int(self.y)), int(self.current_size))


orbital_particles = [OrbitalParticle() for _ in range(100)]

# 计算曼德勃罗集颜色


def mandelbrot_color(x, y, params):
    zx, zy = x, y
    cx, cy = x, y

    if params.mode == 1:  # Julia集
        zx, zy = x, y
        cx, cy = params.julia_cx, params.julia_cy
    elif params.mode == 2:  # Burning Ship
        zx, zy = abs(x), abs(y)
        cx, cy = x, y
    elif params.mode == 3:  # Tricorn
        zx, zy = x, y
        cx, cy = x, y

    iter_count = 0
    while iter_count < params.max_iter:
        if params.mode == 3:  # Tricorn使用共轭复数
            zx, zy = zx*zx - zy*zy + cx, -2*zx*zy + cy
        else:
            zx, zy = zx*zx - zy*zy + cx, 2*zx*zy + cy

        if zx*zx + zy*zy > 4:
            break
        iter_count += 1

    if iter_count == params.max_iter:
        return COLORS[0]

    # 平滑着色
    if iter_count < params.max_iter:
        log_zn = math.log(zx*zx + zy*zy) / 2
        nu = math.log(log_zn / math.log(2)) / math.log(2)
        iter_count = iter_count + 1 - nu

    t = (iter_count + params.time * params.color_speed) / params.max_iter

    # 创建渐变色
    color_count = len(COLORS)
    scaled_t = t * (color_count - 1)

    idx1 = int(scaled_t) % color_count
    idx2 = (idx1 + 1) % color_count

    local_t = scaled_t - int(scaled_t)

    r = COLORS[idx1][0] * (1 - local_t) + COLORS[idx2][0] * local_t
    g = COLORS[idx1][1] * (1 - local_t) + COLORS[idx2][1] * local_t
    b = COLORS[idx1][2] * (1 - local_t) + COLORS[idx2][2] * local_t

    return (int(r), int(g), int(b))

# 绘制分形


def draw_fractal(surface, params):
    for y in range(0, HEIGHT, 2):  # 跳着绘制提高性能
        for x in range(0, WIDTH, 2):
            # 计算复平面坐标
            x0 = (x - WIDTH // 2) / \
                (0.5 * params.zoom * WIDTH) + params.center_x
            y0 = (y - HEIGHT // 2) / \
                (0.5 * params.zoom * HEIGHT) + params.center_y

            # 旋转
            if params.rotation != 0:
                ang = params.rotation
                x0, y0 = x0 * \
                    math.cos(ang) - y0 * math.sin(ang), x0 * \
                    math.sin(ang) + y0 * math.cos(ang)

            color = mandelbrot_color(x0, y0, params)

            # 添加脉动效果
            if params.pulse > 0:
                pulse_factor = 1 + 0.1 * math.sin(params.pulse)
                r = min(255, int(color[0] * pulse_factor))
                g = min(255, int(color[1] * pulse_factor))
                b = min(255, int(color[2] * pulse_factor))
                color = (r, g, b)

            # 绘制像素块
            pygame.draw.rect(surface, color, (x, y, 2, 2))

# 绘制UI


def draw_ui(surface, params, mouse_pos):
    # 半透明背景
    ui_bg = pygame.Surface((WIDTH, 120), pygame.SRCALPHA)
    ui_bg.fill(UI_COLORS['background'])
    surface.blit(ui_bg, (0, 0))

    # 标题
    title = draw_text_with_shadow(surface, "复杂交互式分形宇宙", font_title,
                                  WIDTH // 2 - 200, 20, UI_COLORS['highlight'])

    # 副标题
    subtitle = font_medium.render("探索数学之美 | 腾讯元宝 2026", True, (200, 200, 255))
    surface.blit(subtitle, (WIDTH // 2 - subtitle.get_width() // 2, 70))

    # 模式显示
    mode_names = ["曼德勃罗集", "朱利亚集", "燃烧船分形", "三角分形"]
    mode_text = f"当前模式: {mode_names[params.mode]}"
    if params.mode == 1:  # Julia集显示参数
        mode_text += f"  C=({params.julia_cx:.3f}, {params.julia_cy:.3f})"

    draw_text_with_shadow(surface, mode_text, font_large, 20, 10)

    # 参数显示
    params_text = font_medium.render(f"缩放: {params.zoom:.3f} | 迭代次数: {params.max_iter} | 动画: {'开' if params.animate else '关'}",
                                     True, UI_COLORS['text'])
    surface.blit(params_text, (20, 50))

    # 按钮区域
    button_y = HEIGHT - 120
    button_spacing = 120

    # 模式切换按钮
    mode_button_rect = draw_button(surface, "切换模式 (空格)", font_medium,
                                   WIDTH // 2 - button_spacing * 2, button_y, 200, 40,
                                   pygame.Rect(WIDTH // 2 - button_spacing * 2, button_y, 200, 40).collidepoint(mouse_pos))

    # 动画开关按钮
    anim_text = "关闭动画" if params.animate else "开启动画"
    anim_button_rect = draw_button(surface, f"{anim_text} (A)", font_medium,
                                   WIDTH // 2 - button_spacing, button_y, 200, 40,
                                   pygame.Rect(WIDTH // 2 - button_spacing, button_y, 200, 40).collidepoint(mouse_pos))

    # 重置按钮
    reset_button_rect = draw_button(surface, "重置参数 (R)", font_medium,
                                    WIDTH // 2, button_y, 200, 40,
                                    pygame.Rect(WIDTH // 2, button_y, 200, 40).collidepoint(mouse_pos))

    # 控制说明区域
    controls_bg = pygame.Surface((WIDTH, 100), pygame.SRCALPHA)
    controls_bg.fill((0, 0, 0, 100))
    surface.blit(controls_bg, (0, HEIGHT - 100))

    # 控制说明
    controls = [
        "鼠标控制: 拖动移动视图 | 滚轮缩放",
        "键盘控制: 空格切换模式 | R重置参数 | A动画开关 | 方向键调整参数",
        "ESC退出程序 | 按住Shift加速移动"
    ]

    for i, text in enumerate(controls):
        draw_text_with_shadow(surface, text, font_small,
                              WIDTH // 2 - 300, HEIGHT - 80 + i * 25)

    # 返回按钮区域用于点击检测
    return {
        'mode': mode_button_rect,
        'animate': anim_button_rect,
        'reset': reset_button_rect
    }


# 主循环
clock = pygame.time.Clock()
dragging = False
last_mouse_pos = (0, 0)
mouse_pos = (0, 0)
button_rects = {}
running = True
shift_pressed = False

while running:
    dt = clock.tick(60) / 1000.0
    fractal_params.time += dt
    fractal_params.pulse += dt * 2

    if fractal_params.animate:
        fractal_params.julia_cx = 0.7885 * math.cos(fractal_params.time * 0.2)
        fractal_params.julia_cy = 0.7885 * math.sin(fractal_params.time * 0.2)
        fractal_params.rotation = math.sin(fractal_params.time * 0.1) * 0.1

    # 处理事件
    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:
                fractal_params.mode = (fractal_params.mode + 1) % 4
            elif event.key == K_r:
                fractal_params = FractalParams()
            elif event.key == K_a:
                fractal_params.animate = not fractal_params.animate
            elif event.key == K_UP:
                fractal_params.max_iter += 10
            elif event.key == K_DOWN and fractal_params.max_iter > 20:
                fractal_params.max_iter -= 10
            elif event.key == K_RIGHT:
                fractal_params.julia_cx += 0.01
            elif event.key == K_LEFT:
                fractal_params.julia_cx -= 0.01
            elif event.key == K_LSHIFT or event.key == K_RSHIFT:
                shift_pressed = True

        elif event.type == KEYUP:
            if event.key == K_LSHIFT or event.key == K_RSHIFT:
                shift_pressed = False

        elif event.type == MOUSEBUTTONDOWN:
            if event.button == 1:  # 左键
                # 检查是否点击了按钮
                if 'mode' in button_rects and button_rects['mode'].collidepoint(event.pos):
                    fractal_params.mode = (fractal_params.mode + 1) % 4
                elif 'animate' in button_rects and button_rects['animate'].collidepoint(event.pos):
                    fractal_params.animate = not fractal_params.animate
                elif 'reset' in button_rects and button_rects['reset'].collidepoint(event.pos):
                    fractal_params = FractalParams()
                else:
                    # 不是点击按钮，开始拖动
                    dragging = True
                    last_mouse_pos = event.pos
            elif event.button == 4:  # 滚轮上
                fractal_params.zoom *= 1.1
            elif event.button == 5:  # 滚轮下
                fractal_params.zoom /= 1.1

        elif event.type == MOUSEBUTTONUP:
            if event.button == 1:
                dragging = False

        elif event.type == MOUSEMOTION:
            mouse_pos = event.pos
            if dragging:
                dx = event.pos[0] - last_mouse_pos[0]
                dy = event.pos[1] - last_mouse_pos[1]

                # 按住Shift键时移动更快
                speed_multiplier = 3.0 if shift_pressed else 1.0

                fractal_params.center_x -= dx / \
                    (fractal_params.zoom * WIDTH) * speed_multiplier
                fractal_params.center_y -= dy / \
                    (fractal_params.zoom * HEIGHT) * speed_multiplier

                last_mouse_pos = event.pos

    # 清除表面
    screen.fill((0, 0, 0))
    particle_surface.fill((0, 0, 0, 0))
    ui_surface.fill((0, 0, 0, 0))

    # 绘制分形
    draw_fractal(fractal_surface, fractal_params)

    # 更新和绘制粒子
    for particle in particles:
        particle.update()
        particle.draw(particle_surface)

    # 更新和绘制轨道粒子
    for op in orbital_particles:
        op.update(fractal_params.time)
        op.draw(particle_surface)

    # 绘制UI
    button_rects = draw_ui(ui_surface, fractal_params, mouse_pos)

    # 合并所有表面
    screen.blit(fractal_surface, (0, 0))
    screen.blit(particle_surface, (0, 0))
    screen.blit(ui_surface, (0, 0))

    # 添加一些光效
    glow = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    for _ in range(20):
        x = random.randint(0, WIDTH)
        y = random.randint(0, HEIGHT)
        radius = random.randint(20, 50)
        pygame.draw.circle(glow, (255, 255, 255, 10), (x, y), radius)
    screen.blit(glow, (0, 0), special_flags=pygame.BLEND_ADD)

    # 显示帧率和鼠标位置
    fps_text = f"帧率: {int(clock.get_fps())} FPS | 鼠标: ({mouse_pos[0]}, {mouse_pos[1]})"
    draw_text_with_shadow(screen, fps_text, font_small, 10, HEIGHT - 30)

    # 显示当前坐标
    coord_text = f"坐标: ({fractal_params.center_x:.4f}, {fractal_params.center_y:.4f})"
    draw_text_with_shadow(screen, coord_text, font_small, 10, HEIGHT - 60)

    # 更新显示
    pygame.display.flip()

pygame.quit()
sys.exit()