"""Loading, inspecting and slimming the Qwen3.5-4B checkpoint.

Qwen3.5-4B is a hybrid multimodal model — 24 Gated DeltaNet (linear attention)
layers interleaved with 8 full-attention layers, plus a 24-block vision tower
and a multi-token-prediction head. For a text-only math benchmark the vision
tower and MTP head are never executed, so they are ~455M parameters (9.7% of the
checkpoint) of pure dead weight.

Exact class names and module paths for this architecture vary across
transformers releases, so everything here introspects the loaded module tree
rather than hardcoding paths. ``describe_model`` exists to dump ground truth
before we commit to any of it.
"""

from __future__ import annotations

import os
import re
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import torch

DEFAULT_MODEL = "Qwen/Qwen3.5-4B"

# First matching pattern wins, so order matters: `visual`/`mtp` must precede the
# generic `mlp`/`norm` patterns since those submodules contain MLPs too.
BUDGET_GROUPS: tuple[tuple[str, str], ...] = (
    ("vision", r"visual|vision"),
    ("mtp", r"(^|\.)mtp\."),
    ("embed", r"embed_tokens|lm_head"),
    ("linear_attn", r"linear_attn"),
    ("full_attn", r"self_attn"),
    ("mlp", r"\.mlp\."),
    ("norm", r"norm"),
)

# Tiny but precision-critical: the SSM recurrence dynamics. Keeping every one of
# these in fp32 costs ~3 MB, and quantizing them is what makes low-bit SSM
# quantization collapse.
SSM_SENSITIVE = r"A_log|dt_bias|conv1d|in_proj_a|in_proj_b"


@dataclass
class GroupStats:
    name: str
    num_tensors: int = 0
    num_params: int = 0  # unique storage only
    num_bytes: int = 0  # unique storage only
    raw_params: int = 0  # counting tied aliases twice
    num_aliases: int = 0
    examples: list[str] = field(default_factory=list)

    def add(self, name: str, tensor: torch.Tensor, *, counted: bool) -> None:
        self.num_tensors += 1
        self.raw_params += tensor.numel()
        if counted:
            self.num_params += tensor.numel()
            self.num_bytes += tensor.numel() * tensor.element_size()
        else:
            self.num_aliases += 1
        if len(self.examples) < 3:
            self.examples.append(name + ("  (tied alias)" if not counted else ""))


def _group_for(name: str) -> str:
    for group, pattern in BUDGET_GROUPS:
        if re.search(pattern, name):
            return group
    return "other"


def parameter_budget(state_dict: dict[str, torch.Tensor]) -> dict[str, Any]:
    """Group every tensor by component and report params / bytes per group.

    Tied weights (Qwen3.5 ties ``lm_head`` to ``embed_tokens``) appear twice in a
    ``state_dict`` while sharing one allocation. Counting both would overstate
    the checkpoint by 636M params / 1.27 GB here — and checkpoint bytes is the
    metric we are graded on — so aliases are detected by storage pointer and
    counted once.
    """
    groups: dict[str, GroupStats] = {}
    seen_storage: dict[int, str] = {}
    aliases: list[dict[str, str]] = []

    for name, tensor in state_dict.items():
        if not isinstance(tensor, torch.Tensor):
            continue
        pointer = tensor.data_ptr()
        is_alias = pointer != 0 and pointer in seen_storage
        if is_alias:
            aliases.append({"name": name, "aliases": seen_storage[pointer]})
        else:
            seen_storage[pointer] = name
        groups.setdefault(_group_for(name), GroupStats(_group_for(name))).add(
            name, tensor, counted=not is_alias
        )

    total_params = sum(g.num_params for g in groups.values())
    total_bytes = sum(g.num_bytes for g in groups.values())
    return {
        "total_params": total_params,
        "total_bytes": total_bytes,
        "total_gb": total_bytes / 1e9,
        "raw_params_with_aliases": sum(g.raw_params for g in groups.values()),
        "tied_aliases": aliases,
        "groups": {
            name: {
                "num_tensors": g.num_tensors,
                "num_params": g.num_params,
                "num_bytes": g.num_bytes,
                "gb": g.num_bytes / 1e9,
                "num_aliases": g.num_aliases,
                "pct_params": 100 * g.num_params / total_params if total_params else 0.0,
                "examples": g.examples,
            }
            for name, g in sorted(groups.items(), key=lambda kv: -kv[1].num_params)
        },
    }


