# Copyright and license: follows the repository's existing convention.
"""Download-free model structure & parameter analysis for quantization planning.

Given a HuggingFace model id, this inspects the model's ``config.json`` and the
**safetensors headers only** (no weight download) to produce:

* a per-category parameter distribution (MoE experts, shared experts, attention,
  router/gate, dense MLP, lm_head, embeddings, vision, norms, …), matching the
  format of the mixed-precision reference docs, and
* recommended ``ignore_layers`` / mixed-precision ``layer_config`` presets a user
  can drop straight into the advanced submission fields.

The structure comes from the **safetensors index** (``model.safetensors.index.json``) —
a single small JSON that lists every tensor name and its shard. No weights, shapes, or
dtypes are downloaded, so even 100+ shard / trillion-param models are analyzed in seconds.
Module types are inferred from tensor names (Linear / Embedding / Norm / Bias).
"""
from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass, field

from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
from transformers import AutoConfig

logger = logging.getLogger(__name__)

# Collapse per-layer AND per-expert indices so
# ``...layers.12.mlp.experts.37.gate_proj.weight`` and
# ``...layers.0.mlp.experts.5.gate_proj.weight`` bucket to one module.
_LAYER_IDX_RE = re.compile(r"\.\d+\.")
# Also collapse a trailing ``.<digits>`` (e.g. ``oe_embed_proj0`` style is rare;
# handles names like ``...experts.37`` with no suffix).
_TRAIL_IDX_RE = re.compile(r"\.\d+$")


@dataclass
class ModuleStat:
    """One normalized module (layer/expert indices collapsed to N)."""
    name: str            # e.g. "...layers.N.block_sparse_moe.experts.N.gate_proj"
    category: str = ""
    kind: str = ""       # "Linear" | "Embedding" | "Norm (1D)" | "Bias (1D)"
    count: int = 0       # how many real tensors collapsed into this row


@dataclass
class ModelAnalysis:
    model_id: str
    ok: bool = True
    error: str | None = None
    architectures: list[str] = field(default_factory=list)
    model_type: str = ""
    hidden_size: int | None = None
    num_layers: int | None = None
    num_experts: int | None = None
    vocab_size: int | None = None
    is_moe: bool = False
    has_shared_experts: bool = False
    has_attn_indexer: bool = False
    has_vision: bool = False
    num_tensors: int = 0
    modules: list[ModuleStat] = field(default_factory=list)
    recommended_ignore_layers: str = ""
    recommended_layer_config: str = ""
    modules: list[ModuleStat] = field(default_factory=list)
    recommended_ignore_layers: str = ""
    recommended_layer_config: str = ""


# ── Categorization ────────────────────────────────────────────────────────
# Order matters: the first matching rule wins. Each rule is (category, predicate).
def _categorize(norm_name: str) -> str:
    n = norm_name
    # Routers / gates first (must beat gate_proj which is a normal MLP proj).
    if (".gate." in n or n.endswith(".gate") or n.endswith(".gate.weight")
            or ".router" in n or n.endswith(".router")) and "gate_proj" not in n and "gate_up_proj" not in n:
        return "router_gate"
    if "shared_expert" in n:
        return "shared_experts"
    if ".experts." in n or n.endswith(".experts") or ".block_sparse_moe.experts" in n:
        return "moe_experts"
    if "vision" in n or "visual" in n or "vit" in n:
        return "vision"
    if "lm_head" in n:
        return "lm_head"
    if "embed" in n or "wte" in n or "word_embeddings" in n:
        return "embeddings"
    if "self_attn" in n or ".attn." in n or "attention" in n:
        if "index" in n or "indexer" in n:
            return "attn_indexer"
        if "norm" in n:
            return "norm"
        return "self_attn"
    if "mtp" in n:
        return "mtp"
    if "projector" in n or "patch_merge" in n or "multi_modal" in n:
        return "projector"
    if "mlp" in n or "feed_forward" in n or "ffn" in n:
        return "dense_mlp"
    if "norm" in n or "layernorm" in n or "ln_" in n:
        return "norm"
    return "other"


_CATEGORY_ORDER = [
    "moe_experts", "shared_experts", "dense_mlp", "self_attn", "attn_indexer",
    "router_gate", "mtp", "vision", "projector", "lm_head", "embeddings",
    "norm", "other",
]

