import argparse

import imagesize

import numpy as np

import os


base_path = "data/megadepth"
# Remove the trailing / if need be.
if base_path[-1] in ["/", "\\"]:
    base_path = base_path[:-1]


base_depth_path = os.path.join(base_path, "phoenix/S6/zl548/MegaDepth_v1")
base_undistorted_sfm_path = os.path.join(base_path, "Undistorted_SfM")

scene_ids = os.listdir(base_undistorted_sfm_path)
for scene_id in scene_ids:
    if os.path.exists(
        f"{base_path}/prep_scene_info/detections/detections_{scene_id}.npy"
    ):
        print(f"skipping {scene_id} as it exists")
        continue
    undistorted_sparse_path = os.path.join(
        base_undistorted_sfm_path, scene_id, "sparse-txt"
    )
    if not os.path.exists(undistorted_sparse_path):
        print("sparse path doesnt exist")
        continue

    depths_path = os.path.join(base_depth_path, scene_id, "dense0", "depths")
    if not os.path.exists(depths_path):
        print("depths doesnt exist")

        continue

    images_path = os.path.join(base_undistorted_sfm_path, scene_id, "images")
    if not os.path.exists(images_path):
        print("images path doesnt exist")
        continue

    # Process cameras.txt
    if not os.path.exists(os.path.join(undistorted_sparse_path, "cameras.txt")):
        print("no cameras")
        continue
    with open(os.path.join(undistorted_sparse_path, "cameras.txt"), "r") as f:
        raw = f.readlines()[3:]  # skip the header

    camera_intrinsics = {}
    for camera in raw:
        camera = camera.split(" ")
        camera_intrinsics[int(camera[0])] = [float(elem) for elem in camera[2:]]

    # Process points3D.txt
    with open(os.path.join(undistorted_sparse_path, "points3D.txt"), "r") as f:
        raw = f.readlines()[3:]  # skip the header

    points3D = {}
    for point3D in raw:
        point3D = point3D.split(" ")
        points3D[int(point3D[0])] = np.array(
            [float(point3D[1]), float(point3D[2]), float(point3D[3])]
        )

    # Process images.txt
    with open(os.path.join(undistorted_sparse_path, "images.txt"), "r") as f:
        raw = f.readlines()[4:]  # skip the header

    image_id_to_idx = {}
    image_names = []
    raw_pose = []
    camera = []
    points3D_id_to_2D = []
    n_points3D = []
    id_to_detections = {}
    for idx, (image, points) in enumerate(zip(raw[::2], raw[1::2])):
        image = image.split(" ")
        points = points.split(" ")

        image_id_to_idx[int(image[0])] = idx

        image_name = image[-1].strip("\n")
        image_names.append(image_name)

        raw_pose.append([float(elem) for elem in image[1:-2]])
        camera.append(int(image[-2]))
        points_np = np.array(points).astype(np.float32).reshape(len(points) // 3, 3)
        visible_points = points_np[points_np[:, 2] != -1]
        id_to_detections[idx] = visible_points
    np.save(
        f"{base_path}/prep_scene_info/detections/detections_{scene_id}.npy",
        id_to_detections,
    )
    print(f"{scene_id} done")