"""

coco_to_yolo.py

Converts a COCO-formatted dataset to a YOLO-formatted dataset, flattening
the dataset (to a single folder) in the process.

If the input and output folders are the same, writes .txt files to the input folder,
and neither moves nor modifies images.

Currently ignores segmentation masks, and errors if an annotation has a
segmentation polygon but no bbox.

Has only been tested on a handful of COCO Camera Traps data sets; if you
use it for more general COCO conversion, YMMV.

"""

#%% Imports and constants

import json
import os
import shutil
import sys
import argparse

from collections import defaultdict
from tqdm import tqdm

from megadetector.utils.path_utils import safe_create_link,find_images


#%% Support functions

def write_yolo_dataset_file(yolo_dataset_file,
                            dataset_base_dir,
                            class_list,
                            train_folder_relative=None,
                            val_folder_relative=None,
                            test_folder_relative=None):
    """
    Write a YOLOv5 dataset.yaml file to the absolute path [yolo_dataset_file] (should
    have a .yaml extension, though it's only a warning if it doesn't).

    Args:
        yolo_dataset_file (str): the file, typically ending in .yaml or .yml, to write.
            Does not have to be within dataset_base_dir.
        dataset_base_dir (str): the absolute base path of the YOLO dataset
        class_list (list or str): an ordered list of class names (the first item will be class 0,
            etc.), or the name of a text file containing an ordered list of class names (one per
            line, starting from class zero).
        train_folder_relative (str, optional): train folder name, used only to
            populate dataset.yaml.  Can also be a filename (e.g. a .txt file with image
            files).
        val_folder_relative (str, optional): val folder name, used only to
            populate dataset.yaml.  Can also be a filename (e.g. a .txt file with image
            files).
        test_folder_relative (str, optional): test folder name, used only to
            populate dataset.yaml.  Can also be a filename (e.g. a .txt file with image
            files).
    """

    # Read class names
    if isinstance(class_list,str):
        with open(class_list,'r') as f:
            class_lines = f.readlines()
        class_lines = [s.strip() for s in class_lines]
        class_list = [s for s in class_lines if len(s) > 0]

    if not (yolo_dataset_file.endswith('.yml') or yolo_dataset_file.endswith('.yaml')):
        print('Warning: writing dataset file to a non-yml/yaml extension:\n{}'.format(
            yolo_dataset_file))

    # Write dataset.yaml
    with open(yolo_dataset_file,'w') as f:

        f.write('# Train/val sets\n')
        f.write('path: {}\n'.format(dataset_base_dir))
        if train_folder_relative is not None:
            f.write('train: {}\n'.format(train_folder_relative))
        if val_folder_relative is not None:
            f.write('val: {}\n'.format(val_folder_relative))
        if test_folder_relative is not None:
            f.write('val: {}\n'.format(test_folder_relative))

        f.write('\n')

        f.write('# Classes\n')
        f.write('names:\n')
        for i_class,class_name in enumerate(class_list):
            f.write('  {}: {}\n'.format(i_class,class_name))

# ...def write_yolo_dataset_file(...)


