"""

merge_detections.py

Merge high-confidence detections from one or more results files into another
file.  Typically used to combine results from MDv5b and/or MDv4 and/or MDv5a
into a "primary" results file from MDv5a or MDv1000-redwood.

Detection categories must be the same in both files; if you want to first remap
one file's category mapping to be the same as another's, see remap_detection_categories.

If you want to literally merge two .json files, see combine_batch_outputs.py.

"""

#%% Constants and imports

import argparse
import sys
import json
import os

from tqdm import tqdm

from megadetector.utils.ct_utils import get_iou
from megadetector.utils.ct_utils import write_json

from megadetector.postprocessing.classification_postprocessing import \
    merge_classification_categories


#%% Structs

class MergeDetectionsOptions:
    """
    Class defining options for merge_detections().
    """

    def __init__(self):

        #: Maximum detection size to include in the merged output
        self.max_detection_size = 1.01

        #: Minimum detection size to include in the merged output
        self.min_detection_size = 0

        #: Exclude detections whose confidence in the source file(s) is less
        #: than this.  Should have the same length as the number of source files.
        self.source_confidence_thresholds = [0.05]

        #: Don't bother merging into target images if there is a similar detection
        #: above this threshold (or if there is *any* detection above this threshold,
        #: and merge_empty_only is True)
        self.target_confidence_threshold = 0.2

        #: If you want to merge only certain categories, specify one
        #: (but not both) of these.  These are category IDs, not names.
        self.categories_to_include = None

        #: If you want to merge only certain categories, specify one
        #: (but not both) of these.  These are category IDs, not names.
        self.categories_to_exclude = None

        #: Only merge detections into images that have *no* above-threshold
        #: detections in the target results file.
        self.merge_empty_only = False

        #: IoU threshold above which two detections are considered the same
        self.iou_threshold = 0.65

        #: Error if this is False and the output file exists
        self.overwrite = False

        #: Add a 'merged' field (value True) to all merged detections.
        #:
        #: Only used for debugging.
        self.mark_copied_detections = False


#%% Main function

