# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0

import argparse
import numpy as np
from tqdm import tqdm

from src.tools.utils import (
    get_mscoco_idmap, 
    load_all_gt_games,
    xywh2xyxy
    )


def run(args):
    idmap = get_mscoco_idmap(args.img_dir)
    games = load_all_gt_games(args.src, verbose=True)
    results = []
    # with jsonlines.open(args.src) as reader:
    img_done = set()
    for g in tqdm(games):
        if g.image_info['id'] in img_done:
            continue
        item = dict()

        # remove later: only for missing npy
        if g.image_info['id'] not in idmap: continue

        item['file_name'] = g.image_info['file_name']
        item['file_path'] = idmap[g.image_info['id']]
        item['bbox'] = np.array([
            xywh2xyxy(box) for box in g.bboxs])
        item['num_box'] = len(item['bbox'])
        results.append(item)
        img_done |= set([g.image_info['id']])
    print("[INFO] %d different images used in the games." % len(results))
    np.save(args.out, np.array(results))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Run the training or testing.')
    parser.add_argument('--src', required=True, type=str,
                        help='one of guesswhat.{train/valid/test}.jsonl file.')
    parser.add_argument('--img-dir', required=True, type=str,
                        help='The root directory where you put COCO dataset.')
    parser.add_argument('--out', required=True, type=str,
                        help='Output path.')
    args = parser.parse_args()
    run(args)