def coco_to_yolo(input_image_folder,
                 output_folder,
                 input_file,
                 source_format='coco',
                 overwrite_images=False,
                 create_image_and_label_folders=False,
                 class_file_name='classes.txt',
                 allow_empty_annotations=False,
                 clip_boxes=False,
                 image_id_to_output_image_json_file=None,
                 images_to_exclude=None,
                 path_replacement_char='#',
                 category_names_to_exclude=None,
                 category_names_to_include=None,
                 write_output=True,
                 flatten_paths=False):
    """
    Converts a COCO-formatted dataset to a YOLO-formatted dataset, optionally flattening the
    dataset to a single folder in the process.

    If the input and output folders are the same, writes .txt files to the input folder,
    and neither moves nor modifies images.

    Currently ignores segmentation masks, and errors if an annotation has a
    segmentation polygon but no bbox.

    Args:
        input_image_folder (str): the folder where images live; filenames in the COCO .json
            file [input_file] should be relative to this folder
        output_folder (str): the base folder for the YOLO dataset
        input_file (str): a .json file in COCO format; can be the same as [input_image_folder], in which case
            images are left alone.
        source_format (str, optional): can be 'coco' (default) or 'coco_camera_traps'.  The only difference
            is that when source_format is 'coco_camera_traps', we treat an image with a non-bbox
            annotation as a special case, i.e. that's how an empty image is indicated.  The original
            COCO standard is a little ambiguous on this issue.  If source_format is 'coco', we
            either treat images as empty or error, depending on the value of [allow_empty_annotations].
            [allow_empty_annotations] has no effect if source_format is 'coco_camera_traps'.
        overwrite_images (bool, optional): over-write images in the output folder if they exist
        create_image_and_label_folders (bool, optional): whether to create separate folders called 'images' and
            'labels' in the YOLO output folder.  If create_image_and_label_folders is False,
            a/b/c/image001.jpg will become a#b#c#image001.jpg, and the corresponding text file will
            be a#b#c#image001.txt.  If create_image_and_label_folders is True, a/b/c/image001.jpg will become
            images/a#b#c#image001.jpg, and the corresponding text file will be
            labels/a#b#c#image001.txt.
        class_file_name (str, optional): .txt file (relative to the output folder) that we should
            populate with a list of classes (or None to omit)
        allow_empty_annotations (bool, optional): if this is False and [source_format] is 'coco',
            we'll error on annotations that have no 'bbox' field
        clip_boxes (bool, optional): whether to clip bounding box coordinates to the range [0,1] before
            converting to YOLO xywh format
        image_id_to_output_image_json_file (str, optional): an optional *output* file, to which we will write
            a mapping from image IDs to output file names
        images_to_exclude (list, optional): a list of image files (relative paths in the input folder) that we
            should ignore
        path_replacement_char (str, optional): only relevant if [flatten_paths] is True; this is used to replace
            path separators, e.g. if [path_replacement_char] is '#' and [flatten_paths] is True, a/b/c/d.jpg
            becomes a#b#c#d.jpg
        category_names_to_exclude (str, optional): category names that should not be represented in the
            YOLO output; only impacts annotations, does not prevent copying images.  There's almost no reason
            you would want to specify this and [category_names_to_include].
        category_names_to_include (str, optional): allow-list of category names that should be represented
            in the YOLO output; only impacts annotations, does not prevent copying images.  There's almost
            no reason you would want to specify this and [category_names_to_exclude].
        write_output (bool, optional): determines whether we actually copy images and write annotations;
            setting this to False mostly puts this function in "dry run" "mode.  The class list
            file is written regardless of the value of write_output.
        flatten_paths (bool, optional): replace /'s in image filenames with [path_replacement_char],
            which ensures that the output folder is a single flat folder.

    Returns:
        dict: information about the coco --> yolo mapping, containing at least the fields:

        - class_list_filename: the filename to which we wrote the flat list of class names required
          by the YOLO format.
        - source_image_to_dest_image: a dict mapping source images to destination images
        - coco_id_to_yolo_id: a dict mapping COCO category IDs to YOLO category IDs
    """

    ## Validate input

    if category_names_to_include is not None and category_names_to_exclude is not None:
        raise ValueError('category_names_to_include and category_names_to_exclude are mutually exclusive')

    if output_folder is None:
        output_folder = input_image_folder

    if images_to_exclude is not None:
        images_to_exclude = set(images_to_exclude)

    if category_names_to_exclude is None:
        category_names_to_exclude = {}

    assert os.path.isdir(input_image_folder)
    assert os.path.isfile(input_file)
    os.makedirs(output_folder,exist_ok=True)

    if (output_folder == input_image_folder) and (overwrite_images) and \
        (not create_image_and_label_folders) and (not flatten_paths):
            print('Warning: output folder and input folder are the same, disabling overwrite_images')
            overwrite_images = False

    ## Read input data

    with open(input_file,'r') as f:
        data = json.load(f)


    ## Parse annotations

    image_id_to_annotations = defaultdict(list)

    # i_ann = 0; ann = data['annotations'][0]
    for i_ann,ann in enumerate(data['annotations']):

        # Make sure no annotations have *only* segmentation data
        if ( \
            ('segmentation' in ann.keys()) and \
            (ann['segmentation'] is not None) and \
            (len(ann['segmentation']) > 0) ) \
            and \
            (('bbox' not in ann.keys()) or (ann['bbox'] is None) or (len(ann['bbox'])==0)):
                raise ValueError('Oops: segmentation data present without bbox information, ' + \
                                 'this script isn\'t ready for this dataset')

        image_id_to_annotations[ann['image_id']].append(ann)

    print('Parsed annotations for {} images'.format(len(image_id_to_annotations)))

    # Re-map class IDs to make sure they run from 0...n-classes-1
    #
    # Note: this allows unused categories in the output data set.  This is OK for
    # some training pipelines, not for others.
    next_category_id = 0
    coco_id_to_yolo_id = {}
    coco_id_to_name = {}
    yolo_id_to_name = {}
    coco_category_ids_to_exclude = set()

    for category in data['categories']:
        coco_id_to_name[category['id']] = category['name']
        if (category_names_to_include is not None) and \
            (category['name'] not in category_names_to_include):
            coco_category_ids_to_exclude.add(category['id'])
            continue
        elif (category['name'] in category_names_to_exclude):
            coco_category_ids_to_exclude.add(category['id'])
            continue
        assert category['id'] not in coco_id_to_yolo_id
        coco_id_to_yolo_id[category['id']] = next_category_id
        yolo_id_to_name[next_category_id] = category['name']
        next_category_id += 1


    ## Process images (everything but I/O)

    # List of dictionaries with keys 'source_image','dest_image','bboxes','dest_txt'
    images_to_copy = []

    missing_images = []
    excluded_images = []

    image_names = set()

    typical_image_extensions = set(['.jpg','.jpeg','.png','.gif','.tif','.bmp'])

    printed_empty_annotation_warning = False

    image_id_to_output_image_name = {}

    print('Processing annotations')

    n_clipped_boxes = 0
    n_total_boxes = 0

    # i_image = 0; im = data['images'][i_image]
    for i_image,im in tqdm(enumerate(data['images']),total=len(data['images'])):

        output_info = {}
        source_image = os.path.join(input_image_folder,im['file_name'])
        output_info['source_image'] = source_image

        if images_to_exclude is not None and im['file_name'] in images_to_exclude:
            excluded_images.append(im['file_name'])
            continue

        tokens = os.path.splitext(im['file_name'])
        if tokens[1].lower() not in typical_image_extensions:
            print('Warning: unusual image file name {}'.format(im['file_name']))

        if flatten_paths:
            image_name = tokens[0].replace('\\','/').replace('/',path_replacement_char) + \
                '_' + str(i_image).zfill(6)
        else:
            image_name = tokens[0]

        assert image_name not in image_names, 'Image name collision for {}'.format(image_name)
        image_names.add(image_name)

        assert im['id'] not in image_id_to_output_image_name
        image_id_to_output_image_name[im['id']] = image_name

        dest_image_relative = image_name + tokens[1]
        output_info['dest_image_relative'] = dest_image_relative
        dest_txt_relative = image_name + '.txt'
        output_info['dest_txt_relative'] = dest_txt_relative
        output_info['bboxes'] = []

        # assert os.path.isfile(source_image), 'Could not find image {}'.format(source_image)
        if not os.path.isfile(source_image):
            print('Warning: could not find image {}'.format(source_image))
            missing_images.append(im['file_name'])
            continue

        image_id = im['id']

        image_bboxes = []

        if image_id in image_id_to_annotations:

            for ann in image_id_to_annotations[image_id]:

                # If this annotation has no bounding boxes...
                if 'bbox' not in ann or ann['bbox'] is None or len(ann['bbox']) == 0:

                    if source_format == 'coco':

                        if not allow_empty_annotations:
                            # This is not entirely clear from the COCO spec, but it seems to be consensus
                            # that if you want to specify an image with no objects, you don't include any
                            # annotations for that image.
                            raise ValueError('If an annotation exists, it should have content')
                        else:
                            continue

                    elif source_format == 'coco_camera_traps':

                        # We allow empty bbox lists in COCO camera traps files; this is typically a
                        # negative example in a dataset that has bounding boxes, and 0 is typically
                        # the empty category, which is typically 0.
                        if ann['category_id'] != 0:
                            if not printed_empty_annotation_warning:
                                printed_empty_annotation_warning = True
                                print('Warning: non-bbox annotation found with category {}'.format(
                                    ann['category_id']))
                        continue

                    else:

                        raise ValueError('Unrecognized COCO variant: {}'.format(source_format))

                # ...if this is an empty annotation

                coco_bbox = ann['bbox']

                # This category isn't in our category list.  This typically corresponds to whole sets
                # of images that were excluded from the YOLO set.
                if ann['category_id'] in coco_category_ids_to_exclude:
                    continue

                yolo_category_id = coco_id_to_yolo_id[ann['category_id']]

                # COCO: [x_min, y_min, width, height] in absolute coordinates
                # YOLO: [class, x_center, y_center, width, height] in normalized coordinates

                # Convert from COCO coordinates to YOLO coordinates
                img_w = im['width']
                img_h = im['height']

                if source_format in ('coco','coco_camera_traps'):

                    x_min_absolute = coco_bbox[0]
                    y_min_absolute = coco_bbox[1]
                    box_w_absolute = coco_bbox[2]
                    box_h_absolute = coco_bbox[3]

                    x_center_absolute = (x_min_absolute + (x_min_absolute + box_w_absolute)) / 2
                    y_center_absolute = (y_min_absolute + (y_min_absolute + box_h_absolute)) / 2

                    x_center_relative = x_center_absolute / img_w
                    y_center_relative = y_center_absolute / img_h

                    box_w_relative = box_w_absolute / img_w
                    box_h_relative = box_h_absolute / img_h

                else:

                    raise ValueError('Unrecognized source format {}'.format(source_format))

                if clip_boxes:

                    clipped_box = False

                    box_right = x_center_relative + (box_w_relative / 2.0)
                    if box_right > 1.0:
                        clipped_box = True
                        overhang = box_right - 1.0
                        box_w_relative -= overhang
                        x_center_relative -= (overhang / 2.0)

                    box_bottom = y_center_relative + (box_h_relative / 2.0)
                    if box_bottom > 1.0:
                        clipped_box = True
                        overhang = box_bottom - 1.0
                        box_h_relative -= overhang
                        y_center_relative -= (overhang / 2.0)

                    box_left = x_center_relative - (box_w_relative / 2.0)
                    if box_left < 0.0:
                        clipped_box = True
                        overhang = abs(box_left)
                        box_w_relative -= overhang
                        x_center_relative += (overhang / 2.0)

                    box_top = y_center_relative - (box_h_relative / 2.0)
                    if box_top < 0.0:
                        clipped_box = True
                        overhang = abs(box_top)
                        box_h_relative -= overhang
                        y_center_relative += (overhang / 2.0)

                    if clipped_box:
                        n_clipped_boxes += 1

                yolo_box = [yolo_category_id,
                            x_center_relative, y_center_relative,
                            box_w_relative, box_h_relative]

                image_bboxes.append(yolo_box)
                n_total_boxes += 1

            # ...for each annotation

        # ...if this image has annotations

        output_info['bboxes'] = image_bboxes

        images_to_copy.append(output_info)

    # ...for each image

    print('\nWriting {} boxes ({} clipped) for {} images'.format(n_total_boxes,
                                                               n_clipped_boxes,len(images_to_copy)))
    print('{} missing images (of {})'.format(len(missing_images),len(data['images'])))

    if images_to_exclude is not None:
        print('{} excluded images (of {})'.format(len(excluded_images),len(data['images'])))


    ## Write output

    print('Generating class list')

    if class_file_name is not None:
        class_list_filename = os.path.join(output_folder,class_file_name)
        with open(class_list_filename, 'w') as f:
            print('Writing class list to {}'.format(class_list_filename))
            for i_class in range(0,len(yolo_id_to_name)):
                # Category IDs should range from 0..N-1
                assert i_class in yolo_id_to_name
                f.write(yolo_id_to_name[i_class] + '\n')

    if image_id_to_output_image_json_file is not None:
        print('Writing image ID mapping to {}'.format(image_id_to_output_image_json_file))
        with open(image_id_to_output_image_json_file,'w') as f:
            json.dump(image_id_to_output_image_name,f,indent=1)


    if (output_folder == input_image_folder) and (not create_image_and_label_folders):
        print('Creating annotation files (not copying images, input and output folder are the same)')
    else:
        print('Copying images and creating annotation files')

    if create_image_and_label_folders:
        dest_image_folder = os.path.join(output_folder,'images')
        dest_txt_folder = os.path.join(output_folder,'labels')
    else:
        dest_image_folder = output_folder
        dest_txt_folder = output_folder

    source_image_to_dest_image = {}

    label_files_written = []
    n_boxes_written = 0

    # TODO: parallelize this loop
    #
    # output_info = images_to_copy[0]
    for output_info in tqdm(images_to_copy):

        source_image = output_info['source_image']
        dest_image_relative = output_info['dest_image_relative']
        dest_txt_relative = output_info['dest_txt_relative']

        dest_image = os.path.join(dest_image_folder,dest_image_relative)
        dest_txt = os.path.join(dest_txt_folder,dest_txt_relative)

        source_image_to_dest_image[source_image] = dest_image

        # Copy the image if necessary
        if write_output:

            os.makedirs(os.path.dirname(dest_image),exist_ok=True)
            os.makedirs(os.path.dirname(dest_txt),exist_ok=True)

            if not create_image_and_label_folders:
                assert os.path.dirname(dest_image) == os.path.dirname(dest_txt)

            if (not os.path.isfile(dest_image)) or (overwrite_images):
                shutil.copyfile(source_image,dest_image)

        bboxes = output_info['bboxes']

        # Write the annotation file if necessary
        #
        # Only write an annotation file if there are bounding boxes.  Images with
        # no .txt files are treated as hard negatives, at least by YOLOv5:
        #
        # https://github.com/ultralytics/yolov5/issues/3218
        #
        # I think this is also true for images with empty .txt files, but
        # I'm using the convention suggested on that issue, i.e. hard
        # negatives are expressed as images without .txt files.
        if len(bboxes) > 0:

            n_boxes_written += len(bboxes)
            label_files_written.append(dest_txt)

            if write_output:

                with open(dest_txt,'w') as f:

                    # bbox = bboxes[0]
                    for bbox in bboxes:
                        assert len(bbox) == 5
                        s = '{} {} {} {} {}'.format(bbox[0],bbox[1],bbox[2],bbox[3],bbox[4])
                        f.write(s + '\n')


    # ...for each image

    coco_to_yolo_info = {}
    coco_to_yolo_info['class_list_filename'] = class_list_filename
    coco_to_yolo_info['source_image_to_dest_image'] = source_image_to_dest_image
    coco_to_yolo_info['coco_id_to_yolo_id'] = coco_id_to_yolo_id
    coco_to_yolo_info['label_files_written'] = label_files_written
    coco_to_yolo_info['n_boxes_written'] = n_boxes_written

    return coco_to_yolo_info

