import os
import time
import open3d as o3d
import trimesh
from pathlib import Path


def test_folder_objs(folder_path, output_folder="output_objs"):
    """测试文件夹中所有OBJ文件的读写性能"""
    # 确保输出文件夹存在
    Path(output_folder).mkdir(exist_ok=True)

    # 收集文件夹中所有OBJ文件
    obj_files = [f for f in os.listdir(folder_path) if f.lower().endswith('.obj')]

    if not obj_files:
        print(f"在文件夹 {folder_path} 中未找到OBJ文件")
        return

    print(f"找到 {len(obj_files)} 个OBJ文件，开始测试...")

    # 准备结果记录
    results = []

    for obj_file in obj_files:
        file_path = os.path.join(folder_path, obj_file)
        output_path = os.path.join(output_folder, obj_file)

        print(f"\n测试文件: {obj_file}")

        # 测试open3d
        o3d_write, o3d_read = test_open3d_io(file_path, output_path.replace('.obj', '_o3d.obj'))

        # 测试trimesh
        tm_write, tm_read = test_trimesh_io(file_path, output_path.replace('.obj', '_tm.obj'))

        # 记录结果
        file_stats = {
            'filename': obj_file,
            'o3d_write': o3d_write,
            'o3d_read': o3d_read,
            'tm_write': tm_write,
            'tm_read': tm_read,
            'write_ratio': o3d_write / tm_write if tm_write > 0 else 0,
            'read_ratio': o3d_read / tm_read if tm_read > 0 else 0
        }
        results.append(file_stats)

        # 打印当前文件结果
        print(f"  open3d  | 写入: {o3d_write:.3f}s | 读取: {o3d_read:.3f}s")
        print(f"  trimesh | 写入: {tm_write:.3f}s | 读取: {tm_read:.3f}s")
        print(f"  写入速度比(trimesh/open3d): {file_stats['write_ratio']:.1f}x")
        print(f"  读取速度比(trimesh/open3d): {file_stats['read_ratio']:.1f}x")

    # 打印汇总结果
    print("\n=== 汇总结果 ===")
    avg_write_ratio = sum(r['write_ratio'] for r in results) / len(results)
    avg_read_ratio = sum(r['read_ratio'] for r in results) / len(results)
    print(f"平均写入速度比(trimesh/open3d): {avg_write_ratio:.1f}x")
    print(f"平均读取速度比(trimesh/open3d): {avg_read_ratio:.1f}x")


def test_open3d_io(input_path, output_path):
    """测试open3d的读写性能"""
    # 读取
    start = time.time()
    mesh = o3d.io.read_triangle_mesh(input_path)
    read_time = time.time() - start

    # 写入
    start = time.time()
    o3d.io.write_triangle_mesh(output_path, mesh)
    write_time = time.time() - start

    return write_time, read_time


def test_trimesh_io(input_path, output_path):
    """测试trimesh的读写性能"""
    # 读取
    start = time.time()
    mesh = trimesh.load(input_path)
    read_time = time.time() - start

    # 写入
    start = time.time()
    mesh.export(output_path)
    write_time = time.time() - start

    return write_time, read_time


if __name__ == "__main__":
    # 设置包含OBJ文件的文件夹路径
    obj_folder = "/data/datasets_20t/9_big/"  # 替换为你的OBJ文件夹路径

    # 运行测试
    test_folder_objs(obj_folder)