"""UD (IQ2_M) layer plan - map GGUF tensor name -> quant dtype.

Direct translation of llama.cpp llama-quant.cpp `set_type_for_tensor` for
`LLAMA_FTYPE_MOSTLY_IQ2_M` (lines 450-529), for the dense `qwen35` arch:

  token_embd           -> IQ3_S
  output (lm_head)     -> Q5_K
  attn_v               -> Q4_K   (n_gqa = 24/4 = 6 >= 4)
  attn_output / ssm_out-> IQ3_S
  ffn_down (first n/8) -> IQ3_S, rest -> IQ2_S
  everything else      -> IQ2_S

Small / norm / 1D tensors are handled by the caller as F32/F16 and must NOT
be listed here.
"""

from __future__ import annotations

# GGML dtype ids.
F32 = 0
F16 = 1
Q8_0 = 8
Q2_K = 10
Q5_K = 13
Q4_K = 12
IQ2_S = 22
IQ3_S = 21


def _default_dtype(name: str, i_ffn_down: int, n_ffn_down: int,
                   n_head: int, n_head_kv: int, ftype_iq2_m: bool) -> int:
    """llama-quant.cpp set_type_for_tensor for MOSTLY_IQ2_M."""
    if name == "token_embd.weight":
        return IQ3_S if ftype_iq2_m else Q2_K
    if name == "output.weight":
        return Q4_K if ftype_iq2_m else Q8_0
    if name.endswith(".attn_v.weight"):
        # n_gqa >= 4 -> Q4_K; else IQ3_S (or Q2_K for non-IQ2_M).
        if n_head // n_head_kv >= 4:
            return Q4_K
        return IQ3_S if ftype_iq2_m else Q2_K
    if name.endswith(".attn_output.weight") or name.endswith(".ssm_out.weight"):
        return IQ3_S if ftype_iq2_m else Q2_K
    if name.endswith(".ffn_down.weight"):
        if ftype_iq2_m and i_ffn_down < n_ffn_down // 8:
            return IQ3_S
        return IQ2_S
    return IQ2_S


def make_iq2_m_plan(tensor_names: list[str], *, n_head: int, n_head_kv: int,
                    ftype_iq2_m: bool = True) -> dict[str, int]:
    """Return {gguf_name: gguf_dtype} for the quantizable tensors.

    tensor_names must be sorted in GGUF info-table order (llama.cpp walks
    them in file order) so the `ffn_down` first-n/8 rule applies correctly.
    """
    plan: dict[str, int] = {}
    n_ffn_down = sum(1 for n in tensor_names if n.endswith(".ffn_down.weight"))
    i_ffn_down = 0
    for name in tensor_names:
        if name.endswith(".ffn_down.weight"):
            plan[name] = _default_dtype(name, i_ffn_down, n_ffn_down, n_head, n_head_kv, ftype_iq2_m)
            i_ffn_down += 1
        else:
            plan[name] = _default_dtype(name, 0, n_ffn_down, n_head, n_head_kv, ftype_iq2_m)
    return plan