# ...def coco_to_yolo(...)


def create_yolo_symlinks(source_folder,
                         images_folder,
                         labels_folder,
                         class_list_file=None,
                         class_list_output_name='object.data',
                         force_lowercase_image_extension=False):
    """
    Given a YOLO-formatted folder of images and .txt files, creates a folder
    of symlinks to all the images, and a folder of symlinks to all the labels.
    Used to support preview/editing tools that assume images and labels are in separate
    folders.

    Args:
        source_folder (str): input folder
        images_folder (str): output folder with links to images
        labels_folder (str): output folder with links to labels
        class_list_file (str, optional): list to classes.txt file
        class_list_output_name (str, optional): output file to write with class information
        force_lowercase_image_extension (bool, False): create symlinks with, e.g., .jpg, even
            if the input image is, e.g., .JPG

    :meta private:
    """

    assert source_folder != images_folder and source_folder != labels_folder

    os.makedirs(images_folder,exist_ok=True)
    os.makedirs(labels_folder,exist_ok=True)

    image_files_relative = find_images(source_folder,recursive=True,return_relative_paths=True)

    # image_fn_relative = image_files_relative[0]=
    for image_fn_relative in tqdm(image_files_relative):

        source_file_abs = os.path.join(source_folder,image_fn_relative)
        target_file_abs = os.path.join(images_folder,image_fn_relative)

        if force_lowercase_image_extension:
            tokens = os.path.splitext(target_file_abs)
            target_file_abs = tokens[0] + tokens[1].lower()

        os.makedirs(os.path.dirname(target_file_abs),exist_ok=True)
        safe_create_link(source_file_abs,target_file_abs)
        source_annotation_file_abs = os.path.splitext(source_file_abs)[0] + '.txt'
        if os.path.isfile(source_annotation_file_abs):
            target_annotation_file_abs = \
                os.path.splitext(os.path.join(labels_folder,image_fn_relative))[0] + '.txt'
            os.makedirs(os.path.dirname(target_annotation_file_abs),exist_ok=True)
            safe_create_link(source_annotation_file_abs,target_annotation_file_abs)

    # ...for each image

    if class_list_file is not None:
        target_class_list_file = os.path.join(labels_folder,class_list_output_name)
        safe_create_link(class_list_file,target_class_list_file)

