import argparse
from glob import glob
import json
import os

import tqdm as tq

# TODO(jfemiani) super categories
aliases = {
    '_ignore_': dict(label='__ignore__', occluded=False),
    #  Convert ROI to ___ignore__ regions.
    'ROI': dict(label='roi', occluded=False),
    'Balcony': dict(label='balcony', occluded=False),
    'Bay Window': dict(label='bay window', occluded=False),
    'Blind': dict(label='blind', occluded=False),
    'Blind/Window': dict(label='blind', occluded=False),
    'Building': dict(label='building', occluded=False),
    'Facade': dict(label='building', occluded=False),  # TODO(jfemiani) Need subategory
    'Chimney': dict(label='chimney', occluded=False),
    'Column': dict(label='column', occluded=False),
    'Deco': dict(label='deco', occluded=False),
    'Door': dict(label='door', occluded=False),
    'Cornice': dict(label='cornice', occluded=False),
    'Ground': dict(label='ground', occluded=False),
    'Molding': dict(label='molding', occluded=False),
    'Roof': dict(label='roof', occluded=False),
    'Shop/billboards': dict(label='shop/sign', occluded=False),
    'Sill': dict(label='sill', occluded=False),
    'Sky': dict(label='sky', occluded=False),
    'Vegetation': dict(label='vegetation', occluded=False),
    'Vehicle': dict(label='vehicle', occluded=False),
    'Person': dict(label='person', occluded=False),
    'Windows': dict(label='window', occluded=False),
    #
    # Datapure uses separate labels to mark occluded objects
    'Occluded/W': dict(label='window', occluded=True),
    'Occluded/V': dict(label='vehicle', occluded=True),
    'Occluded/B': dict(label='building', occluded=True),
    'Occluded/D': dict(label='door', occluded=True),
    'Occluded/C': dict(label='column', occluded=True),
}


def convert(folder):
    """Convert labels from a folder.

    This will output labelm': dict(label='', occluded=False), and labels files

    Args:
        folder (str): The root folder for a batch of annotations (contains 'images.json' and 'classes.json' )
    """
    pass


def read_json(path):
    with open(path) as f:
        return json.load(f)


def run_list(args):
    """List all of the labels and annotations we have

    Args:
        args (Namesppace/CfgNode/Easydict): The options.
    """
    all_classes = {}

    annotation_files = glob(f'{args.folder}/**/annotations.json', recursive=True)
    folders = [os.path.dirname(f) for f in annotation_files]
    classes_files = [os.path.join(f, 'classes.json') for f in folders]

    for i, folder in enumerate(tq.tqdm(folders, leave=False)):
        classes_path = classes_files[i]
        annotation_path = annotation_files[i]

        try:
            c = read_json(classes_path)
        except FileNotFoundError:
            tq.tqdm.write(f"Missing {classes_path}")
            continue

        for cls in c:
            name = cls['name']
            if name in all_classes:
                all_classes[name].append(cls)
            else:
                all_classes[name] = [cls]

    for name in all_classes:
        print(f'{name:20} => '
              f'{aliases[name]["label"]:15}: '
              f'{len(all_classes[name])} batches')


p = argparse.ArgumentParser()

p.add_argument('--folder', type=str, default=os.curdir,
               help='Root folder(s) for labelled batches. Each has a annotations.json and classes.json')
sps = p.add_subparsers()
sp_list = sps.add_parser('list', help="List all of the classes used in every batch/subfolder")
sp_list.set_defaults(action=run_list)

args = p.parse_args()


# TODO: Replace ROI with 4 __ignore__ regions

args.action(args)