# Source sentence structure
src = ["The", "agreement", "on", "the", "European", "Economic",
       "Area", "was", "signed", "in", "August", "1992", "."]

# French translation target
tgt = ["L'", "accord", "sur", "la", "zone", "économique",
       "européenne", "a", "été", "signé", "en", "août", "1992", "."]

# The decoder generating "accord" (position 1) needs "agreement" (position 1)
# The decoder generating "zone"  (position 4) needs "Area" (position 6)
# The decoder generating "économique" needs "Economic" (position 5)

# Problem: in vanilla Seq2Seq, at decoder step t, the decoder ONLY sees
# the single final encoder hidden state h_T — it cannot directly access
# h_1 (the state that best captured "agreement") or h_5 (captured "Economic")

# The encoder has ALL this information in its hidden states h_1, h_2, ..., h_13
# But vanilla Seq2Seq throws away h_1 through h_12 and only keeps h_13!

print("Information the decoder NEEDS vs. what it CAN ACCESS (vanilla Seq2Seq):")
print()
for dec_step, (tgt_word, src_focus) in enumerate([
    ("L'",         "The"),
    ("accord",     "agreement"),
    ("sur",        "on"),
    ("la",         "the"),
    ("zone",       "Area"),
    ("économique", "Economic"),
    ("européenne", "European"),
]):
    src_idx = src.index(src_focus) if src_focus in src else -1
    print(f"  Decoder step {dec_step+1} (generating '{tgt_word}'):")
    print(f"    Needs to focus on: '{src_focus}' (encoder state h_{src_idx+1})")
    print(f"    Vanilla Seq2Seq: only has h_13 (final state)  ← BOTTLENECK")
    print(f"    Attention:       directly reads h_{src_idx+1} ← SOLVED")

import torch
import torch.nn as nn
import torch.nn.functional as F

class BahdanauAttention(nn.Module):
    """
    Bahdanau (additive) attention.
    Computes a soft alignment between decoder query and encoder hidden states.
    """
    def __init__(self, query_dim, key_dim, attn_dim):
        """
        query_dim: dimension of decoder hidden state (s_{t-1})
        key_dim:   dimension of encoder hidden states (h_s)
        attn_dim:  dimension of the internal alignment model
        """
        super().__init__()
        # W1 projects decoder query to attention dim
        self.W1 = nn.Linear(query_dim, attn_dim, bias=False)
        # W2 projects encoder keys to attention dim (can be precomputed!)
        self.W2 = nn.Linear(key_dim, attn_dim, bias=False)
        # V projects to scalar score
        self.V  = nn.Linear(attn_dim, 1, bias=False)

    def forward(self, query, keys):
        """
        query: (batch, query_dim)          — decoder hidden state s_{t-1}
        keys:  (batch, src_len, key_dim)   — all encoder hidden states

        Returns:
            context:      (batch, key_dim)         — weighted sum of encoder states
            attn_weights: (batch, src_len)          — attention distribution α_{t,s}
        """
        # Project query: (batch, 1, attn_dim) — add time dimension for broadcasting
        q_proj = self.W1(query).unsqueeze(1)      # (B, 1, A)

        # Project keys: (batch, src_len, attn_dim)
        k_proj = self.W2(keys)                    # (B, T, A)

        # Compute attention energies e_{t,s} = V^T · tanh(W1·q + W2·k)
        # Broadcasting: q_proj (B,1,A) + k_proj (B,T,A) → (B,T,A)
        energy = self.V(torch.tanh(q_proj + k_proj)).squeeze(2)  # (B, T)

        # Normalize to attention weights via softmax
        attn_weights = F.softmax(energy, dim=1)   # (B, T) — sums to 1 along src_len

        # Compute context vector: weighted sum of encoder states
        # attn_weights: (B, T) → (B, 1, T) for bmm
        # keys:         (B, T, key_dim)
        context = torch.bmm(attn_weights.unsqueeze(1), keys)  # (B, 1, key_dim)
        context = context.squeeze(1)                            # (B, key_dim)

        return context, attn_weights


