import os
import numpy as np
import trimesh  # 计算normal

def load_off_file(file_path):
    """加载单个 .off 文件，返回顶点坐标和normal"""
    try:
        mesh = trimesh.load(file_path)
        points = mesh.vertices  # x y z
        normals = mesh.vertex_normals if hasattr(mesh, 'vertex_normals') and mesh.vertex_normals is not None else np.zeros((len(points), 3))  # 如果无normal，用0填充
        return points, normals
    except Exception as e:
        print(f"Error loading {file_path}: {e}")
        return None, None

def convert_off_to_txt(input_dir, output_dir, num_points=1024):
    """将整个目录下的 .off 文件转换为 .txt 文件，支持递归train/test"""
    os.makedirs(output_dir, exist_ok=True)

    # 递归遍历input_dir下的所有.off文件
    for root, dirs, files in os.walk(input_dir):
        for off_file in files:
            if not off_file.endswith('.off'):
                continue

            off_path = os.path.join(root, off_file)
            # 相对路径保存，保持train/test结构
            rel_path = os.path.relpath(root, input_dir)
            class_name = rel_path.split(os.sep)[0] if rel_path != '.' else os.path.basename(root)
            txt_file = off_file.replace('.off', '.txt')
            txt_path = os.path.join(output_dir, rel_path, txt_file)
            os.makedirs(os.path.dirname(txt_path), exist_ok=True)

            print(f"Processing: {off_path}")

            points, normals = load_off_file(off_path)
            if points is None:
                continue

            # 随机采样
            n_points = len(points)
            if n_points < num_points:
                indices = np.random.choice(n_points, num_points, replace=True)
            else:
                indices = np.random.choice(n_points, num_points, replace=False)
            sampled_points = points[indices]
            sampled_normals = normals[indices]

            # 保存 .txt (x y z nx ny nz，空格分隔)
            data_to_save = np.hstack((sampled_points, sampled_normals))
            np.savetxt(txt_path, data_to_save, delimiter=' ', fmt='%.6f')
            print(f"  -> Saved: {txt_path}")

if __name__ == "__main__":
    INPUT_DIR = "/home/lab/LAD/bitpointV3/data/raw_modelnet400915/modelnet40_off"  # 调整到你的根
    OUTPUT_DIR = "/home/lab/LAD/bitpointV3/data/modelnet40_normal_resampled"

    convert_off_to_txt(INPUT_DIR, OUTPUT_DIR, num_points=1024)
    print("✅ Conversion completed!")