import os
import trimesh
import numpy as np
from scipy.spatial import KDTree
from tqdm import tqdm
import mesh2sdf

# 加载BREP模型并转换为三角网格
def load_brep_to_mesh(file_path: str) -> trimesh.Trimesh:
    mesh = trimesh.load_mesh(file_path)
    return mesh

def process_mesh_to_sdf(
    mesh: trimesh.Trimesh,
    filename:str, 
    save_dir:str,
    idx: str,
    mesh_scale:int = 0.8,
    size:int= 128,
    level: int = None, 
) -> str:
    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 create_grid(x_range, y_range, z_range, resolution):
    x = np.linspace(x_range[0], x_range[1], resolution)
    y = np.linspace(y_range[0], y_range[1], resolution)
    z = np.linspace(z_range[0], z_range[1], resolution)
    grid = np.meshgrid(x, y, z)
    points = np.vstack(list(map(np.ravel, grid))).T
    return points
'''



def test():
    # 参数设置
    input_file = "/mnt/disk2/dataset/furniture/step/furniture_dataset_step/train/bathtub_0004.step"
    output_file = "tmp_data"

    # 加载BREP模型并转换为三角网格
    mesh = load_brep_to_mesh(input_file)

    process_mesh_to_sdf(
        mesh=mesh,
        filename=input_file,
        save_dir=output_file,
        idx="0000000",
    )






# 真的执行转换
def main():

    INPUT = '/mnt/disk2/dataset/furniture/step/furniture_dataset_step' # 下面train，val和test
    OUTPUT = '/app/data/furniture_sdf' 
    valid = 0
    for set_ in ['train', 'val', 'test']:   
        with ProcessPoolExecutor(max_workers=os.cpu_count() // 2) as executor:
            futures = {}
            for step_folder in glob.glob(os.path.join(input_folder, set_, '*.step')):
                future = executor.submit(process, step_folder, timeout=300)
                futures[future] = step_folder

            for future in tqdm(as_completed(futures), total=len(step_dirs)):
                try:
                    status = future.result(timeout=300)
                    valid += status
                except TimeoutError:
                    print(f"Timeout occurred while processing {futures[future]}")
                except Exception as e:
                    print(f"An error occurred while processing {futures[future]}: {e}")

    print(f'Done... Data Converted Ratio {100.0*valid/len(step_dirs)}%')



if __name__ == "__main__":
    test()