import pygame
import sys
import math

# pygame初始化
pygame.init()
W, H = 800, 600
screen = pygame.display.set_mode((W, H))
pygame.display.set_caption("3D打印机｜鼠标拖拽视角 | 滚轮缩放｜1234切换模型")
clock = pygame.time.Clock()

# 颜色常量
WHITE = (255, 255, 255)
BLACK = (30, 30, 30)
ORANGE = (255, 130, 0)
GRAY = (80, 80, 80)
RED = (220, 40, 40)
GREEN = (80, 220, 80)

# 斜二测基础参数
iso_scale = 1.2

# ===== 视角控制参数 =====
cam_offset_x = 0    # 视角横向偏移
cam_offset_y = 0    # 视角纵向偏移
cam_zoom = 1.0      # 缩放系数
mouse_drag = False
drag_start_x = 0
drag_start_y = 0

# 喷头初始参数
head_x, head_y, head_z = 40, 40, 0
speed = 1.2
print_lines = []
path_index = 0
last_head_screen_pos = None

# =========== 模型定义 ===========
def get_cube():
    return [
        (40, 40, 0), (160, 40, 0), (160, 160, 0), (40, 160, 0), (40, 40, 0),
        (40, 40, 40), (160, 40, 40), (160, 160, 40), (40, 160, 40), (40, 40, 40),
    ]

def get_circle():
    pts = []
    radius = 60
    cx, cy = 100, 100
    for angle in range(0, 361, 6):
        rad = math.radians(angle)
        px = cx + radius * math.cos(rad)
        py = cy + radius * math.sin(rad)
        pts.append((px, py, 0))
    for angle in range(0, 361, 6):
        rad = math.radians(angle)
        px = cx + radius * math.cos(rad)
        py = cy + radius * math.sin(rad)
        pts.append((px, py, 35))
    return pts

def get_pyramid():
    return [
        (40, 40, 0), (160, 40, 0), (160, 160, 0), (40, 160, 0), (40, 40, 0),
        (100, 100, 50),
        (160, 40, 0), (100, 100, 50),
        (160, 160, 0), (100, 100, 50),
        (40, 160, 0), (100, 100, 50),
    ]

def get_star():
    pts = []
    center_x, center_y = 100, 100
    outer_r = 70
    inner_r = 32
    for i in range(10):
        angle = math.radians(i * 36 - 90)
        r = outer_r if i % 2 == 0 else inner_r
        px = center_x + r * math.cos(angle)
        py = center_y + r * math.sin(angle)
        pts.append((px, py, 0))
    for i in range(10):
        angle = math.radians(i * 36 - 90)
        r = outer_r if i % 2 == 0 else inner_r
        px = center_x + r * math.cos(angle)
        py = center_y + r * math.sin(angle)
        pts.append((px, py, 38))
    return pts

# 模型映射
model_list = {
    pygame.K_1: {"name": "双层立方体", "func": get_cube},
    pygame.K_2: {"name": "双层圆环", "func": get_circle},
    pygame.K_3: {"name": "金字塔", "func": get_pyramid},
    pygame.K_4: {"name": "双层五角星", "func": get_star},
}
current_model_name = ""
path_points = get_cube()

# 字体兼容
try:
    font = pygame.font.SysFont("simhei", 22)
except:
    font = pygame.font.Font(None, 22)


def iso_proj(x, y, z):
    """逻辑坐标 → 屏幕坐标，叠加视角偏移与缩放"""
    screen_x = (x - y) * iso_scale + W // 2
    screen_y = (x + y) * 0.6 - z * 1.2 + H // 3

    # 缩放
    screen_x = (screen_x - W/2) * cam_zoom + W/2
    screen_y = (screen_y - H/2) * cam_zoom + H/2

    # 平移
    screen_x += cam_offset_x
    screen_y += cam_offset_y
    return int(screen_x), int(screen_y)


def switch_model(key):
    global path_points, path_index, print_lines, head_x, head_y, head_z, current_model_name, last_head_screen_pos
    if key in model_list:
        info = model_list[key]
        path_points = info["func"]()
        current_model_name = info["name"]
        path_index = 0
        print_lines.clear()
        head_x, head_y, head_z = path_points[0]
        last_head_screen_pos = None


switch_model(pygame.K_1)

running = True
while running:
    clock.tick(60)
    screen.fill(BLACK)

    # 事件处理
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

        # 按键切换模型
        if event.type == pygame.KEYDOWN:
            switch_model(event.key)

        # 鼠标按下 开始拖拽
        if event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 1:
                mouse_drag = True
                drag_start_x, drag_start_y = pygame.mouse.get_pos()
            # 滚轮缩放
            if event.button == 4:  # 上滚放大
                cam_zoom += 0.08
            if event.button == 5:  # 下滚缩小
                cam_zoom -= 0.08
            # 限制缩放范围，防止过大过小
            cam_zoom = max(0.4, min(2.5, cam_zoom))

        # 鼠标松开 停止拖拽
        if event.type == pygame.MOUSEBUTTONUP:
            if event.button == 1:
                mouse_drag = False

        # 鼠标拖拽移动视角
        if event.type == pygame.MOUSEMOTION and mouse_drag:
            mx, my = pygame.mouse.get_pos()
            dx = mx - drag_start_x
            dy = my - drag_start_y
            cam_offset_x += dx
            cam_offset_y += dy
            drag_start_x, drag_start_y = mx, my

    # 喷头自动移动逻辑
    target_x, target_y, target_z = path_points[path_index]
    dx = target_x - head_x
    dy = target_y - head_y
    dz = target_z - head_z

    head_x += dx * 0.02 * speed
    head_y += dy * 0.02 * speed
    head_z += dz * 0.02 * speed

    distance = abs(dx) + abs(dy) + abs(dz)
    if distance < 1.5:
        path_index += 1
        if path_index >= len(path_points):
            path_index = 0
            print_lines.clear()

    # 绘制打印轨迹
    current_screen_pos = iso_proj(head_x, head_y, head_z)
    if last_head_screen_pos is not None:
        print_lines.append([last_head_screen_pos, current_screen_pos])
    last_head_screen_pos = current_screen_pos

    # 打印平台
    p1 = iso_proj(0, 0, 0)
    p2 = iso_proj(200, 0, 0)
    p3 = iso_proj(200, 200, 0)
    p4 = iso_proj(0, 200, 0)
    pygame.draw.polygon(screen, GRAY, [p1, p2, p3, p4], 2)

    # 耗材线条
    for line in print_lines:
        pygame.draw.line(screen, ORANGE, line[0], line[1], 3)

    # 红色喷头
    pygame.draw.circle(screen, RED, current_screen_pos, 6)
    base_pos = iso_proj(head_x, head_y, 0)
    pygame.draw.line(screen, WHITE, base_pos, current_screen_pos, 1)

    # UI文字
    tip1 = font.render("【1/2/3/4切换模型 | 鼠标左键拖拽视角 | 滚轮缩放】", True, WHITE)
    tip2 = font.render(f"当前模型：{current_model_name}", True, GREEN)
    screen.blit(tip1, (10, 10))
    screen.blit(tip2, (10, 38))

    pygame.display.flip()

pygame.quit()
sys.exit()