"""List every tensor in a safetensors checkpoint and count the parameters.

Purpose: open the safetensors file(s) of a downloaded model, print each tensor's name,
         shape and dtype, total the parameters and bytes, and compare the file with the
         arithmetic that config.json predicts, so that the numbers on a model card can be
         checked against the checkpoint on your own disk.
Platform: all (this reads file headers only; no accelerator and no model loading)
Minimum memory: 8 GB (the script itself needs a few megabytes)
Assumes: a model directory downloaded with `hf download ... --local-dir <dir>` containing
         one or more *.safetensors files and a config.json; the course environment from
         Part 1. Only the standard library is used.

Usage: python list-tensors.py --model ~/llm-course/models/qwen3-1.7b [--labbook labbook.md]
           [--limit 12] [--layer 0]
"""
import argparse
import json
import struct
from pathlib import Path

# Bytes per element for the dtype strings the safetensors header uses.
DTYPE_BYTES = {
    "BOOL": 1, "U8": 1, "I8": 1, "F8_E4M3": 1, "F8_E5M2": 1,
    "I16": 2, "U16": 2, "F16": 2, "BF16": 2,
    "I32": 4, "U32": 4, "F32": 4,
    "I64": 8, "U64": 8, "F64": 8,
}


def read_header(path: Path) -> tuple[dict, int]:
    """Read a safetensors header without touching the weights.

    The format specification says the file starts with 8 bytes holding an unsigned
    little-endian 64-bit integer N, followed by N bytes of JSON, followed by the byte
    buffer. Reading the first 8 + N bytes is therefore enough to list everything.
    Returns the parsed header and N.
    """
    with path.open("rb") as fh:
        (header_len,) = struct.unpack("<Q", fh.read(8))
        header = json.loads(fh.read(header_len).decode("utf-8"))
    return header, header_len


def numel(shape: list[int]) -> int:
    n = 1
    for dim in shape:
        n *= dim
    return n


