#!/usr/bin/env python3

import argparse, os, shutil
from tqdm import tqdm

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 = [
    "test.tar",
    "train.tar",
]

for file in files:
    if not os.path.isfile(os.path.join(args.path, file)):
        print("Please download from the Lyft level-5 website and place into the folder specified by --path:")
        print("    Perception Dataset > Training dataset")
        print("    Perception Dataset > Testing and Validation dataset")
        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(os.path.join(args.path, "test_images"), (args.height, args.width))
    georeg.data.prepare.resize(os.path.join(args.path, "train_images"), (args.height, args.width))

# Change folder structure
for split in ["train", "test"]:
    split_path = os.path.join(args.path, split)
    os.makedirs(split_path)
    for type in ["data", "images", "lidar", "maps"]:
        src = os.path.join(args.path, f"{split}_{type}")
        dest = os.path.join(args.path, split, type)
        shutil.move(src, dest)
