from __future__ import annotations

import warnings


def resolve_dtype(torch, device, precision: str):
    """Map a precision string to a torch dtype, gated on hardware support.

    bf16 is only worth it on hardware with native kernels (any modern GPU, or a
    CPU with AVX512-BF16/AMX). Without them PyTorch silently routes bf16 matmuls
    onto a fallback that is orders of magnitude slower than fp32, so this gate
    falls back to fp32 with a warning instead. The CPU check is the same one
    PyTorch uses to route bf16 matmuls onto oneDNN's native kernels.
    """
    if precision == "fp32":
        return torch.float32
    if precision != "bf16":
        raise RuntimeError(f"Unsupported precision '{precision}'. Use 'fp32' or 'bf16'.")
    if device.type == "cuda":
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
        reason = "this GPU has no bf16 support"
    else:
        try:
            native = torch.ops.mkldnn._is_mkldnn_bf16_supported()
        except (AttributeError, RuntimeError):
            native = False
        if native:
            return torch.bfloat16
        reason = "this CPU has no native bf16 support (needs AVX512-BF16/AMX)"
    warnings.warn(f"precision 'bf16' requested but {reason}; running fp32 instead.")
    return torch.float32


# Keep the Library object alive: dropping it would un-register the impl.
_deform_conv2d_shim: object | None = None


def ensure_deform_conv2d_autocast_cpu(torch) -> None:
    """Give torchvision's deform_conv2d an AutocastCPU wrapper (fp32 fallback).

    torchvision (<= 0.27 at least) ships no AutocastCPU registration and no
    BFloat16 CPU kernel for ``torchvision::deform_conv2d``, so running a model
    that uses it (BiRefNet) under ``torch.autocast("cpu", bf16)`` crashes with
    "deformable_im2col not implemented for 'BFloat16'". This registers a
    dispatcher-level wrapper that casts the op's inputs to fp32 and runs it
    outside autocast — upstream/downstream convs keep their bf16 speed, only
    this one op pays a cast. No-op if called twice or if a future torchvision
    registers its own wrapper.
    """
    global _deform_conv2d_shim
    if _deform_conv2d_shim is not None:
        return
    import torchvision  # noqa: F401  (registers torchvision::deform_conv2d)

    lib = torch.library.Library("torchvision", "IMPL", "AutocastCPU")

    def _deform_conv2d_fp32(input, weight, offset, mask, bias,
                            stride_h, stride_w, pad_h, pad_w,
                            dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask):
        with torch.autocast("cpu", enabled=False):
            return torch.ops.torchvision.deform_conv2d(
                input.float(), weight.float(), offset.float(), mask.float(),
                bias.float(), stride_h, stride_w, pad_h, pad_w,
                dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask,
            )

    try:
        lib.impl("deform_conv2d", _deform_conv2d_fp32)
    except RuntimeError:
        # torchvision grew its own AutocastCPU wrapper; ours is unneeded.
        return
    _deform_conv2d_shim = lib