# Test the attention module
BATCH     = 4
SRC_LEN   = 12
DEC_DIM   = 512
ENC_DIM   = 1024   # bidirectional → 2 * 512
ATTN_DIM  = 256

attn = BahdanauAttention(query_dim=DEC_DIM, key_dim=ENC_DIM, attn_dim=ATTN_DIM)

query = torch.randn(BATCH, DEC_DIM)                 # decoder hidden state
keys  = torch.randn(BATCH, SRC_LEN, ENC_DIM)        # all encoder states

context, weights = attn(query, keys)

print(f"Query shape:      {query.shape}")       # (4, 512)
print(f"Keys shape:       {keys.shape}")        # (4, 12, 1024)
print(f"Context shape:    {context.shape}")     # (4, 1024)
print(f"Weights shape:    {weights.shape}")     # (4, 12) — attention over 12 src tokens
print(f"Weights sum to 1: {weights.sum(1).allclose(torch.ones(BATCH))}")  # True

# What attention weights look like:
print(f"\nSample attention weights (first example in batch):")
print(f"  {[f'{w:.3f}' for w in weights[0].detach().numpy()]}")
# Should sum to 1.0; initially uniform since untrained

class AttentionDecoder(nn.Module):
    """
    LSTM decoder augmented with Bahdanau attention.
    At each step: computes attention-weighted context → concatenates with embedding.
    """
    def __init__(self, vocab_size, embed_dim, dec_hidden, enc_hidden, attn_dim, dropout=0.3):
        super().__init__()
        self.embedding  = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.attention  = BahdanauAttention(dec_hidden, enc_hidden, attn_dim)

        # Decoder LSTM: input = embedding + context vector
        self.lstm = nn.LSTMCell(embed_dim + enc_hidden, dec_hidden)

        self.fc_out  = nn.Linear(dec_hidden + enc_hidden + embed_dim, vocab_size)
        self.dropout = nn.Dropout(dropout)

    def forward_step(self, token_ids, h, c, encoder_outputs):
        """
        token_ids:       (batch,)               — current input token
        h, c:            (batch, dec_hidden)     — decoder state
        encoder_outputs: (batch, src_len, enc_H) — all encoder hidden states

        Returns: logits (batch, vocab), h_new, c_new, attn_weights (batch, src_len)
        """
        # 1. Embed current token
        embedded = self.dropout(self.embedding(token_ids))  # (B, E)

        # 2. Compute attention using PREVIOUS hidden state as query
        context, attn_weights = self.attention(h, encoder_outputs)  # (B, enc_H)

        # 3. Concatenate embedding + context → LSTM input
        lstm_input = torch.cat([embedded, context], dim=1)  # (B, E + enc_H)

        # 4. LSTM step
        h_new, c_new = self.lstm(lstm_input, (h, c))

        # 5. Predict next token from all available information
        prediction_input = torch.cat([h_new, context, embedded], dim=1)
        logits = self.fc_out(self.dropout(prediction_input))

        return logits, h_new, c_new, attn_weights


print("Attention-augmented decoder flow:")
print("  1. Embed current token")
print("  2. Query all encoder states with current decoder state")
print("  3. Form context = weighted sum of relevant encoder states")
print("  4. Feed [embedding + context] to LSTM")
print("  5. Predict next token from [hidden state + context + embedding]")
print()
print("Key insight: the context vector is DIFFERENT at each decoder step!")
print("  Step t=1 (generating 'accord'):     context ≈ h_encoder['agreement']")
print("  Step t=4 (generating 'zone'):       context ≈ h_encoder['Area']")
print("  Step t=5 (generating 'économique'): context ≈ h_encoder['Economic']")

import torch
import torch.nn.functional as F

# SQL analogy (hard lookup):
# SELECT value FROM database WHERE key = query  ← returns exactly one result

# Soft attention analogy:
# RETURN SUM(similarity(query, key_i) * value_i for all i)  ← weighted blend