_CATEGORY_LABELS = {
    "moe_experts": "MoE routed experts",
    "shared_experts": "MoE shared experts",
    "dense_mlp": "Dense MLP",
    "self_attn": "Attention (q/k/v/o)",
    "attn_indexer": "Attention indexer (sparse)",
    "router_gate": "Router / gate",
    "mtp": "MTP module",
    "vision": "Vision tower",
    "projector": "Projector / merger",
    "lm_head": "lm_head",
    "embeddings": "Embeddings",
    "norm": "Norms (1D, auto-skipped)",
    "other": "Other",
}


def _cfg_get(cfg, *attrs, default=None):
    """Read the first present attribute from a config or its nested text_config."""
    sources = [cfg]
    if hasattr(cfg, "text_config") and cfg.text_config is not None:
        sources.append(cfg.text_config)
    for src in sources:
        for a in attrs:
            v = getattr(src, a, None)
            if v is not None:
                return v
    return default


def analyze_model_structure(model_id: str, revision: str = "main", token: str | None = None) -> ModelAnalysis:
    """Analyze *model_id* without downloading weights. Never raises."""
    res = ModelAnalysis(model_id=model_id)

    # 1. Config (best-effort; some fields inform MoE detection & the report).
    try:
        cfg = AutoConfig.from_pretrained(model_id, revision=revision, token=token, trust_remote_code=True)
        arch = getattr(cfg, "architectures", None) or []
        res.architectures = list(arch)
        res.model_type = getattr(cfg, "model_type", "") or ""
        res.hidden_size = _cfg_get(cfg, "hidden_size", "n_embd", "d_model")
        res.num_layers = _cfg_get(cfg, "num_hidden_layers", "n_layer", "num_layers")
        res.num_experts = _cfg_get(cfg, "num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts")
        res.vocab_size = _cfg_get(cfg, "vocab_size")
    except Exception as e:
        logger.warning("[analyze] config load failed for %s: %s", model_id, e)

    # 2. Tensor NAMES only — from the safetensors index (one small JSON), so even
    #    194-shard / trillion-param models are analyzed in seconds. No weights, no
    #    shapes/dtypes downloaded. Falls back to a single-file header for unsharded models.
    names, err = _list_tensor_names(model_id, revision, token)
    if err:
        res.ok = False
        res.error = err
        return res

    res.num_tensors = len(names)
    mods: dict[str, ModuleStat] = {}
    cat_present: set[str] = set()
    for tname in names:
        # Collapse layer + expert indices: any ".<digits>." and a trailing ".<digits>".
        norm = _TRAIL_IDX_RE.sub(".N", _LAYER_IDX_RE.sub(".N.", tname))
        cat = _categorize(norm)
        cat_present.add(cat)
        mkey = _module_key(norm)
        ms = mods.get(mkey)
        if ms is None:
            ms = ModuleStat(name=mkey, category=cat, kind=_kind_from_name(mkey))
            mods[mkey] = ms
        ms.count += 1

    # Order modules by category (structural grouping), then name — reads like the
    # reference "model structure" table rather than a params ranking.
    cat_rank = {c: i for i, c in enumerate(_CATEGORY_ORDER)}
    res.modules = sorted(mods.values(), key=lambda m: (cat_rank.get(m.category, 99), m.name))

    # 3. Derived traits (from names + config).
    res.is_moe = ("moe_experts" in cat_present) or bool(res.num_experts)
    res.has_shared_experts = "shared_experts" in cat_present
    res.has_attn_indexer = "attn_indexer" in cat_present
    res.has_vision = "vision" in cat_present

    # 4. Suggested controls (optional; derived from REAL module names → precise substrings).
    res.recommended_ignore_layers, res.recommended_layer_config = _recommend(res, cat_present, mods)
    return res


def _list_tensor_names(model_id: str, revision: str, token: str | None):
    """Return ``(names, error)`` — the model's tensor names without downloading weights.

    Strategy (fast → fallback):
      1. ``model.safetensors.index.json`` (sharded) → ``weight_map`` keys. One file.
      2. a single ``*.safetensors`` file → read just its header via the HfApi.
    Never downloads weight bytes.
    """
    api = HfApi(token=token)
    # 1. Sharded index (covers the vast majority of large models).
    for index_name in ("model.safetensors.index.json", "pytorch_model.bin.index.json"):
        try:
            path = hf_hub_download(model_id, index_name, revision=revision, token=token)
            with open(path) as f:
                data = json.load(f)
            wm = data.get("weight_map") or {}
            if wm:
                return list(wm.keys()), None
        except EntryNotFoundError:
            continue
        except Exception as e:
            logger.warning("[analyze] index read failed for %s (%s): %s", model_id, index_name, e)

    # 2. Single-file safetensors → header only.
    try:
        files = api.list_repo_files(model_id, revision=revision)
        st_files = [f for f in files if f.endswith(".safetensors")]
        if st_files:
            meta = api.parse_safetensors_file_metadata(model_id, st_files[0], revision=revision)
            return list((meta.tensors or {}).keys()), None
        # No safetensors at all.
        return [], (
            "No safetensors found (model may be GGUF / pytorch_model.bin only). "
            "Structure analysis needs a safetensors checkpoint."
        )
    except Exception as e:
        return [], (
            f"Could not read model index/metadata: {e}. "
            "The model may be gated/private (log in) or lack safetensors."
        )