def checkpoint_size_bytes(state_dict: dict[str, torch.Tensor]) -> int:
    """Unique bytes in a state_dict — the number the leaderboard scores."""
    return parameter_budget(state_dict)["total_bytes"]


def format_budget(budget: dict[str, Any]) -> str:
    lines = [
        f"{'component':<16}{'params':>16}{'% params':>11}{'GB':>9}{'tensors':>9}",
        "-" * 61,
    ]
    for name, stats in budget["groups"].items():
        suffix = f"  ({stats['num_aliases']} tied)" if stats["num_aliases"] else ""
        lines.append(
            f"{name:<16}{stats['num_params']:>16,}{stats['pct_params']:>10.1f}%"
            f"{stats['gb']:>9.3f}{stats['num_tensors']:>9}{suffix}"
        )
    lines.append("-" * 61)
    lines.append(
        f"{'TOTAL':<16}{budget['total_params']:>16,}{100.0:>10.1f}%"
        f"{budget['total_gb']:>9.3f}"
    )
    if budget.get("tied_aliases"):
        excess = budget["raw_params_with_aliases"] - budget["total_params"]
        lines.append(
            f"  ({len(budget['tied_aliases'])} tied alias tensor(s) counted once; "
            f"naive state_dict sum would overstate by {excess:,} params)"
        )
    return "\n".join(lines)


def load_model(
    model_id: str = DEFAULT_MODEL,
    dtype: str = "bfloat16",
    device_map: str | None = "auto",
    cache_dir: str | None = None,
    device: str | None = None,
    multimodal: bool = False,
):
    """Load the model, trying each plausible auto-class for this architecture.

    On transformers >= 5, ``AutoModelForCausalLM`` resolves Qwen3.5 to
    ``Qwen3_5ForCausalLM`` and materializes only the language model — the vision
    tower and MTP head are never allocated, so the 455M-param text-only saving
    happens for free at load time.

    Pass ``device`` to pin the whole model to one GPU. At 8.4 GB it fits on any
    of ours, and ``device_map="auto"`` across several GPUs would pipeline-shard
    it instead, adding cross-device hops on every layer for no benefit.
    """
    import transformers

    torch_dtype = {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }[dtype]

    kwargs: dict[str, Any] = {
        "dtype": torch_dtype,
        "trust_remote_code": True,
    }
    if device is not None:
        kwargs["device_map"] = {"": device}
    elif device_map is not None:
        kwargs["device_map"] = device_map
    if cache_dir:
        kwargs["cache_dir"] = cache_dir

    candidates = [
        "AutoModelForCausalLM",
        "AutoModelForImageTextToText",
        "AutoModelForVision2Seq",
        "AutoModel",
    ]
    if multimodal:
        # vLLM only registers Qwen3_5ForConditionalGeneration and rejects a
        # text-only Qwen3_5TextConfig (vllm#39231). Any checkpoint we intend to
        # evaluate with vLLM must therefore keep the multimodal wrapper, even
        # though the vision tower is never executed for a math prompt.
        candidates = [
            "Qwen3_5ForConditionalGeneration",
            "AutoModelForImageTextToText",
            "AutoModelForVision2Seq",
            "AutoModelForCausalLM",
        ]
    errors: list[str] = []
    for class_name in candidates:
        auto_class = getattr(transformers, class_name, None)
        if auto_class is None:
            continue
        try:
            model = auto_class.from_pretrained(model_id, **kwargs)
            model.eval()
            print(f"[model] loaded {model_id} via {class_name}")
            return model
        except Exception as exc:  # noqa: BLE001 - we want the full error list
            errors.append(f"  {class_name}: {type(exc).__name__}: {exc}")

    raise RuntimeError(
        f"Could not load {model_id} with any auto-class.\n" + "\n".join(errors)
    )


