"""
=============================================================================
Lecture 6 - Attention & the Transformer Architecture
Practical Session: Visualizing Attention in BERT

# Methods Seminar: Multimodal Computational Methods for Political Science
# Computational Social Science Program, LMU Munich
=============================================================================

In this practical we will:
  1. Load a pre-trained BERT model from Hugging Face
  2. Feed it UK parliamentary sentences
  3. Extract the attention weights from every layer and head
  4. Visualize them as heatmaps
  5. Compare attention patterns across layers and heads
  6. Probe politically interesting patterns (negation, coreference, contrast)

We do NOT train anything here. We "dissect" a pre-trained model to see
what the attention mechanism from the lecture actually does on real text.

This is best run in Google Colab (free, no setup). A GPU is NOT required
for this practical -- we only run inference on single sentences.

Setup:
  pip install transformers torch matplotlib seaborn
=============================================================================
"""

# -----------------------------------------------------------------------------
# 0. SETUP
# -----------------------------------------------------------------------------

# In Colab, run once:
# !pip install transformers torch matplotlib seaborn -q

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from transformers import AutoTokenizer, AutoModel

# For reproducibility (attention is deterministic, but good habit)
torch.manual_seed(42)

# BERT-base is small enough to run on CPU for single sentences
MODEL_NAME = "bert-base-uncased"

