#!/usr/bin/env python3

import argparse, os, pyntcloud, tqdm
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument("--path", type=str, required=True)
parser.add_argument("--width", type=int, default=-1)
parser.add_argument("--height", type=int, default=-1)
args = parser.parse_args()
assert (args.width > 0) == (args.height > 0)

import georeg

urls = [
    "https://s3.amazonaws.com/argoai-argoverse/tracking_train1_v1.1.tar.gz",
    "https://s3.amazonaws.com/argoai-argoverse/tracking_train2_v1.1.tar.gz",
    "https://s3.amazonaws.com/argoai-argoverse/tracking_train3_v1.1.tar.gz",
    "https://s3.amazonaws.com/argoai-argoverse/tracking_train4_v1.1.tar.gz",
    "https://s3.amazonaws.com/argoai-argoverse/tracking_val_v1.1.tar.gz",
    "https://s3.amazonaws.com/argoai-argoverse/tracking_test_v1.1.tar.gz"
]

for url in urls:
    file = os.path.join(args.path, url.split("/")[-1])
    georeg.data.prepare.download(url, file)
    georeg.data.prepare.extract(file, args.path)

# TODO: rm -rf */*/*/.ipynb_checkpoints

if args.height > 0:
    georeg.data.prepare.resize(args.path, (args.height, args.width))

tasks = []
path = os.path.join(args.path, "argoverse-tracking")
for split in os.listdir(path):
    split_path = os.path.join(path, split)
    if os.path.isdir(split_path):
        for scene in os.listdir(split_path):
            scene_path = os.path.join(split_path, scene)
            if os.path.isdir(scene_path):
                lidar_path = os.path.join(scene_path, "lidar")
                for lidar_file in [os.path.join(lidar_path, f) for f in os.listdir(lidar_path) if f.endswith(".ply")]:
                    tasks.append(lidar_file)
for lidar_file in tqdm.tqdm(tasks):
    data = pyntcloud.PyntCloud.from_file(lidar_file)
    points = np.stack((np.array(data.points.x), np.array(data.points.y), np.array(data.points.z)), axis=-1) # In ego space
    np.savez_compressed(lidar_file[:-3] + "npz", points)
    os.remove(lidar_file)