def load_tokenizer(model_id: str = DEFAULT_MODEL, cache_dir: str | None = None):
    from transformers import AutoTokenizer

    kwargs: dict[str, Any] = {"trust_remote_code": True}
    if cache_dir:
        kwargs["cache_dir"] = cache_dir
    tokenizer = AutoTokenizer.from_pretrained(model_id, **kwargs)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"  # required for correct batched generation
    return tokenizer


def estimate_memory(
    model: torch.nn.Module, batch_size: int, seq_len: int
) -> dict[str, float]:
    """Predict peak VRAM for a generation run, so OOM is caught before it costs an hour.

    Only the 8 full-attention layers hold a growing KV cache; the 24 Gated
    DeltaNet layers keep a fixed-size recurrent state regardless of length,
    which is why this model's cache is far cheaper than a dense transformer's.
    """
    config = getattr(model, "config", None)
    config = getattr(config, "text_config", config)

    layer_types = getattr(config, "layer_types", None)
    num_layers = getattr(config, "num_hidden_layers", 32)
    if layer_types:
        num_full = sum(1 for t in layer_types if "full" in str(t))
    else:
        interval = getattr(config, "full_attention_interval", 4)
        num_full = num_layers // interval

    kv_heads = getattr(config, "num_key_value_heads", 4)
    head_dim = getattr(config, "head_dim", 256)
    dtype_size = 2

    kv_bytes_per_token = num_full * kv_heads * head_dim * 2 * dtype_size
    kv_bytes = kv_bytes_per_token * seq_len * batch_size

    # DynamicCache grows by torch.cat on every decode step: allocate (n+1),
    # copy, free n. Peak is therefore ~2x the steady-state cache, which is
    # exactly the "ladder" you see climbing in nvtop until it OOMs.
    # expandable_segments lets the allocator grow a segment in place, cutting
    # the transient sharply -- but we still budget for it.
    growth_factor = 1.4 if "expandable_segments" in os.environ.get(
        "PYTORCH_CUDA_ALLOC_CONF", ""
    ) else 2.0
    kv_peak_bytes = kv_bytes * growth_factor

    # Gated DeltaNet recurrent state: [v_heads, k_dim, v_dim] per layer, length-independent.
    v_heads = getattr(config, "linear_num_value_heads", 32)
    k_dim = getattr(config, "linear_key_head_dim", 128)
    v_dim = getattr(config, "linear_value_head_dim", 128)
    state_bytes = (num_layers - num_full) * v_heads * k_dim * v_dim * 4 * batch_size

    weight_bytes = sum(p.numel() * p.element_size() for p in model.parameters())

    return {
        "weights_gb": weight_bytes / 1e9,
        "kv_cache_gb": kv_bytes / 1e9,
        "kv_peak_gb": kv_peak_bytes / 1e9,
        "ssm_state_gb": state_bytes / 1e9,
        "steady_gb": (weight_bytes + kv_bytes + state_bytes) / 1e9,
        # What the preflight must check against: peak, not steady state.
        "total_gb": (weight_bytes + kv_peak_bytes + state_bytes) / 1e9,
        "kv_bytes_per_token": kv_bytes_per_token,
        "growth_factor": growth_factor,
    }


def free_vram_gb(device: str | torch.device) -> tuple[float, float]:
    """(free, total) GB on ``device`` — reflects other users' jobs on a shared box."""
    free, total = torch.cuda.mem_get_info(torch.device(device))
    return free / 1e9, total / 1e9


def describe_model(model: torch.nn.Module, max_depth: int = 3) -> dict[str, Any]:
    """Summarize the module tree — the ground truth we build stripping on."""
    tree: list[dict[str, Any]] = []
    for name, module in model.named_modules():
        depth = name.count(".")
        if name and depth <= max_depth:
            n_params = sum(p.numel() for p in module.parameters(recurse=True))
            tree.append(
                {
                    "path": name,
                    "class": type(module).__name__,
                    "depth": depth,
                    "num_params": n_params,
                }
            )

    state_dict = model.state_dict()
    sensitive = [n for n in state_dict if re.search(SSM_SENSITIVE, n)]
    return {
        "model_class": type(model).__name__,
        "num_parameters": sum(p.numel() for p in model.parameters()),
        "module_tree": tree,
        "budget": parameter_budget(state_dict),
        "ssm_sensitive_tensors": {
            "count": len(sensitive),
            "num_params": sum(state_dict[n].numel() for n in sensitive),
            "names": sensitive[:20],
        },
        "state_dict_prefixes": sorted(
            {".".join(n.split(".")[:2]) for n in state_dict}
        ),
    }


