# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

"""ONNX Runtime inference micro-benchmark."""

import time
import statistics
from pathlib import Path

from packaging import version
import torch
import torchvision.models
import numpy as np

from superbench.common.utils import logger
from superbench.benchmarks import BenchmarkRegistry, Platform, Precision, ReturnCode
from superbench.benchmarks.micro_benchmarks import MicroBenchmark


class ORTInferenceBenchmark(MicroBenchmark):
    """ONNXRuntime inference micro-benchmark class."""
    _execution_provider_map = {
        'cuda': 'CUDAExecutionProvider',
        'rocm': 'ROCMExecutionProvider',
        'migraphx': 'MIGraphXExecutionProvider',
        'cpu': 'CPUExecutionProvider',
    }
    _execution_provider_preference = [
        'CUDAExecutionProvider',
        'ROCMExecutionProvider',
        'MIGraphXExecutionProvider',
        'CPUExecutionProvider',
    ]

    def __init__(self, name, parameters=''):
        """Constructor.

        Args:
            name (str): benchmark name.
            parameters (str): benchmark parameters.
        """
        super().__init__(name, parameters)

        self._pytorch_models = [
            'resnet50',
            'resnet101',
            'resnet152',
            'densenet169',
            'densenet201',
            'vgg11',
            'vgg13',
            'vgg16',
            'vgg19',
        ]
        self.__graph_opt_level = None
        self.__model_cache_path = Path(torch.hub.get_dir()) / 'checkpoints'
        self.__execution_provider = None

    def add_parser_arguments(self):
        """Add the specified arguments."""
        super().add_parser_arguments()

        self._parser.add_argument(
            '--pytorch_models',
            type=str,
            nargs='+',
            default=self._pytorch_models,
            help='ONNX models for ONNX Runtime inference benchmark, e.g., {}.'.format(', '.join(self._pytorch_models)),
        )

        self._parser.add_argument(
            '--precision',
            type=Precision,
            choices=[Precision.FLOAT32, Precision.FLOAT16, Precision.INT8],
            default=Precision.FLOAT16,
            required=False,
            help='Precision for inference, allow int8, float16, or float32 only.',
        )

        self._parser.add_argument(
            '--graph_opt_level',
            type=int,
            default=3,
            choices=[0, 1, 2, 3],
            required=False,
            help='ONNXRuntime graph optimization level, 0 for ORT_DISABLE_ALL, 1 for ORT_ENABLE_BASIC, '
            '2 for ORT_ENABLE_EXTENDED, 3 for ORT_ENABLE_ALL.',
        )

        self._parser.add_argument(
            '--batch_size',
            type=int,
            default=32,
            required=False,
            help='Set batch size for inference.',
        )

        self._parser.add_argument(
            '--num_warmup',
            type=int,
            default=64,
            required=False,
            help='The number of warmup step before the benchmarking.',
        )

        self._parser.add_argument(
            '--num_steps',
            type=int,
            default=256,
            required=False,
            help='The number of test step for benchmarking.',
        )

        self._parser.add_argument(
            '--execution_provider',
            type=str,
            choices=['auto'] + list(self._execution_provider_map.keys()) + list(self._execution_provider_map.values()),
            default='auto',
            required=False,
            help='ONNX Runtime execution provider. Use auto, cuda, rocm, migraphx, cpu, or the full provider name.',
        )

        self._parser.add_argument(
            '--pretrained',
            action='store_true',
            default=False,
            required=False,
            help='Whether to use pretrained torchvision model weights when exporting ONNX models.',
        )

    def __select_execution_provider(self, available_providers):
        """Select ONNX Runtime execution provider.

        Args:
            available_providers (List[str]): available ONNX Runtime execution providers.

        Return:
            str: selected execution provider.
        """
        provider = self._args.execution_provider
        if provider != 'auto':
            return self._execution_provider_map.get(provider, provider)

        for preferred_provider in self._execution_provider_preference:
            if preferred_provider in available_providers:
                return preferred_provider

        return ''

    def _preprocess(self):
        """Preprocess/preparation operations before the benchmarking.

        Return:
            True if _preprocess() succeed.
        """
        if not super()._preprocess():
            return False

        import onnxruntime as ort
        self.__graph_opt_level = {
            0: ort.GraphOptimizationLevel.ORT_DISABLE_ALL,
            1: ort.GraphOptimizationLevel.ORT_ENABLE_BASIC,
            2: ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED,
            3: ort.GraphOptimizationLevel.ORT_ENABLE_ALL,
        }

        available_providers = ort.get_available_providers()
        self.__execution_provider = self.__select_execution_provider(available_providers)
        if self.__execution_provider not in available_providers:
            self._result.set_return_code(ReturnCode.MICROBENCHMARK_UNSUPPORTED_ARCHITECTURE)
            logger.error(
                'Unsupported ONNX Runtime execution provider - benchmark: {}, provider: {}, available providers: {}.'.
                format(self._name, self._args.execution_provider, available_providers)
            )
            return False
        logger.info('Using ONNX Runtime execution provider: %s.', self.__execution_provider)

        for model in self._args.pytorch_models:
            if hasattr(torchvision.models, model):
                data_type = Precision.FLOAT16.value if self._args.precision == Precision.FLOAT16 \
                    else Precision.FLOAT32.value
                model_path = f'{self.__model_cache_path / (model + "." + data_type + ".onnx")}'
                if version.parse(torchvision.__version__) < version.parse('0.13'):
                    model_args = {'pretrained': self._args.pretrained}
                else:
                    model_args = {'pretrained': True} if self._args.pretrained else {'weights': None}
                torch.onnx.export(
                    getattr(torchvision.models, model)(**model_args).to(dtype=getattr(torch, data_type)).cuda(),
                    torch.randn(self._args.batch_size, 3, 224, 224, device='cuda', dtype=getattr(torch, data_type)),
                    model_path,
                    input_names=['input'],
                )
                if self._args.precision == Precision.INT8:
                    file_name = '{model}.{precision}.onnx'.format(model=model, precision=self._args.precision)
                    # For quantization of ONNXRuntime, refer
                    # https://onnxruntime.ai/docs/performance/quantization.html#quantization-overview
                    from onnxruntime.quantization import quantize_dynamic
                    quantize_dynamic(model_path, f'{self.__model_cache_path / file_name}')
            else:
                logger.error('Cannot find PyTorch model %s.', model)
                return False

        return True

    def _benchmark(self):
        """Implementation for benchmarking."""
        import onnxruntime as ort
        precision_metric = {'float16': 'fp16', 'float32': 'fp32', 'int8': 'int8'}

        for model in self._args.pytorch_models:
            sess_options = ort.SessionOptions()
            sess_options.graph_optimization_level = self.__graph_opt_level[self._args.graph_opt_level]
            file_name = '{model}.{precision}.onnx'.format(model=model, precision=self._args.precision)
            ort_sess = ort.InferenceSession(
                f'{self.__model_cache_path / file_name}', sess_options, providers=[self.__execution_provider]
            )

            elapse_times = self.__inference(ort_sess)

            if self._args.precision.value in precision_metric:
                precision = precision_metric[self._args.precision.value]
            else:
                precision = self._args.precision.value
            metric = '{}_{}_time'.format(precision, model)
            if not self._process_numeric_result(metric, elapse_times, cal_percentile=True):
                return False

            logger.info(
                'ORT Inference - round: {}, name: {}, model: {}, precision: {}, latency: {} ms'.format(
                    self._curr_run_index, self._name, model, self._args.precision, statistics.mean(elapse_times)
                )
            )

        return True

    def __inference(self, ort_sess):
        """Do inference given the ORT inference session.

        Args:
            ort_sess (InferenceSession): inference session for ORT.

        Return:
            elapse_times (List[float]): latency of every iterations.
        """
        precision = np.float16 if self._args.precision == Precision.FLOAT16 else np.float32
        input_tensor = np.random.randn(self._args.batch_size, 3, 224, 224).astype(dtype=precision)

        for i in range(self._args.num_warmup):
            ort_sess.run(None, {'input': input_tensor})

        elapse_times = list()
        for i in range(self._args.num_steps):
            start = time.time()
            ort_sess.run(None, {'input': input_tensor})
            end = time.time()
            elapse_times.append((end - start) * 1000)

        return elapse_times


BenchmarkRegistry.register_benchmark(
    'ort-inference',
    ORTInferenceBenchmark,
    platform=Platform.CUDA,
)
BenchmarkRegistry.register_benchmark(
    'ort-inference',
    ORTInferenceBenchmark,
    platform=Platform.ROCM,
)
BenchmarkRegistry.register_benchmark(
    'ort-inference',
    ORTInferenceBenchmark,
    platform=Platform.DTK,
)
