IamHere/ball-demo.py

import pygame
import math
import time
import os
try:
    import pyperclip
    clipboard_available = True
except ImportError:
    clipboard_available = False
    print("警告: pyperclip 模块未安装，剪贴板功能不可用。请运行 'pip install pyperclip' 安装。")

from pygame.locals import *
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *

# 初始化字体
pygame.font.init()
# 使用默认字体，增大字体大小
font = pygame.font.Font(None, 24)

# 初始化 GLUT
glutInit()

# 模拟模式开关（True: 使用模拟数据，False: 扫描真实目录）
# USE_MOCK_DATA = True
USE_MOCK_DATA = False

# 初始化 Pygame
pygame.init()

# 获取屏幕尺寸
info = pygame.display.Info()
screen_width = info.current_w
screen_height = info.current_h

# 设置窗口大小（高度为屏幕高度的80%，保持原始宽高比）
height = int(screen_height * 0.8)
width = int(height * (800 / 600))  # 保持原始宽高比 800:600
display = pygame.display.set_mode((width, height), DOUBLEBUF | OPENGL)
pygame.display.set_caption('3D Yellow Ball')

# 启用深度测试
glEnable(GL_DEPTH_TEST)

# 设置背景色
glClearColor(0.1, 0.1, 0.1, 1.0)
# 禁用混合，确保完全不透明
glDisable(GL_BLEND)

# 设置光照
glEnable(GL_LIGHTING)
glEnable(GL_LIGHT0)

# 设置光源位置
light_position = [2.0, 2.0, 2.0, 1.0]  # 点光源
glLightfv(GL_LIGHT0, GL_POSITION, light_position)

# 设置光源颜色
light_ambient = [0.3, 0.3, 0.3, 1.0]
light_diffuse = [1.0, 1.0, 1.0, 1.0]
light_specular = [1.0, 1.0, 1.0, 1.0]
glLightfv(GL_LIGHT0, GL_AMBIENT, light_ambient)
glLightfv(GL_LIGHT0, GL_DIFFUSE, light_diffuse)
glLightfv(GL_LIGHT0, GL_SPECULAR, light_specular)

# 设置投影矩阵（透视投影）
glMatrixMode(GL_PROJECTION)
glLoadIdentity()
gluPerspective(45, (width/height), 0.1, 50.0)

# 设置模型视图矩阵
glMatrixMode(GL_MODELVIEW)
glLoadIdentity()
# 移动相机（现在在主循环中每帧设置）
# glTranslatef(0.0, 0.0, -5)

def generate_mock_tree():
    # 生成模拟目录树数据
    nodes = [
        ("root", True, -1, 0, "/mock/root"),
        ("dir1", True, 0, 1, "/mock/root/dir1"),
        ("dir2", True, 0, 1, "/mock/root/dir2"),
        ("file1.txt", False, 0, 1, "/mock/root/file1.txt"),
        ("file2.py", False, 0, 1, "/mock/root/file2.py"),
        ("subdir1", True, 1, 2, "/mock/root/dir1/subdir1"),
        ("file3.js", False, 1, 2, "/mock/root/dir1/file3.js"),
        ("file4.md", False, 2, 2, "/mock/root/dir2/file4.md"),
    ]
    
    edges = [
        (0, 1), (0, 2), (0, 3), (0, 4),
        (1, 5), (1, 6),
        (2, 7),
    ]
    
    # 构建子节点列表
    children = [[] for _ in range(len(nodes))]
    for parent_id, child_id in edges:
        children[parent_id].append(child_id)
    
    # 计算节点位置（使用原有算法）
    positions = []
    for i in range(len(nodes)):
        positions.append([0.0, 0.0, 0.0])
    
    # 计算最大深度
    max_depth_found = 0
    for node in nodes:
        max_depth_found = max(max_depth_found, node[3])
    
    # 递归计算位置函数
    def calculate_positions(node_id, start_angle, end_angle, depth):
        if depth > max_depth_found:
            return
        
        if node_id == 0:  # 根节点
            positions[node_id] = [0.0, max_depth_found * 0.4, 0.0]
        else:
            parent_id = nodes[node_id][2]
            if parent_id >= 0:
                child_index = children[parent_id].index(node_id)
                num_siblings = len(children[parent_id])
                angle = start_angle + (end_angle - start_angle) * (child_index / max(num_siblings, 1))
                radius = 0.75 * (0.75 ** depth)
                px, py, pz = positions[parent_id]
                x = px + radius * math.cos(angle)
                y = py - 0.65
                z = pz + radius * math.sin(angle)
                positions[node_id] = [x, y, z]
        
        node_children = children[node_id]
        if node_children:
            angle_range = math.pi * 1.5
            angle_start = -angle_range / 2
            angle_step = angle_range / len(node_children)
            for i, child_id in enumerate(node_children):
                child_angle_start = angle_start + i * angle_step
                child_angle_end = angle_start + (i + 1) * angle_step
                calculate_positions(child_id, child_angle_start, child_angle_end, depth + 1)
    
    calculate_positions(0, -math.pi, math.pi, 0)
    return nodes, edges, positions, children

