"""
Transformers & Attention Mechanism — comprehensive topic content with interactive graphs.
"""

import json
import numpy as np


def _attention_heatmap():
    """Show a self-attention matrix for a sentence."""
    words = ["The", "cat", "sat", "on", "the", "mat"]
    n = len(words)
    np.random.seed(42)
    attn = np.random.dirichlet(np.ones(n) * 0.5, size=n)
    # Make "cat" attend to "sat" and "mat"
    attn[1] = [0.05, 0.1, 0.35, 0.05, 0.05, 0.40]
    attn[2] = [0.05, 0.30, 0.2, 0.15, 0.05, 0.25]
    attn[5] = [0.05, 0.25, 0.15, 0.20, 0.10, 0.25]

    data = [{
        "z": attn.tolist(), "x": words, "y": words,
        "type": "heatmap",
        "colorscale": [[0, "#0a0a1f"], [0.5, "#6C63FF"], [1, "#FF6584"]],
        "text": [[f"{attn[i][j]:.2f}" for j in range(n)] for i in range(n)],
        "texttemplate": "%{text}", "textfont": {"size": 12, "color": "#fff"},
        "showscale": True, "colorbar": {"title": "Attention", "titlefont": {"size": 12}},
    }]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Self-Attention Matrix", "font": {"size": 18}},
        "xaxis": {"title": "Key (attending to)", "side": "bottom"},
        "yaxis": {"title": "Query (attending from)", "autorange": "reversed"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 60},
    }
    return json.dumps({"data": data, "layout": layout})


def _multi_head_attention():
    """Show different attention heads focusing on different things."""
    heads = ["Head 1\n(Syntax)", "Head 2\n(Semantic)", "Head 3\n(Position)", "Head 4\n(Coreference)"]
    focus = ["Subject-Verb", "Noun-Adjective", "Adjacent Words", "Pronoun-Entity"]
    importance = [0.35, 0.28, 0.20, 0.17]
    colors = ["#6C63FF", "#FF6584", "#34d399", "#fbbf24"]

    data = [{"x": heads, "y": importance, "type": "bar",
             "marker": {"color": colors, "line": {"width": 1.5, "color": "#fff"}},
             "text": focus, "textposition": "outside",
             "textfont": {"color": "#e0e0e0", "size": 12}}]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Multi-Head Attention — What Each Head Learns", "font": {"size": 18}},
        "yaxis": {"title": "Attention Weight", "gridcolor": "rgba(255,255,255,0.08)"},
        "margin": {"l": 60, "r": 30, "t": 50, "b": 80},
    }
    return json.dumps({"data": data, "layout": layout})


def _transformer_vs_rnn():
    """Compare Transformer vs RNN performance."""
    tasks = ["Translation", "Summarization", "Question\nAnswering", "Sentiment", "Text\nGeneration"]
    rnn_scores = [35.2, 28.1, 72.3, 88.5, 18.4]
    transformer_scores = [41.8, 39.5, 88.7, 94.2, 34.1]

    data = [
        {"x": tasks, "y": rnn_scores, "name": "RNN/LSTM", "type": "bar",
         "marker": {"color": "#FF6584", "line": {"width": 1.5, "color": "#fff"}}},
        {"x": tasks, "y": transformer_scores, "name": "Transformer", "type": "bar",
         "marker": {"color": "#34d399", "line": {"width": 1.5, "color": "#fff"}}},
    ]
    layout = {
        "template": "plotly_dark", "paper_bgcolor": "rgba(0,0,0,0)", "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": "Transformer vs RNN on NLP Benchmarks", "font": {"size": 18}},
        "yaxis": {"title": "Score / BLEU", "gridcolor": "rgba(255,255,255,0.08)"},
        "barmode": "group",
        "margin": {"l": 60, "r": 30, "t": 50, "b": 80},
    }
    return json.dumps({"data": data, "layout": layout})