def expected_from_config(config: dict) -> dict | None:
    """Parameter arithmetic for a dense Qwen3-style block, from config.json alone.

    Per block: q_proj (heads*head_dim x hidden), k_proj and v_proj (kv_heads*head_dim x
    hidden), o_proj (hidden x heads*head_dim), q_norm and k_norm (head_dim each),
    gate_proj and up_proj (intermediate x hidden), down_proj (hidden x intermediate),
    two RMSNorm scales (hidden each). Then the final norm and the embedding matrix.
    Returns None when the config lacks a field or declares attention biases.
    """
    fields = ("hidden_size", "intermediate_size", "num_hidden_layers", "num_attention_heads",
              "num_key_value_heads", "head_dim", "vocab_size")
    if any(f not in config for f in fields) or config.get("attention_bias"):
        return None
    h, inter, layers = config["hidden_size"], config["intermediate_size"], config["num_hidden_layers"]
    q_width = config["num_attention_heads"] * config["head_dim"]
    kv_width = config["num_key_value_heads"] * config["head_dim"]
    attention = 2 * h * q_width + 2 * h * kv_width + 2 * config["head_dim"]
    mlp = 3 * h * inter
    norms = 2 * h
    per_layer = attention + mlp + norms
    embedding = config["vocab_size"] * h
    non_embedding = per_layer * layers + h
    return {
        "attention_per_layer": attention, "mlp_per_layer": mlp, "norms_per_layer": norms,
        "per_layer": per_layer, "non_embedding": non_embedding, "embedding": embedding,
        "total_tied": non_embedding + embedding, "total_untied": non_embedding + 2 * embedding,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", required=True, help="directory holding the checkpoint")
    parser.add_argument("--limit", type=int, default=12, help="tensors to print in full before summarising")
    parser.add_argument("--layer", type=int, default=None, help="also print every tensor of this layer index")
    parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file")
    args = parser.parse_args()

    model_dir = Path(args.model).expanduser()
    if not model_dir.is_dir():
        raise SystemExit(f"{model_dir} is not a directory; check the --model path")
    files = sorted(model_dir.glob("*.safetensors"))
    if not files:
        raise SystemExit(f"no *.safetensors files in {model_dir}; the download did not finish, or the path is wrong")

    config_path = model_dir / "config.json"
    config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.exists() else {}
    if not config:
        print(f"warning: no config.json in {model_dir}; the config comparison is skipped")

    tensors: list[tuple[str, list[int], str, int, str]] = []
    metadata: dict[str, str] = {}
    per_file: list[dict] = []
    for path in files:
        header, header_len = read_header(path)
        meta = header.pop("__metadata__", None)
        if isinstance(meta, dict):
            metadata.update({str(k): str(v) for k, v in meta.items()})
        count = 0
        for name, info in header.items():
            shape = [int(d) for d in info["shape"]]
            tensors.append((name, shape, info["dtype"], numel(shape), path.name))
            count += 1
        per_file.append({"file": path.name, "bytes": path.stat().st_size, "header_bytes": header_len, "tensors": count})

    tensors.sort(key=lambda t: t[0])
    total_params = sum(t[3] for t in tensors)
    total_bytes = sum(t[3] * DTYPE_BYTES.get(t[2], 0) for t in tensors)
    dtypes = sorted({t[2] for t in tensors})
    names = {t[0] for t in tensors}

    # The embedding matrix is the tensor whose first dimension is the vocabulary; an
    # output head, if stored, has the same shape under a different name.
    vocab_size = config.get("vocab_size")
    vocab_sized = [t for t in tensors if vocab_size and t[1] and t[1][0] == vocab_size]
    embedding_params = sum(t[3] for t in vocab_sized)
    tied = bool(config.get("tie_word_embeddings"))
    lm_head_stored = "lm_head.weight" in names
    lm_head_params = next((t[3] for t in tensors if t[0] == "lm_head.weight"), 0)
    # What the model occupies once loaded: a tied head is one matrix in memory even
    # when the file stores it twice.
    unique_params = total_params - (lm_head_params if tied and lm_head_stored else 0)
    bytes_per_param = DTYPE_BYTES.get(dtypes[0], 0) if len(dtypes) == 1 else None

    print(f"directory:  {model_dir}")
    for f in per_file:
        print(f"file:       {f['file']:36s} {f['bytes']:>15,} bytes on disk   "
              f"header {f['header_bytes']:>7,} bytes   {f['tensors']} tensors")
    print(f"tensors:    {len(tensors)}")
    print(f"dtypes:     {', '.join(dtypes)}")
    if metadata:
        print(f"metadata:   {metadata}")
    print()
    for name, shape, dtype, n, _ in tensors[: args.limit]:
        print(f"  {name:52s} {str(tuple(shape)):>22s}  {dtype:5s}  {n:>13,}")
    if len(tensors) > args.limit:
        print(f"  ... and {len(tensors) - args.limit} more tensors")

    layer_params = None
    if args.layer is not None:
        prefix = f"model.layers.{args.layer}."
        block = [t for t in tensors if t[0].startswith(prefix)]
        if not block:
            raise SystemExit(f"no tensors named {prefix}*; the model has {config.get('num_hidden_layers', '?')} layers")
        layer_params = sum(t[3] for t in block)
        print(f"\nevery tensor of layer {args.layer} ({len(block)} tensors, {layer_params:,} parameters):")
        for name, shape, dtype, n, _ in block:
            print(f"  {name[len(prefix):]:40s} {str(tuple(shape)):>16s}  {dtype:5s}  {n:>13,}")

    print()
    print(f"parameters in the file:        {total_params:>15,}")
    print(f"  in vocabulary-sized tensors: {embedding_params:>15,}   "
          f"({len(vocab_sized)} tensor(s): {', '.join(t[0] for t in vocab_sized) or 'none'})")
    print(f"  everything else:             {total_params - embedding_params:>15,}")
    print(f"tie_word_embeddings: {tied}   lm_head.weight stored in the file: {lm_head_stored}")
    if tied and lm_head_stored:
        print(f"  the output head is a second copy of the embedding matrix; it is loaded once, so")
        print(f"  parameters once loaded:      {unique_params:>15,}")
    print(f"weight bytes on disk:          {total_bytes:>15,}  ({total_bytes / 1e9:.2f} GB)")
    if bytes_per_param:
        print(f"weight bytes once loaded:      {unique_params * bytes_per_param:>15,}  "
              f"({unique_params * bytes_per_param / 1e9:.2f} GB at {bytes_per_param} bytes per parameter)")

    expected = expected_from_config(config) if config else None
    if config:
        print()
        for field in ("model_type", "num_hidden_layers", "hidden_size", "intermediate_size",
                      "num_attention_heads", "num_key_value_heads", "head_dim",
                      "vocab_size", "max_position_embeddings", "tie_word_embeddings", "torch_dtype"):
            if field in config:
                print(f"  config {field:26s} {config[field]}")
    if expected:
        print()
        print("from config.json alone (dense Qwen3-style block):")
        print(f"  attention per layer:  {expected['attention_per_layer']:>15,}")
        print(f"  feed-forward per layer:{expected['mlp_per_layer']:>14,}")
        print(f"  norms per layer:      {expected['norms_per_layer']:>15,}")
        print(f"  per layer:            {expected['per_layer']:>15,}"
              + (f"   file says {layer_params:,}  match: {layer_params == expected['per_layer']}" if layer_params else ""))
        print(f"  non-embedding:        {expected['non_embedding']:>15,}   file says {total_params - embedding_params:,}"
              f"  match: {expected['non_embedding'] == total_params - embedding_params}")
        print(f"  embedding matrix:     {expected['embedding']:>15,}")
        print(f"  total, tied head:     {expected['total_tied']:>15,}")
        print(f"  total, untied head:   {expected['total_untied']:>15,}")

    if args.labbook:
        record = {
            "lab": "part-02/list-tensors",
            "model_dir": str(model_dir),
            "files": per_file,
            "tensors": len(tensors),
            "dtypes": dtypes,
            "parameters_in_file": total_params,
            "parameters_loaded": unique_params,
            "embedding_parameters": embedding_params,
            "lm_head_stored": lm_head_stored,
            "weight_bytes_on_disk": total_bytes,
            "config": {k: config.get(k) for k in ("num_hidden_layers", "hidden_size", "intermediate_size",
                                                  "num_attention_heads", "num_key_value_heads",
                                                  "head_dim", "vocab_size", "tie_word_embeddings")},
            "expected_from_config": expected,
        }
        with Path(args.labbook).expanduser().open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(record) + "\n")
        print(f"\nrecorded in {args.labbook}")


if __name__ == "__main__":
    main()