#!/usr/bin/env python3

import argparse, os, shutil, tqdm
import pandas as pd
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

files = [
    "pandaset_0.zip",
    "pandaset_1.zip",
    "pandaset_2.zip",
]

for file in files:
    if not os.path.isfile(os.path.join(args.path, file)):
        print("Please download and place into the folder specified by --path:")
        print("    Part 1: pandaset_0.zip")
        print("    Part 2: pandaset_1.zip")
        print("    Part 3: pandaset_2.zip")
        sys.exit(-1)

for file in files:
    georeg.data.prepare.extract(os.path.join(args.path, file), args.path)

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

for scene_id in tqdm.tqdm(sorted(os.listdir(args.path))):
    lidar_path = os.path.join(args.path, scene_id, "lidar")
    if os.path.isdir(lidar_path):
        for lidar_file in sorted([os.path.join(lidar_path, f) for f in os.listdir(lidar_path) if f.endswith(".pkl.gz")]):
            points = pd.read_pickle(lidar_file).to_numpy()
            np.savez_compressed(lidar_file[:-6] + "npz", points)
            os.remove(lidar_file)
