import open3d as o3d
import numpy as np
import os 
import time 

#specify the location from where we'll read the point cloud data generated by APIC/FLIP/PIC 
dirfiles = os.listdir('./outpic_pic_50_mod/')
sorted_files = sorted(dirfiles, key=lambda x: int(x.split('frameparticles')[1]))
correct_file_path = ['./outpic_pic_50_mod/' + x for x in sorted_files]

#specify the output directory where the final meshes, images will stored
output_direcory_point_cloud = os.getcwd() + "/PIC_Particles_UPDATED"
output_direcory_normal = os.getcwd() + "/PIC_NORMALS_UPDATED"
output_directory_mesh = os.getcwd() + "/PIC_MESHES_UPDATED"
output_directory_mesh_img = os.getcwd() + "/PIC_MESHES_IMAGES"

if not os.path.isdir(output_direcory_point_cloud):
    os.mkdir(output_direcory_point_cloud)
if not os.path.isdir(output_direcory_normal):
    os.mkdir(output_direcory_normal)
if not os.path.isdir(output_directory_mesh):
    os.mkdir(output_directory_mesh)
if not os.path.isdir(output_directory_mesh_img):
    os.mkdir(output_directory_mesh_img)

#point cloud visualization code
vis = o3d.visualization.Visualizer()
vis.create_window(width=1920, height=1080,visible = True) 
print('visualizing point cloud')
for point_cloud_file in correct_file_path: 
    print(point_cloud_file)
    pcd = o3d.io.read_point_cloud(point_cloud_file,format="xyz")
    pcd.paint_uniform_color([0,0.5,0.75])
    vis.add_geometry(pcd)
    vis.get_render_option().load_from_json("renderoption.json")
    vis.poll_events()
    vis.update_renderer()
    time.sleep(1)
    vis.capture_screen_image(output_direcory_point_cloud +'/' + point_cloud_file.split('/')[-1] + '.png')
    vis.clear_geometries()



#normal estimation visualization code
print('estimating normals')
vis = o3d.visualization.Visualizer()
vis.create_window(visible = True) 
for point_cloud_file in correct_file_path: 
    print(point_cloud_file)
    pcd = o3d.io.read_point_cloud(point_cloud_file,format="xyz")
    pcd.paint_uniform_color([0,0.5,0.75])
    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(
        radius=0.1, max_nn=30))
    #once the normal has been estimated compute the corresponding mesh 
    # o3d.visualization.draw_geometries([triangle_mesh])
    vis.add_geometry(pcd)
    vis.get_render_option().load_from_json("renderoption.json")
    vis.poll_events()
    vis.update_renderer()
    time.sleep(1)
    #save the normal
    vis.capture_screen_image(output_direcory_normal +'/' + point_cloud_file.split('/')[-1] + '.png')
    vis.clear_geometries()

#mesh generation and visualization script
print("estimate the surface and save ply")
vis = o3d.visualization.Visualizer()
vis.create_window(visible = False) 
#generate the .ply files and the meshes
for point_cloud_file in correct_file_path: 
    print(point_cloud_file)
    pcd = o3d.io.read_point_cloud(point_cloud_file,format="xyz")
    pcd.paint_uniform_color([0,0.5,0.75])
    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(
        radius=0.1, max_nn=30))
    #once the normal has been estimated compute the corresponding mesh 
    triangle_mesh, double_vector = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd)
    vis.add_geometry(triangle_mesh)
    vis.get_render_option().load_from_json("renderoption.json")
    vis.poll_events()
    vis.update_renderer()
    time.sleep(1)
    #save the normal
    vis.capture_screen_image(output_directory_mesh_img +'/' + point_cloud_file.split('/')[-1] + '.png')
    o3d.io.write_triangle_mesh(output_directory_mesh +'/' + point_cloud_file.split('/')[-1] + '.ply', triangle_mesh)
    vis.clear_geometries()