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

# This mapping is extracted from the official LVIS mapping:
# https://github.com/lvis-dataset/lvis-api/blob/master/data/coco_to_synset.json
COCO_SYNSET_CATEGORIES = [
    {"synset": "person.n.01", "coco_cat_id": 1},
    {"synset": "bicycle.n.01", "coco_cat_id": 2},
    {"synset": "car.n.01", "coco_cat_id": 3},
    {"synset": "motorcycle.n.01", "coco_cat_id": 4},
    {"synset": "airplane.n.01", "coco_cat_id": 5},
    {"synset": "bus.n.01", "coco_cat_id": 6},
    {"synset": "train.n.01", "coco_cat_id": 7},
    {"synset": "truck.n.01", "coco_cat_id": 8},
    {"synset": "boat.n.01", "coco_cat_id": 9},
    {"synset": "traffic_light.n.01", "coco_cat_id": 10},
    {"synset": "fireplug.n.01", "coco_cat_id": 11},
    {"synset": "stop_sign.n.01", "coco_cat_id": 13},
    {"synset": "parking_meter.n.01", "coco_cat_id": 14},
    {"synset": "bench.n.01", "coco_cat_id": 15},
    {"synset": "bird.n.01", "coco_cat_id": 16},
    {"synset": "cat.n.01", "coco_cat_id": 17},
    {"synset": "dog.n.01", "coco_cat_id": 18},
    {"synset": "horse.n.01", "coco_cat_id": 19},
    {"synset": "sheep.n.01", "coco_cat_id": 20},
    {"synset": "beef.n.01", "coco_cat_id": 21},
    {"synset": "elephant.n.01", "coco_cat_id": 22},
    {"synset": "bear.n.01", "coco_cat_id": 23},
    {"synset": "zebra.n.01", "coco_cat_id": 24},
    {"synset": "giraffe.n.01", "coco_cat_id": 25},
    {"synset": "backpack.n.01", "coco_cat_id": 27},
    {"synset": "umbrella.n.01", "coco_cat_id": 28},
    {"synset": "bag.n.04", "coco_cat_id": 31},
    {"synset": "necktie.n.01", "coco_cat_id": 32},
    {"synset": "bag.n.06", "coco_cat_id": 33},
    {"synset": "frisbee.n.01", "coco_cat_id": 34},
    {"synset": "ski.n.01", "coco_cat_id": 35},
    {"synset": "snowboard.n.01", "coco_cat_id": 36},
    {"synset": "ball.n.06", "coco_cat_id": 37},
    {"synset": "kite.n.03", "coco_cat_id": 38},
    {"synset": "baseball_bat.n.01", "coco_cat_id": 39},
    {"synset": "baseball_glove.n.01", "coco_cat_id": 40},
    {"synset": "skateboard.n.01", "coco_cat_id": 41},
    {"synset": "surfboard.n.01", "coco_cat_id": 42},
    {"synset": "tennis_racket.n.01", "coco_cat_id": 43},
    {"synset": "bottle.n.01", "coco_cat_id": 44},
    {"synset": "wineglass.n.01", "coco_cat_id": 46},
    {"synset": "cup.n.01", "coco_cat_id": 47},
    {"synset": "fork.n.01", "coco_cat_id": 48},
    {"synset": "knife.n.01", "coco_cat_id": 49},
    {"synset": "spoon.n.01", "coco_cat_id": 50},
    {"synset": "bowl.n.03", "coco_cat_id": 51},
    {"synset": "banana.n.02", "coco_cat_id": 52},
    {"synset": "apple.n.01", "coco_cat_id": 53},
    {"synset": "sandwich.n.01", "coco_cat_id": 54},
    {"synset": "orange.n.01", "coco_cat_id": 55},
    {"synset": "broccoli.n.01", "coco_cat_id": 56},
    {"synset": "carrot.n.01", "coco_cat_id": 57},
    # {"synset": "frank.n.02", "coco_cat_id": 58},
    {"synset": "sausage.n.01", "coco_cat_id": 58},
    {"synset": "pizza.n.01", "coco_cat_id": 59},
    {"synset": "doughnut.n.02", "coco_cat_id": 60},
    {"synset": "cake.n.03", "coco_cat_id": 61},
    {"synset": "chair.n.01", "coco_cat_id": 62},
    {"synset": "sofa.n.01", "coco_cat_id": 63},
    {"synset": "pot.n.04", "coco_cat_id": 64},
    {"synset": "bed.n.01", "coco_cat_id": 65},
    {"synset": "dining_table.n.01", "coco_cat_id": 67},
    {"synset": "toilet.n.02", "coco_cat_id": 70},
    {"synset": "television_receiver.n.01", "coco_cat_id": 72},
    {"synset": "laptop.n.01", "coco_cat_id": 73},
    {"synset": "mouse.n.04", "coco_cat_id": 74},
    {"synset": "remote_control.n.01", "coco_cat_id": 75},
    {"synset": "computer_keyboard.n.01", "coco_cat_id": 76},
    {"synset": "cellular_telephone.n.01", "coco_cat_id": 77},
    {"synset": "microwave.n.02", "coco_cat_id": 78},
    {"synset": "oven.n.01", "coco_cat_id": 79},
    {"synset": "toaster.n.02", "coco_cat_id": 80},
    {"synset": "sink.n.01", "coco_cat_id": 81},
    {"synset": "electric_refrigerator.n.01", "coco_cat_id": 82},
    {"synset": "book.n.01", "coco_cat_id": 84},
    {"synset": "clock.n.01", "coco_cat_id": 85},
    {"synset": "vase.n.01", "coco_cat_id": 86},
    {"synset": "scissors.n.01", "coco_cat_id": 87},
    {"synset": "teddy.n.01", "coco_cat_id": 88},
    {"synset": "hand_blower.n.01", "coco_cat_id": 89},
    {"synset": "toothbrush.n.01", "coco_cat_id": 90},
]

