"""点云加载模块"""
import numpy as np
from pathlib import Path
import struct


class PointCloudLoader:
    """加载点云文件"""

    @staticmethod
    def load_ply(file_path: str) -> np.ndarray:
        """
        加载PLY格式点云

        Args:
            file_path: PLY文件路径

        Returns:
            np.ndarray: 点云数组 (N, 3)
        """
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"点云文件不存在: {file_path}")

        print(f"正在加载点云: {file_path}")

        with open(path, 'rb') as f:
            # 读取头部
            header_lines = []
            while True:
                line = f.readline().decode('ascii').strip()
                header_lines.append(line)
                if line == 'end_header':
                    break

            # 解析头部
            vertex_count = 0
            format_type = None
            properties = []

            for line in header_lines:
                if line.startswith('element vertex'):
                    vertex_count = int(line.split()[-1])
                elif line.startswith('format'):
                    format_type = line.split()[1]
                elif line.startswith('property'):
                    parts = line.split()
                    properties.append({
                        'type': parts[1],
                        'name': parts[2]
                    })

            if vertex_count == 0:
                raise ValueError("PLY文件中没有顶点数据")

            # 读取顶点数据
            if format_type == 'binary_little_endian':
                points = PointCloudLoader._read_binary_ply(f, vertex_count, properties)
            elif format_type == 'ascii':
                points = PointCloudLoader._read_ascii_ply(f, vertex_count)
            else:
                raise ValueError(f"不支持的PLY格式: {format_type}")

        print(f"  加载完成: {len(points)} 个点")
        return points

    @staticmethod
    def _read_binary_ply(f, vertex_count: int, properties: list) -> np.ndarray:
        """读取二进制PLY数据（优化版：批量读取）"""
        # 构建numpy dtype
        type_map = {
            'float': np.float32,
            'double': np.float64,
            'uchar': np.uint8,
            'uint': np.uint32,
            'int': np.int32
        }

        # 构建结构化数组的dtype
        dtype_list = []
        for prop in properties:
            dtype_list.append((prop['name'], type_map.get(prop['type'], np.float32)))

        dtype = np.dtype(dtype_list)

        # 一次性读取所有数据
        raw_data = np.fromfile(f, dtype=dtype, count=vertex_count)

        # 提取xyz坐标（假设前3个属性是x, y, z）
        # 如果属性名不是x,y,z，可以通过索引访问
        if 'x' in raw_data.dtype.names and 'y' in raw_data.dtype.names and 'z' in raw_data.dtype.names:
            points = np.column_stack([raw_data['x'], raw_data['y'], raw_data['z']])
        else:
            # 假设前3个字段是xyz
            points = np.column_stack([raw_data[raw_data.dtype.names[0]],
                                     raw_data[raw_data.dtype.names[1]],
                                     raw_data[raw_data.dtype.names[2]]])

        return points.astype(np.float64)

    @staticmethod
    def _read_ascii_ply(f, vertex_count: int) -> np.ndarray:
        """读取ASCII PLY数据（优化版：批量读取）"""
        # 读取剩余所有行
        remaining = f.read().decode('ascii')

        # 使用numpy的快速文本解析
        from io import StringIO
        points = np.loadtxt(StringIO(remaining), dtype=np.float64, max_rows=vertex_count, usecols=(0, 1, 2))

        return points