def merge_detections(source_files,
                     target_file,
                     output_file=None,
                     options=None):
    """
    Merge high-confidence detections from one or more results files into another
    file.  Typically used to combine results from, e.g., MDv5a and/or MDv5b into a
    "primary" results file from MDv1000-redwood.

    [source_files] (a list of files or a single filename) specifies the set of
    results files that will be merged into [target_file].  The difference between a
    "source file" and the "target file" is that if no merging is necessary, either because
    two boxes are nearly identical or because merge_only_empty is True and the target
    file already has above-threshold detection for an image+category, the output file gets
    the results of the "target" file.  I.e., the "target" file wins all ties.

    Merges classification categories based on category names (ignoring descriptions).

    The results are written to [output_file].

    Args:
        source_files (list of str): list of files to merge into the results in [target_file]
        target_file (str): filename that is treated as the primary source of results
        output_file (str, optional): file to which we should write merged results
        options (MergeDetectionsOptions, optional): see MergeDetectionsOptions

    Returns:
        dict in MegaDetector format, containing merged results
    """

    if isinstance(source_files,str):
        source_files = [source_files]

    if options is None:
        options = MergeDetectionsOptions()

    if (output_file is not None) and (not options.overwrite) and (os.path.isfile(output_file)):
        print('File {} exists, bypassing merge'.format(output_file))
        return

    assert not ((options.categories_to_exclude is not None) and \
                (options.categories_to_include is not None)), \
                'categories_to_include and categories_to_exclude are mutually exclusive'

    if options.categories_to_exclude is not None:
        options.categories_to_exclude = [int(c) for c in options.categories_to_exclude]

    if options.categories_to_include is not None:
        options.categories_to_include = [int(c) for c in options.categories_to_include]

    assert len(source_files) == len(options.source_confidence_thresholds), \
        '{} source files provided, but {} source confidence thresholds provided'.format(
            len(source_files),len(options.source_confidence_thresholds))

    for fn in source_files:
        assert os.path.isfile(fn), 'Could not find source file {}'.format(fn)

    assert os.path.isfile(target_file)

    with open(target_file,'r') as f:
        output_data = json.load(f)

    print('Loaded results for {} images'.format(len(output_data['images'])))

    fn_to_image = {}

    # im = output_data['images'][0]
    for im in output_data['images']:
        fn_to_image[im['file']] = im

    if 'detections_transferred_from' not in output_data['info']:
        output_data['info']['detections_transferred_from'] = []

    if 'detector' not in output_data['info']:
        output_data['info']['detector'] = 'MDv4 (assumed)'

    detection_categories_raw = output_data['detection_categories'].keys()

    # Determine whether we should be processing all categories, or just a subset
    # of categories.
    detection_categories = []

    if options.categories_to_exclude is not None:
        for c in detection_categories_raw:
            if int(c) not in options.categories_to_exclude:
                detection_categories.append(c)
            else:
                print('Excluding category {}'.format(c))
    elif options.categories_to_include is not None:
        for c in detection_categories_raw:
            if int(c) in options.categories_to_include:
                print('Including category {}'.format(c))
                detection_categories.append(c)
    else:
        detection_categories = detection_categories_raw

    # For each source file...
    #
    # i_source_file = 0; source_file = source_files[i_source_file]
    for i_source_file,source_file in enumerate(source_files):

        print('Processing detections from file {}'.format(source_file))

        with open(source_file,'r') as f:
            source_data = json.load(f)

        # If necessary, map all the classification categories in the source file
        # into the space of the target file
        #
        # This will no-op if neither file has classification categories.
        source_data = \
            merge_classification_categories(target_file=output_data,
                                            source_file=source_data,
                                            output_file=None)

        # If classification data is present in "target_file", merge_classification_categories
        # guarantees compatible classification categories in source_data
        if 'classification_categories' in output_data:
            assert 'classification_categories' in source_data, 'Category merge error A'
            assert len(source_data['classification_categories']) >= \
                   len(output_data['classification_categories']), 'Category merge error B'
            for category_id in output_data['classification_categories']:
                assert category_id in source_data['classification_categories']
                assert source_data['classification_categories'][category_id] == \
                       output_data['classification_categories'][category_id], 'Category merge error C'
            output_data['classification_categories'] = source_data['classification_categories']
            if 'classification_category_descriptions' in output_data:
                output_data['classification_category_descriptions'] = \
                    source_data['classification_category_descriptions']

        if 'detector' in source_data['info']:
            source_detector_name = source_data['info']['detector']
        else:
            source_detector_name = os.path.basename(source_file)

        output_data['info']['detections_transferred_from'].append(os.path.basename(source_file))
        output_data['info']['detector'] = output_data['info']['detector'] + ' + ' + source_detector_name

        assert source_data['detection_categories'] == output_data['detection_categories'], \
            'Cannot merge files with different detection category maps'

        source_confidence_threshold = options.source_confidence_thresholds[i_source_file]

        # For each image...
        #
        # source_im = source_data['images'][0]
        for source_im in tqdm(source_data['images']):

            image_filename = source_im['file']

            assert image_filename in fn_to_image, \
                'Image {} not in target image set'.format(image_filename)
            target_im = fn_to_image[image_filename]

            if 'detections' not in source_im or source_im['detections'] is None:
                continue

            if 'detections' not in target_im or target_im['detections'] is None:
                continue

            source_detections_this_image = source_im['detections']
            target_detections_this_image = target_im['detections']

            detections_to_transfer = []

            # For each detection category...
            #
            # detection_category = list(detection_categories)[0]
            for detection_category in detection_categories:

                target_detections_this_category = \
                    [det for det in target_detections_this_image if det['category'] == \
                     detection_category]

                max_target_confidence_this_category = 0.0

                if len(target_detections_this_category) > 0:
                    max_target_confidence_this_category = max([det['conf'] for \
                      det in target_detections_this_category])

                # If we have a valid detection in the target file, and we're only merging
                # into images that have no detections at all, we don't need to review the individual
                # detections in the source file.
                if options.merge_empty_only and \
                    (max_target_confidence_this_category >= options.target_confidence_threshold):
                    continue

                source_detections_this_category_raw = [det for det in \
                  source_detections_this_image if det['category'] == detection_category]

                # Boxes are x/y/w/h
                # source_sizes = [det['bbox'][2]*det['bbox'][3] for det in source_detections_this_category_raw]

                # Only look at source boxes within the size range
                source_detections_this_category_filtered = [
                    det for det in source_detections_this_category_raw if \
                        (det['bbox'][2]*det['bbox'][3] <= options.max_detection_size) and \
                        (det['bbox'][2]*det['bbox'][3] >= options.min_detection_size) \
                        ]

                # For each detection in this category in the source image...
                #
                # det = source_detections_this_category_filtered[0]
                for det in source_detections_this_category_filtered:

                    if det['conf'] >= source_confidence_threshold:

                        # Check only whole images
                        if options.merge_empty_only:

                            # We verified this above, asserting here for clarity
                            assert max_target_confidence_this_category < options.target_confidence_threshold
                            det['transferred_from'] = source_detector_name
                            detections_to_transfer.append(det)

                        # Check individual detections
                        else:

                            # Does this source detection match any existing above-threshold
                            # target category detections?
                            matches_existing_box = False

                            # target_detection = target_detections_this_category[0]
                            for target_detection in target_detections_this_category:

                                if (target_detection['conf'] >= options.target_confidence_threshold) \
                                   and \
                                   (get_iou(det['bbox'],target_detection['bbox']) >= options.iou_threshold):
                                    matches_existing_box = True
                                    break

                            if (not matches_existing_box):
                                det['transferred_from'] = source_detector_name
                                detections_to_transfer.append(det)

                    # ...if this source detection is above the confidence threshold

                # ...for each source detection within category

            # ...for each detection category

            if len(detections_to_transfer) > 0:

                # print('Adding {} detections to image {}'.format(len(detections_to_transfer),image_filename))
                detections = fn_to_image[image_filename]['detections']
                detections.extend(detections_to_transfer)

                if options.mark_copied_detections:
                    for det in detections_to_transfer:
                        det['merged'] = True

                # Update the max_detection_conf field (if present)
                if 'max_detection_conf' in fn_to_image[image_filename]:
                    fn_to_image[image_filename]['max_detection_conf'] = \
                        max([d['conf'] for d in detections])

            # ...if we have any detections to transfer

        # ...for each image

    # ...for each source file

    if output_file is not None:
        write_json(output_file,output_data)
        print('Saved merged results to {}'.format(output_file))

    return output_data

