import os
import sys
import trimesh
import mesh2sdf
import numpy as np
import time
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor




def process_mesh_to_sdf(
    filename:str, 
    save_dir:str,
    idx: str,
    mesh_scale:int = 0.8
    size:int= 128
    level: int = None 
):
    if not level:
        level = 2 / size

    mesh = trimesh.load(filename, force='mesh')

    # normalize mesh
    vertices = mesh.vertices
    bbmin = vertices.min(0)
    bbmax = vertices.max(0)
    center = (bbmin + bbmax) * 0.5
    scale = 2.0 * mesh_scale / (bbmax - bbmin).max()
    vertices = (vertices - center) * scale

    # fix mesh
    t0 = time.time()
    sdf, mesh = mesh2sdf.compute(
        vertices, mesh.faces, size, fix=True, level=level, return_mesh=True)
    t1 = time.time()

    # sdf to x,y,z,sdf
    xyzsdf = np.zeros((size, size, size, 4))
    for x in range(size):
        for y in range(size):
            for z in range(size):
                xyzsdf[x, y, z] = [x, y, z, sdf[x, y, z]]
    # output
    mesh.vertices = mesh.vertices / scale + center
    name = filename.split('/')[-1][:-4]
    mesh.export(os.path.join(save_dir, filename[:-4] + '.fixed.obj'))
    np.save(os.path.join(save_dir, filename[:-4] + '.xyzsdf.npy'), xyzsdf)
    #mesh.export(filename[:-4] + '.fixed.obj')
    #np.save(filename[:-4] + '.npy', sdf)
    #print('It takes %.4f seconds to process %s' % (t1-t0, filename))
    return idx


def mesh_to_sdf():
    dataset_dir = "/home/wch/data/abc"
    valid_id = []
    '''
    for i in tqdm(range(4025,10000)):
        internal_dir_name = f"{i:08d}"
        data_dir = os.path.join(dataset_dir, internal_dir_name)
        for name in os.listdir(data_dir):
            file = os.path.join(data_dir,name)
            if os.path.isfile(file) and name.endswith(".obj") and not name.endswith(".fixed.obj"):
                #process(file,data_dir)
                executor.submit(process, file, data_dir)
                valid_id.append(internal_dir_name)
                '''
    # 创建一个进程池
    with ProcessPoolExecutor(max_workers=5) as executor:
        futures = []
        for i in tqdm(range(4025, 10000)):
            internal_dir_name = f"{i:08d}"
            data_dir = os.path.join(dataset_dir, internal_dir_name)
            for name in os.listdir(data_dir):
                file = os.path.join(data_dir, name)
                if os.path.isfile(file) and name.endswith(".obj") and not name.endswith(".fixed.obj"):
                    future = executor.submit(process, file, data_dir,internal_dir_name)
                    futures.append(future)

        # 等待所有任务完成，并收集结果
        for future in tqdm(futures):
            valid_id.append(future.result())

    with open("/home/wch/data/abc/valid.txt", "w") as f:
        for idx in valid_id:
            f.write(idx+"\n")