def _module_key(norm_name: str) -> str:
    """Normalized tensor name → module key (drop trailing .weight/.bias/.scale)."""
    for suf in (".weight", ".bias", ".weight_scale", ".weight_packed", ".scale", ".g_idx", ".qweight", ".qzeros", ".scales"):
        if norm_name.endswith(suf):
            return norm_name[: -len(suf)]
    return norm_name


def _kind_from_name(tname: str) -> str:
    """Infer the module type from its name only (no shape needed).

    Mirrors the human 'type' column in the reference structure tables.
    """
    n = tname.lower()
    if "lm_head" in n:
        return "Linear (head)"
    if "embed" in n or "wte" in n or "word_embeddings" in n:
        return "Embedding"
    if n.endswith("_bias") or n.endswith(".bias") or "correction_bias" in n:
        return "Bias (1D)"
    if "norm" in n or "layernorm" in n or "ln_f" in n or "ln_1" in n or "ln_2" in n:
        return "Norm (1D)"
    # Everything else that carries weights is a Linear/projection in these models
    # (q/k/v/o_proj, gate, experts w1/w2/w3, fc1/fc2, router classifier, …).
    return "Linear"


def _rel_name(module_key: str) -> str:
    """Strip the model wrapper + the ``layers.N.`` prefix to get a substring a user
    can paste into ignore_layers. e.g.
    ``language_model.model.layers.N.block_sparse_moe.gate`` → ``block_sparse_moe.gate``.
    """
    key = module_key
    for pre in ("language_model.model.", "language_model.", "model.model.", "model.", "transformer."):
        if key.startswith(pre):
            key = key[len(pre):]
            break
    # Drop everything up to and including the first ``layers.N.``
    key = re.sub(r"^.*?layers\.N\.", "", key)
    return key


def _char_lcp(strings: list[str]) -> str:
    """Character-level longest common prefix (used to fold sibling leaves into one
    precise substring, e.g. self_attn.index_q_proj + self_attn.index_k_proj →
    ``self_attn.index_``)."""
    if not strings:
        return ""
    s1, s2 = min(strings), max(strings)
    i = 0
    while i < len(s1) and i < len(s2) and s1[i] == s2[i]:
        i += 1
    return s1[:i]


def _ignore_token_for(category_keys: list[str]) -> str | None:
    """Turn a category's real module keys into ONE precise ignore substring.

    Uses the relative names (after ``layers.N.``); if there are several siblings,
    a character-level common prefix yields a safe substring (never the bare last
    segment like ``gate`` which would also hit ``gate_proj``).
    """
    rels = sorted({_rel_name(k) for k in category_keys})
    rels = [r for r in rels if r]
    if not rels:
        return None
    if len(rels) == 1:
        return rels[0]
    lcp = _char_lcp(rels)
    # Require a reasonably specific prefix; otherwise just return the shortest name.
    if len(lcp) >= 4:
        return lcp
    return min(rels, key=len)


def _experts_layer_config_key(expert_keys: list[str]) -> str:
    """Derive the precise ``layer_config`` key for routed experts from real names.

    e.g. ``block_sparse_moe.experts.N.gate_proj`` → ``block_sparse_moe.experts``;
    ``mlp.experts.N.gate_proj`` → ``mlp.experts``. Using ``<parent>.experts``
    (not bare ``experts``) keeps it precise and never touches ``shared_experts``.
    """
    for k in expert_keys:
        rel = _rel_name(k)
        segs = rel.split(".")
        for i, s in enumerate(segs):
            if s == "experts":
                return ".".join(segs[max(0, i - 1):i + 1]) if i > 0 else "experts"
    return "experts"


def _keys_in(mods: dict, category: str) -> list[str]:
    return [k for k, m in mods.items() if m.category == category]


