import argparse
import json
import os

import tqdm
from visual_nuscenes import NuScenes


def parse_args():
    parser = argparse.ArgumentParser(description="parse args")
    parser.add_argument("result_json", type=str)
    parser.add_argument("out_dir", default="./result_vis/", type=str)
    parser.add_argument("--use_gt", action="store_true")
    parser.add_argument("--data_root", default="data/nuscenes", type=str)
    parser.add_argument("--score_thresh", default=0.25, type=float)
    parser.add_argument("--num_vis", default=100, type=int)

    args = parser.parse_args()
    return args


if __name__ == "__main__":

    args = parse_args()

    use_gt = args.use_gt
    result_json = args.result_json
    out_dir = args.out_dir
    dataroot = args.data_root
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    if use_gt:
        nusc = NuScenes(
            version="v1.0-trainval",
            dataroot=dataroot,
            verbose=True,
            pred=False,
            annotations="sample_annotation",
        )
    else:
        nusc = NuScenes(
            version="v1.0-trainval",
            dataroot=dataroot,
            verbose=True,
            pred=True,
            annotations=result_json,
            score_thr=args.score_thresh,
        )

    with open("{}.json".format(result_json)) as f:
        table = json.load(f)
    tokens = list(table["results"].keys())

    num_vis = min(args.num_vis, len(tokens))
    for token in tqdm.tqdm(tokens[:num_vis]):
        if use_gt:
            out_path = os.path.join(out_dir, token + "_gt.png")
            nusc.render_sample(token, out_path=out_path, verbose=False)
        else:
            out_path = os.path.join(out_dir, token + "_pred.png")
            nusc.render_sample(token, out_path=out_path, verbose=False)
