import importlib
from typing import Any, Optional
import torch


ops = ["torch.Tensor.__matmul__", "torch.addbmm", "torch.addmm", "torch.addmv", "torch.addr", "torch.baddbmm", "torch.bmm", "torch.chain_matmul", "torch.linalg.multi_dot", "torch.nn.functional.conv1d", "torch.nn.functional.conv2d", "torch.nn.functional.conv3d", "torch.nn.functional.conv_transpose1d", "torch.nn.functional.conv_transpose2d", "torch.nn.functional.conv_transpose3d", "torch.nn.GRUCell", "torch.nn.functional.linear", "torch.nn.LSTMCell", "torch.matmul", "torch.mm", "torch.mv", "torch.prelu", "torch.nn.RNNCell", "torch.embedding"]
supported_cast_pairs = {
    torch.float16: (torch.float32,),
    torch.float32: (torch.float16,),
}


def forward(op, args: tuple, kwargs: dict):
    if not torch.dml.is_autocast_enabled:
        return op(*args, **kwargs)
    args = list(map(cast, args))
    for kwarg in kwargs:
        kwargs[kwarg] = cast(kwargs[kwarg])
    return op(*args, **kwargs)


def cast(tensor: torch.Tensor):
    if not torch.is_tensor(tensor):
        return tensor
    dtype: torch.dtype = tensor.dtype
    if dtype not in supported_cast_pairs or (torch.dml.autocast_gpu_dtype != dtype and torch.dml.autocast_gpu_dtype not in supported_cast_pairs[dtype]):
        return tensor
    return tensor.type(torch.dml.autocast_gpu_dtype)


def cond(op: str):
    if isinstance(op, str):
        func_path = op.split('.')
        for i in range(len(func_path)-1, -1, -1):
            try:
                resolved_obj = importlib.import_module('.'.join(func_path[:i]))
                break
            except ImportError:
                pass
        for attr_name in func_path[i:-1]:
            resolved_obj = getattr(resolved_obj, attr_name)
        op = getattr(resolved_obj, func_path[-1])
        setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: forward(op, args, kwargs))


for o in ops:
    cond(o)


class autocast:
    prev: bool

    fast_dtype: torch.dtype = torch.float16
    prev_fast_dtype: torch.dtype
    def __init__(self, dtype: Optional[torch.dtype] = torch.float16):
        self.fast_dtype = dtype

    def __enter__(self):
        self.prev = torch.dml.is_autocast_enabled
        self.prev_fast_dtype = torch.dml.autocast_gpu_dtype
        torch.dml.is_autocast_enabled = True
        torch.dml.autocast_gpu_dtype = self.fast_dtype

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
        torch.dml.is_autocast_enabled = self.prev
        torch.dml.autocast_gpu_dtype = self.prev_fast_dtype