from collections.abc import Sequence
from typing import Any, Union

import onnxruntime


def torch_device_to_onnx_providers(
    device: str,
) -> Sequence[Union[str, tuple[str, dict[str, Any]]]]:
    """
    PyTorch のデバイス種別を ONNX の ExecutionProvider に変換する

    Args:
        device (str): PyTorch のデバイス種別

    Returns:
        Sequence[Union[str, tuple[str, dict[str, Any]]]]: ExecutionProvider のリスト
    """

    if device.startswith("cuda"):
        return [
            # cudnn_conv_algo_search を DEFAULT にすると推論速度が大幅に向上する
            ## ref: https://medium.com/neuml/debug-onnx-gpu-performance-c9290fe07459
            ("CUDAExecutionProvider", {"arena_extend_strategy": "kSameAsRequested", "cudnn_conv_algo_search": "DEFAULT"}),
            # CUDA が利用できない場合、可能であれば DirectML を利用する (明示的な device_id 指定が必要)
            ## device_id: 0 は、システムにインストールされているプライマリディスプレイ用 GPU に対応する
            ## プライマリディスプレイ用 GPU (GPU 0) よりも性能の高い GPU が接続されている環境では、 適宜 device_id を変更する必要がある
            ## ref: https://github.com/w-okada/voice-changer/issues/410#issuecomment-1627994911
            ("DmlExecutionProvider", {"device_id": 0}),
            # arena_extend_strategy を kSameAsRequested にすると、推論セッションによって作成される
            # メモリアリーナが、実際に推論に必要な容量以上にメモリを確保する問題を防ぐことができる
            ## ref: https://github.com/microsoft/onnxruntime/issues/11627#issuecomment-1137668551
            ## ref: https://skottmckay.github.io/onnxruntime/docs/reference/api/c-api.html
            ("CPUExecutionProvider", {"arena_extend_strategy": "kSameAsRequested"}),
        ]  # fmt: skip
    else:
        return [
            ("CPUExecutionProvider", {"arena_extend_strategy": "kSameAsRequested"}),
        ]


def get_onnx_device_options(
    onnx_session: onnxruntime.InferenceSession,
    onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
) -> tuple[str, int, onnxruntime.RunOptions]:
    """
    ONNX 推論時のデバイス関連のオプションを取得する

    Args:
        onnx_session (onnxruntime.InferenceSession): 初期化済みの ONNX セッション
        onnx_providers (Sequence[Union[str, tuple[str, dict[str, Any]]]]): ExecutionProvider のリスト

    Returns:
        tuple[str, int, onnxruntime.RunOptions]: 入力テンソルの転送に使用するデバイス種別, デバイス ID, 実行オプション
    """

    # ONNX セッションに対応する SessionOptions を取得
    sess_options = onnx_session.get_session_options()

    # 実際に推論に用いられる ExecutionProvider を取得
    first_provider = onnx_session.get_providers()[0]

    # 入力テンソルを転送する推論デバイスを取得
    ## GPU への I/O Binding は CUDA と DirectML 以外はサポートされていないので、
    ## それ以外の ExecutionProvider が指定された場合は CPU に入力テンソルを転送する
    if first_provider == "CUDAExecutionProvider":
        device_type = "cuda"
    elif first_provider == "DmlExecutionProvider":
        device_type = "dml"
    else:
        device_type = "cpu"

    # 入力テンソルを転送する GPU デバイスの ID を取得
    ## ExecutionProvider に指定したオプションの中から device_id を取得し、入力テンソルの転送先として指定する
    ## InferenceSession で利用するデバイス ID と入力テンソルの転送先デバイス ID は一致している必要がある
    ## 本来は ExecutionProvider に指定したオプションは InferenceSession.get_provider_options() で取得できるはずだが、
    ## 手元環境では DmlExecutionProvider のみ常に空の辞書が返されることがあったため、当面 onnx_providers から直接オプションを取り出している
    ## CPU 推論時の device_id は 0 で固定
    device_id = 0
    onnx_providers_dict: dict[str, dict[str, Any]] = {}
    if device_type != "cpu":
        for provider in onnx_providers:
            if isinstance(provider, tuple):
                provider_name, options = provider
                onnx_providers_dict[provider_name] = options
            else:
                onnx_providers_dict[provider] = {}
        first_provider_options = onnx_providers_dict[first_provider]
        if "device_id" in first_provider_options:
            device_id = int(first_provider_options["device_id"])

    # 推論後にメモリアリーナを縮小し、確保されていたメモリを解放する
    ## この設定により、多量の推論処理を行った際に漸進的にメモリが消費される現象を回避できる
    ## 特に CPU 推論時に有効な設定で、これによりアイドル時のメモリ消費量が大幅に削減される
    ## 音声合成モデルの CPU 推論時に SessionOptions の enable_cpu_mem_arena を無効化すると 推論速度が 1 秒以上低下してしまうが、
    ## enable_cpu_mem_arena は有効のまま下記設定を行うと、推論速度を維持しながらメモリ消費量を削減できる
    ## SessionOptions の enable_cpu_mem_arena が False の時は実行しない
    ## ref: https://onnxruntime.ai/docs/get-started/with-c.html
    ## ref: https://github.com/microsoft/onnxruntime/issues/9313#issuecomment-2182919186
    ## ref: https://github.com/microsoft/onnxruntime/issues/11627
    ## ref: https://github.com/microsoft/onnxruntime/issues/22297#issuecomment-2438629814
    ## ref: https://github.com/microsoft/onnxruntime/blob/v1.20.1/onnxruntime/test/python/onnxruntime_test_python.py#L1626-L1647
    ## ref: https://github.com/microsoft/onnxruntime/blob/v1.20.1/include/onnxruntime/core/session/onnxruntime_run_options_config_keys.h#L19-L27
    run_options = onnxruntime.RunOptions()
    if sess_options.enable_cpu_mem_arena is True:
        if first_provider == "CPUExecutionProvider":
            # CPU 推論時は cpu:0 を指定
            run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu:0")  # fmt: skip
        elif first_provider == "DmlExecutionProvider":
            # DirectML 推論時はこのオプションはサポートされていないようなので、何も指定しない
            # "The registered allocator for device-id combination is not an arena based allocator: gpu:0" のようなエラーが出る…
            pass
        elif first_provider == "CUDAExecutionProvider":
            # CUDA 推論時は cpu:0;gpu:(device_id) を指定
            ## 公式テストコードを読む限り、CUDA だけでなく CPU のメモリも明示的に解放した方がよいらしい
            run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", f"cpu:0;gpu:{device_id}")  # fmt: skip

    return device_type, device_id, run_options