"""Quantization -> bits-per-weight table and weight memory math.

Values are the effective bits per weight (including block scales/min) as used
by llama.cpp's k_quants / imatrix quants. Sourced from the llama.cpp docs and
the widely-cited community tables. These are estimates; real file sizes vary a
little due to alignment and metadata.
"""

from __future__ import annotations

from dataclasses import dataclass


# quant name -> effective bits per weight
QUANT_BPW: dict[str, float] = {
    # 2-bit
    "Q2_K": 2.5625,
    "IQ2_XXS": 2.0625,
    "IQ2_XS": 2.3125,
    "IQ2_S": 2.5,
    "IQ2_M": 2.7,
    # 3-bit
    "Q3_K_S": 3.4375,
    "Q3_K_M": 3.84375,
    "Q3_K_L": 4.125,
    "IQ3_XXS": 3.0625,
    "IQ3_XS": 3.25,
    "IQ3_S": 3.5,
    "IQ3_M": 3.7,
    # 4-bit
    "Q4_0": 4.5,
    "Q4_1": 5.0,
    "Q4_K_S": 4.5,
    "Q4_K_M": 4.84375,
    "IQ4_NL": 4.5,
    "IQ4_XS": 4.25,
    # 5-bit
    "Q5_0": 5.5,
    "Q5_1": 6.0,
    "Q5_K_S": 5.5,
    "Q5_K_M": 5.6875,
    # 6-bit
    "Q6_K": 6.5625,
    # 8-bit
    "Q8_0": 8.5,
    "Q8_1": 9.0,
    # unquantized
    "F16": 16.0,
    "BF16": 16.0,
    "F32": 32.0,
    # AMD ROCm floating-point block quants (ciru-ai/ROCmFPX fork types).
    # These are *effective* bits-per-weight = real_file_size * 8 / total_params,
    # measured from the jcbtc/chadrock3.6 ROCmFP4/FPX GGUFs. They are higher than
    # the nominal 4.0/8.0 because the typical model is a hybrid attention+SSM
    # (Mamba) qwen35 with co-stored F32 state/conv tensors and mixed tensor types
    # (the 35B FPX file is mostly Q6_K experts + ROCmFPX attention). Using the
    # effective BPW reproduces the real file size, which is what a VRAM estimate
    # needs. NOTE: stock llama.cpp cannot run these — they require the pinned
    # ciru-ai/ROCmFPX runner; see command_preview for the warning.
    "ROCmFP4": 4.34,   # 27B Strix Lean: 14.82 GB / 27.32B params
    "ROCmFPX": 7.08,   # 35B A3B MoEQuality: model card states 7.08 BPW
    # nvfp4 (NVIDIA FP4, GGML_TYPE_NVFP4=40; llama.cpp PR #22196, merged Apr
    # 2026). Effective ~4.5 bpw from the block_nvfp4 struct: 64 weights in
    # 36 bytes = 4 bytes E4M3 sub-block scales + 32 bytes packed FP4 →
    # 36*8/64 = 4.5 bpw. Reproduces full-param file sizes (the stripped
    # Qwen3.6-27B NVFP4 ~14 GB file corresponds to ~4.15 bpw once you
    # account for the missing vision/MTP tensors). Blackwell GPUs get the
    # FP4 tensor-core speedup; older cards (Ada/Ampere/Hopper) get memory
    # savings only. See command_preview for the experimental caveat.
    "nvfp4": 4.5,
}


@dataclass(frozen=True)
class QuantInfo:
    name: str
    bpw: float


# Unsloth Dynamic (UD-*) quant base mapping. UD variants re-allocate bits
# per-tensor adaptively, so they have no single canonical bpw and aren't in
# QUANT_BPW. For dropdown sync we map them to the closest standard quant; the
# user can override the dropdown afterward. Add entries as Unsloth ships new
# UD variant suffixes.
_UD_BASE: dict[str, str] = {
    "Q2_K_XL": "Q2_K", "Q2_K_XS": "Q2_K", "Q2_K_L": "Q2_K",
    "Q3_K_XL": "Q3_K_M", "Q3_K_XS": "Q3_K_M", "Q3_K_L": "Q3_K_L",
    "Q4_K_XL": "Q4_K_M", "Q4_K_XS": "Q4_K_M", "Q4_K_L": "Q4_K_M",
    "Q5_K_XL": "Q5_K_M", "Q5_K_XS": "Q5_K_M",
    "Q6_K_XL": "Q6_K", "Q6_K_XS": "Q6_K",
    "Q8_K_XL": "Q8_0", "Q8_K_XS": "Q8_0",
}


def quant_from_filename(filename: str) -> str | None:
    """Detect a quant name from a GGUF filename like `Model-Q4_K_M.gguf`.

    Returns the canonical quant name (matching a key in QUANT_BPW) or None.
    Matching is case-insensitive on the quant token. Unsloth Dynamic (UD-*)
    variants (e.g. ``UD-Q4_K_XL``) have no canonical bpw; they map to their
    closest standard quant via ``_UD_BASE`` so the dropdown still syncs.
    """
    import re

    upper = filename.upper()
    # Try longer / more specific names first so e.g. Q4_K_S isn't caught by Q4_K.
    ordered = sorted(QUANT_BPW.keys(), key=len, reverse=True)
    for q in ordered:
        # match as a token: surrounded by non-alphanumeric or string edges
        if re.search(rf"(^|[^A-Z0-9]){re.escape(q.upper())}([^A-Z0-9]|$)", upper):
            return q

    # UD-* fallback: Unsloth Dynamic quants (UD-Q4_K_XL …) don't token-match a
    # QUANT_BPW key (the _XL/_XS suffix isn't a key), so check the known UD
    # bases directly and return the mapped standard quant.
    for ud_key, base in _UD_BASE.items():
        if re.search(rf"(^|[^A-Z0-9]){re.escape(ud_key)}([^A-Z0-9]|$)", upper):
            return base
    return None


def weight_bytes(params: int, quant: str) -> float:
    """Estimated weight memory in bytes for `params` parameters at `quant`.

    `params` is the total parameter count (including all MoE experts; the GGUF
    already encodes them so this is the natural input).
    """
    if quant not in QUANT_BPW:
        raise ValueError(f"Unknown quant type: {quant!r}")
    return params * QUANT_BPW[quant] / 8.0