def _recommend(res: ModelAnalysis, cat_present: set, mods: dict) -> tuple[str, str]:
    """Produce suggested ignore_layers + mixed-precision layer_config using the
    model's REAL module names, so the substrings are precise and safe.

    Critical: never emit a bare last segment like ``gate`` — under auto-round's
    substring matching that would also hit ``gate_proj`` (a dense-MLP projection).
    We derive ``<parent>.gate`` / ``self_attn.index_`` etc. from actual names.

    Heuristic (from the reference mixed-precision docs):
      * ignore: lm_head, router/gate, vision tower, projectors, attention indexer.
      * embeddings + norms are auto-skipped by AutoRound (not listed).
      * MoE routed experts → MXFP4 via layer_config; the rest stay at the global scheme.
    """
    ignore: list[str] = []
    for cat in ("lm_head", "router_gate", "vision", "projector", "attn_indexer"):
        if cat in cat_present:
            tok = _ignore_token_for(_keys_in(mods, cat))
            if tok:
                ignore.append(tok)
    # Deduplicate, preserve order.
    seen: set[str] = set()
    ignore = [x for x in ignore if not (x in seen or seen.add(x))]

    layer_config = ""
    if res.is_moe and "moe_experts" in cat_present:
        key = _experts_layer_config_key(_keys_in(mods, "moe_experts"))
        layer_config = f"{{{key}:{{bits:4,data_type:mx_fp}}}}"

    return ",".join(ignore), layer_config


# ── Markdown rendering ─────────────────────────────────────────────────────
def render_analysis_markdown(res: ModelAnalysis) -> str:
    if not res.ok:
        return f"### ⚠️ Model structure analysis failed\n\n{res.error}"

    lines: list[str] = []
    lines.append(f"### 🔍 Model structure — `{res.model_id}`\n")

    # Overview
    lines.append("| Field | Value |")
    lines.append("|---|---|")
    if res.architectures:
        lines.append(f"| Architecture | `{', '.join(res.architectures)}` |")
    if res.model_type:
        lines.append(f"| model_type | `{res.model_type}` |")
    if res.num_layers is not None:
        lines.append(f"| Layers | {res.num_layers} |")
    if res.num_experts:
        lines.append(f"| Experts | {res.num_experts} |")
    if res.hidden_size is not None:
        lines.append(f"| hidden_size | {res.hidden_size} |")
    if res.vocab_size is not None:
        lines.append(f"| vocab_size | {res.vocab_size} |")
    lines.append(f"| Tensors (total) | {res.num_tensors} |")
    traits = []
    if res.is_moe:
        traits.append("MoE")
    if res.has_shared_experts:
        traits.append("shared-experts")
    if res.has_attn_indexer:
        traits.append("sparse-attn-indexer")
    if res.has_vision:
        traits.append("vision")
    if traits:
        lines.append(f"| Traits | {', '.join(traits)} |")
    lines.append("")

    # Module schema — the model's composition, grouped by category. Layer & expert
    # indices are shown as `N`. These are the exact substrings a user pastes into
    # Ignore Layers / Layer Config.
    lines.append("#### Model structure — modules (normalized: layer & expert indices = `N`)\n")
    lines.append("Use these names to craft **Ignore Layers** / **Layer Config** substrings.\n")
    lines.append("| Module | Type | Category | Count |")
    lines.append("|---|---|---|---:|")
    _MAX_ROWS = 80
    shown = res.modules[:_MAX_ROWS]
    for ms in shown:
        label = _CATEGORY_LABELS.get(ms.category, ms.category)
        lines.append(f"| `{ms.name}` | {ms.kind} | {label} | {ms.count} |")
    if len(res.modules) > _MAX_ROWS:
        lines.append(f"| … | | | *(+{len(res.modules) - _MAX_ROWS} more)* |")
    lines.append("")

    # Suggested controls (optional; users can define their own from the table above)
    lines.append("#### Suggested quantization controls (optional)\n")
    if res.recommended_ignore_layers:
        lines.append(f"- **Ignore Layers:** `{res.recommended_ignore_layers}`")
    else:
        lines.append("- **Ignore Layers:** *(nothing extra suggested — defaults are fine)*")
    if res.recommended_layer_config:
        lines.append(f"- **Layer Config (mixed precision):** `{res.recommended_layer_config}`")
        lines.append("  - routes MoE routed experts to MXFP4 while the rest stay at the global scheme.")
    lines.append("")
    lines.append("> Suggestions only — review against the module table above before submitting. "
                 "Norms (1D) and embeddings are skipped automatically by AutoRound.")
    return "\n".join(lines)