"""
Lab 08b: See attention for yourself
Chapter 8 - Applied Cybersecurity: Foundations to Mastery

Visualise BERT attention weights with bertviz.
Run inside a Jupyter notebook for the interactive head_view / model_view output.

Toolchain (pinned to 11 May 2026 access date):
  transformers==5.8.0
  bertviz==1.4.1
  torch (CPU)
"""

# ── Step 1: Load a small model with output_attentions=True ────────────────────
from transformers import BertTokenizer, BertModel
from bertviz import head_view, model_view
import torch

print("Loading bert-base-uncased (attention outputs enabled)...")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased", output_attentions=True)
model.eval()

# ── Step 2: Run a single forward pass ─────────────────────────────────────────
# The book's suggested sentence: coreference that "it" -> "server"
sentence = "The server refused the request because it was overloaded."

inputs = tokenizer(sentence, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

attention = outputs.attentions   # tuple of (batch, heads, seq, seq) tensors
tokens    = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

print(f"\nTokens: {tokens}")
print(f"Number of layers: {len(attention)}")
print(f"Attention shape (layer 0): {attention[0].shape}")
print("  = (batch, num_heads, seq_len, seq_len)")

# ── Step 3: Render the attention heat-maps ────────────────────────────────────
# These calls produce interactive HTML widgets - run in Jupyter for full effect.
# When run as a plain script they still confirm the shapes are correct.
try:
    print("\nRendering head_view (interactive - run in Jupyter for full widget)...")
    head_view(attention, tokens)
    print("Rendering model_view...")
    model_view(attention, tokens)
except Exception as e:
    print(f"Widget render skipped in script mode ({e}).")
    print("Run this file as a Jupyter notebook for the interactive visualisation.")

# ── Step 4: Inspect coreference by hand ───────────────────────────────────────
# Find which token "it" attends to most strongly in each layer
it_pos = tokens.index("it") if "it" in tokens else None
if it_pos is not None:
    print(f"\n'it' is at token position {it_pos}")
    print("Layer-by-layer: top attended token from 'it':")
    for layer_idx, layer_attn in enumerate(attention):
        # Average across all heads
        mean_attn = layer_attn[0].mean(dim=0)   # (seq, seq)
        it_attn   = mean_attn[it_pos]            # attention FROM 'it'
        top_pos   = it_attn.argmax().item()
        print(f"  Layer {layer_idx:2d}: 'it' -> '{tokens[top_pos]}' "
              f"(weight={it_attn[top_pos]:.3f})")
    print("\nWhen 'it' attends most strongly to 'server', the Q.K^T grid is visible.")
else:
    print("'it' not found as a single token - try a different tokenisation.")