# ...def create_yolo_symlinks(...)


#%% Interactive driver

if False:

    pass

    #%% Options

    input_file = os.path.expanduser('~/data/md-test-coco.json')
    image_folder = os.path.expanduser('~/data/md-test')
    output_folder = os.path.expanduser('~/data/md-test-yolo')
    create_image_and_label_folders=False
    class_file_name='classes.txt'
    allow_empty_annotations=False
    clip_boxes=False
    image_id_to_output_image_json_file=None
    images_to_exclude=None
    path_replacement_char='#'
    category_names_to_exclude=None


    #%% Programmatic execution

    coco_to_yolo_results = coco_to_yolo(image_folder,output_folder,input_file,
                     source_format='coco',
                     overwrite_images=False,
                     create_image_and_label_folders=create_image_and_label_folders,
                     class_file_name=class_file_name,
                     allow_empty_annotations=allow_empty_annotations,
                     clip_boxes=clip_boxes)

    create_yolo_symlinks(source_folder=output_folder,
                         images_folder=output_folder + '/images',
                         labels_folder=output_folder + '/labels',
                         class_list_file=coco_to_yolo_results['class_list_filename'],
                         class_list_output_name='object.data',
                         force_lowercase_image_extension=True)


    #%% Prepare command-line example

    s = 'python coco_to_yolo.py {} {} {} --create_bounding_box_editor_symlinks'.format(
        image_folder,output_folder,input_file)
    print(s)
    import clipboard; clipboard.copy(s)


