import argparse
import numpy as np
import open3d as o3d

parser = argparse.ArgumentParser()
parser.add_argument("--npz-path", type=str, required=True)
args = parser.parse_args()

data = np.load(args.npz_path)
print(data.files)
points = data["vertices"]
normals = data["vertex_normals"]
num_viewpoints = int(data["num_viewpoints"])
print(points.shape)
print(num_viewpoints)

pcd_list = []
view_index = 10

pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)
pcd.normals = o3d.utility.Vector3dVector(normals)
pcd.colors = o3d.utility.Vector3dVector(
    np.repeat(np.array((0, 1, 0))[None, :], len(points), axis=0))
pcd_list.append(pcd)

partial_point_indices = data[f"partial_point_indices_{view_index}"]
partial_points = points[partial_point_indices]
partial_normals = normals[partial_point_indices]
color = (1, 0, 0)
partial_colors = np.repeat(np.array(color)[None, :],
                           len(partial_points),
                           axis=0)
partial_pcd = o3d.geometry.PointCloud()
partial_pcd.points = o3d.utility.Vector3dVector(partial_points)
partial_pcd.normals = o3d.utility.Vector3dVector(partial_normals)
partial_pcd.colors = o3d.utility.Vector3dVector(partial_colors)
pcd_list.append(partial_pcd)


def change_background_to_black(vis):
    opt = vis.get_render_option()
    opt.background_color = np.asarray([0, 0, 0])
    return False


key_to_callback = {}
key_to_callback[ord("K")] = change_background_to_black
o3d.visualization.draw_geometries_with_key_callbacks(pcd_list, key_to_callback)
