# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Run evaluation for a model exported to ONNX"""

import os
import time
from tempfile import TemporaryDirectory

import numpy as np
import onnxruntime as ort
from pycocotools.coco import COCO

from src.dataset import create_maskrcnn_dataset
from src.model_utils.config import config
from src.util import coco_eval, bbox2result_1image, results2json, get_seg_masks


def create_session(checkpoint_path, target_device):
    """Load ONNX model and create ORT session"""
    if target_device == 'GPU':
        providers = ['CUDAExecutionProvider']
    elif target_device == 'CPU':
        providers = ['CPUExecutionProvider']
    else:
        raise ValueError(f"Unsupported target device '{target_device}'. Expected one of: 'CPU', 'GPU'")
    session = ort.InferenceSession(checkpoint_path, providers=providers)
    input_names = [x.name for x in session.get_inputs()]
    return session, input_names


def run_eval(onnx_checkpoint_path, mindrecord_path, batch_size, num_classes, target_device, input_type,
             max_predictions_count=128):
    """Run ONNX model evaluation

    Args:
        onnx_checkpoint_path (str): path to ONNX checkpoint
        mindrecord_path (str): path to MindRecord, including name prefix
        batch_size (int): batch size
        num_classes (int): number of classes
        target_device (str): 'GPU' / 'CPU'
        input_type (str): 'float16' / 'float32'
        max_predictions_count (int, optional): maximum number of predictions. Defaults to 128.

    Returns:
        List[Tuple]: predictions - (boxes, masks) tuples for each image
    """
    session, (images_input_name, shapes_input_name) = create_session(onnx_checkpoint_path, target_device)

    dataset = create_maskrcnn_dataset(mindrecord_path, batch_size, is_training=False)
    it = dataset.create_dict_iterator(output_numpy=True, num_epochs=1)

    outputs = []
    for eval_iter, batch in enumerate(it):
        start_time = time.time()

        inputs = {
            images_input_name: batch['image'].astype(input_type),
            shapes_input_name: batch['image_shape'].astype(input_type)
        }
        pred_boxes, pred_labels, valid_predictions_masks, pred_masks = session.run(None, inputs)
        pred_labels = np.squeeze(pred_labels, -1)
        valid_predictions_masks = np.squeeze(valid_predictions_masks, -1)

        it = zip(valid_predictions_masks, pred_boxes, pred_labels, pred_masks, batch['image_shape'])
        for valid_mask, boxes_batch, labels_batch, masks_batch, metas_batch in it:
            boxes = boxes_batch[valid_mask]
            labels = labels_batch[valid_mask]
            masks = masks_batch[valid_mask]

            if len(boxes) > max_predictions_count:
                indices = boxes[:, -1].argsort()[::-1]
                indices = indices[:max_predictions_count]
                boxes = boxes[indices]
                labels = labels[indices]
                masks = masks[indices]

            result_boxes = bbox2result_1image(boxes, labels, num_classes)
            result_masks = get_seg_masks(masks, boxes, labels, metas_batch,
                                         rescale=True, num_classes=num_classes)
            outputs.append((result_boxes, result_masks))

        end_time = time.time()
        # flush must be true to avoid losing output when this script is launched with "&> log_file.txt"
        print(f"Iter {eval_iter} / {dataset.get_dataset_size()} took {end_time - start_time:.2f} s", flush=True)

    return outputs


def report_metrics(outputs, coco_annotation_path, eval_types=('bbox', 'segm')):
    """Print COCO metrics"""
    coco_dataset = COCO(coco_annotation_path)
    with TemporaryDirectory() as tempdir:
        prefix = os.path.join(tempdir, 'results.pkl')
        result_files = results2json(coco_dataset, outputs, prefix)
        coco_eval(result_files, eval_types, coco_dataset, single_result=False)


def main():
    """Run ONNX eval from command line"""
    mindrecord_path = os.path.join(config.coco_root, config.mindrecord_dir, 'MaskRcnn_eval.mindrecord')
    coco_annotation_path = os.path.join(config.coco_root, config.ann_file)

    outputs = run_eval(config.file_name, mindrecord_path, config.test_batch_size,
                       config.num_classes, config.device_target, config.export_input_type)
    report_metrics(outputs, coco_annotation_path)


if __name__ == '__main__':
    main()