import os
import argparse
import json
from collections import defaultdict


def uied2coco(uied_dir, img_id_map, cat_id_map, output_file = None):
    pred = []
    for file in os.listdir(uied_dir):
        if not file.endswith('.json'):
            continue
        img_id = img_id_map[file.split('.')[-2]]
        if img_id is None:
            continue
        with open(os.path.join(uied_dir, file), 'r') as f:
            data = json.load(f)['compos']
        for i in range(len(data)):
            cat_id = cat_id_map[data[i]['class']]
            x = data[i]['position']['column_min']
            y = data[i]['position']['row_min']
            width = data[i]['width']
            height = data[i]['height']
            bbox = [x, y, width, height]
            score = 1
            pred.append({'image_id': img_id, 'category_id': cat_id, 'bbox': bbox, 'score': score})
    if output_file is not None:
        with open(output_file, 'w') as f:
            json.dump(pred, f)
    return pred


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--uied_dir', type=str, required=True)
    parser.add_argument('--ann', type=str, required=True)
    parser.add_argument('--output_file', type=str, required=True)
    args = parser.parse_args()
    with open(args.ann, 'r') as f:
        ann = json.load(f)
    img_id_map = defaultdict(lambda: None)
    for img in ann['images']:
        img_id_map[img['file_name'].split('.')[-2]] = img['id']
    cat_id_map = defaultdict(lambda: 1)
    for cat in ann['categories']:
        cat_id_map[cat['name']] = cat['id']
    uied2coco(args.uied_dir, img_id_map, cat_id_map, args.output_file)