# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.  See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------

import argparse
import csv
import logging
import os
import sys
import time
import timeit
from datetime import datetime
from enum import Enum

import coloredlogs
import numpy
import onnx
import torch
from packaging import version

logger = logging.getLogger(__name__)


class Precision(Enum):
    FLOAT32 = "fp32"
    FLOAT16 = "fp16"
    INT8 = "int8"

    def __str__(self):
        return self.value


class OptimizerInfo(Enum):
    # no_opt means using the raw ONNX model, but OnnxRuntime might still apply optimization as long as
    # graph optimization level is not 0 (disable all).
    NOOPT = "no_opt"
    BYORT = "by_ort"
    BYSCRIPT = "by_script"

    def __str__(self):
        return self.value


class ConfigModifier:
    def __init__(self, num_layers):
        self.num_layers = num_layers

    def modify(self, config):
        if self.num_layers is None:
            return
        if hasattr(config, "num_hidden_layers"):
            config.num_hidden_layers = self.num_layers
            logger.info(f"Modifying pytorch model's number of hidden layers to: {self.num_layers}")
        if hasattr(config, "encoder_layers"):
            config.encoder_layers = self.num_layers
            logger.info(f"Modifying pytorch model's number of encoder layers to: {self.num_layers}")
        if hasattr(config, "decoder_layers "):
            config.decoder_layers = self.num_layers
            logger.info(f"Modifying pytorch model's number of decoder layers to: {self.num_layers}")

    def get_layer_num(self):
        return self.num_layers


IO_BINDING_DATA_TYPE_MAP = {
    "float32": numpy.float32,
    # TODO: Add more.
}


def create_onnxruntime_session(
    onnx_model_path,
    use_gpu,
    provider=None,
    enable_all_optimization=True,
    num_threads=-1,
    enable_profiling=False,
    verbose=False,
    provider_options={},  # map execution provider name to its option
):
    session = None
    try:
        from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions

        sess_options = SessionOptions()

        if enable_all_optimization:
            sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
        else:
            sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC

        if enable_profiling:
            sess_options.enable_profiling = True

        if num_threads > 0:
            sess_options.intra_op_num_threads = num_threads
            logger.debug(f"Session option: intra_op_num_threads={sess_options.intra_op_num_threads}")

        if verbose:
            sess_options.log_severity_level = 0
        else:
            sess_options.log_severity_level = 4

        logger.debug(f"Create session for onnx model: {onnx_model_path}")
        if use_gpu:
            if provider == "dml":
                providers = ["DmlExecutionProvider", "CPUExecutionProvider"]
            elif provider == "rocm":
                providers = ["ROCMExecutionProvider", "CPUExecutionProvider"]
            elif provider == "migraphx":
                providers = [
                    "MIGraphXExecutionProvider",
                    "ROCMExecutionProvider",
                    "CPUExecutionProvider",
                ]
            elif provider == "cuda":
                providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
            elif provider == "tensorrt":
                providers = [
                    "TensorrtExecutionProvider",
                    "CUDAExecutionProvider",
                    "CPUExecutionProvider",
                ]
            else:
                providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        else:
            providers = ["CPUExecutionProvider"]

        if provider_options:
            providers = [(name, provider_options[name]) if name in provider_options else name for name in providers]

        session = InferenceSession(onnx_model_path, sess_options, providers=providers)
    except:
        logger.error(f"Exception", exc_info=True)

    return session


def setup_logger(verbose=True):
    if verbose:
        coloredlogs.install(
            level="DEBUG",
            fmt="[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s",
        )
    else:
        coloredlogs.install(fmt="%(message)s")
        logging.getLogger("transformers").setLevel(logging.WARNING)


def prepare_environment(cache_dir, output_dir, use_gpu, provider=None):
    if cache_dir and not os.path.exists(cache_dir):
        os.makedirs(cache_dir)

    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    import onnxruntime

    if use_gpu:
        if provider == "dml":
            assert (
                "DmlExecutionProvider" in onnxruntime.get_available_providers()
            ), "Please install onnxruntime-directml package to test GPU inference."

        else:
            assert (
                "CUDAExecutionProvider" in onnxruntime.get_available_providers()
            ), "Please install onnxruntime-gpu package to test GPU inference."

    import transformers

    logger.info(f"PyTorch Version:{torch.__version__}")
    logger.info(f"Transformers Version:{transformers.__version__}")
    logger.info(f"Onnxruntime Version:{onnxruntime.__version__}")

    # Support three major versions of PyTorch and OnnxRuntime, and up to 6 months of transformers.
    from packaging import version

    assert version.parse(torch.__version__) >= version.parse("1.5.0")
    assert version.parse(transformers.__version__) >= version.parse("3.0.0")
    assert version.parse(onnxruntime.__version__) >= version.parse("1.4.0")


