import torch
import intel_extension_for_pytorch._C as core
import warnings
from typing import Any, Optional
from torch.types import _dtype


# Expand torch.amp.autocast_mode.autocast to support both torch.bfloat16 and torch.float16 on cpu.
class _mode_autocast(torch.amp.autocast_mode.autocast):
    def __init__(
        self,
        device_type: str,
        dtype: Optional[_dtype] = None,
        enabled: bool = True,
        cache_enabled: Optional[bool] = None,
    ):
        if torch._jit_internal.is_scripting():
            self._enabled = enabled
            self.device = device_type
            self.fast_dtype = dtype
            # TODO: support get_autocast_gpu/cpu_dtype
            assert dtype is not None
            return
        self.device = device_type
        if self.device == "cuda":
            self.fast_dtype = torch.get_autocast_gpu_dtype()
        elif self.device == "cpu":
            self.fast_dtype = torch.get_autocast_cpu_dtype()
        elif self.device == "xpu":
            self.fast_dtype = torch.xpu.get_autocast_xpu_dtype()  # type: ignore[attr-defined]
        else:
            raise RuntimeError(
                "User specified autocast device_type must be 'cuda' or 'cpu'"
            )
        self._cache_enabled = torch.is_autocast_cache_enabled()
        if (
            enabled
            and self.device == "cuda"
            and torch.cuda.amp.common.amp_definitely_not_available()
        ):
            warnings.warn(
                "User provided device_type of 'cuda', but CUDA is not available. Disabling"
            )
            enabled = False
        if dtype is not None:
            self.fast_dtype = dtype
        if cache_enabled is not None:
            self._cache_enabled = cache_enabled

        if self.device == "cpu":
            supported_dtype = [torch.bfloat16, torch.float16]
            if self.fast_dtype not in supported_dtype:
                error_message = "In CPU autocast, but the target dtype is not supported. Disabling autocast.\n"
                error_message += "CPU Autocast only supports dtype of torch.bfloat16 and torch.float16 currently."
                warnings.warn(error_message)
                enabled = False
        if self.device == "xpu":
            supported_dtype = [torch.bfloat16, torch.float16]
            if self.fast_dtype not in supported_dtype:
                error_message = "In XPU autocast, but the target dtype is not supported. Disabling autocast.\n"
                error_message += (
                    "XPU Autocast only supports dtype of torch.bfloat16 currently."
                )
                warnings.warn(error_message)
                enabled = False
        if self.device == "cuda":
            if self.fast_dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
                raise RuntimeError(
                    "Current CUDA Device does not support bfloat16. Please switch dtype to float16."
                )
        self._enabled = enabled


# same as torch.cpu.amp.autocast
class autocast_cpu(_mode_autocast):
    r"""
    See :class:`torch.autocast`.
    ``torch.cpu.amp.autocast(args...)`` is equivalent to ``torch.autocast("cpu", args...)``
    """

    def __init__(
        self,
        enabled: bool = True,
        dtype: torch.dtype = torch.bfloat16,
        cache_enabled: bool = True,
    ):
        if torch._jit_internal.is_scripting():
            self._enabled = enabled
            self.device = "cpu"
            self.fast_dtype = dtype
            return
        super().__init__(
            "cpu", enabled=enabled, dtype=dtype, cache_enabled=cache_enabled
        )

    def __enter__(self):
        if torch._jit_internal.is_scripting():
            return self
        return super().__enter__()

    # TODO: discuss a unified TorchScript-friendly API for autocast
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):  # type: ignore[override]
        if torch._jit_internal.is_scripting():
            return
        return super().__exit__(exc_type, exc_val, exc_tb)

    def __call__(self, func):
        if torch._jit_internal.is_scripting():
            return func
        return super().__call__(func)


# Expand torch.cpu.amp.autocast to support both torch.bfloat16 & torch.half
# and support the disabling of cache_enabled for autocast within jit.trace.
class _autocast(autocast_cpu):
    def __enter__(self):
        self.prev_cache_enabled = torch.is_autocast_cache_enabled()
        self.prev = torch.is_autocast_cpu_enabled()
        self.prev_fast_dtype = core.get_autocast_dtype()
        torch.set_autocast_cpu_enabled(self._enabled)
        core.set_autocast_dtype(self.fast_dtype)
        torch.autocast_increment_nesting()
        torch.set_autocast_cache_enabled(self._cache_enabled)

    def __exit__(self, *args):
        # Drop the cache when we exit to a nesting level that's outside any instance of autocast.
        if torch.autocast_decrement_nesting() == 0:
            core.clear_autocast_cache()
            torch.clear_autocast_cache()
        torch.set_autocast_cpu_enabled(self.prev)
        core.set_autocast_dtype(self.prev_fast_dtype)
        torch.set_autocast_cache_enabled(self.prev_cache_enabled)
        return False


if core._has_cpu():
    torch.cpu.amp.autocast = _autocast