def build_directory_tree(root_path, max_depth=3, max_children_per_node=30):
    # 如果启用模拟模式，返回模拟数据
    if USE_MOCK_DATA:
        return generate_mock_tree()
    
    # 递归扫描目录树
    nodes = []  # 每个节点: (name, is_dir, parent_id, depth)
    edges = []  # (parent_id, child_id)
    
    # 递归扫描函数
    def scan_directory(current_path, parent_id, depth):
        if depth >= max_depth:
            return
        
        try:
            entries = os.listdir(current_path)
        except Exception:
            return  # 忽略所有访问错误
        
        # 过滤掉一些不需要的目录
        filtered_entries = []
        for entry in entries:
            # 过滤常见的版本控制和IDE目录
            if entry in ['.git', '.vscode', '.idea', 'node_modules', '__pycache__']:
                continue
            filtered_entries.append(entry)
        
        # 限制每个节点的最大子项数量
        filtered_entries = filtered_entries[:max_children_per_node]
        
        # 当前目录的节点ID：根节点为0，其他目录为parent_id
        current_node_id = 0 if parent_id == -1 else parent_id
        
        for entry in filtered_entries:
            entry_path = os.path.join(current_path, entry)
            is_dir = os.path.isdir(entry_path)
            
            # 添加子节点
            child_id = len(nodes)
            nodes.append((entry, is_dir, current_node_id, depth + 1, entry_path))
            edges.append((current_node_id, child_id))
            
            # 如果是目录，递归扫描
            if is_dir:
                scan_directory(entry_path, child_id, depth + 1)
    
    # 添加根节点
    root_name = os.path.basename(root_path.rstrip('\\'))
    nodes.append((root_name, True, -1, 0, root_path))
    
    # 从根节点开始递归扫描
    scan_directory(root_path, -1, 0)
    
    # 计算节点位置
    positions = []
    for i in range(len(nodes)):
        positions.append([0.0, 0.0, 0.0])  # 初始位置
    
    # 构建子节点列表
    children = [[] for _ in range(len(nodes))]
    for edge in edges:
        parent_id, child_id = edge
        children[parent_id].append(child_id)
    
    # 计算最大深度
    max_depth_found = 0
    for node in nodes:
        max_depth_found = max(max_depth_found, node[3])
    
    # 递归计算位置
    def calculate_positions(node_id, start_angle, end_angle, depth):
        if depth > max_depth_found:
            return
        
        # 计算当前节点位置
        if node_id == 0:  # 根节点
            positions[node_id] = [0.0, max_depth_found * 0.4, 0.0]  # 顶部，高度减半
        else:
            parent_id = nodes[node_id][2]
            if parent_id >= 0:
                # 子节点围绕父节点在圆上分布
                child_index = children[parent_id].index(node_id)
                num_siblings = len(children[parent_id])
                
                # 计算角度（在父节点周围的圆上）
                angle = start_angle + (end_angle - start_angle) * (child_index / max(num_siblings, 1))
                
                # 半径随深度减小
                radius = 0.75 * (0.75 ** depth)  # 深度越大，半径越小，基础半径减半
                
                # 计算位置
                px, py, pz = positions[parent_id]
                x = px + radius * math.cos(angle)
                y = py - 0.65  # 每层向下移动，垂直距离减半
                z = pz + radius * math.sin(angle)
                
                positions[node_id] = [x, y, z]
        
        # 为子节点递归计算位置
        node_children = children[node_id]
        if node_children:
            # 计算子节点的角度范围
            angle_range = math.pi * 1.5  # 270度范围
            angle_start = -angle_range / 2
            angle_step = angle_range / len(node_children)
            
            for i, child_id in enumerate(node_children):
                child_angle_start = angle_start + i * angle_step
                child_angle_end = angle_start + (i + 1) * angle_step
                calculate_positions(child_id, child_angle_start, child_angle_end, depth + 1)
    
    # 从根节点开始计算位置
    calculate_positions(0, -math.pi, math.pi, 0)
    
    return nodes, edges, positions, children

# 构建目录树（使用当前目录）
# root_path = "."
root_path = "E:\\agricultural_research_platform"

tree_nodes, tree_edges, node_positions, tree_children = build_directory_tree(root_path)

