# SPDX-License-Identifier: Apache-2.0
"""Data type mapping between ONNX and PyTorch."""

from typing import Dict, Optional

import onnx
import torch

# ONNX TensorProto data type to PyTorch dtype mapping
DTYPE_MAP: Dict[int, Optional[torch.dtype]] = {
    onnx.TensorProto.FLOAT: torch.float32,
    onnx.TensorProto.FLOAT16: torch.float16,
    onnx.TensorProto.BFLOAT16: torch.bfloat16,
    onnx.TensorProto.DOUBLE: torch.float64,
    onnx.TensorProto.INT8: torch.int8,
    onnx.TensorProto.INT16: torch.int16,
    onnx.TensorProto.INT32: torch.int32,
    onnx.TensorProto.INT64: torch.int64,
    onnx.TensorProto.UINT8: torch.uint8,
    onnx.TensorProto.UINT16: torch.uint16,
    onnx.TensorProto.UINT32: torch.uint32,
    onnx.TensorProto.UINT64: torch.uint64,
    onnx.TensorProto.BOOL: torch.bool,
    onnx.TensorProto.COMPLEX64: torch.complex64,
    onnx.TensorProto.COMPLEX128: torch.complex128,
    onnx.TensorProto.STRING: None,
}

# Reverse mapping: PyTorch dtype to ONNX TensorProto data type
TORCH_TO_ONNX_DTYPE: Dict[torch.dtype, int] = {
    torch_dtype: onnx_dtype
    for onnx_dtype, torch_dtype in DTYPE_MAP.items()
    if torch_dtype is not None
}


def onnx_dtype_to_torch(onnx_dtype: int) -> Optional[torch.dtype]:
    """Convert ONNX TensorProto data type to PyTorch dtype.

    Parameters
    ----------
    onnx_dtype : int
        ONNX TensorProto data type enum value.

    Returns
    -------
    Optional[torch.dtype]
        Corresponding PyTorch dtype, or None if not supported.
    """
    return DTYPE_MAP.get(onnx_dtype)


def torch_dtype_to_onnx(torch_dtype: torch.dtype) -> Optional[int]:
    """Convert PyTorch dtype to ONNX TensorProto data type.

    Parameters
    ----------
    torch_dtype : torch.dtype
        PyTorch dtype.

    Returns
    -------
    Optional[int]
        Corresponding ONNX TensorProto data type, or None if not supported.
    """
    return TORCH_TO_ONNX_DTYPE.get(torch_dtype)


# ONNX stash_type to PyTorch dtype mapping (used in normalization ops)
_STASH_TYPE_MAP: Dict[int, torch.dtype] = {
    1: torch.float32,
    10: torch.float16,
    11: torch.float64,
    16: torch.bfloat16,
}


def stash_type_to_torch_dtype(stash_type: int) -> torch.dtype:
    """Convert ONNX stash_type attribute to PyTorch dtype.

    The stash_type attribute specifies the floating-point precision
    for intermediate computations in normalization operations.

    Parameters
    ----------
    stash_type : int
        ONNX stash_type value:
        - 1: float32
        - 10: float16
        - 11: float64
        - 16: bfloat16

    Returns
    -------
    torch.dtype
        Corresponding PyTorch dtype. Defaults to float32 for unknown values.
    """
    return _STASH_TYPE_MAP.get(stash_type, torch.float32)


__all__ = [
    "DTYPE_MAP",
    "TORCH_TO_ONNX_DTYPE",
    "onnx_dtype_to_torch",
    "stash_type_to_torch_dtype",
    "torch_dtype_to_onnx",
]