def get_content() -> dict:
    return {
        "title": "Transformers & Attention",
        "icon": "⚡",
        "subtitle": "The architecture behind GPT, BERT, and modern AI",
        "sections": [
            {"id": "introduction", "heading": "What are Transformers?", "type": "text",
             "content": ("The Transformer is a neural network architecture that processes sequences using "
                         "<strong>self-attention</strong> instead of recurrence. It can look at the <strong>entire input "
                         "at once</strong>, making it massively parallelizable and far more powerful than RNNs.\n\n"
                         "Introduced in the 2017 paper '<em>Attention Is All You Need</em>', Transformers power GPT, BERT, "
                         "T5, and virtually all modern language models and increasingly vision models too.")},
            {"id": "attention", "heading": "Self-Attention Visualized", "type": "graph",
             "description": "Each word attends to every other word in the sequence. The heatmap shows how much each word (row) pays attention to each other word (column). Notice how 'cat' strongly attends to 'sat' and 'mat'.",
             "graph_json": _attention_heatmap()},
            {"id": "math", "heading": "The Mathematics", "type": "math",
             "content": [
                 {"label": "Scaled Dot-Product Attention", "formula": "\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V",
                  "explanation": "Q (query), K (key), V (value) are linear projections of the input. Scaling by √d_k prevents vanishing gradients in softmax."},
                 {"label": "Multi-Head Attention", "formula": "\\text{MultiHead}(Q,K,V) = \\text{Concat}(\\text{head}_1, ..., \\text{head}_h)W^O",
                  "explanation": "Multiple attention heads run in parallel, each learning different relationships. Outputs are concatenated and projected."},
                 {"label": "Positional Encoding", "formula": "PE_{(pos, 2i)} = \\sin(pos / 10000^{2i/d}), \\quad PE_{(pos, 2i+1)} = \\cos(pos / 10000^{2i/d})",
                  "explanation": "Since there's no recurrence, position info is injected via sinusoidal encoding. Each position gets a unique pattern."},
             ]},
            {"id": "multi-head", "heading": "Multi-Head Attention", "type": "graph",
             "description": "Different attention heads specialize in different linguistic relationships: syntax, semantics, spatial proximity, and coreference resolution.",
             "graph_json": _multi_head_attention()},
            {"id": "architecture", "heading": "Transformer Architecture", "type": "list",
             "entries": [
                 {"title": "Encoder", "detail": "Processes the input sequence. Stack of layers, each with: Multi-Head Attention → Add & Norm → Feed-Forward → Add & Norm."},
                 {"title": "Decoder", "detail": "Generates output autoregressively. Same as encoder but with masked attention (can't peek at future tokens) and cross-attention to encoder."},
                 {"title": "Add & Norm", "detail": "Residual connection (skip) + Layer Normalization after each sub-layer. Enables training very deep models."},
                 {"title": "Feed-Forward Network", "detail": "Simple 2-layer MLP applied to each position independently. Typically 4× the hidden dimension."},
                 {"title": "Positional Encoding", "detail": "Injects position information since attention is order-agnostic. Fixed sinusoidal or learned embeddings."},
             ]},
            {"id": "comparison", "heading": "Transformer vs RNN", "type": "graph",
             "description": "Transformers consistently outperform RNNs across NLP tasks. Key advantages: parallel training, better long-range dependencies, and scalability to billions of parameters.",
             "graph_json": _transformer_vs_rnn()},
            {"id": "models", "heading": "Famous Transformer Models", "type": "list",
             "entries": [
                 {"title": "BERT (2018)", "detail": "Encoder-only. Bidirectional understanding. Fine-tune for classification, QA, NER. 110M-340M parameters."},
                 {"title": "GPT-2/3/4 (2019-2023)", "detail": "Decoder-only. Autoregressive text generation. The 'G' stands for Generative. GPT-4: ~1.7T parameters."},
                 {"title": "T5 (2019)", "detail": "Encoder-decoder. 'Text-to-Text' framework treats every NLP task as text generation."},
                 {"title": "Vision Transformer (ViT)", "detail": "Applies transformers to images by splitting them into patches. Competitive with CNNs on ImageNet."},
                 {"title": "Whisper", "detail": "Transformer for speech recognition. Trained on 680K hours of audio. Multilingual and robust."},
             ]},
            {"id": "implementation", "heading": "Python Implementation", "type": "code", "language": "python",
             "content": ("import torch\nimport torch.nn as nn\nimport math\n\n"
                         "class SelfAttention(nn.Module):\n    def __init__(self, d_model, n_heads):\n"
                         "        super().__init__()\n        self.n_heads = n_heads\n"
                         "        self.d_k = d_model // n_heads\n"
                         "        self.W_q = nn.Linear(d_model, d_model)\n"
                         "        self.W_k = nn.Linear(d_model, d_model)\n"
                         "        self.W_v = nn.Linear(d_model, d_model)\n"
                         "        self.W_o = nn.Linear(d_model, d_model)\n\n"
                         "    def forward(self, x):\n"
                         "        B, T, C = x.shape\n"
                         "        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)\n"
                         "        k = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)\n"
                         "        v = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)\n\n"
                         "        # Scaled dot-product attention\n"
                         "        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)\n"
                         "        attn = torch.softmax(attn, dim=-1)\n"
                         "        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)\n"
                         "        return self.W_o(out)\n\n"
                         "# Usage\nattn = SelfAttention(d_model=256, n_heads=8)\n"
                         "x = torch.randn(2, 10, 256)  # batch=2, seq_len=10, dim=256\n"
                         "print(f'Output shape: {attn(x).shape}')  # [2, 10, 256]\n")},
        ],
    }