def draw_text(x, y, z, text):
    """在3D空间中绘制文本"""
    glDisable(GL_LIGHTING)
    glColor3f(1.0, 1.0, 1.0)  # 白色文本
    
    # 保存当前矩阵
    glPushMatrix()
    
    # 移动到指定位置
    glTranslatef(x, y, z)
    
    # 旋转文本，使其面向相机
    glRotatef(rotation_angle, 0.0, 1.0, 0.0)
    
    # 缩小文本，使其适合3D空间
    scale = 0.001
    glScalef(scale, scale, scale)
    
    # 计算文本宽度，以便居中显示
    text_width = 0
    for char in text:
        text_width += glutBitmapWidth(GLUT_BITMAP_HELVETICA_12, ord(char))
    
    # 设置文本位置，使其居中
    glRasterPos2f(-text_width / 2, 0)
    
    # 绘制文本
    for char in text:
        glutBitmapCharacter(GLUT_BITMAP_HELVETICA_12, ord(char))
    
    # 恢复矩阵
    glPopMatrix()
    glEnable(GL_LIGHTING)

# 初始化选中节点（根节点）
selected_node_index = 0

# 剪贴板相关变量
clipboard_content = ""
clipboard_check_time = 0
blinking_nodes = set()
file_content_cache = {}
blink_start_time = pygame.time.get_ticks()

# 旋转控制变量
rotation_angle = 0.0

# 相机控制变量
camera_z = -5  # 相机的 Z 轴位置，初始值为 -5