def map_name(x):
    x = x.replace('_', ' ')
    if '(' in x:
        x = x[:x.find('(')]
    return x.lower().strip()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cc_ann', default='datasets/cc3m/train_image_info.json')
    parser.add_argument('--out_path', default='datasets/cc3m/train_image_info_tags.json')
    parser.add_argument('--keep_images', action='store_true')
    parser.add_argument('--allcaps', action='store_true')
    parser.add_argument('--cat_path', default='')
    parser.add_argument('--convert_caption', action='store_true')
    # parser.add_argument('--lvis_ann', default='datasets/lvis/lvis_v1_val.json')
    args = parser.parse_args()

    # lvis_data = json.load(open(args.lvis_ann, 'r'))
    cc_data = json.load(open(args.cc_ann, 'r'))
    if args.convert_caption:
        num_caps = 0
        caps = defaultdict(list)
        for x in cc_data['annotations']:
            caps[x['image_id']].append(x['caption'])
        for x in cc_data['images']:
            x['captions'] = caps[x['id']]
            num_caps += len(x['captions'])
        print('# captions', num_caps)

    if args.cat_path != '':
        print('Loading', args.cat_path)
        cats = json.load(open(args.cat_path))['categories']
        if 'synonyms' not in cats[0]:
            cocoid2synset = {x['coco_cat_id']: x['synset'] \
                for x in COCO_SYNSET_CATEGORIES}
            synset2synonyms = {x['synset']: x['synonyms'] \
                for x in cc_data['categories']}
            for x in cats:
                synonyms = synset2synonyms[cocoid2synset[x['id']]]
                x['synonyms'] = synonyms
                x['frequency'] = 'f'
        cc_data['categories'] = cats

    id2cat = {x['id']: x for x in cc_data['categories']}
    class_count = {x['id']: 0 for x in cc_data['categories']}
    class_data = {x['id']: [' ' + map_name(xx) + ' ' for xx in x['synonyms']] \
            for x in cc_data['categories']}
    num_examples = 5
    examples = {x['id']: [] for x in cc_data['categories']}

    print('class_data', class_data)

    images = []
    for i, x in enumerate(cc_data['images']):
        if i % 10000 == 0:
            print(i, len(cc_data['images']))
        if args.allcaps:
            caption = (' '.join(x['captions'])).lower()
        else:
            caption = x['captions'][0].lower()
        x['pos_category_ids'] = []
        for cat_id, cat_names in class_data.items():
            find = False
            for c in cat_names:
                if c in caption or caption.startswith(c[1:]) \
                    or caption.endswith(c[:-1]):
                    find = True
                    break
            if find:
                x['pos_category_ids'].append(cat_id)
                class_count[cat_id] += 1
                if len(examples[cat_id]) < num_examples:
                    examples[cat_id].append(caption)
        if len(x['pos_category_ids']) > 0 or args.keep_images:
            images.append(x)

    zero_class = []
    for cat_id, count in class_count.items():
        print(id2cat[cat_id]['name'], count, end=', ')
        if count == 0:
            zero_class.append(id2cat[cat_id])
    print('==')
    print('zero class', zero_class)

    # for freq in ['r', 'c', 'f']:
    #     print('#cats', freq, len([x for x in cc_data['categories'] \
    #         if x['frequency'] == freq] and class_count[x['id']] > 0))

    for freq in ['r', 'c', 'f']:
        print('#Images', freq, sum([v for k, v in class_count.items() \
        if id2cat[k]['frequency'] == freq]))

    try:
        out_data = {'images': images, 'categories': cc_data['categories'], \
            'annotations': []}
        for k, v in out_data.items():
            print(k, len(v))
        if args.keep_images and not args.out_path.endswith('_full.json'):
            args.out_path = args.out_path[:-5] + '_full.json'
        print('Writing to', args.out_path)
        json.dump(out_data, open(args.out_path, 'w'))
    except:
        pass