# Copyright (c) Facebook, Inc. and its affiliates.
import argparse
import json

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--ann", default='datasets/lvis/lvis_v1_train.json')
    parser.add_argument("--add_freq", action='store_true')
    parser.add_argument("--r_thresh", type=int, default=10)
    parser.add_argument("--c_thresh", type=int, default=100)
    args = parser.parse_args()

    print('Loading', args.ann)
    data = json.load(open(args.ann, 'r'))
    cats = data['categories']
    image_count = {x['id']: set() for x in cats}
    ann_count = {x['id']: 0 for x in cats}
    for x in data['annotations']:
        image_count[x['category_id']].add(x['image_id'])
        ann_count[x['category_id']] += 1
    num_freqs = {x: 0 for x in ['r', 'f', 'c']}
    for x in cats:
        x['image_count'] = len(image_count[x['id']])
        x['instance_count'] = ann_count[x['id']]
        if args.add_freq:
            freq = 'f'
            if x['image_count'] < args.c_thresh:
                freq = 'c'
            if x['image_count'] < args.r_thresh:
                freq = 'r'
            x['frequency'] = freq
            num_freqs[freq] += 1
    print(cats)
    image_counts = sorted([x['image_count'] for x in cats])
    # print('image count', image_counts)
    # import pdb; pdb.set_trace()
    if args.add_freq:
        for x in ['r', 'c', 'f']:
            print(x, num_freqs[x])
    out = cats # {'categories': cats}
    out_path = args.ann[:-5] + '_cat_info.json'
    print('Saving to', out_path)
    json.dump(out, open(out_path, 'w'))