# 主循环
running = True
while running:
    # 获取当前时间（用于闪烁效果）
    current_time = pygame.time.get_ticks()
    
    # 更新旋转角度（每20秒旋转一圈）
    rotation_angle = (current_time * 0.018) % 360.0
    
    # 检查剪贴板内容是否变化（如果剪贴板功能可用）
    if clipboard_available:
        if current_time - clipboard_check_time > 1000:  # 每秒检查一次
            clipboard_check_time = current_time
            try:
                new_clipboard_content = pyperclip.paste()
                if new_clipboard_content != clipboard_content:
                    clipboard_content = new_clipboard_content
                    # 剪贴板内容变化，重新检查所有文件节点
                    blinking_nodes.clear()
                    if clipboard_content.strip():  # 剪贴板非空
                        for i, (name, is_dir, parent_id, depth, full_path) in enumerate(tree_nodes):
                            if not is_dir:  # 文件节点
                                # 检查文件内容是否包含剪贴板文本
                                try:
                                    # 从缓存读取或加载文件内容
                                    if i not in file_content_cache:
                                        with open(full_path, 'r', encoding='utf-8', errors='ignore') as f:
                                            file_content_cache[i] = f.read()
                                    file_content = file_content_cache[i]
                                    if clipboard_content in file_content:
                                        blinking_nodes.add(i)
                                except Exception as e:
                                    pass  # 忽略读取错误
            except Exception as e:
                pass  # 忽略剪贴板访问错误
    
    # 处理事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            # 获取当前节点信息
            current_node = tree_nodes[selected_node_index]
            current_parent_id = current_node[2]
            current_depth = current_node[3]
            
            if event.key == pygame.K_UP:  # 上键：移动到父节点
                if current_parent_id >= 0:  # 有父节点
                    selected_node_index = current_parent_id
            elif event.key == pygame.K_DOWN:  # 下键：移动到第一个子节点
                if tree_children[selected_node_index]:  # 有子节点
                    selected_node_index = tree_children[selected_node_index][0]
            elif event.key == pygame.K_LEFT:  # 左键：移动到前一个兄弟节点
                if current_parent_id >= 0:  # 有父节点
                    siblings = tree_children[current_parent_id]
                    if len(siblings) > 1:
                        current_index_in_siblings = siblings.index(selected_node_index)
                        # 移动到前一个兄弟，如果当前是第一个则移动到最后一个
                        new_index = (current_index_in_siblings - 1) % len(siblings)
                        selected_node_index = siblings[new_index]
            elif event.key == pygame.K_RIGHT:  # 右键：移动到后一个兄弟节点
                if current_parent_id >= 0:  # 有父节点
                    siblings = tree_children[current_parent_id]
                    if len(siblings) > 1:
                        current_index_in_siblings = siblings.index(selected_node_index)
                        # 移动到后一个兄弟，如果当前是最后一个则移动到第一个
                        new_index = (current_index_in_siblings + 1) % len(siblings)
                        selected_node_index = siblings[new_index]
        elif event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 4:  # 鼠标滚轮向上滚动，放大
                camera_z += 0.5  # 相机向Z轴正方向移动，放大视图
            elif event.button == 5:  # 鼠标滚轮向下滚动，缩小
                camera_z -= 0.5  # 相机向Z轴负方向移动，缩小视图
            
            # 限制相机的Z轴范围，防止过度放大或缩小
            camera_z = max(-10, min(-2, camera_z))
    
    # 清除屏幕
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    
    # 重置模型视图矩阵并设置相机位置
    glLoadIdentity()
    glTranslatef(0.0, 0.0, camera_z)  # 相机位置，可通过鼠标滚轮调整
    glRotatef(rotation_angle, 0.0, 1.0, 0.0)  # 绕竖直中线/Y轴旋转
    
    # 目录树可视化
    ball_radius = 0.0667  # 节点球的半径（默认大小）
    
    # 绘制所有连线（目录树边）
    glDisable(GL_LIGHTING)  # 禁用光照以使用纯色
    glColor3f(1.0, 1.0, 1.0)  # 白色连线
    glLineWidth(2.0)
    
    glBegin(GL_LINES)
    for edge in tree_edges:
        node1_idx, node2_idx = edge
        x1, y1, z1 = node_positions[node1_idx]
        x2, y2, z2 = node_positions[node2_idx]
        glVertex3f(x1, y1, z1)
        glVertex3f(x2, y2, z2)
    glEnd()
    
    glEnable(GL_LIGHTING)  # 重新启用光照
    
    # 绘制所有节点球（根据类型使用不同颜色）
    for i, (name, is_dir, parent_id, depth, full_path) in enumerate(tree_nodes):
        x, y, z = node_positions[i]
        
        # 根据节点类型设置材质颜色（选中节点显示为红色）
        if i in blinking_nodes:  # 闪烁节点
            # 计算闪烁因子（正弦波，周期约2秒）
            blink_factor = (math.sin((current_time - blink_start_time) * 0.005) + 1) * 0.5
            # 基础颜色 #c31c1f
            base_r, base_g, base_b = 0.7647, 0.1098, 0.1216
            # 亮度在0.7到1.0之间变化
            brightness = 0.7 + 0.3 * blink_factor
            color = [base_r * brightness, base_g * brightness, base_b * brightness, 1.0]
            ambient = [base_r * 0.5, base_g * 0.5, base_b * 0.5, 1.0]
            specular = [base_r * 0.8, base_g * 0.8, base_b * 0.8, 1.0]
        elif i == selected_node_index:  # 选中节点
            color = [1.0, 0.0, 0.0, 1.0]  # 红色
            ambient = [0.5, 0.0, 0.0, 1.0]
            specular = [0.8, 0.3, 0.3, 1.0]
        elif i == 0:  # 根节点
            color = [1.0, 1.0, 0.0, 1.0]  # 黄色
            ambient = [0.5, 0.5, 0.0, 1.0]
            specular = [0.8, 0.8, 0.4, 1.0]
        elif is_dir:  # 目录
            color = [0.2902, 0.5961, 0.3098, 1.0]  # #4a984f
            ambient = [0.2, 0.4, 0.2, 1.0]
            specular = [0.4, 0.8, 0.4, 1.0]
        else:  # 文件
            color = [0.7059, 0.7176, 0.2549, 1.0]  # #b4b741
            ambient = [0.4, 0.4, 0.15, 1.0]
            specular = [0.8, 0.8, 0.4, 1.0]
        
        shininess = [30.0]
        
        glMaterialfv(GL_FRONT, GL_AMBIENT, ambient)
        glMaterialfv(GL_FRONT, GL_DIFFUSE, color)
        glMaterialfv(GL_FRONT, GL_SPECULAR, specular)
        glMaterialfv(GL_FRONT, GL_SHININESS, shininess)
        
        glPushMatrix()
        glTranslatef(x, y, z)
        quad = gluNewQuadric()
        gluQuadricNormals(quad, GLU_SMOOTH)
        # 根据深度和节点类型调整球体大小：深度越大，球体越小；文件节点的半径缩小为原来的1/3，文件夹节点的半径缩小为原来的1/2
        if is_dir:
            current_radius = (ball_radius / 2) * (0.85 ** depth)  # 文件夹节点的半径缩小为原来的1/2
        else:
            current_radius = (ball_radius / 3) * (0.85 ** depth)  # 文件节点的半径缩小为原来的1/3
        gluSphere(quad, current_radius, 32, 32)
        gluDeleteQuadric(quad)
        glPopMatrix()
        
        # 绘制文件名
        draw_text(x, y + current_radius + 0.05, z, name)
    
    # 交换缓冲区
    pygame.display.flip()
    # 控制帧率
    pygame.time.wait(10)

# 退出 Pygame
pygame.quit()