########
#
# get_lila_image_counts.py
#
# Count the number of images and bounding boxes with each label in one or more LILA datasets.
#
# This script doesn't write these counts out anywhere other than the console, it's just intended
# as a template for doing operations like this on LILA data.  get_lila_annotation_counts.py writes 
# information out to a .json file, but it counts *annotations*, not *images*, for each category.
#
########

#%% Constants and imports

import json
import os

from collections import defaultdict

from data_management.lila.lila_common import read_lila_metadata, read_metadata_file_for_dataset

# If None, will use all datasets
datasets_of_interest = None

# We'll write images, metadata downloads, and temporary files here
lila_local_base = os.path.expanduser('~/lila')

metadata_dir = os.path.join(lila_local_base,'metadata')
os.makedirs(metadata_dir,exist_ok=True)


#%% Download and parse the metadata file

metadata_table = read_lila_metadata(metadata_dir)


#%% Download and extract metadata for the datasets we're interested in

if datasets_of_interest is None:
    datasets_of_interest = list(metadata_table.keys())

for ds_name in datasets_of_interest:    
    metadata_table[ds_name]['json_filename'] = read_metadata_file_for_dataset(ds_name=ds_name,
                                                                         metadata_dir=metadata_dir,
                                                                         metadata_table=metadata_table)
    
    
#%% Count categories

ds_name_to_category_counts = {}

# ds_name = datasets_of_interest[0]
for ds_name in datasets_of_interest:
    
    category_to_image_count = {}
    category_to_bbox_count = {}
    
    print('Counting categories in: ' + ds_name)
    
    json_filename = metadata_table[ds_name]['json_filename']
    with open(json_filename, 'r') as f:
        data = json.load(f)
    
    categories = data['categories']
    category_ids = [c['id'] for c in categories]
    for c in categories:
        category_id_to_name = {c['id']:c['name'] for c in categories}
    annotations = data['annotations']
    images = data['images']
    
    for category_id in category_ids:        
        category_name = category_id_to_name[category_id]        
        category_to_image_count[category_name] = 0
        category_to_bbox_count[category_name] = 0
        
    image_id_to_category_names = defaultdict(set)
    
    # Go through annotations, marking each image with the categories that are present
    #
    # ann = annotations[0]
    for ann in annotations:
        
        category_name = category_id_to_name[ann['category_id']]
        image_id_to_category_names[ann['image_id']].add(category_name)

    # Now go through images and count categories
    category_to_count = defaultdict(int)
    
    # im = images[0]
    for im in images:
        categories_this_image = image_id_to_category_names[im['id']]
        for category_name in categories_this_image:
            category_to_count[category_name] += 1

    ds_name_to_category_counts[ds_name] = category_to_count
    
# ...for each dataset
    

#%% Print the results

for ds_name in ds_name_to_category_counts:
    
    print('\n** Category counts for {} **\n'.format(ds_name))
    
    category_to_count = ds_name_to_category_counts[ds_name]
    category_to_count = {k: v for k, v in sorted(category_to_count.items(), reverse=True, 
                                                 key=lambda item: item[1])}
    
    for category_name in category_to_count.keys():        
        print('{}: {}'.format(category_name,category_to_count[category_name]))
