"""
Shape-Based Architecture Inference — derive model parameters from weight shapes alone.

Covers all major transformer and CNN architectures that use SafeTensors format:

  Decoder-only LLMs:    Llama, Mistral, Qwen, Falcon, GPT-2, MPT, Phi, OPT, …
  Encoder-only LLMs:    BERT, RoBERTa, DistilBERT, ALBERT, DeBERTa, …
  Encoder-decoder:      T5, mT5, BART, mBART, …
  Vision-Language:      LLaVA, InternVL, Qwen2-VL, Phi-3-Vision, Idefics, MiniCPM-V, …
  MoE:                  Mixtral, DeepSeek-V2/V3, …
  MLA attention:        DeepSeek
  Vision Transformers:  CLIP, SigLIP, DINOv2, EVA-CLIP, … (standalone or in VLMs)
  CNNs:                 ResNet, EfficientNet, ConvNeXt, VGG, MobileNet, …

All inference reads only SafeTensors headers — zero weight loading cost.

Usage::

    from dokodemo_ai.graph.shape_inference import ShapeInference

    inferred = ShapeInference.from_index(safetensors_index, param_meta)
    print(inferred.summary())
    # ShapeInferredConfig(arch=decoder, hidden=4096, heads=32/32, head_dim=128, layers=32, ...)
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from dokodemo_ai.graph.compiler import TensorRef

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Result dataclass
# ---------------------------------------------------------------------------


@dataclass
class ShapeInferredConfig:
    """
    Architecture parameters derived purely from weight tensor shapes.

    Fields that could not be determined are left at their zero/empty defaults.
    Use ``merge_with_config(config)`` to blend these into an HF config dict.
    """

    # Transformer / LLM
    hidden_size: int = 0
    vocab_size: int = 0
    num_layers: int = 0
    num_q_heads: int = 0
    num_kv_heads: int = 0
    head_dim: int = 0
    is_moe: bool = False
    num_experts: int = 0
    is_mla: bool = False
    is_vlm: bool = False
    vision_arch: str = ""

    # Broad classification
    arch_type: str = "unknown"  # decoder | encoder | encoder_decoder | vit | cnn | unknown

    # Encoder-decoder specific (T5 / BART)
    num_encoder_layers: int = 0
    num_decoder_layers: int = 0

    # CNN / ViT specific
    num_classes: int = 0    # from classifier head weight
    in_channels: int = 0   # from first conv weight

    # Tracks which parameters were successfully inferred (for debug logging)
    _sources: dict[str, str] = field(default_factory=dict, repr=False)

    def merge_with_config(self, config: dict) -> dict:
        """
        Return a new config dict where missing/zero values are filled from
        weight-inferred parameters.

        The HF config always takes priority; inferred values only fill gaps.
        """
        merged = dict(config)

        def _fill(config_key: str, inferred_value: int | str | bool) -> None:
            if not inferred_value:
                return
            if not merged.get(config_key):
                merged[config_key] = inferred_value

        _fill("hidden_size", self.hidden_size)
        _fill("vocab_size", self.vocab_size)
        _fill("num_attention_heads", self.num_q_heads)
        _fill("num_key_value_heads", self.num_kv_heads)
        # NOTE: head_dim is intentionally NOT merged.
        # InferenceEngine computes head_dim = hidden_size // num_attention_heads,
        # which is always correct.  Merging shape-inferred head_dim risks injecting
        # a wrong value (e.g. 128 instead of 64 for Qwen2-0.5B) when the HF
        # config.json lacks an explicit head_dim key.
        _fill("num_hidden_layers", self.num_layers)

        # Encoder-decoder (T5 uses "num_layers" for encoder count)
        if self.num_encoder_layers:
            _fill("num_layers", self.num_encoder_layers)
        if self.num_decoder_layers:
            _fill("num_decoder_layers", self.num_decoder_layers)

        # CNN / ViT
        _fill("num_labels", self.num_classes)
        _fill("num_channels", self.in_channels)

        # Also fill into text_config sub-dict for VLMs
        if "text_config" in merged and isinstance(merged["text_config"], dict):
            tc = dict(merged["text_config"])
            for k, v in [
                ("hidden_size", self.hidden_size),
                ("num_attention_heads", self.num_q_heads),
                ("num_key_value_heads", self.num_kv_heads),
                # head_dim intentionally omitted — see note above.
                ("num_hidden_layers", self.num_layers),
            ]:
                if v and not tc.get(k):
                    tc[k] = v
            merged["text_config"] = tc

        return merged

    def summary(self) -> str:
        parts = [f"arch={self.arch_type}"]
        if self.hidden_size:
            parts.append(f"hidden={self.hidden_size}")
        if self.vocab_size:
            parts.append(f"vocab={self.vocab_size}")
        if self.num_layers:
            parts.append(f"layers={self.num_layers}")
        if self.num_encoder_layers or self.num_decoder_layers:
            parts.append(f"enc/dec={self.num_encoder_layers}/{self.num_decoder_layers}")
        if self.num_q_heads:
            parts.append(f"heads={self.num_q_heads}/{self.num_kv_heads}")
        if self.head_dim:
            parts.append(f"head_dim={self.head_dim}")
        if self.num_classes:
            parts.append(f"classes={self.num_classes}")
        if self.in_channels:
            parts.append(f"in_ch={self.in_channels}")
        if self.is_moe:
            parts.append(f"MoE(experts={self.num_experts})")
        if self.is_mla:
            parts.append("MLA")
        if self.is_vlm:
            parts.append(f"VLM/{self.vision_arch or '?'}")
        return "ShapeInferredConfig(" + ", ".join(parts) + ")"


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

# Common head dimensions, ordered most→least frequent.
_COMMON_HEAD_DIMS = (128, 64, 256, 96, 80, 48, 40, 32)

# ---------------------------------------------------------------------------
# Layer index regexes
# ---------------------------------------------------------------------------

# Primary pattern — covers decoder-only and most encoders.
# Groups:   layers            (Llama / Mistral / Qwen / DeepSeek)
#           h                 (GPT-2, Falcon)
#           blocks            (Bloom, misc)
#           decoder_layers /
#           encoder_layers    (some seq2seq models)
#           resblocks         (OpenAI-format CLIP / EVA-CLIP)
#           block             (T5 — singular: encoder.block.0)
#           layer             (BERT / RoBERTa — singular: encoder.layer.0)
#           transformer_blocks (Stable-Diffusion UNet text encoder)
_LAYER_IDX_RE = re.compile(
    r"(?:^|\.)"
    r"(?:layers|h|blocks|decoder_layers|encoder_layers"
    r"|resblocks|block|layer|transformer_blocks)"
    r"\.(\d+)\."
)

# T5-style encoder / decoder block patterns — tracked separately for enc-dec split.
_T5_ENCODER_RE = re.compile(r"(?:^|\.)encoder\.block\.(\d+)\.")
_T5_DECODER_RE = re.compile(r"(?:^|\.)decoder\.block\.(\d+)\.")

# BERT-style encoder layer pattern — used for arch_type classification.
_BERT_ENCODER_RE = re.compile(r"(?:^|\.)encoder\.layer\.(\d+)\.")

# ---------------------------------------------------------------------------
# Projection suffixes
# ---------------------------------------------------------------------------

# Separate Q projections
_Q_PROJ_SUFFIXES = (
    ".q_proj.weight",
    ".query.weight",
    ".query_proj.weight",
    ".Wq.weight",
    ".attn_q.weight",
    ".self.query.weight",           # BERT: attention.self.query.weight
    ".attention.query.weight",      # some BERT variants
    ".SelfAttention.q.weight",      # T5
    ".EncDecAttention.q.weight",    # T5 cross-attention
)

# Separate K projections
_K_PROJ_SUFFIXES = (
    ".k_proj.weight",
    ".key.weight",
    ".key_proj.weight",
    ".Wk.weight",
    ".attn_k.weight",
    ".self.key.weight",             # BERT
    ".attention.key.weight",        # some BERT variants
    ".SelfAttention.k.weight",      # T5
    ".EncDecAttention.k.weight",    # T5 cross-attention
)

# Combined QKV projections — shape encodes Q, K, and V together.
_QKV_COMBINED_SUFFIXES = (
    ".c_attn.weight",               # GPT-2: [hidden, 3*hidden]
    ".qkv_proj.weight",             # some ViTs
    "in_proj_weight",               # PyTorch MultiheadAttention (no leading dot)
    ".Wqkv.weight",                 # MPT: [3*hidden, hidden]
    ".query_key_value.weight",      # Falcon GQA: [hidden + 2*kv, hidden]
    ".qkv.weight",                  # generic ViT
    ".attention.qkv.weight",        # some HF transformers
    ".attn.qkv.weight",             # some ViTs
)

# ---------------------------------------------------------------------------
# Embedding weight suffixes (suffix-matched to tolerate arbitrary prefixes)
# ---------------------------------------------------------------------------
_EMBED_SUFFIXES = (
    "embed_tokens.weight",
    "wte.weight",
    "word_embeddings.weight",
    "embeddings.word_embeddings.weight",  # BERT / RoBERTa
    "shared.weight",                       # T5 tied embed
    "token_embedding.weight",
    "embed.weight",
    "tok_embeddings.weight",               # early LLaMA
    "decoder.embed_tokens.weight",         # some encoder-decoder models
    "embedding.weight",                    # generic
    "word_embedding.weight",               # some older models
)

# ---------------------------------------------------------------------------
# VLM vision tower / projector prefixes
# ---------------------------------------------------------------------------
_VLM_TOWER_PREFIXES = (
    "vision_tower.",
    "visual.",
    "multi_modal_projector.",
    "mm_projector.",
    "model.vision_tower.",
    "image_encoder.",           # Idefics / Flamingo
    "vision_model.",            # standalone vision encoders used in VLMs
    "img_projection.",          # Phi-3-Vision
    "perceiver.",               # Flamingo perceiver resampler
    "vpm.",                     # MiniCPM-V (InternViT-based)
    "resampler.",               # Qwen-VL resampler
)

# ---------------------------------------------------------------------------
# CNN / ViT specific
# ---------------------------------------------------------------------------

# Classifier head weight name suffixes → num_classes
_CLASSIFIER_SUFFIXES = (
    "classifier.weight",
    "fc.weight",
    "head.weight",
    "head.fc.weight",
    "classifier.1.weight",
    "output.weight",
    "linear.weight",
)

# Common first-conv weight names → in_channels
_FIRST_CONV_NAMES = (
    "conv1.weight",
    "patch_embed.proj.weight",          # ViT patch embedding
    "patch_embed.projection.weight",
    "features.0.0.weight",              # EfficientNet / MobileNet (torchvision)
    "stem.conv.weight",                 # ConvNeXt stem
    "stem.0.weight",
    "backbone.conv1.weight",
)


# ---------------------------------------------------------------------------
# Inference engine
# ---------------------------------------------------------------------------


class ShapeInference:
    """
    Derive model architecture parameters from SafeTensors weight shapes.

    All methods are static/classmethods — this class is a namespace, not an instance.
    """

    @classmethod
    def from_index(
        cls,
        safetensors_index: dict[str, str],
        param_meta: dict[str, "TensorRef"],
    ) -> ShapeInferredConfig:
        """
        Build a ShapeInferredConfig from the SafeTensors parameter index.

        Args:
            safetensors_index: {param_name → filename} mapping (all shards).
            param_meta: {param_name → TensorRef} with shape info from headers.

        Returns:
            ShapeInferredConfig with as many fields filled as possible.
        """
        cfg = ShapeInferredConfig()
        sources: dict[str, str] = {}

        # ── 1. Classify broad architecture type ───────────────────────────
        arch_type, num_enc, num_dec = cls._classify_arch(safetensors_index)
        cfg.arch_type = arch_type
        cfg.num_encoder_layers = num_enc
        cfg.num_decoder_layers = num_dec
        if num_enc or num_dec:
            sources["arch_type"] = "encoder/decoder block scan"

        # ── 2. Embedding → hidden_size, vocab_size ────────────────────────
        hidden_size, vocab_size = cls._infer_embedding(param_meta)
        if hidden_size:
            cfg.hidden_size = hidden_size
            cfg.vocab_size = vocab_size
            sources["hidden_size"] = "embedding weight"
            sources["vocab_size"] = "embedding weight"

        # ── 3. Num layers ─────────────────────────────────────────────────
        num_layers = cls._infer_num_layers(safetensors_index)
        if num_layers:
            cfg.num_layers = num_layers
            sources["num_layers"] = "layer index scan"

        # ── 4. Boolean flags ──────────────────────────────────────────────
        cfg.is_mla = cls._detect_mla(safetensors_index)
        cfg.is_moe, cfg.num_experts = cls._detect_moe(safetensors_index)
        cfg.is_vlm, cfg.vision_arch = cls._detect_vlm_arch(safetensors_index)
        if cfg.is_vlm:
            cfg.arch_type = "decoder"   # VLMs are decoder-based

        # ── 5. Head structure (skip for pure MLA — config has explicit dims)
        if not cfg.is_mla:
            head_dim, q_heads, kv_heads = cls._infer_head_structure(
                safetensors_index, param_meta, hidden_size, cfg.arch_type
            )
            if head_dim:
                cfg.head_dim = head_dim
                cfg.num_q_heads = q_heads
                cfg.num_kv_heads = kv_heads
                sources["head_dim"] = "proj weight shape"
                sources["num_q_heads"] = "q_proj/qkv shape"
                sources["num_kv_heads"] = "k_proj/qkv shape"

        # ── 6. CNN / ViT specific ─────────────────────────────────────────
        if cfg.arch_type in ("cnn", "vit"):
            num_classes, in_channels = cls._infer_cnn_params(
                safetensors_index, param_meta
            )
            if num_classes:
                cfg.num_classes = num_classes
                sources["num_classes"] = "classifier weight"
            if in_channels:
                cfg.in_channels = in_channels
                sources["in_channels"] = "first conv weight"

        cfg._sources = sources
        logger.info("Weight-inferred: %s", cfg.summary())
        return cfg

    # ------------------------------------------------------------------
    # Architecture classification
    # ------------------------------------------------------------------

    @staticmethod
    def _classify_arch(
        safetensors_index: dict[str, str],
    ) -> tuple[str, int, int]:
        """
        Return (arch_type, num_encoder_layers, num_decoder_layers).

        arch_type ∈ {decoder, encoder, encoder_decoder, vit, cnn, unknown}
        """
        names = list(safetensors_index.keys())

        # ── T5-style encoder-decoder (encoder.block.N / decoder.block.N) ──
        enc_max = dec_max = -1
        for name in names:
            m = _T5_ENCODER_RE.search(name)
            if m:
                enc_max = max(enc_max, int(m.group(1)))
            m = _T5_DECODER_RE.search(name)
            if m:
                dec_max = max(dec_max, int(m.group(1)))
        if enc_max >= 0 and dec_max >= 0:
            return "encoder_decoder", enc_max + 1, dec_max + 1

        # ── BART-style encoder-decoder (.encoder.layers.N and .decoder.layers.N) ──
        has_encoder_layers = any(
            ".encoder.layers." in n or ".encoder.layer." in n for n in names
        )
        has_decoder_layers = any(
            ".decoder.layers." in n or ".decoder.layer." in n for n in names
        )
        if has_encoder_layers and has_decoder_layers:
            enc_max = dec_max = -1
            for name in names:
                if ".encoder.layers." in name or ".encoder.layer." in name:
                    m = _LAYER_IDX_RE.search(name)
                    if m:
                        enc_max = max(enc_max, int(m.group(1)))
                if ".decoder.layers." in name or ".decoder.layer." in name:
                    m = _LAYER_IDX_RE.search(name)
                    if m:
                        dec_max = max(dec_max, int(m.group(1)))
            return (
                "encoder_decoder",
                enc_max + 1 if enc_max >= 0 else 0,
                dec_max + 1 if dec_max >= 0 else 0,
            )

        # ── ViT markers (standalone vision transformer, no text embedding) ──
        has_cls_token = any("cls_token" in n for n in names)
        has_patch_embed = any("patch_embed" in n for n in names)
        has_text_embed = any(any(n.endswith(s) for s in _EMBED_SUFFIXES) for n in names)
        if (has_cls_token or has_patch_embed) and not has_text_embed:
            return "vit", 0, 0

        # ── CNN (conv layers, no text embedding, no attention Q/K) ─────────
        all_suffixes = _Q_PROJ_SUFFIXES + _K_PROJ_SUFFIXES + _QKV_COMBINED_SUFFIXES
        has_attn = any(any(sfx in n for sfx in all_suffixes) for n in names)
        has_conv = any("conv" in n.lower() and n.endswith(".weight") for n in names)
        if has_conv and not has_text_embed and not has_attn:
            return "cnn", 0, 0

        # ── BERT-style encoder-only ────────────────────────────────────────
        has_bert_layers = any(_BERT_ENCODER_RE.search(n) for n in names)
        has_causal = any("causal" in n.lower() for n in names)
        if has_bert_layers and not has_causal and has_text_embed:
            return "encoder", 0, 0

        # ── Default: decoder-only ──────────────────────────────────────────
        if has_text_embed or has_attn:
            return "decoder", 0, 0

        return "unknown", 0, 0

    # ------------------------------------------------------------------
    # Individual inference helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _infer_embedding(
        param_meta: dict[str, "TensorRef"],
    ) -> tuple[int, int]:
        """Return (hidden_size, vocab_size) from the embedding weight shape.

        For vision-only models (ViTs, CNNs) that have no token embedding, falls back
        to the patch embedding output channels as hidden_size (vocab_size stays 0).
        """
        for name, ref in param_meta.items():
            if not ref.shape or len(ref.shape) != 2:
                continue
            for suffix in _EMBED_SUFFIXES:
                if name.endswith(suffix):
                    vocab_size, hidden_size = ref.shape[0], ref.shape[1]
                    # Sanity: vocab should be >> hidden for any real model
                    if vocab_size > hidden_size and hidden_size > 0:
                        return hidden_size, vocab_size

        # Fallback for ViT: patch embedding projection gives hidden_size
        # Shape: [out_channels, in_channels, patch_H, patch_W] → out_channels = hidden_size
        for name, ref in param_meta.items():
            if ref.shape and len(ref.shape) == 4:
                if any(name.endswith(s) for s in (
                    "patch_embed.proj.weight",
                    "patch_embed.projection.weight",
                    "patch_embedding.projection.weight",
                )):
                    return ref.shape[0], 0  # hidden_size, vocab_size=0

        return 0, 0

    @staticmethod
    def _infer_num_layers(safetensors_index: dict[str, str]) -> int:
        """Return num_layers from the maximum layer index across all param names."""
        max_idx = -1
        for name in safetensors_index:
            m = _LAYER_IDX_RE.search(name)
            if m:
                idx = int(m.group(1))
                if idx > max_idx:
                    max_idx = idx
        return max_idx + 1 if max_idx >= 0 else 0

    @classmethod
    def _infer_head_structure(
        cls,
        safetensors_index: dict[str, str],
        param_meta: dict[str, "TensorRef"],
        hidden_size: int,
        arch_type: str = "unknown",
    ) -> tuple[int, int, int]:
        """
        Return (head_dim, num_q_heads, num_kv_heads).

        Tries separate Q/K projections first, then falls back to combined QKV.

        Square-projection disambiguation:
          When k_proj rows == hidden_size (BERT, T5, GPT-2, early models), num_heads
          cannot be determined from shape alone.  We break the tie with two signals:

            1. Attention biases present (query.bias / c_attn.bias / …):
               → older-style model (BERT, GPT-2) → prefer head_dim=64
            2. arch_type is "encoder" or "encoder_decoder":
               → T5 (no biases but clearly seq2seq) → prefer head_dim=64

          Modern decoder LLMs (Llama, Mistral, Qwen…) never have attention biases
          and are "decoder" arch → head_dim=128 stays first.
        """
        # Detect attention projection biases — reliable signal for pre-LLaMA models.
        # Exclude VLM vision tower params: CLIP/SigLIP attention layers have biases
        # but the LLM backbone (Llama/Vicuna) does not; mixing them gives a false positive.
        _ATTN_BIAS_SUFFIXES = (
            ".query.bias", ".q_proj.bias", ".c_attn.bias",
            ".key.bias", ".k_proj.bias",
            ".SelfAttention.q.bias", ".attention.query.bias",
        )
        has_attn_bias = any(
            any(n.endswith(sfx) for sfx in _ATTN_BIAS_SUFFIXES)
            and not any(n.startswith(p) for p in _VLM_TOWER_PREFIXES)
            for n in safetensors_index
        )
        prefer_small = has_attn_bias or arch_type in ("encoder", "encoder_decoder", "vit")

        result = cls._infer_head_from_qk_proj(
            safetensors_index, param_meta, hidden_size, prefer_small
        )
        if result[0]:
            return result
        return cls._infer_head_from_combined_qkv(
            safetensors_index, param_meta, hidden_size, prefer_small
        )

    @staticmethod
    def _infer_head_from_qk_proj(
        safetensors_index: dict[str, str],
        param_meta: dict[str, "TensorRef"],
        hidden_size: int,
        prefer_small_head_dim: bool = False,
    ) -> tuple[int, int, int]:
        """
        Return (head_dim, num_q_heads, num_kv_heads) from separate Q/K projections.

        Bootstrap strategy:
          - Find k_proj.weight for the smallest layer index within the LLM domain
          - Try each common head_dim: accept the first that divides k_proj rows evenly
          - Derive num_kv_heads = k_proj_rows / head_dim
          - Derive num_q_heads  = q_proj_rows / head_dim

        When k_rows == hidden_size (square projection), the model uses standard MHA
        and num_heads is otherwise ambiguous.  prefer_small_head_dim=True shifts the
        search order to try head_dim=64 first (BERT / T5 / GPT-2 style).
        """
        smallest_layer = None
        k_ref: "TensorRef | None" = None
        k_name_prefix = ""  # text before layer path — keeps us within the same domain

        for name in safetensors_index:
            for suffix in _K_PROJ_SUFFIXES:
                if suffix in name:
                    m = _LAYER_IDX_RE.search(name)
                    if m:
                        idx = int(m.group(1))
                        if smallest_layer is None or idx < smallest_layer:
                            if name in param_meta and param_meta[name].shape:
                                smallest_layer = idx
                                k_ref = param_meta[name]
                                prefix_end = m.start()
                                k_name_prefix = name[: prefix_end + 1]

        if k_ref is None or len(k_ref.shape) < 2:
            return 0, 0, 0

        k_rows = k_ref.shape[0]

        # Bootstrap head_dim.
        # For square projections (k_rows == hidden_size), shift search order when
        # the caller signals an older-style model (BERT / T5 / GPT-2).
        is_square = hidden_size > 0 and k_rows == hidden_size
        if is_square and prefer_small_head_dim:
            candidate_dims = (64, 96, 48, 32, 128, 256, 80, 40)
        else:
            candidate_dims = _COMMON_HEAD_DIMS

        head_dim = 0
        for candidate in candidate_dims:
            if k_rows % candidate == 0:
                if hidden_size and hidden_size % candidate != 0:
                    continue
                head_dim = candidate
                break

        if head_dim == 0:
            return 0, 0, 0

        num_kv_heads = k_rows // head_dim

        # Find q_proj within same domain prefix to avoid VLM/LLM confusion
        num_q_heads = num_kv_heads  # default: GQA ratio = 1
        if smallest_layer is not None:
            found_q = False
            for name in safetensors_index:
                if not name.startswith(k_name_prefix):
                    continue
                for suffix in _Q_PROJ_SUFFIXES:
                    if suffix in name:
                        m = _LAYER_IDX_RE.search(name)
                        if m and int(m.group(1)) == smallest_layer:
                            ref = param_meta.get(name)
                            if ref and ref.shape and len(ref.shape) >= 2:
                                q_rows = ref.shape[0]
                                if q_rows % head_dim == 0:
                                    num_q_heads = q_rows // head_dim
                                    found_q = True
                                    break
                if found_q:
                    break

        return head_dim, num_q_heads, num_kv_heads

    @staticmethod
    def _infer_head_from_combined_qkv(
        safetensors_index: dict[str, str],
        param_meta: dict[str, "TensorRef"],
        hidden_size: int,
        prefer_small_head_dim: bool = False,
    ) -> tuple[int, int, int]:
        """
        Return (head_dim, num_q_heads, num_kv_heads) from combined QKV projections.

        Handles:
          - Symmetric QKV (c_attn, Wqkv, qkv_proj): combined_dim = 3 * q_dim
          - Falcon GQA (query_key_value):            rows = hidden + 2 * kv_size

        prefer_small_head_dim applies the same head_dim=64-first search order used by
        _infer_head_from_qk_proj when q_dim == hidden_size.
        """
        for name in safetensors_index:
            matched = any(
                name.endswith(sfx) or name.endswith(sfx.lstrip("."))
                for sfx in _QKV_COMBINED_SUFFIXES
            )
            if not matched:
                continue

            m = _LAYER_IDX_RE.search(name)
            if not m:
                continue

            ref = param_meta.get(name)
            if not ref or not ref.shape or len(ref.shape) < 2:
                continue

            rows, cols = ref.shape[0], ref.shape[1]
            combined_dim = max(rows, cols)

            # Case 1 — symmetric QKV: combined_dim = 3 * q_dim
            if combined_dim % 3 == 0:
                q_dim = combined_dim // 3
                is_square = hidden_size > 0 and q_dim == hidden_size
                dims = (
                    (64, 96, 48, 32, 128, 256, 80, 40)
                    if (is_square and prefer_small_head_dim)
                    else _COMMON_HEAD_DIMS
                )
                for candidate in dims:
                    if q_dim % candidate == 0:
                        if hidden_size and hidden_size % candidate != 0:
                            continue
                        num_heads = q_dim // candidate
                        return candidate, num_heads, num_heads

            # Case 2 — Falcon GQA: rows = hidden + 2 * kv_size
            if hidden_size and rows > hidden_size and "query_key_value" in name:
                kv_combined = rows - hidden_size
                if kv_combined % 2 == 0:
                    kv_size = kv_combined // 2
                    for candidate in _COMMON_HEAD_DIMS:
                        if kv_size % candidate == 0 and hidden_size % candidate == 0:
                            num_kv = kv_size // candidate
                            num_q = hidden_size // candidate
                            return candidate, num_q, num_kv

        return 0, 0, 0

    @staticmethod
    def _infer_cnn_params(
        safetensors_index: dict[str, str],
        param_meta: dict[str, "TensorRef"],
    ) -> tuple[int, int]:
        """
        Return (num_classes, in_channels) for CNN and ViT models.

        num_classes  — from classifier / head weight (final linear layer)
        in_channels  — from the first conv / patch-embedding weight (4-D tensor)
        """
        num_classes = 0
        in_channels = 0

        # ── Classifier head → num_classes ─────────────────────────────────
        for name in safetensors_index:
            for suffix in _CLASSIFIER_SUFFIXES:
                if name.endswith(suffix):
                    ref = param_meta.get(name)
                    if ref and ref.shape and len(ref.shape) == 2:
                        r, c = ref.shape
                        # PyTorch Linear stores weight as [out_features, in_features].
                        # out_features = num_classes; in_features = feature_dim (larger).
                        # Sanity: both dims must be plausible class / feature counts.
                        if 2 <= r <= 200_000 and 2 <= c <= 200_000:
                            num_classes = r  # shape[0] is always out_features
                        break
            if num_classes:
                break

        # ── First conv / patch embed → in_channels ────────────────────────
        for first_conv in _FIRST_CONV_NAMES:
            for name in safetensors_index:
                if name.endswith(first_conv):
                    ref = param_meta.get(name)
                    if ref and ref.shape and len(ref.shape) == 4:
                        in_channels = ref.shape[1]  # [out_ch, in_ch, kH, kW]
                        break
            if in_channels:
                break

        # Fall back: scan any 4-D conv weight (lexicographic order → early layers first)
        if not in_channels:
            for name in sorted(safetensors_index.keys()):
                if "conv" in name.lower() and name.endswith(".weight"):
                    ref = param_meta.get(name)
                    if ref and ref.shape and len(ref.shape) == 4:
                        in_channels = ref.shape[1]
                        break

        return num_classes, in_channels

    # ------------------------------------------------------------------
    # Boolean flag detectors
    # ------------------------------------------------------------------

    @staticmethod
    def _detect_moe(safetensors_index: dict[str, str]) -> tuple[bool, int]:
        """Return (is_moe, num_experts) by scanning for expert param names."""
        expert_re = re.compile(r"\.experts\.(\d+)\.")
        # Find the first (smallest) layer index that actually has experts
        first_moe_layer: int | None = None
        is_moe = False
        for name in safetensors_index:
            if ".experts." in name:
                is_moe = True
                m = _LAYER_IDX_RE.search(name)
                if m:
                    idx = int(m.group(1))
                    if first_moe_layer is None or idx < first_moe_layer:
                        first_moe_layer = idx

        if not is_moe or first_moe_layer is None:
            return False, 0

        # Count unique expert indices in the first MoE layer
        expert_indices: set[int] = set()
        first_str = str(first_moe_layer)
        for name in safetensors_index:
            if ".experts." not in name:
                continue
            m = _LAYER_IDX_RE.search(name)
            if m and m.group(1) == first_str:
                em = expert_re.search(name)
                if em:
                    expert_indices.add(int(em.group(1)))
        return True, len(expert_indices)

    @staticmethod
    def _detect_mla(safetensors_index: dict[str, str]) -> bool:
        """Return True if the model uses DeepSeek MLA (Multi-head Latent Attention)."""
        return any("kv_a_proj_with_mqa" in name for name in safetensors_index)

    @staticmethod
    def _detect_vlm_arch(safetensors_index: dict[str, str]) -> tuple[bool, str]:
        """Return (is_vlm, vision_arch) by scanning for vision tower params."""
        has_vlm_params = any(
            any(name.startswith(p) for p in _VLM_TOWER_PREFIXES)
            for name in safetensors_index
        )
        if not has_vlm_params:
            return False, ""

        # Identify vision backbone from specific param name patterns
        for name in safetensors_index:
            if name.startswith("vision_tower."):
                # SigLIP models have "siglip" anywhere in their param paths
                if any("siglip" in n for n in safetensors_index):
                    return True, "siglip"
                return True, "clip"

            if name.startswith("visual.patch_embed"):
                return True, "intern"       # InternViT (InternVL)

            if name.startswith("vpm."):
                return True, "intern"       # MiniCPM-V uses InternViT

            if name.startswith("visual."):
                return True, "siglip"       # Qwen2-VL uses SigLIP

            if name.startswith("image_encoder."):
                return True, "clip"         # Idefics / Flamingo

            if name.startswith("vision_model."):
                return True, "clip"         # generic CLIP-style

            if name.startswith("img_projection."):
                return True, "clip"         # Phi-3-Vision

            if name.startswith("perceiver."):
                return True, "clip"         # Flamingo perceiver resampler

        # Projector-only present but vision arch not determinable
        return True, ""