#%% Command-line driver

def main(): # noqa

    parser = argparse.ArgumentParser(
        description='Convert COCO-formatted data to YOLO format, flattening the image structure')

    # input_image_folder,output_folder,input_file

    parser.add_argument(
        'input_folder',
        type=str,
        help='Path to input images')

    parser.add_argument(
        'output_folder',
        type=str,
        help='Path to flat, YOLO-formatted dataset')

    parser.add_argument(
        'input_file',
        type=str,
        help='Path to COCO dataset file (.json)')

    parser.add_argument(
        '--create_bounding_box_editor_symlinks',
        action='store_true',
        help='Prepare symlinks so the whole folder appears to contain "images" and "labels" folderss')

    if len(sys.argv[1:]) == 0:
        parser.print_help()
        parser.exit()

    args = parser.parse_args()

    coco_to_yolo_results = coco_to_yolo(args.input_folder,args.output_folder,args.input_file)

    if args.create_bounding_box_editor_symlinks:
        create_yolo_symlinks(source_folder=args.output_folder,
                             images_folder=args.output_folder + '/images',
                             labels_folder=args.output_folder + '/labels',
                             class_list_file=coco_to_yolo_results['class_list_filename'],
                             class_list_output_name='object.data',
                             force_lowercase_image_extension=True)

if __name__ == '__main__':
    main()