def find_submodule(model: torch.nn.Module, pattern: str) -> list[str]:
    """Paths of modules whose name matches ``pattern`` and whose parent does not."""
    regex = re.compile(pattern)
    hits = [name for name, _ in model.named_modules() if name and regex.search(name)]
    # Keep only the shallowest match on each branch.
    return [h for h in hits if not any(h.startswith(o + ".") for o in hits)]


def strip_unused_modules(
    model: torch.nn.Module,
    drop_vision: bool = True,
    drop_mtp: bool = True,
) -> dict[str, Any]:
    """Delete text-irrelevant submodules in place; return what was removed.

    The vision tower and MTP head together are ~455M params. A text-only math
    eval never executes either: the ViT has no image inputs, and the MTP head is
    a speculative-decoding draft head that HF ``generate()`` does not call.
    """
    removed: list[dict[str, Any]] = []
    targets: list[str] = []
    if drop_vision:
        targets.append(r"(^|\.)(visual|vision_tower)$")
    if drop_mtp:
        targets.append(r"(^|\.)mtp$")

    for pattern in targets:
        for path in find_submodule(model, pattern):
            parent = model
            parts = path.split(".")
            for part in parts[:-1]:
                parent = getattr(parent, part)
            child = getattr(parent, parts[-1], None)
            if child is None:
                continue
            n_params = sum(p.numel() for p in child.parameters(recurse=True))
            setattr(parent, parts[-1], None)
            removed.append({"path": path, "num_params": n_params})

    # transformers >= 5 loads Qwen3.5 through Qwen3_5ForCausalLM, which never
    # allocates the vision tower or MTP head. Finding nothing to remove is the
    # expected outcome there, not a failure.
    already_text_only = not removed and not find_submodule(model, r"visual|vision_tower|mtp")

    return {
        "removed": removed,
        "num_params_removed": sum(r["num_params"] for r in removed),
        "bytes_removed_bf16": 2 * sum(r["num_params"] for r in removed),
        "already_text_only": already_text_only,
    }


# Files the model repo carries that `save_pretrained` does not reproduce.
# vLLM's multimodal path builds an image processor even for a text-only prompt,
# so a checkpoint missing `preprocessor_config.json` fails to load outright.
AUXILIARY_FILES = (
    "preprocessor_config.json",
    "video_preprocessor_config.json",
    "chat_template.jinja",
)


def save_checkpoint(
    model: torch.nn.Module,
    out_dir: str | Path,
    source_model: str = DEFAULT_MODEL,
    cache_dir: str | None = None,
) -> list[str]:
    """Write a checkpoint that vLLM and transformers can both load.

    ``model.save_pretrained`` emits weights + config, and the tokenizer covers
    vocab/merges, but the processor and chat-template files come from the source
    repo and must be copied across explicitly.
    """
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(out)
    load_tokenizer(source_model, cache_dir=cache_dir).save_pretrained(out)

    copied: list[str] = []
    source_dir = Path(source_model)
    for filename in AUXILIARY_FILES:
        target = out / filename
        if target.exists():
            copied.append(filename)
            continue
        try:
            if source_dir.is_dir():
                candidate = source_dir / filename
                if not candidate.is_file():
                    continue
                shutil.copy(candidate, target)
            else:
                from huggingface_hub import hf_hub_download

                downloaded = hf_hub_download(
                    source_model, filename, cache_dir=cache_dir
                )
                shutil.copy(downloaded, target)
            copied.append(filename)
        except Exception as exc:  # noqa: BLE001 - optional files; report and continue
            print(f"[save_checkpoint] could not copy {filename}: {exc}")
    return copied