# ...def merge_detections(...)


#%% Command-line driver

def main():
    """
    Command-line driver for merge_detections.py
    """

    default_options = MergeDetectionsOptions()

    parser = argparse.ArgumentParser(
        description='Merge detections from one or more MegaDetector results files into an existing results file')
    parser.add_argument(
        'source_files',
        nargs='+',
        help='Path to source .json file(s) to merge from')
    parser.add_argument(
        'target_file',
        help='Path to a .json file to merge detections into')
    parser.add_argument(
        'output_file',
        help='Path to output .json results file')
    parser.add_argument(
        '--max_detection_size',
        type=float,
        default=default_options.max_detection_size,
        help='Ignore detections with an area larger than this (as a fraction of ' + \
             'image size) (default {})'.format(
             default_options.max_detection_size))
    parser.add_argument(
        '--min_detection_size',
        default=default_options.min_detection_size,
        type=float,
        help='Ignore detections with an area smaller than this (as a fraction of ' + \
             'image size) (default {})'.format(
             default_options.min_detection_size))
    parser.add_argument(
        '--source_confidence_thresholds',
        nargs='+',
        type=float,
        default=default_options.source_confidence_thresholds,
        help='List of thresholds for each source file (default {}). '.format(
            default_options.source_confidence_thresholds) + \
            'Merge only if the source file\'s detection confidence is higher than its ' + \
            'corresponding threshold.  Should be the same length as the number of source files.')
    parser.add_argument(
        '--target_confidence_threshold',
        type=float,
        default=default_options.target_confidence_threshold,
        help="Do not merge if target file detection confidence is already higher " + \
             "than this (default {})".format(default_options.target_confidence_threshold))
    parser.add_argument(
        '--categories_to_include',
        type=int,
        nargs='+',
        default=None,
        help='List of numeric detection category IDs to include')
    parser.add_argument(
        '--categories_to_exclude',
        type=int,
        nargs='+',
        default=None,
        help='List of numeric detection categories to exclude')
    parser.add_argument(
        '--merge_empty_only',
        action='store_true',
        help='Ignore individual detections and only merge images for which the target ' + \
             'file contains no detections')
    parser.add_argument(
        '--iou_threshold',
        type=float,
        default=default_options.iou_threshold,
        help='Sets the minimum IoU for a source detection to be considered the same as ' + \
             'a target detection (default {})'.format(default_options.iou_threshold))

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

    args = parser.parse_args()

    options = MergeDetectionsOptions()
    options.max_detection_size = args.max_detection_size
    options.min_detection_size = args.min_detection_size
    options.source_confidence_thresholds = args.source_confidence_thresholds
    options.target_confidence_threshold = args.target_confidence_threshold
    options.categories_to_include = args.categories_to_include
    options.categories_to_exclude = args.categories_to_exclude
    options.merge_empty_only = args.merge_empty_only
    options.iou_threshold = args.iou_threshold

    merge_detections(args.source_files, args.target_file, args.output_file, options)

if __name__ == '__main__':
    main()