def get_latency_result(runtimes, batch_size):
    latency_ms = sum(runtimes) / float(len(runtimes)) * 1000.0
    latency_variance = numpy.var(runtimes, dtype=numpy.float64) * 1000.0
    throughput = batch_size * (1000.0 / latency_ms)

    return {
        "test_times": len(runtimes),
        "latency_variance": "{:.2f}".format(latency_variance),
        "latency_90_percentile": "{:.2f}".format(numpy.percentile(runtimes, 90) * 1000.0),
        "latency_95_percentile": "{:.2f}".format(numpy.percentile(runtimes, 95) * 1000.0),
        "latency_99_percentile": "{:.2f}".format(numpy.percentile(runtimes, 99) * 1000.0),
        "average_latency_ms": "{:.2f}".format(latency_ms),
        "QPS": "{:.2f}".format(throughput),
    }


def output_details(results, csv_filename):
    with open(csv_filename, mode="a", newline="") as csv_file:
        column_names = [
            "engine",
            "version",
            "providers",
            "device",
            "precision",
            "optimizer",
            "io_binding",
            "model_name",
            "inputs",
            "threads",
            "batch_size",
            "sequence_length",
            "custom_layer_num",
            "datetime",
            "test_times",
            "QPS",
            "average_latency_ms",
            "latency_variance",
            "latency_90_percentile",
            "latency_95_percentile",
            "latency_99_percentile",
        ]

        csv_writer = csv.DictWriter(csv_file, fieldnames=column_names)
        csv_writer.writeheader()
        for result in results:
            csv_writer.writerow(result)

    logger.info(f"Detail results are saved to csv file: {csv_filename}")


def output_summary(results, csv_filename, args):
    with open(csv_filename, mode="a", newline="") as csv_file:
        header_names = [
            "model_name",
            "inputs",
            "custom_layer_num",
            "engine",
            "version",
            "providers",
            "device",
            "precision",
            "optimizer",
            "io_binding",
            "threads",
        ]
        data_names = []
        for batch_size in args.batch_sizes:
            for sequence_length in args.sequence_lengths:
                data_names.append(f"b{batch_size}_s{sequence_length}")

        csv_writer = csv.DictWriter(csv_file, fieldnames=header_names + data_names)
        csv_writer.writeheader()
        for model_name in args.models:
            for input_count in [1, 2, 3]:
                for engine_name in args.engines:
                    for io_binding in [True, False, ""]:
                        for threads in args.num_threads:
                            row = {}
                            for result in results:
                                if (
                                    result["model_name"] == model_name
                                    and result["inputs"] == input_count
                                    and result["engine"] == engine_name
                                    and result["io_binding"] == io_binding
                                    and result["threads"] == threads
                                ):
                                    headers = {k: v for k, v in result.items() if k in header_names}
                                    if not row:
                                        row.update(headers)
                                        row.update({k: "" for k in data_names})
                                    else:
                                        for k in header_names:
                                            assert row[k] == headers[k]
                                    b = result["batch_size"]
                                    s = result["sequence_length"]
                                    row[f"b{b}_s{s}"] = result["average_latency_ms"]
                            if row:
                                csv_writer.writerow(row)

    logger.info(f"Summary results are saved to csv file: {csv_filename}")


def output_fusion_statistics(model_fusion_statistics, csv_filename):
    from transformers import __version__ as transformers_version

    with open(csv_filename, mode="a", newline="") as csv_file:
        column_names = ["model_filename", "datetime", "transformers", "torch"] + list(
            next(iter(model_fusion_statistics.values())).keys()
        )
        csv_writer = csv.DictWriter(csv_file, fieldnames=column_names)
        csv_writer.writeheader()
        for key in model_fusion_statistics.keys():
            model_fusion_statistics[key]["datetime"] = str(datetime.now())
            model_fusion_statistics[key]["transformers"] = transformers_version
            model_fusion_statistics[key]["torch"] = torch.__version__
            model_fusion_statistics[key]["model_filename"] = key
            csv_writer.writerow(model_fusion_statistics[key])
    logger.info(f"Fusion statistics is saved to csv file: {csv_filename}")


def inference_ort(ort_session, ort_inputs, result_template, repeat_times, batch_size, warm_up_repeat=0):
    result = {}
    timeit.repeat(lambda: ort_session.run(None, ort_inputs), number=1, repeat=warm_up_repeat)  # Dry run
    runtimes = timeit.repeat(lambda: ort_session.run(None, ort_inputs), number=1, repeat=repeat_times)
    result.update(result_template)
    result.update({"io_binding": False})
    result.update(get_latency_result(runtimes, batch_size))
    return result