# Illustrative example: 3 "database entries" (encoder states)
# Keys and values are the same here (common in cross-attention)

keys   = torch.tensor([[1.0, 0.0, 0.0],   # h_agreement: "noun, formal, contract"
                        [0.0, 1.0, 0.0],   # h_Economic:  "adj, financial, trade"
                        [0.0, 0.0, 1.0]])  # h_signed:    "verb, past, action"

values = keys.clone()  # in cross-attention, keys ≈ values

# Query: decoder wants to find "a noun about formal agreement"
query_for_accord  = torch.tensor([[0.9, 0.05, 0.05]])  # mostly key-0-like
query_for_zone    = torch.tensor([[0.05, 0.9, 0.05]])  # mostly key-1-like

for q_name, q in [("decoder wants 'accord' → looks for agreement", query_for_accord),
                  ("decoder wants 'zone'   → looks for Economic",  query_for_zone)]:
    scores  = (q @ keys.T)                      # (1, 3) — raw dot products
    weights = F.softmax(scores, dim=-1)          # (1, 3) — attention weights
    context = (weights @ values)                 # (1, 3) — weighted sum

    print(f"\n{q_name}:")
    print(f"  Query:   {q.numpy()}")
    print(f"  Scores:  {scores.detach().numpy()}")
    print(f"  Weights: {weights.detach().numpy()} (sums to {weights.sum():.1f})")
    print(f"  Context: {context.detach().numpy()}")
    print(f"  ↑ Context is mostly key_{weights.argmax().item()} — correct alignment!")

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Scaled Dot-Product Attention — the fundamental building block of Transformers.

    Q: (batch, heads, seq_len_q, d_k)   — queries
    K: (batch, heads, seq_len_k, d_k)   — keys
    V: (batch, heads, seq_len_v, d_v)   — values  (seq_len_k == seq_len_v)
    mask: (batch, 1, seq_len_q, seq_len_k) — optional mask (e.g., causal or padding)

    Returns:
        output:  (batch, heads, seq_len_q, d_v)  — attention-weighted values
        weights: (batch, heads, seq_len_q, seq_len_k)  — attention weights
    """
    d_k = Q.size(-1)

    # Step 1: Compute raw attention scores QK^T
    # Q: (B, H, T_q, d_k) × K^T: (B, H, d_k, T_k) → (B, H, T_q, T_k)
    scores = torch.matmul(Q, K.transpose(-2, -1))   # (B, H, T_q, T_k)

    # Step 2: Scale by 1/√d_k to prevent softmax saturation
    scores = scores / math.sqrt(d_k)

    # Step 3: Apply mask (set masked positions to -∞ → 0 after softmax)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Step 4: Softmax to get attention weights
    weights = F.softmax(scores, dim=-1)   # (B, H, T_q, T_k) — rows sum to 1

    # Step 5: Weighted sum of values
    output = torch.matmul(weights, V)     # (B, H, T_q, d_v)

    return output, weights


# Demonstrate the scaling effect
print("Scaling effect on softmax distribution:")
print()
d_k = 64
q = torch.randn(1, 1, 1, d_k)
k = torch.randn(1, 1, 5, d_k)

scores_unscaled = torch.matmul(q, k.transpose(-2, -1))
scores_scaled   = scores_unscaled / math.sqrt(d_k)

weights_unscaled = F.softmax(scores_unscaled, dim=-1)
weights_scaled   = F.softmax(scores_scaled,   dim=-1)

print(f"Unscaled scores magnitude: {scores_unscaled.abs().mean():.2f}")
print(f"Scaled   scores magnitude: {scores_scaled.abs().mean():.2f}")
print()
print(f"Unscaled weights: {weights_unscaled[0,0,0].detach().numpy().round(3)}")
print(f"Scaled   weights: {weights_scaled[0,0,0].detach().numpy().round(3)}")
print()
print("Unscaled: one weight ≈ 1.0 (softmax peaked → vanishing gradients)")
print("Scaled:   weights more distributed (healthier gradients)")

# Full example: attention on a 5-word sentence
batch_size, num_heads, seq_len, d_k = 2, 1, 5, 64
d_v = 64

Q = torch.randn(batch_size, num_heads, seq_len, d_k)
K = torch.randn(batch_size, num_heads, seq_len, d_k)
V = torch.randn(batch_size, num_heads, seq_len, d_v)

output, weights = scaled_dot_product_attention(Q, K, V)

print("Scaled Dot-Product Attention shapes:")
print(f"  Q, K: {Q.shape}")           # (2, 1, 5, 64)
print(f"  V:    {V.shape}")           # (2, 1, 5, 64)
print(f"  Output:  {output.shape}")   # (2, 1, 5, 64)
print(f"  Weights: {weights.shape}")  # (2, 1, 5, 5) — 5×5 attention matrix

print(f"\nAttention weights for first example (5×5 matrix):")
print(f"  {weights[0,0].detach().numpy().round(3)}")
print(f"\nRow sums (should all be 1.0): {weights[0,0].sum(1).detach().numpy().round(3)}")

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SelfAttention(nn.Module):
    """
    Single-head self-attention layer.
    Q, K, V all come from the same input sequence.
    """
    def __init__(self, embed_dim, d_k=None):
        super().__init__()
        self.d_k = d_k or embed_dim

        # Three separate learned linear projections
        self.W_Q = nn.Linear(embed_dim, self.d_k, bias=False)
        self.W_K = nn.Linear(embed_dim, self.d_k, bias=False)
        self.W_V = nn.Linear(embed_dim, self.d_k, bias=False)

        self.out_proj = nn.Linear(self.d_k, embed_dim, bias=False)

    def forward(self, x, mask=None):
        """
        x: (batch, seq_len, embed_dim) — the INPUT sequence
        All of Q, K, V derived from x.
        """
        # Project x to Q, K, V — three DIFFERENT linear projections
        Q = self.W_Q(x)   # (B, T, d_k) — "what am I looking for?"
        K = self.W_K(x)   # (B, T, d_k) — "what do I contain?"
        V = self.W_V(x)   # (B, T, d_k) — "what will I contribute if attended to?"

        # Add head dimension (single head here)
        Q = Q.unsqueeze(1)   # (B, 1, T, d_k)
        K = K.unsqueeze(1)
        V = V.unsqueeze(1)

        # Scaled dot-product attention
        output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)

        # Remove head dimension, project back to embed_dim
        output = output.squeeze(1)           # (B, T, d_k)
        output = self.out_proj(output)       # (B, T, embed_dim)

        return output, attn_weights.squeeze(1)   # attn_weights: (B, T, T)


# Example: self-attention on a 6-word sentence
BATCH, SEQ_LEN, EMBED_DIM = 2, 6, 64

x = torch.randn(BATCH, SEQ_LEN, EMBED_DIM)  # (2, 6, 64)

self_attn = SelfAttention(embed_dim=EMBED_DIM, d_k=64)
output, weights = self_attn(x)

print(f"Input shape:   {x.shape}")        # (2, 6, 64)
print(f"Output shape:  {output.shape}")   # (2, 6, 64) — same!
print(f"Weights shape: {weights.shape}")  # (2, 6, 6) — each position attends to all others
print()

# Interpret the attention matrix:
# weights[b, i, j] = how much position i attends to position j
print("Attention weight matrix for batch 0 (6×6):")
print("  Rows = query positions, Columns = key positions")
print("  Each row sums to 1.0")
print(f"  {weights[0].detach().numpy().round(2)}")

print("Self-attention vs. RNN for long-range dependencies:")
print()
print("RNN (LSTM):")
print("  To connect position 1 to position T: must pass through T-1 hidden states")
print("  Information can decay even in LSTM for very long sequences")
print("  Maximum path length: O(T)")
print()
print("Self-Attention:")
print("  Any position can directly attend to any other position in ONE step")
print("  Position 1 can attend to position T directly: one dot product!")
print("  Maximum path length: O(1) — constant regardless of sequence length")
print()
print("This is why transformers MASSIVELY outperform LSTMs on long documents,")
print("code understanding, and cross-sentence reasoning tasks.")
print()

# Sentence: "The animal didn't cross the street because it was too tired."
# Indices:   0     1      2      3     4   5       7      8  9   10  11
# When computing representation of "it" (position 7):
# Self-attention can directly attend to "animal" (position 1) in ONE operation.
# LSTM would need to pass the "animal" information forward through 6 hidden states.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    """
    Multi-Head Attention as in 'Attention is All You Need' (Vaswani et al., 2017).

    Runs h independent attention heads in parallel, each with different
    linear projections of Q, K, V. Concatenates the outputs and projects back.
    """
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super().__init__()
        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.d_k = embed_dim // num_heads   # dimension per head

        # Single large linear layers — more efficient than separate per head
        self.W_Q = nn.Linear(embed_dim, embed_dim, bias=False)
        self.W_K = nn.Linear(embed_dim, embed_dim, bias=False)
        self.W_V = nn.Linear(embed_dim, embed_dim, bias=False)

        # Final projection after concatenating all heads
        self.W_O = nn.Linear(embed_dim, embed_dim)

        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        """
        query: (batch, T_q, embed_dim)
        key:   (batch, T_k, embed_dim)
        value: (batch, T_v, embed_dim)   [T_k == T_v]
        mask:  (batch, 1, T_q, T_k)     optional

        When query==key==value: this is self-attention.
        When query from decoder, key/value from encoder: this is cross-attention.
        """
        B = query.size(0)
        H = self.num_heads
        d_k = self.d_k

        # Project Q, K, V all at once
        Q = self.W_Q(query)   # (B, T_q, embed_dim)
        K = self.W_K(key)     # (B, T_k, embed_dim)
        V = self.W_V(value)   # (B, T_v, embed_dim)

        # Reshape into (B, H, T, d_k) — split embed_dim into H heads
        Q = Q.view(B, -1, H, d_k).transpose(1, 2)  # (B, H, T_q, d_k)
        K = K.view(B, -1, H, d_k).transpose(1, 2)  # (B, H, T_k, d_k)
        V = V.view(B, -1, H, d_k).transpose(1, 2)  # (B, H, T_v, d_k)

        # Compute attention for all heads in parallel
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
        # attn_output: (B, H, T_q, d_k)
        # attn_weights: (B, H, T_q, T_k)

        # Concatenate heads: (B, H, T_q, d_k) → (B, T_q, H*d_k) = (B, T_q, embed_dim)
        attn_output = attn_output.transpose(1, 2).contiguous()   # (B, T_q, H, d_k)
        attn_output = attn_output.view(B, -1, self.embed_dim)     # (B, T_q, embed_dim)

        # Final linear projection
        output = self.W_O(self.dropout(attn_output))  # (B, T_q, embed_dim)

        return output, attn_weights


# Test with standard Transformer hyperparameters
EMBED_DIM = 512
NUM_HEADS = 8       # d_k = 512 / 8 = 64 per head
SEQ_LEN   = 10
BATCH     = 4

mha = MultiHeadAttention(embed_dim=EMBED_DIM, num_heads=NUM_HEADS)

x = torch.randn(BATCH, SEQ_LEN, EMBED_DIM)

# Self-attention: query = key = value = x
output, weights = mha(x, x, x)

print(f"Multi-Head Attention:")
print(f"  embed_dim={EMBED_DIM}, num_heads={NUM_HEADS}, d_k={EMBED_DIM//NUM_HEADS}")
print(f"  Input:   {x.shape}")       # (4, 10, 512)
print(f"  Output:  {output.shape}")  # (4, 10, 512) — same shape as input!
print(f"  Weights: {weights.shape}") # (4, 8, 10, 10) — 8 heads, 10×10 attention each
print(f"\nParameter count:")
total = sum(p.numel() for p in mha.parameters())
print(f"  {total:,} parameters")

# Using PyTorch's built-in MultiheadAttention (preferred for production)
import torch
import torch.nn as nn

mha_builtin = nn.MultiheadAttention(
    embed_dim=512,
    num_heads=8,
    dropout=0.1,
    batch_first=True   # input shape: (batch, seq, embed_dim)
)

x = torch.randn(4, 10, 512)

# Self-attention
output, attn_weights = mha_builtin(x, x, x)   # (query, key, value)
print(f"PyTorch MHA output:  {output.shape}")    # (4, 10, 512)
print(f"PyTorch MHA weights: {attn_weights.shape}")  # (4, 10, 10) — averaged over heads

# Cross-attention (decoder attending to encoder)
decoder_query = torch.randn(4, 7, 512)   # decoder: 7 tokens
encoder_kvs   = torch.randn(4, 12, 512) # encoder: 12 tokens
cross_out, cross_w = mha_builtin(decoder_query, encoder_kvs, encoder_kvs)
print(f"\nCross-attention output:  {cross_out.shape}")   # (4, 7, 512)
print(f"Cross-attention weights: {cross_w.shape}")      # (4, 7, 12) — 7 decoder × 12 encoder

import torch

def create_causal_mask(seq_len, device='cpu'):
    """
    Creates a causal (lower-triangular) mask for decoder self-attention.
    Position i can attend to positions 0, 1, ..., i (but NOT i+1, i+2, ...)
    """
    # Lower triangular matrix: 1 where attending is allowed, 0 where not
    mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
    # Shape: (seq_len, seq_len)
    # We'll expand to (1, 1, seq_len, seq_len) for broadcast with (B, H, T, T)
    return mask.unsqueeze(0).unsqueeze(0)

# Visualize the causal mask for a 5-token sequence
mask = create_causal_mask(5)
print("Causal mask for seq_len=5:")
print("  (1 = can attend, 0 = cannot attend / blocked)")
print()
for i, row in enumerate(mask[0, 0].numpy().astype(int)):
    attention_to = [f"pos{j}" for j, v in enumerate(row) if v == 1]
    blocked      = [f"pos{j}" for j, v in enumerate(row) if v == 0]
    print(f"  pos{i} can attend to: {attention_to}")
    if blocked:
        print(f"  pos{i}    BLOCKED from: {blocked}")
    print()

print(mask[0, 0].numpy().astype(int))

import torch
import torch.nn.functional as F
import math

def masked_self_attention(x, mask=None):
    """
    Demonstrate causal masking in action.
    Shows how -inf scores become 0 after softmax (blocked positions contribute nothing).
    """
    B, T, D = x.shape
    d_k = D

    # Simplified: Q = K = V = x for illustration
    scores = (x @ x.transpose(-2, -1)) / math.sqrt(d_k)  # (B, T, T)

    print(f"Attention scores (before masking), seq_len={T}:")
    print(scores[0].detach().numpy().round(2))
    print()

    if mask is not None:
        # Set masked positions to -inf → after softmax they become 0
        scores = scores.masked_fill(mask[0, 0] == 0, float('-inf'))
        print(f"Attention scores (after causal masking):")
        print(scores[0].detach().numpy().round(2))
        print("  (-inf becomes 0 after softmax — future tokens completely blocked)")
        print()

    weights = F.softmax(scores, dim=-1)  # NaN-safe: -inf → 0 via softmax
    print(f"Attention weights (after softmax):")
    print(weights[0].detach().numpy().round(3))

    return weights

x = torch.randn(1, 4, 8)  # batch=1, seq_len=4, embed_dim=8
mask = create_causal_mask(4)
_ = masked_self_attention(x, mask)

print("Transformer Decoder has THREE attention sub-layers:")
print()
print("1. Masked Self-Attention:")
print("   Q, K, V = decoder input (target sequence)")
print("   Mask: causal — can only see past positions")
print("   Purpose: each target position attends to all previous target positions")
print()
print("2. Cross-Attention (Encoder-Decoder Attention):")
print("   Q  = decoder hidden states")
print("   K, V = encoder output (the encoded source sequence)")
print("   Mask: padding mask only (no causal mask — full encoder is known)")
print("   Purpose: decoder attends to relevant parts of the source")
print("   This is where the translation alignment happens!")
print()
print("3. Feed-Forward Network:")
print("   Applied independently to each position")
print("   No attention — just two linear transformations + ReLU")
print("   Purpose: additional non-linear transformation of attended representations")
print()
print("All three sub-layers have residual connections and layer normalization.")

import matplotlib.pyplot as plt
import numpy as np

def plot_attention_heatmap(attention_weights, src_tokens, tgt_tokens,
                            title="Attention Alignment", figsize=(8, 6)):
    """
    Plot attention weights as a heatmap (alignment matrix).
    Rows = target tokens, Columns = source tokens.
    Bright = high attention, dark = low attention.
    """
    fig, ax = plt.subplots(figsize=figsize)

    im = ax.imshow(attention_weights, cmap='Blues', aspect='auto', vmin=0, vmax=1)
    plt.colorbar(im, ax=ax, label='Attention weight')

    ax.set_xticks(range(len(src_tokens)))
    ax.set_yticks(range(len(tgt_tokens)))
    ax.set_xticklabels(src_tokens, rotation=45, ha='right')
    ax.set_yticklabels(tgt_tokens)
    ax.set_xlabel('Source (encoder positions)')
    ax.set_ylabel('Target (decoder steps)')
    ax.set_title(title)

    # Add weight values in each cell
    for i in range(len(tgt_tokens)):
        for j in range(len(src_tokens)):
            text = ax.text(j, i, f'{attention_weights[i, j]:.2f}',
                          ha='center', va='center', fontsize=7,
                          color='white' if attention_weights[i, j] > 0.5 else 'black')

    plt.tight_layout()
    plt.show()


# Simulated attention weights for English → French translation
# In practice: save attn_weights from BahdanauAttention during inference

src_tokens = ["The", "agreement", "was", "signed", "in", "August", "1992", "."]
tgt_tokens = ["L'", "accord",     "a",   "été",    "signé", "en",  "août", "1992", "."]

# This is what a well-trained attention model actually produces!
# Near-diagonal pattern for well-aligned language pairs
synthetic_attn = np.array([
    [0.85, 0.05, 0.04, 0.02, 0.01, 0.01, 0.01, 0.01],  # L' → The
    [0.03, 0.90, 0.03, 0.02, 0.01, 0.00, 0.00, 0.01],  # accord → agreement
    [0.02, 0.02, 0.82, 0.06, 0.04, 0.02, 0.01, 0.01],  # a → was
    [0.01, 0.02, 0.05, 0.82, 0.05, 0.03, 0.01, 0.01],  # été → (part of "was signed")
    [0.01, 0.01, 0.04, 0.05, 0.80, 0.05, 0.03, 0.01],  # signé → signed
    [0.01, 0.01, 0.02, 0.03, 0.04, 0.82, 0.05, 0.02],  # en → in
    [0.01, 0.01, 0.01, 0.02, 0.02, 0.04, 0.85, 0.04],  # août → August
    [0.01, 0.01, 0.01, 0.01, 0.02, 0.02, 0.04, 0.88],  # 1992 → 1992
    [0.03, 0.01, 0.01, 0.01, 0.02, 0.01, 0.02, 0.89],  # . → .
])

plot_attention_heatmap(synthetic_attn, src_tokens, tgt_tokens,
                        "English → French Alignment (Bahdanau Attention)")

import numpy as np

# Raw alignment scores e_{t,s} for the decoder generating "accord",
# BEFORE softmax — these come from the Bahdanau/QKV scoring function.
# "agreement" scores far higher because it is the true aligned word.
src_tokens = ["The", "agreement", "on", "the", "European", "Economic", "Area"]
raw_scores = np.array([0.8, 4.5, 0.3, 0.5, 1.2, 1.6, 0.6])

def softmax_with_temperature(scores, T):
    """softmax(scores / T) — lower T sharpens the distribution, higher T flattens it."""
    scaled = scores / T
    scaled = scaled - scaled.max()          # numerical stability
    exp = np.exp(scaled)
    return exp / exp.sum()

for T in [0.2, 1.0, 3.0]:
    weights = softmax_with_temperature(raw_scores, T)
    print(f"T={T:.1f}: {np.round(weights, 3)}  (max weight: {weights.max():.2f})")

# T=0.2: attention collapses almost entirely onto "agreement" (near one-hot)
# T=1.0: "agreement" dominates but other words retain some weight
# T=3.0: distribution is nearly uniform across all source words

from transformers import BertTokenizer, BertModel
import torch

# BERT Attention Visualization with BertViz (pip install bertviz)
# from bertviz import head_view, model_view

# BERT has 12 layers, each with 12 attention heads = 144 total attention matrices
# Each head specialises in different linguistic phenomena

print("Attention head specialization in BERT (from research papers):")
print()
print("Layer 1, Head 1:  Direct object tracking (verb → its object)")
print("Layer 2, Head 3:  Separator token [SEP] — mostly attends to [SEP]")
print("Layer 3, Head 7:  Coreference — pronoun to antecedent")
print("Layer 5, Head 2:  Subject-verb agreement")
print("Layer 9, Head 6:  Preposition attachment")
print("Layer 11, Head 8: Semantic role labeling")
print()
print("This analysis shows attention has learned real linguistic structure!")

# Load BERT and extract attention weights
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased', output_attentions=True)

text = "The animal didn't cross the street because it was too tired"
inputs = tokenizer(text, return_tensors='pt')
with torch.no_grad():
    outputs = model(**inputs)

all_attention = outputs.attentions  # Tuple of 12 tensors, one per layer
print(f"\nBERT attention structure:")
print(f"  Number of layers: {len(all_attention)}")
print(f"  Shape per layer:  {all_attention[0].shape}")
# Each layer: (batch=1, num_heads=12, seq_len, seq_len)

# Find which head in layer 3 attends from "it" to "animal"
layer3_attn = all_attention[3][0]  # (12, seq_len, seq_len)
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
print(f"  Tokens: {tokens}")

it_idx = tokens.index('it')
animal_idx = tokens.index('animal')
print(f"\n  Head attention scores from 'it' to 'animal' (layer 4):")
for head_idx in range(12):
    score = layer3_attn[head_idx, it_idx, animal_idx].item()
    if score > 0.05:
        print(f"    Head {head_idx+1:2d}: {score:.4f} {'← HIGH' if score > 0.2 else ''}")

print("Key insight: attention is strictly more powerful than RNNs for sequences")
print()
print("RNN limitations that attention solves:")
print()
print("  1. Sequential computation:")
print("     RNN:       must process step 1, then step 2, ..., then step T (serial)")
print("     Attention: computes all positions simultaneously (parallel on GPU)")
print(f"     Speed gain: O(T) → O(1) time depth (with O(T^2) memory)")
print()
print("  2. Long-range dependencies:")
print("     RNN:       path length between positions i and j = |i - j| steps")
print("     Attention: path length = 1 step (direct attention)")
print()
print("  3. Variable input handling:")
print("     RNN:       handles variable length natively, but with bounded memory")
print("     Attention: handles variable length, all positions equally accessible")
print()
print("What the Transformer keeps from RNNs:")
print("  - Residual connections (originally from ResNet)")
print("  - Layer normalization")
print("  - Dropout")
print()
print("What the Transformer adds:")
print("  - Positional encoding (since attention has no inherent order)")
print("  - Multiple stacked encoder layers (each with self-attention + FFN)")
print("  - Multiple stacked decoder layers (self-attn + cross-attn + FFN)")
print()
print("Preview: Lesson 53 covers the full Transformer architecture in detail.")