# https://onnxruntime.ai/docs/api/python/auto_examples/plot_profiling.html
from time import time

import numpy as np
import onnxruntime as ort
from tqdm import tqdm


def Benchmark(onnx_file: str, steps: int = 10000, providers=['CUDAExecutionProvider'], provider_options=None) -> float:
    """Benckmark with Onnxruntime

    * Quantized Model of Onnxruntime - TensorrtExecutionProvider and Onnxruntime - CUDAExecutionProvider has different format.

    * Onnx that generated by PPQ is not supportable with TensorrtExecutionProvider.

    * Set providers=CUDAExecutionProvider before benchmark this file.

    """
    sess = ort.InferenceSession(path_or_bytes=onnx_file, providers=providers, provider_options=provider_options)

    feed_dict, output_names = {}, []
    for input_meta in sess.get_inputs():
        name, dtype, shape = input_meta.name, input_meta.type, input_meta.shape

        for element in shape:
            if element is None or type(element) == str:
                raise TypeError('Dynamic input is not supported by this function.')

        if dtype == 'tensor(float)':
            feed_dict[name] = np.random.random(size=shape).astype(np.float32)
        else:
            raise Exception(f'Input {name} has unexpected data type.')

    for output_meta in sess.get_outputs():
        output_names.append(output_meta.name)

    tick = time()
    for _ in tqdm(range(steps)):
        sess.run(output_names=output_names, input_feed=feed_dict)
    tok = time()

    print(f'Time span: {tok - tick: .4f} sec')
    return tick - tok


def Profile(onnx_file: str, steps: int = 1, providers=['CUDAExecutionProvider'], provider_options=None):
    """Profile with Onnxruntime

    * Quantized Model of Onnxruntime - TensorrtExecutionProvider and Onnxruntime - CUDAExecutionProvider has different format.

    * Onnx that generated by PPQ is not supportable with TensorrtExecutionProvider.

    * Set providers=CUDAExecutionProvider before benchmark this file.

    """
    options = ort.SessionOptions()
    options.enable_profiling = True
    sess = ort.InferenceSession(
        path_or_bytes=onnx_file, providers=providers, provider_options=provider_options, sess_options=options
    )

    feed_dict, output_names = {}, []
    for input_meta in sess.get_inputs():
        name, dtype, shape = input_meta.name, input_meta.type, input_meta.shape

        for element in shape:
            if element is None or type(element) == str:
                raise TypeError('Dynamic input is not supported by this function.')

        if dtype == 'tensor(float)':
            feed_dict[name] = np.random.random(size=shape).astype(np.float32)
        else:
            raise Exception(f'Input {name} has unexpected data type.')

    for output_meta in sess.get_outputs():
        output_names.append(output_meta.name)

    for _ in tqdm(range(steps)):
        sess.run(output_names=output_names, input_feed=feed_dict)

    prof_file = sess.end_profiling()
    print(f'Profile file is generated at {prof_file}, open it with your web browser chrome://tracing/')