def inference_ort_with_io_binding(
    ort_session,
    ort_inputs,
    result_template,
    repeat_times,
    ort_output_names,
    ort_outputs,
    output_buffers,
    output_buffer_max_sizes,
    batch_size,
    device,
    data_type=numpy.longlong,
    warm_up_repeat=0,
):
    result = {}

    # Bind inputs and outputs to onnxruntime session
    io_binding = ort_session.io_binding()
    # Bind inputs to device
    for name in ort_inputs.keys():
        np_input = torch.from_numpy(ort_inputs[name]).to(device)
        input_type = (
            IO_BINDING_DATA_TYPE_MAP[str(ort_inputs[name].dtype)]
            if str(ort_inputs[name].dtype) in IO_BINDING_DATA_TYPE_MAP
            else data_type
        )
        io_binding.bind_input(
            name,
            np_input.device.type,
            0,
            input_type,
            np_input.shape,
            np_input.data_ptr(),
        )
    # Bind outputs buffers with the sizes needed if not allocated already
    if len(output_buffers) == 0:
        allocateOutputBuffers(output_buffers, output_buffer_max_sizes, device)

    for i in range(len(ort_output_names)):
        io_binding.bind_output(
            ort_output_names[i],
            output_buffers[i].device.type,
            0,
            numpy.float32,
            ort_outputs[i].shape,
            output_buffers[i].data_ptr(),
        )
    timeit.repeat(
        lambda: ort_session.run_with_iobinding(io_binding),
        number=1,
        repeat=warm_up_repeat,
    )  # Dry run
    runtimes = timeit.repeat(
        lambda: ort_session.run_with_iobinding(io_binding),
        number=1,
        repeat=repeat_times,
    )
    result.update(result_template)
    result.update({"io_binding": True})
    result.update(get_latency_result(runtimes, batch_size))
    return result


def allocateOutputBuffers(output_buffers, output_buffer_max_sizes, device):
    # Allocate output tensors with the largest test size needed. So the allocated memory can be reused
    # for each test run.

    for i in output_buffer_max_sizes:
        output_buffers.append(torch.empty(i, dtype=torch.float32, device=device))


def set_random_seed(seed=123):
    """Set random seed manully to get deterministic results"""
    import random

    random.seed(seed)
    numpy.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # torch.backends.cudnn.enabled = False
    # torch.backends.cudnn.benchmark = False
    # torch.backends.cudnn.deterministic = True


def measure_memory(is_gpu, func):
    import os
    from time import sleep

    import psutil

    class MemoryMonitor:
        def __init__(self, keep_measuring=True):
            self.keep_measuring = keep_measuring

        def measure_cpu_usage(self):
            max_usage = 0
            while True:
                max_usage = max(max_usage, psutil.Process(os.getpid()).memory_info().rss / 1024**2)
                sleep(0.005)  # 5ms
                if not self.keep_measuring:
                    break
            return max_usage

        def measure_gpu_usage(self):
            from py3nvml.py3nvml import (
                NVMLError,
                nvmlDeviceGetCount,
                nvmlDeviceGetHandleByIndex,
                nvmlDeviceGetMemoryInfo,
                nvmlDeviceGetName,
                nvmlInit,
                nvmlShutdown,
            )

            max_gpu_usage = []
            gpu_name = []
            try:
                nvmlInit()
                deviceCount = nvmlDeviceGetCount()
                max_gpu_usage = [0 for i in range(deviceCount)]
                gpu_name = [nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(deviceCount)]
                while True:
                    for i in range(deviceCount):
                        info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i))
                        max_gpu_usage[i] = max(max_gpu_usage[i], info.used / 1024**2)
                    sleep(0.005)  # 5ms
                    if not self.keep_measuring:
                        break
                nvmlShutdown()
                return [
                    {
                        "device_id": i,
                        "name": gpu_name[i],
                        "max_used_MB": max_gpu_usage[i],
                    }
                    for i in range(deviceCount)
                ]
            except NVMLError as error:
                if not self.silent:
                    self.logger.error("Error fetching GPU information using nvml: %s", error)
                return None

    monitor = MemoryMonitor(False)

    memory_before_test = monitor.measure_gpu_usage() if is_gpu else monitor.measure_cpu_usage()

    from concurrent.futures import ThreadPoolExecutor

    with ThreadPoolExecutor() as executor:
        monitor = MemoryMonitor()
        mem_thread = executor.submit(monitor.measure_gpu_usage if is_gpu else monitor.measure_cpu_usage)
        try:
            fn_thread = executor.submit(func)
            result = fn_thread.result()
        finally:
            monitor.keep_measuring = False
            max_usage = mem_thread.result()

        if is_gpu:
            print(f"GPU memory usage: before={memory_before_test}  peak={max_usage}")
            if len(memory_before_test) >= 1 and len(max_usage) >= 1:
                before = memory_before_test[0]["max_used_MB"]
                after = max_usage[0]["max_used_MB"]
                return after - before
            else:
                return None
        else:
            print(f"CPU memory usage: before={memory_before_test:.1f} MB, peak={max_usage:.1f} MB")
            return max_usage - memory_before_test
