"""Run one forward pass: print the next-token distribution and one head's attention.

Purpose: show the two things the lessons describe but cannot show on a page - the
         probability the model assigns to every token in its vocabulary for the position
         after the prompt (with the effect of temperature on that distribution), and the
         attention weights of one head in one layer, checked for the causal mask.
Platform: all (cuda on Tracks S and N, ROCm on Track X, mps on Track M, or the CPU;
          the device is chosen automatically and printed)
Minimum memory: 8 GB (about 3.4 GB of weights for Qwen3-1.7B in BF16)
Assumes: a model directory downloaded with `hf download ... --local-dir <dir>`, and
         torch, transformers, safetensors and accelerate installed in the course
         environment from Part 1.

Usage: python next-token-and-attention.py --model ~/llm-course/models/qwen3-1.7b
           [--prompt "The capital of France is"] [--layer -1] [--head 0]
           [--top 10] [--dtype auto|float32|bfloat16|float16] [--labbook labbook.md]
"""
import argparse
import json
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

TEMPERATURES = (0.5, 1.0, 1.5)


def pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", required=True, help="model directory or Hub id")
    parser.add_argument("--prompt", default="The capital of France is")
    parser.add_argument("--layer", type=int, default=-1, help="layer index; -1 is the last")
    parser.add_argument("--head", type=int, default=0, help="attention head index within that layer")
    parser.add_argument("--top", type=int, default=10, help="how many next tokens to print")
    parser.add_argument("--dtype", default="auto", choices=["auto", "float32", "bfloat16", "float16"])
    parser.add_argument("--max-print", type=int, default=24, help="print the full matrix only up to this many tokens")
    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 model_dir.is_dir() and not (model_dir / "config.json").exists():
        raise SystemExit(f"{model_dir} has no config.json; the download did not finish, or the path is wrong")

    device = pick_device()
    if args.dtype == "auto":
        # bfloat16 is the precision the checkpoint was saved in and halves the weight
        # memory, which is what keeps this lab inside the 8 GB tier. If an operator is
        # missing for bfloat16 on your CPU, rerun with --dtype float32 and twice the memory.
        dtype = torch.bfloat16
    else:
        dtype = getattr(torch, args.dtype)
    print(f"device: {device}   dtype: {dtype}   torch {torch.__version__}")

    tokeniser = AutoTokenizer.from_pretrained(args.model)
    # The eager implementation computes the attention matrix explicitly, which is what
    # output_attentions=True returns; the optimised backends never build it.
    model = AutoModelForCausalLM.from_pretrained(
        args.model, dtype=dtype, attn_implementation="eager",
    ).to(device)
    model.eval()

    config = model.config
    groups = config.num_attention_heads // config.num_key_value_heads
    print(f"layers: {config.num_hidden_layers}   query heads: {config.num_attention_heads}   "
          f"key-value heads: {config.num_key_value_heads}   head_dim: {config.head_dim}   "
          f"vocab: {config.vocab_size:,}")
    print(f"weights once loaded: {sum(p.numel() for p in model.parameters()):,} parameters")

    inputs = tokeniser(args.prompt, return_tensors="pt").to(device)
    ids = inputs["input_ids"][0].tolist()
    pieces = tokeniser.convert_ids_to_tokens(ids)
    print(f"\nprompt: {args.prompt!r}")
    print(f"tokens ({len(ids)}): {pieces}")
    print(f"ids: {ids}")

    with torch.no_grad():
        out = model(**inputs, output_attentions=True)

    # ---- the next-token distribution -------------------------------------------------
    print(f"\nlogits: tensor shape {tuple(out.logits.shape)} = (batch, prompt tokens, vocabulary)")
    logits = out.logits[0, -1].float()
    probs = torch.softmax(logits, dim=-1)
    top = torch.topk(probs, args.top)
    print(f"\ntop {args.top} next tokens after {pieces[-1]!r} (softmax of the last row, temperature 1):")
    print(f"  {'rank':>4s} {'id':>8s} {'logit':>9s} {'probability':>12s}  token")
    top_records = []
    for rank, (p, token_id) in enumerate(zip(top.values.tolist(), top.indices.tolist()), start=1):
        token = tokeniser.convert_ids_to_tokens(token_id)
        print(f"  {rank:4d} {token_id:8d} {float(logits[token_id]):9.3f} {p:12.6f}  {token!r}")
        top_records.append({"rank": rank, "id": token_id, "token": token,
                            "logit": round(float(logits[token_id]), 3), "probability": round(p, 6)})
    covered = float(top.values.sum())
    print(f"  the other {config.vocab_size - args.top:,} tokens share {1 - covered:.6f} of the probability")

    # The same logits under three temperatures: softmax(logits / T).
    print(f"\ntemperature on the same logits (probability of the top token, and of the top {args.top} together):")
    temperature_records = []
    for temp in TEMPERATURES:
        p_t = torch.softmax(logits / temp, dim=-1)
        top_t = torch.topk(p_t, args.top)
        entropy = float(-(p_t * torch.log(p_t.clamp_min(1e-12))).sum())
        print(f"  T = {temp:3.1f}   top-1 {top_t.values[0].item():.4f}   top-{args.top} {top_t.values.sum().item():.4f}"
              f"   entropy {entropy:5.2f} nats")
        temperature_records.append({"temperature": temp, "top1": round(top_t.values[0].item(), 4),
                                    "topk": round(top_t.values.sum().item(), 4), "entropy_nats": round(entropy, 2)})

    # ---- one head's attention --------------------------------------------------------
    layer_index = args.layer if args.layer >= 0 else len(out.attentions) + args.layer
    if not 0 <= layer_index < len(out.attentions):
        raise SystemExit(f"--layer {args.layer} is outside 0..{len(out.attentions) - 1}")
    heads = out.attentions[layer_index].shape[1]
    if not 0 <= args.head < heads:
        raise SystemExit(f"--head {args.head} is outside 0..{heads - 1}")

    print(f"\nattentions: {len(out.attentions)} tensors (one per layer), each of shape "
          f"{tuple(out.attentions[layer_index].shape)} = (batch, heads, query position, key position)")
    attention = out.attentions[layer_index][0, args.head].float()
    print(f"layer {layer_index}, query head {args.head} (it reads key-value head {args.head // groups}, "
          f"shared by {groups} query heads): tensor shape {tuple(attention.shape)}")

    row_sums = attention.sum(dim=-1)
    above = attention.triu(diagonal=1)
    print(f"  row sums: min {row_sums.min().item():.4f}  max {row_sums.max().item():.4f}  "
          f"(one, within the rounding of {dtype})")
    print(f"  largest entry above the diagonal: {above.max().item():.6f}  (the causal mask; must be exactly 0)")

    last_row = attention[-1]
    order = torch.argsort(last_row, descending=True)
    print(f"\n  what the last token {pieces[-1]!r} attended to:")
    attended = []
    for pos in order[: min(5, len(ids))].tolist():
        weight = float(last_row[pos])
        print(f"    position {pos:3d}  weight {weight:7.4f}  {pieces[pos]!r}")
        attended.append({"position": pos, "weight": round(weight, 4), "token": pieces[pos]})

    if len(ids) <= args.max_print:
        print("\n  full matrix, rounded to two decimals (rows: query position; columns: key position):")
        header = "        " + " ".join(f"{i:5d}" for i in range(len(ids)))
        print(header)
        for i in range(len(ids)):
            row = " ".join(f"{float(attention[i, j]):5.2f}" for j in range(len(ids)))
            print(f"    {i:3d} {row}")
    else:
        print(f"\n  (matrix not printed: {len(ids)} tokens is more than --max-print {args.max_print})")

    if args.labbook:
        record = {
            "lab": "part-02/next-token-and-attention",
            "model": args.model, "device": str(device), "dtype": str(dtype),
            "torch": torch.__version__,
            "prompt": args.prompt, "prompt_tokens": len(ids),
            "logits_shape": list(out.logits.shape),
            "top_tokens": top_records,
            "temperature": temperature_records,
            "layer": layer_index, "head": args.head, "kv_head": args.head // groups,
            "attention_shape": list(attention.shape),
            "row_sum_min": round(row_sums.min().item(), 4), "row_sum_max": round(row_sums.max().item(), 4),
            "max_above_diagonal": above.max().item(),
            "last_token_attended": attended,
        }
        with Path(args.labbook).expanduser().open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(record, ensure_ascii=False) + "\n")
        print(f"\nrecorded in {args.labbook}")


if __name__ == "__main__":
    main()