print("Loading BERT (first run downloads ~440 MB)...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME, output_attentions=True)
model.eval()   # inference mode -- no training

print(f"Model loaded: {MODEL_NAME}")
print(f"Layers: {model.config.num_hidden_layers}, "
      f"Heads per layer: {model.config.num_attention_heads}")


# -----------------------------------------------------------------------------
# 1. HELPER FUNCTIONS
# -----------------------------------------------------------------------------

def get_attention(sentence):
    """Run BERT on a sentence, return tokens and attention tensors.

    Returns:
        tokens: list of token strings (incl. [CLS] and [SEP])
        attentions: tuple of tensors, one per layer,
                    each shape [1, n_heads, seq_len, seq_len]
    """
    inputs = tokenizer(sentence, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    return tokens, outputs.attentions


def plot_attention_head(sentence, layer, head, figsize=(8, 7)):
    """Visualize one attention head in one layer as a heatmap."""
    tokens, attentions = get_attention(sentence)

    # attentions[layer] has shape [1, n_heads, seq_len, seq_len]
    attn = attentions[layer][0, head].numpy()

    plt.figure(figsize=figsize)
    sns.heatmap(attn,
                xticklabels=tokens,
                yticklabels=tokens,
                cmap="viridis",
                cbar_kws={"label": "attention weight"},
                square=True)
    plt.title(f"BERT attention -- layer {layer}, head {head}\n"
              f'"{sentence}"', fontsize=10)
    plt.xlabel("Key (attended TO)")
    plt.ylabel("Query (attending FROM)")
    plt.xticks(rotation=90)
    plt.yticks(rotation=0)
    plt.tight_layout()
    plt.show()


def plot_attention_grid(sentence, layer, figsize=(16, 12)):
    """Show all heads of one layer in a grid (12 heads for bert-base)."""
    tokens, attentions = get_attention(sentence)
    n_heads = attentions[layer].shape[1]

    n_cols = 4
    n_rows = int(np.ceil(n_heads / n_cols))
    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
    axes = axes.flatten()

    for head in range(n_heads):
        attn = attentions[layer][0, head].numpy()
        ax = axes[head]
        sns.heatmap(attn, ax=ax, cmap="viridis", cbar=False,
                    xticklabels=tokens, yticklabels=tokens)
        ax.set_title(f"Head {head}", fontsize=9)
        ax.tick_params(labelsize=6)
        ax.set_xticklabels(tokens, rotation=90, fontsize=5)
        ax.set_yticklabels(tokens, rotation=0, fontsize=5)

    # Hide any unused subplots
    for i in range(n_heads, len(axes)):
        axes[i].axis("off")

    fig.suptitle(f'All heads in layer {layer}\n"{sentence}"', fontsize=12)
    plt.tight_layout()
    plt.show()


def attention_from_word(sentence, layer, head, query_word):
    """Print which words a given query word attends to most."""
    tokens, attentions = get_attention(sentence)

    # Find the query word's position (first match)
    query_word = query_word.lower()
    if query_word not in tokens:
        print(f"'{query_word}' not found. Tokens are: {tokens}")
        return

    q_idx = tokens.index(query_word)
    attn = attentions[layer][0, head].numpy()
    weights = attn[q_idx]   # row = attention FROM the query word

    # Sort tokens by attention weight
    ranked = sorted(zip(tokens, weights), key=lambda x: -x[1])

    print(f"\nLayer {layer}, Head {head}: what does '{query_word}' attend to?")
    print("-" * 45)
    for tok, w in ranked:
        bar = "#" * int(w * 40)
        print(f"  {tok:15s} {w:.3f} {bar}")


# -----------------------------------------------------------------------------
# 2. FIRST LOOK: TOKENIZATION
# -----------------------------------------------------------------------------

print("\n" + "="*60)
print("  TOKENIZATION")
print("="*60)

sentence = ("The Prime Minister announced new climate measures "
            "despite vocal opposition from her own party.")

tokens, attentions = get_attention(sentence)
print(f"\nSentence: {sentence}")
print(f"\nTokens ({len(tokens)}):")
print(tokens)
print(f"\nNumber of layers: {len(attentions)}")
print(f"Shape of each layer's attention: {tuple(attentions[0].shape)}")
print("  = [batch=1, heads=12, seq_len, seq_len]")

# Note the special tokens [CLS] (start) and [SEP] (end).
# Note also how rare words may be split into subwords (## prefix).


# -----------------------------------------------------------------------------
# 3. VISUALIZE A SINGLE HEAD
# -----------------------------------------------------------------------------

print("\n" + "="*60)
print("  SINGLE HEAD HEATMAP")
print("="*60)
print("Plotting layer 4, head 7 (try changing these!)")

plot_attention_head(sentence, layer=4, head=7)

# The heatmap reads: each ROW is a query word (attending FROM),
# each COLUMN is a key word (attended TO). Bright = high attention.


# -----------------------------------------------------------------------------
# 4. VISUALIZE ALL HEADS IN A LAYER
# -----------------------------------------------------------------------------

print("\n" + "="*60)
print("  ALL HEADS IN A LAYER")
print("="*60)
print("Layer 0 (early -- expect positional/local patterns)")
plot_attention_grid(sentence, layer=0)

print("Layer 6 (middle)")
plot_attention_grid(sentence, layer=6)

print("Layer 11 (late -- expect semantic/abstract patterns)")
plot_attention_grid(sentence, layer=11)

# Compare! Early layers often show diagonal / next-word patterns.
# Later layers show more diffuse, content-driven patterns.


# -----------------------------------------------------------------------------
# 5. PROBE SPECIFIC WORDS
# -----------------------------------------------------------------------------

print("\n" + "="*60)
print("  WHAT DOES A SPECIFIC WORD ATTEND TO?")
print("="*60)

# Where does "announced" look? (subject? object?)
attention_from_word(sentence, layer=8, head=5, query_word="announced")

# Where does "her" look? (coreference -- should point to "minister"?)
attention_from_word(sentence, layer=8, head=5, query_word="her")

# Where does "measures" look? (should it attend to "climate"?)
attention_from_word(sentence, layer=8, head=5, query_word="measures")

# Try different (layer, head) combinations -- patterns vary a lot!


# -----------------------------------------------------------------------------
# 6. THE NEGATION TEST
# -----------------------------------------------------------------------------

print("\n" + "="*60)
print("  NEGATION TEST")
print("="*60)

neg_sentence = "The bill did not pass in the House of Commons."
print(f"Sentence: {neg_sentence}")

# Does any head make "pass" attend strongly to "not"?
# This matters: negation flips meaning, and bag-of-words misses it.
tokens_neg, _ = get_attention(neg_sentence)
print(f"Tokens: {tokens_neg}")

# Scan several layers/heads for "pass" attending to "not"
print("\nScanning for heads where 'pass' attends to 'not'...")
_, attn_neg = get_attention(neg_sentence)
if "pass" in tokens_neg and "not" in tokens_neg:
    pass_idx = tokens_neg.index("pass")
    not_idx = tokens_neg.index("not")
    best = []
    for L in range(len(attn_neg)):
        for H in range(attn_neg[L].shape[1]):
            w = attn_neg[L][0, H, pass_idx, not_idx].item()
            best.append((w, L, H))
    best.sort(reverse=True)
    print("Top 5 (layer, head) where 'pass' -> 'not':")
    for w, L, H in best[:5]:
        print(f"  Layer {L:2d}, Head {H:2d}: weight = {w:.3f}")

    # Visualize the strongest one
    top_w, top_L, top_H = best[0]
    print(f"\nVisualizing the strongest: layer {top_L}, head {top_H}")
    plot_attention_head(neg_sentence, layer=top_L, head=top_H)


# -----------------------------------------------------------------------------
# 7. THE TWO-PARTY TEST
# -----------------------------------------------------------------------------

print("\n" + "="*60)
print("  TWO-PARTY / CONTRAST TEST")
print("="*60)

party_sentence = "Labour rejected the Conservative proposal on immigration."
print(f"Sentence: {party_sentence}")

tokens_party, _ = get_attention(party_sentence)
print(f"Tokens: {tokens_party}")

# Where does "rejected" attend? Subject (Labour)? Object (proposal)?
attention_from_word(party_sentence, layer=8, head=5, query_word="rejected")

# Try a few layers to see how the pattern evolves
for L in [2, 6, 10]:
    attention_from_word(party_sentence, layer=L, head=5,
                        query_word="rejected")


# -----------------------------------------------------------------------------
# 8. AVERAGE ATTENTION ACROSS ALL HEADS
# -----------------------------------------------------------------------------

print("\n" + "="*60)
print("  AVERAGED ATTENTION (all heads, one layer)")
print("="*60)

def plot_averaged_attention(sentence, layer):
    """Average attention over all heads in a layer -- a summary view."""
    tokens, attentions = get_attention(sentence)
    attn = attentions[layer][0].mean(dim=0).numpy()   # mean over heads

    plt.figure(figsize=(8, 7))
    sns.heatmap(attn, xticklabels=tokens, yticklabels=tokens,
                cmap="magma", square=True,
                cbar_kws={"label": "avg attention weight"})
    plt.title(f"Average attention over all heads -- layer {layer}")
    plt.xlabel("Key (attended TO)")
    plt.ylabel("Query (attending FROM)")
    plt.xticks(rotation=90)
    plt.yticks(rotation=0)
    plt.tight_layout()
    plt.show()

plot_averaged_attention(sentence, layer=6)

# Averaging smooths out head-specific quirks and shows the
# overall "flow" of attention in that layer.


# =============================================================================
# DISCUSSION QUESTIONS
# =============================================================================
#
# 1. Did you find a head whose pattern is interpretable as a linguistic
#    relationship (subject-verb, coreference, next-word)? Which layer/head?
#
# 2. How do attention patterns change between EARLY layers (0-3) and
#    LATE layers (9-11)? Use the grid plots from section 4.
#
# 3. The [CLS] and [SEP] tokens attract a lot of attention in many heads.
#    Why might that be? (Hint: [CLS] is used for classification; heads
#    "park" attention there when nothing else is informative.)
#
# 4. In the negation test (section 6): did any head make "pass" attend
#    strongly to "not"? What does this tell you about whether BERT can
#    represent negation (which bag-of-words cannot)?
#
# 5. In the two-party test (section 7): does "rejected" attend more to
#    the subject (Labour) or the object (proposal)? Does this change
#    across layers?
#
# 6. LIMIT OF INTERPRETATION: the model doesn't "mean" to attend in any
#    particular way. Attention patterns are emergent statistics, often
#    SUGGESTIVE of structure but not definitive explanations. Where might
#    you be over-interpreting?
#
# =============================================================================