from PyQt5.QtCore import QObject, pyqtSlot
import open3d as o3d
import numpy as np
import os
from OCC.Core.TopoDS import TopoDS_Face


class Datahandle(QObject):
    def __init__(self):
        super().__init__()
        self.qml_item = None
        self.vertices = np.array([])
        self.colors = np.array([])
        self.triangles = np.array([])  # 新增：三角面索引
        self.normals = np.array([])  # 新增：法线


    def load_data(self):
        """
        加载 3D 数据（这里用随机点云代替实际文件）
        返回: vertices, colors (numpy arrays)
        """
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(np.random.rand(5000, 3))
        pcd.colors = o3d.utility.Vector3dVector(np.random.rand(5000, 3))
        vertices = np.asarray(pcd.points)
        colors = np.asarray(pcd.colors)
        return vertices, colors

    @pyqtSlot('QVariant')
    def set3DItem(self, item):
        """接收 QML 中的 QML3DItem 实例"""
        self.qml_item = item
        if self.qml_item:
            self.qml_item.set_data(self.vertices, self.colors)
            print("✅ 数据已传递给 QML3DItem")

    @pyqtSlot(float, float, float)
    def onPointPicked(self, x, y, z):
        """接收从 QML3DItem 发出的拾取信号"""
        print(f"✅ 拾取到点: ({x:.3f}, {y:.3f}, {z:.3f})")
        # 这里可以扩展：保存坐标、发送到其他模块等

class Dataload(QObject):
    def __init__(self):
        super().__init__()
        self.qml_item = None
        self.vertices = np.array([])
        self.colors = np.array([])
        self.triangles = np.array([])  # 新增：三角面索引
        self.normals = np.array([])  # 新增：法线
        self.picking_color_map = {}

    def load_data_from_file(self, file_path: str):
        """
        使用 numpy-stl 自动识别 STL 格式（兼容旧版本）
        """
        print(f"📁 正在加载: {file_path}")

        if not os.path.exists(file_path):
            print(f"❌ 文件不存在: {file_path}")
            return np.array([]), np.array([]), np.array([]), np.array([])

        ext = os.path.splitext(file_path)[1].lower()
        if ext != '.stl':
            print(f"❌ 不支持的格式: {ext}")
            return np.array([]), np.array([]), np.array([]), np.array([])

        try:
            from stl import mesh as stl_mesh

            # ✅ 自动识别格式（旧版本也支持）
            stl_data = stl_mesh.Mesh.from_file(file_path)

            # 获取三角面 (N, 3, 3)
            vectors = stl_data.vectors.astype(np.float32)
            print(f"✅ 加载 {len(vectors)} 个三角面")

            # 展平顶点
            vertices = vectors.reshape(-1, 3)
            print(f"原始顶点数: {len(vertices)}")

            # 三角面索引
            num_triangles = len(vectors)
            triangles = np.arange(num_triangles * 3, dtype=np.int32).reshape(-1, 3)
            print(f"三角面数量: {num_triangles}")

            # 法线（每个三角面对应一个法线）
            normals = stl_data.normals.astype(np.float32)
            vertex_normals = np.repeat(normals, 3, axis=0)  # 每个顶点复制一次
            print(f"法线数量: {len(vertex_normals)}")

            # 颜色：使用法线生成
            colors = (vertex_normals + 1.0) / 2.0
            if np.any(np.isnan(colors)) or np.any(np.isinf(colors)):
                colors = np.ones_like(vertex_normals) * 0.8
                self.colors = colors

            # --- 归一化 ---
            if len(vertices) > 0:
                min_coords = np.min(vertices, axis=0)
                max_coords = np.max(vertices, axis=0)
                center = (min_coords + max_coords) / 2.0
                scale = np.max(max_coords - min_coords)
                if scale > 0:
                    vertices = (vertices - center) / scale * 2.0

                    z_range = np.max(vertices[:, 2]) - np.min(vertices[:, 2])
                    if z_range < 0.1:
                        z_center = (np.min(vertices[:, 2]) + np.max(vertices[:, 2])) / 2.0
                        vertices[:, 2] = (vertices[:, 2] - z_center) * (0.2 / z_range) + z_center
                        self.vertices = vertices
                        print(f"调整Z轴范围: {z_range:.6f} -> 0.2")

            print(f"✅ 加载成功: {file_path}")
            self._build_picking_color_map()
            return vertices, colors, triangles, vertex_normals

        except Exception as e:
            print(f"❌ 加载失败: {e}")
            return np.array([]), np.array([]), np.array([]), np.array([])

    def _build_picking_color_map(self):
        self.picking_color_map.clear()
        for i in range(len(self.vertices)):
            r, g, b = self._get_picking_color(i)
            key = (round(r, 5), round(g, 5), round(b, 5))
            self.picking_color_map[key] = i
        print("✅ picking_color_map 构建完成")
        # 调试：打印前几个
        print("📊 前5个映射:", list(self.picking_color_map.items())[:5])

    def _get_picking_color(self, index):
        """根据顶点索引生成唯一颜色 (r,g,b 归一化到 0~1)"""
        idx = index + 1  # 避免 0,0,0 黑色
        r = (idx) & 0xFF
        g = (idx >> 8) & 0xFF
        b = (idx >> 16) & 0xFF
        return (r / 255.0, g / 255.0, b / 255.0)