from importlib.metadata import metadata

from packaging import version
import torch

try:
    import triton  # noqa: F401
    import triton.language as tl  # noqa: F401

    triton_available = True
except ImportError:
    triton_available = False


_NF4_QUANT_TABLE = torch.tensor(
    [
        -1.0,
        -0.6961928009986877,
        -0.5250730514526367,
        -0.39491748809814453,
        -0.28444138169288635,
        -0.18477343022823334,
        -0.09105003625154495,
        0.0,
        0.07958029955625534,
        0.16093020141124725,
        0.24611230194568634,
        0.33791524171829224,
        0.44070982933044434,
        0.5626170039176941,
        0.7229568362236023,
        1.0,
    ],
    dtype=torch.float32,
    device="xpu"
    if hasattr(torch, "xpu") and torch.xpu.is_available()
    else "cpu",  # Only cpu/xpu use this table for now.
)
_FP4_QUANT_TABLE = torch.tensor(
    [
        0.0000,
        0.0052,
        0.6667,
        1.0000,
        0.3333,
        0.5000,
        0.1667,
        0.2500,
        0.0000,
        -0.0052,
        -0.6667,
        -1.0000,
        -0.3333,
        -0.5000,
        -0.1667,
        -0.2500,
    ],
    dtype=torch.float32,
    device="xpu"
    if hasattr(torch, "xpu") and torch.xpu.is_available()
    else "cpu",  # Only cpu/xpu use this table for now.
)
CODE = {"nf4": _NF4_QUANT_TABLE, "fp4": _FP4_QUANT_TABLE}

# Cache 4-bit dequantization code tensors per (quant_type, device).
_code_4bit_cache: dict[tuple[str, torch.device], torch.Tensor] = {}


def _get_4bit_code(quant_type: str, device: torch.device) -> torch.Tensor:
    key = (quant_type, device)
    if key not in _code_4bit_cache:
        from bitsandbytes.functional import get_4bit_type

        _code_4bit_cache[key] = get_4bit_type(quant_type, device=device)
    return _code_4bit_cache[key]


def get_gaudi_sw_version():
    """
    Returns the installed version of Gaudi SW.
    """
    try:
        # if we find the spec, examine the installed version
        plugin_metadata = metadata("habana-torch-plugin")
        plugin_version = plugin_metadata.get("Version")
        if plugin_version:
            gaudi_version = version.parse(plugin_version)
    except Exception:
        gaudi_version = None

    return gaudi_version


GAUDI_SW_VER = get_gaudi_sw_version()