import argparse
from pathlib import Path
import open3d as o3d
from tqdm.contrib.concurrent import thread_map, process_map
import numpy as np

def process_scene(data):
    scene_dir, output_dir = data
    point_cloud = o3d.io.read_point_cloud(scene_dir)
    xyz = np.asarray(point_cloud.points)

    rgb = np.array(point_cloud.colors)

    rgb = np.clip(rgb, 0, 255)[:, :3]

    # if rgb [0, 1] then change to [0, 255]
    if not len(rgb):
        return None
        
    if rgb.max() <= 1:
        rgb = (rgb * 255)

    points = np.concatenate([xyz, rgb], axis=1).astype(np.float32)
    output_path = output_dir / f"{scene_dir.parent.name}_point.bin"
    print(f"saving {scene_dir} to {output_path}")
    points.tofile(output_path)

    # print(points, points.shape)

    return output_path

def load_scan(pcd_path):
    pcd_data = np.fromfile(pcd_path, dtype=np.float32).reshape(-1, 6)
    return pcd_data    

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input",  type=str, required=True)
    parser.add_argument("-o", "--output", type=str, required=True)
    args = parser.parse_args()

    input_dir = Path(args.input) / "processed"
    output_dir = Path(args.output)

    output_dir.mkdir(parents=True, exist_ok=True)

    input_files = list(input_dir.glob("*/*.ply"))
    data = [*zip(input_files, [output_dir] * len(input_files))]

    process_map(process_scene, data, max_workers=16)
    

if __name__ == "__main__":
    main()