import argparse
import os
import json
import numpy as np
import imageio

from create_input import depth_to_world_coords_points, camera_list


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--folder_path", type=str)
    parser.add_argument("--video_path", type=str)
    parser.add_argument("--frame_id", type=int, default=0)
    parser.add_argument("--max_depth", type=float, default=25)
    return parser.parse_args()


def save_ply(points: np.ndarray, colors: np.ndarray, out_path: str):
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    n = points.shape[0]
    colors = np.clip(colors, 0, 255).astype(np.uint8)
    header = (
        "ply\n"
        "format ascii 1.0\n"
        f"element vertex {n}\n"
        "property float x\n"
        "property float y\n"
        "property float z\n"
        "property uchar red\n"
        "property uchar green\n"
        "property uchar blue\n"
        "end_header\n"
    )
    with open(out_path, "w") as f:
        f.write(header)
        for p, c in zip(points, colors):
            f.write(f"{float(p[0])} {float(p[1])} {float(p[2])} {int(c[0])} {int(c[1])} {int(c[2])}\n")


if __name__ == "__main__":
    args = parse_args()
    folder_path = args.folder_path
    video_path = args.video_path
    frame_id = args.frame_id
    max_depth = args.max_depth

    reader = imageio.v2.get_reader(video_path)
    for i, frame in enumerate(reader):
        if i == frame_id:
            frame = frame.astype(np.uint8)
            break

    with open(os.path.join(folder_path, "depth_range.json"), "r") as f:
        depth_range = json.load(f)[frame_id]

    rgb = frame[:512]
    depth = frame[512:, :, 0] / 255.0
    depth = depth * (depth_range[1] - depth_range[0]) + depth_range[0]
    depth = 1 / (depth + 1e-6)
    valid_mask = np.logical_and(depth > 0, depth < max_depth)

    intrinsics, extrinsics = camera_list(
        num_frames=1, type="forward", Width=512, Height=512, fx=256, fy=256
    )
    point_map = depth_to_world_coords_points(depth, extrinsics[0], intrinsics[0])
    points = point_map[valid_mask].reshape(-1, 3)
    colors = rgb[valid_mask].reshape(-1, 3)

    out_ply = os.path.join(folder_path, f"frame_{frame_id:06d}.ply")
    save_ply(points, colors, out_ply)
    print(f"Saved point cloud: {out_ply}, number of points: {points.shape[0]}")