def draw_graph(
    edge_index,
    pos,
    nodes,
    color,
    title="",
    size=320,
    width=380,
    height=280,
    show_ids=True,
):
    """Draw a graph with Altair.

    edge_index: [2, E] tensor (directed pairs; each undirected edge is drawn once)
    pos:        {node_id: (x, y)} layout, e.g. from networkx.spring_layout
    nodes:      DataFrame with a `node` column plus anything used for color/tooltip
    color:      an Altair color encoding (string shorthand or alt.Color)
    """
    pairs = np.unique(np.sort(edge_index.cpu().numpy().T, axis=1), axis=0)
    pairs = pairs[
        pairs[:, 0] != pairs[:, 1]
    ]  # drop self-loops from the drawing
    edges = pd.DataFrame(
        {
            "x": [pos[s][0] for s, _ in pairs],
            "y": [pos[s][1] for s, _ in pairs],
            "x2": [pos[t][0] for _, t in pairs],
            "y2": [pos[t][1] for _, t in pairs],
        }
    )
    nodes = nodes.assign(
        x=[pos[n][0] for n in nodes["node"]],
        y=[pos[n][1] for n in nodes["node"]],
    )
    x_enc = alt.X(
        "x:Q", axis=None, scale=alt.Scale(zero=False, padding=20)
    )
    y_enc = alt.Y(
        "y:Q", axis=None, scale=alt.Scale(zero=False, padding=20)
    )
    layers = [
        alt.Chart(edges)
        .mark_rule(color="gray", opacity=0.45)
        .encode(x=x_enc, y=y_enc, x2="x2:Q", y2="y2:Q"),
        alt.Chart(nodes)
        .mark_circle(size=size, opacity=1, stroke="white", strokeWidth=1.5)
        .encode(
            x=x_enc,
            y=y_enc,
            color=color,
            tooltip=[c for c in nodes.columns if c not in ("x", "y")],
        ),
    ]
    if show_ids:
        layers.append(
            alt.Chart(nodes)
            .mark_text(fontSize=11, fontWeight="bold", color="white")
            .encode(x=x_enc, y=y_enc, text="node:N")
        )
    return alt.layer(*layers).properties(
        title=title, width=width, height=height
    )

mo.vstack(
    [
        mo.md("### PyG's building blocks"),
        mo.mermaid(
            """
            graph LR
                A["datasets<br/>Planetoid, TUDataset, OGB..."] --> B["transforms<br/>NormalizeFeatures, ToUndirected..."]
                B --> C["data.Data / HeteroData<br/>x, edge_index, y, masks"]
                C --> D["loader<br/>DataLoader, NeighborLoader"]
                D --> E["nn<br/>GCNConv, GATConv, SAGEConv...<br/>(all built on MessagePassing)"]
                E --> F["loss + optimizer<br/>plain PyTorch"]
                C -. helpers .- G["utils<br/>degree, k_hop_subgraph,<br/>to_networkx, scatter..."]
            """
        ),
        mo.md(
            r"""
            | module | what it gives you | used in this notebook |
            |---|---|---|
            | `torch_geometric.data` | `Data` (one graph), `HeteroData` (several node/edge types), `Batch` | toy graph, Cora |
            | `torch_geometric.datasets` | 100+ benchmark datasets that download themselves | `Planetoid("Cora")` |
            | `torch_geometric.transforms` | preprocessing applied when a graph is loaded | `NormalizeFeatures()` |
            | `torch_geometric.nn` | 60+ conv layers, pooling, normalization, full models | `GCNConv`, `GATConv`, `SAGEConv`, `MessagePassing` |
            | `torch_geometric.loader` | mini-batching many small graphs, or sampling subgraphs of one huge graph | (explained in §3) |
            | `torch_geometric.utils` | graph helpers: degrees, self-loops, subgraphs, conversions | `degree`, `k_hop_subgraph`, `homophily` |
            | `torch_geometric.explain` | explain predictions (GNNExplainer, ...) | (next steps) |

            Every conv layer follows the same **message passing** template, which you'll implement yourself in §1:

            $$
            \mathbf{x}_i' \;=\; \gamma\Big(\mathbf{x}_i,\; \bigoplus_{j \in \mathcal{N}(i)} \phi\big(\mathbf{x}_i, \mathbf{x}_j, \mathbf{e}_{ji}\big)\Big)
            $$

            $\phi$ = **message** function, $\bigoplus$ = **aggregation** (sum / mean / max), $\gamma$ = **update** function.
            """
        ),
    ]
)

_node = mo.md(
    r"""
    **Predict a label for each node.** *Example:* the topic of each paper in a citation network, or whether a
    social-media account is a bot. Usually *semi-supervised*: only some nodes are labelled.

    ```python
    out = model(data.x, data.edge_index)          # [num_nodes, num_classes]
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    ```
    **This notebook does node classification on Cora.**
    """
)
_link = mo.md(
    r"""
    **Predict whether an edge exists (or will).** *Example:* friend recommendation, drug–target interaction,
    "users who bought X also bought Y".

    ```python
    z = model.encode(data.x, data.edge_index)      # node embeddings
    score = (z[src] * z[dst]).sum(dim=-1)          # dot product per candidate edge
    loss = F.binary_cross_entropy_with_logits(score, edge_label)
    ```
    Helpers: `transforms.RandomLinkSplit`, `utils.negative_sampling`.
    """
)
_graph = mo.md(
    r"""
    **Predict one label for a whole graph.** *Example:* is this molecule toxic? Which community does this
    sub-network belong to?

    ```python
    from torch_geometric.loader import DataLoader
    from torch_geometric.nn import global_mean_pool
    for batch in DataLoader(dataset, batch_size=32):   # many small graphs -> one big disconnected graph
        h = model(batch.x, batch.edge_index)
        g = global_mean_pool(h, batch.batch)            # [num_graphs, hidden]
    ```
    Datasets: `TUDataset("MUTAG")`, `MoleculeNet`, OGB.
    """
)
mo.vstack(
    [
        mo.md("### Three kinds of graph tasks"),
        mo.ui.tabs(
            {
                "Node classification": _node,
                "Link prediction": _link,
                "Graph classification": _graph,
            }
        ),
    ]
)

toy_edges_text = mo.ui.text_area(
    value="0-1, 1-2, 2-0, 2-3, 3-4, 4-5, 5-3",
    label="Undirected edges (`a-b`, separated by commas or new lines)",
    full_width=True,
    rows=2,
)
toy_edges_text

def parse_edges(text: str) -> list[tuple[int, int]]:
    pairs = []
    for token in text.replace("\n", ",").split(","):
        if "-" in token:
            a, b = token.split("-", 1)
            pairs.append((int(a.strip()), int(b.strip())))
    return pairs

try:
    _pairs = parse_edges(toy_edges_text.value)
except ValueError:
    _pairs = []
mo.stop(
    not _pairs,
    mo.callout(
        mo.md("Couldn't parse any edges. Use the form `0-1, 1-2`."),
        kind="danger",
    ),
)

# Store each undirected edge in both directions
toy_edge_index = torch.tensor(
    [
        [a for a, b in _pairs] + [b for a, b in _pairs],
        [b for a, b in _pairs] + [a for a, b in _pairs],
    ]
)
_n = int(toy_edge_index.max()) + 1
# Small integer features so the arithmetic below is easy to follow
toy_x = torch.randint(
    0, 6, (_n, 2), generator=torch.Generator().manual_seed(1)
).float()
toy_graph = Data(x=toy_x, edge_index=toy_edge_index).to(device)
toy_pos = nx.spring_layout(nx.Graph(_pairs), seed=7)
for _i in range(
    _n
):  # isolated node ids (gaps in numbering) still need a position
    toy_pos.setdefault(_i, (0.0, 0.0))

_deg = degree(toy_graph.edge_index[1], num_nodes=_n).int().cpu()
_nodes = pd.DataFrame(
    {
        "node": range(_n),
        "degree": _deg.numpy(),
        "x": [str(v.tolist()) for v in toy_x],
    }
)

mo.hstack(
    [
        draw_graph(
            toy_graph.edge_index,
            toy_pos,
            _nodes.rename(columns={"x": "features"}),
            color=alt.Color(
                "degree:O",
                scale=alt.Scale(scheme="blues"),
                legend=alt.Legend(title="degree"),
            ),
            title="Your graph (color = degree)",
        ),
        mo.md(
            f"""
            `{toy_graph}`

            **`x`** (features, one row per node):
            ```
            {toy_x.int().tolist()}
            ```
            **`edge_index`** (`{toy_graph.num_edges}` directed edges = 2 × {len(_pairs)}):
            ```
            sources: {toy_edge_index[0].tolist()}
            targets: {toy_edge_index[1].tolist()}
            ```
            **degree**: `{_deg.tolist()}` · undirected: **{toy_graph.is_undirected()}** ·
            isolated nodes: **{toy_graph.has_isolated_nodes()}** · device: `{toy_graph.x.device}`
            """
        ),
    ],
    widths=[1, 1],
    align="start",
)

mp_node = mo.ui.dropdown(
    {str(i): i for i in range(toy_graph.num_nodes)},
    value="2",
    label="target node",
)
mp_aggr = mo.ui.radio(
    ["sum", "mean", "max"], value="mean", label="aggregation", inline=True
)
mp_self_loops = mo.ui.switch(value=False, label="add self-loops")
mo.hstack([mp_node, mp_aggr, mp_self_loops], justify="start", gap=2)

_ei = toy_graph.edge_index.cpu()
if mp_self_loops.value:
    _ei, _ = add_self_loops(_ei, num_nodes=toy_graph.num_nodes)
_x = toy_x
_i = mp_node.value
_src = _ei[0][_ei[1] == _i]  # neighbours j that send a message j -> i

_msgs = _x[_src]
if len(_src) == 0:
    _agg = torch.zeros(_x.size(1))
elif mp_aggr.value == "sum":
    _agg = _msgs.sum(0)
elif mp_aggr.value == "mean":
    _agg = _msgs.mean(0)
else:
    _agg = _msgs.max(0).values

# The same thing for ALL nodes at once, vectorised: this is what PyG does on the GPU
_all = scatter(
    _x[_ei[0]],
    _ei[1],
    dim=0,
    dim_size=toy_graph.num_nodes,
    reduce=mp_aggr.value,
)

_role = [
    "target"
    if n == _i
    else ("neighbour" if n in _src.tolist() else "other")
    for n in range(toy_graph.num_nodes)
]
_nodes = pd.DataFrame(
    {
        "node": range(toy_graph.num_nodes),
        "role": _role,
        "features": [str(v.int().tolist()) for v in _x],
    }
)
_chart = draw_graph(
    toy_graph.edge_index,
    toy_pos,
    _nodes,
    color=alt.Color(
        "role:N",
        scale=alt.Scale(
            domain=["target", "neighbour", "other"],
            range=["#e45756", "#4c78a8", "#bab0ac"],
        ),
    ),
    title=f"Messages flowing into node {_i}",
)
_table = pd.DataFrame(
    {
        "from node j": _src.tolist(),
        "message x_j": [str(m.int().tolist()) for m in _msgs],
    }
)

mo.hstack(
    [
        _chart,
        mo.vstack(
            [
                mo.md(f"**Incoming messages** ({len(_src)}):"),
                mo.ui.table(
                    _table,
                    selection=None,
                    show_column_summaries=False,
                    pagination=False,
                ),
                mo.md(
                    rf"""
                    **{mp_aggr.value}** of messages → $\mathbf{{x}}_{_i}' =$ `{[round(v, 2) for v in _agg.tolist()]}`

                    Vectorised for every node with `torch_geometric.utils.scatter(x[src], dst, reduce="{mp_aggr.value}")`:
                    ```
                    {[[round(v, 2) for v in row] for row in _all.tolist()]}
                    ```
                    Row {_i} matches the hand computation. One `scatter` call replaces a Python loop over nodes.
                    """
                ),
            ]
        ),
    ],
    widths=[1, 1],
    align="start",
)

_n = toy_graph.num_nodes
_A = to_dense_adj(toy_graph.edge_index.cpu(), max_num_nodes=_n)[0]
_A_tilde = _A + torch.eye(_n)
_d = _A_tilde.sum(1)
_A_hat = _A_tilde / torch.sqrt(_d[:, None] * _d[None, :])

def _heatmap(mat: torch.Tensor, title: str, fmt: str):
    _df = pd.DataFrame(
        [(i, j, float(mat[i, j])) for i in range(_n) for j in range(_n)],
        columns=["row i", "col j", "value"],
    )
    _base = alt.Chart(_df).encode(x=alt.X("col j:O"), y=alt.Y("row i:O"))
    return (
        _base.mark_rect().encode(
            color=alt.Color(
                "value:Q", scale=alt.Scale(scheme="blues"), legend=None
            ),
            tooltip=["row i", "col j", alt.Tooltip("value:Q", format=fmt)],
        )
        + _base.mark_text(fontSize=10).encode(
            text=alt.Text("value:Q", format=fmt),
            color=alt.condition(
                "datum.value > 0.6", alt.value("white"), alt.value("black")
            ),
        )
    ).properties(
        width=min(40 * _n, 320), height=min(40 * _n, 320), title=title
    )

mo.ui.tabs(
    {
        "A (adjacency)": _heatmap(_A, "A: 1 if edge i–j", ".0f"),
        "A + I (self-loops)": _heatmap(_A_tilde, "A + I", ".0f"),
        "Â (GCN-normalized)": _heatmap(
            _A_hat, "Â = D^-1/2 (A+I) D^-1/2", ".2f"
        ),
    }
)

torch.manual_seed(0)
_ref = GCNConv(dataset.num_features, 16).to(device)
_mine = MyGCNConv(dataset.num_features, 16).to(device)
with (
    torch.no_grad()
):  # copy the weights so both layers compute the same function
    _mine.lin.weight.copy_(_ref.lin.weight)
    _mine.bias.copy_(_ref.bias)
    _out_ref = _ref(graph.x, graph.edge_index)
    _out_mine = _mine(graph.x, graph.edge_index)
_max_diff = (_out_ref - _out_mine).abs().max().item()

mo.callout(
    mo.md(
        f"""
        **MyGCNConv vs PyG GCNConv on Cora** (loaded in §2: {graph.num_nodes:,} nodes, on `{device}`):
        output shape `{tuple(_out_mine.shape)}`, max absolute difference **{_max_diff:.2e}**,
        identical: **{torch.allclose(_out_ref, _out_mine, atol=1e-5)}**
        """
    ),
    kind="success" if _max_diff < 1e-5 else "danger",
)

_counts = torch.bincount(
    graph.y.cpu(), minlength=dataset.num_classes
).numpy()
alt.Chart(
    pd.DataFrame({"topic": CLASS_NAMES, "papers": _counts})
).mark_bar().encode(
    x=alt.X("papers:Q", title="papers"),
    y=alt.Y("topic:N", sort="-x", title=None),
    tooltip=["topic", "papers"],
).properties(title="Class balance in Cora", width=480, height=200)

_split = pd.DataFrame(
    {
        "split": [
            "train",
            "validation",
            "test",
            "unlabelled (still used for message passing!)",
        ],
        "nodes": [
            int(graph.train_mask.sum()),
            int(graph.val_mask.sum()),
            int(graph.test_mask.sum()),
            int(
                (
                    ~(graph.train_mask | graph.val_mask | graph.test_mask)
                ).sum()
            ),
        ],
        "used for": [
            "computing the loss & gradients",
            "choosing the best epoch / hyper-parameters",
            "final, untouched evaluation",
            "only their features & edges: this is the *transductive* setting",
        ],
    }
)
_row = graph.x[0].cpu()
mo.accordion(
    {
        "Train / validation / test masks": mo.vstack(
            [
                mo.md(
                    "PyG stores splits as **boolean masks** over nodes (`graph.train_mask`, ...). "
                    "The whole graph goes through the model every time; the mask only decides *which "
                    "nodes' predictions count* for the loss or the metric."
                ),
                mo.ui.table(
                    _split,
                    selection=None,
                    show_column_summaries=False,
                    pagination=False,
                ),
            ]
        ),
        "What does a feature vector look like?": mo.md(
            f"""
            Paper 0 contains **{int((_row > 0).sum())}** of the {dataset.num_features:,} dictionary words.
            After `NormalizeFeatures` each non-zero entry is `1 / {int((_row > 0).sum())}` ≈ **{_row.max():.4f}**
            and the row sums to **{_row.sum():.2f}**.
            Overall `x` is **{(graph.x == 0).float().mean().item():.1%} zeros**: very sparse, like most text features.
            """
        ),
    },
    multiple=True,
)

_h_edge = homophily(graph.edge_index, graph.y, method="edge")
_h_node = homophily(graph.edge_index, graph.y, method="node")
_deg = degree(graph.edge_index[1], graph.num_nodes).cpu()
_deg_df = pd.DataFrame({"degree": _deg.clamp(max=20).int().numpy()})

mo.vstack(
    [
        mo.md(
            rf"""
            ### Why do GNNs work here? Homophily

            *"Birds of a feather flock together"*: papers mostly cite papers **on the same topic**.
            **Edge homophily** is the fraction of edges joining two nodes of the same class:

            $$h = \frac{{|\{{(i,j) \in E : y_i = y_j\}}|}}{{|E|}}$$

            For Cora **h = {_h_edge:.2f}** (node-averaged: {_h_node:.2f}). With 7 classes, random edges would give
            about 0.14. So a neighbour's topic is a strong hint about a node's own topic, and message passing exploits
            exactly that. On *heterophilous* graphs (e.g. fraudsters connecting to normal users), plain GCN struggles.
            """
        ),
        alt.Chart(_deg_df)
        .mark_bar()
        .encode(
            x=alt.X("degree:O", title="node degree (20 = 20 or more)"),
            y=alt.Y("count():Q", title="papers"),
        )
        .properties(
            width=520,
            height=160,
            title=f"Degree distribution: median {int(_deg.median())}, max {int(_deg.max())}, a long tail of hubs",
        ),
    ]
)

ego_node = mo.ui.number(
    start=0, stop=graph.num_nodes - 1, value=1701, label="paper id"
)
ego_hops = mo.ui.slider(
    1, 3, value=2, label="hops (= GNN layers)", show_value=True
)
mo.hstack([ego_node, ego_hops], justify="start", gap=2)

_center = int(ego_node.value)
_subset, _sub_ei, _mapping, _ = k_hop_subgraph(
    _center,
    ego_hops.value,
    graph.edge_index.cpu(),
    relabel_nodes=True,
    num_nodes=graph.num_nodes,
)
_y = graph.y.cpu()
_n_sub = _subset.numel()
_same = (_y[_subset] == _y[_center]).float().mean().item()
_max_draw = 350

_stats = mo.vstack(
    [
        mo.stat(
            f"{_n_sub:,}",
            label=f"papers within {ego_hops.value} hop(s)",
            bordered=True,
        ),
        mo.stat(
            f"{_n_sub / graph.num_nodes:.1%}",
            label="of the whole graph",
            bordered=True,
        ),
        mo.stat(
            f"{_same:.0%}",
            label=f"share the topic of #{_center}",
            bordered=True,
        ),
        mo.md(
            f"True topic of #{_center}: **{CLASS_NAMES[int(_y[_center])]}**"
        ),
    ]
)

if _n_sub > _max_draw:
    _view = mo.callout(
        mo.md(
            f"{_n_sub:,} nodes is too many to draw nicely (limit {_max_draw}). Try fewer hops or another paper. "
            "This explosion of the receptive field is why large graphs use **neighbour sampling** (§3)."
        ),
        kind="info",
    )
else:
    _g = nx.Graph(_sub_ei.t().tolist())
    _g.add_nodes_from(range(_n_sub))
    _pos = nx.spring_layout(_g, seed=0, k=1.5 / max(_n_sub, 1) ** 0.5)
    _hop = nx.single_source_shortest_path_length(_g, int(_mapping))
    _nodes = pd.DataFrame(
        {
            "node": range(_n_sub),
            "paper": _subset.tolist(),
            "topic": [CLASS_NAMES[int(c)] for c in _y[_subset]],
            "hop": [_hop.get(i, -1) for i in range(_n_sub)],
        }
    )
    _view = draw_graph(
        _sub_ei,
        _pos,
        _nodes,
        color=alt.Color("topic:N", scale=alt.Scale(domain=CLASS_NAMES)),
        size=[0, 260, 90, 40][ego_hops.value],
        show_ids=False,
        width=460,
        height=360,
        title=f"{ego_hops.value}-hop neighbourhood of paper #{_center} (big = closer)",
    )
    _center_xy = pd.DataFrame(
        {"x": [_pos[int(_mapping)][0]], "y": [_pos[int(_mapping)][1]]}
    )
    _view = _view + alt.Chart(_center_xy).mark_point(
        shape="diamond",
        size=500,
        filled=False,
        color="black",
        strokeWidth=2,
    ).encode(x="x:Q", y="y:Q")

mo.hstack([_view, _stats], widths=[3, 1], align="start")

_gcn = mo.md(
    r"""
    #### GCN: Graph Convolutional Network (Kipf & Welling, ICLR 2017)

    $$
    \mathbf{h}_i' = \sigma\Big( \sum_{j \in \mathcal{N}(i) \cup \{i\}} \frac{1}{\sqrt{\tilde d_i \tilde d_j}}\; W \mathbf{h}_j \Big)
    \qquad\text{(matrix form: } H' = \sigma(\hat{A} H W)\text{)}
    $$

    **Intuition:** a degree-normalized *average* of yourself and your neighbours, then a linear layer.
    It came from approximating *spectral graph convolutions* (filters on the graph Laplacian's eigenvectors)
    with a cheap first-order polynomial, which is why it's called a "convolution".

    - ✅ simple, fast, few parameters, a very strong baseline on homophilous graphs
    - ❌ all neighbours get **fixed** weights (from degrees only), so it can't learn that some citations matter more
    - ❌ transductive in its classic form; deep stacks **over-smooth** (§6)

    You implemented exactly this layer in §1 (`MyGCNConv`).
    """
)
_gat = mo.md(
    r"""
    #### GAT: Graph Attention Network (Veličković et al., ICLR 2018)

    Instead of fixed weights, GAT **learns** how much node $i$ should attend to each neighbour $j$:

    $$
    e_{ij} = \text{LeakyReLU}\big(\mathbf{a}^\top [\,W\mathbf{h}_i \,\|\, W\mathbf{h}_j\,]\big), \qquad
    \alpha_{ij} = \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}(i) \cup \{i\}} \exp(e_{ik})}, \qquad
    \mathbf{h}_i' = \sigma\Big(\sum_{j} \alpha_{ij}\, W \mathbf{h}_j\Big)
    $$

    The attention scores $\alpha_{ij}$ are a **softmax over each node's neighbourhood**, so they sum to 1.
    **Multi-head attention:** run $K$ independent attention mechanisms and concatenate their outputs, as in
    Transformers. This stabilises training. Here hidden size = heads × per-head size.

    - ✅ adaptive neighbour weights, often the most accurate of the three on citation graphs
    - ✅ attention weights are inspectable (see §5!)
    - ❌ more parameters and compute per edge; dropout on attention is important against overfitting

    *Fun fact:* a Transformer is a GAT on a fully-connected graph.
    """
)
_sage = mo.md(
    r"""
    #### GraphSAGE: SAmple and aggreGatE (Hamilton et al., NeurIPS 2017)

    $$
    \mathbf{h}_i' = \sigma\Big( W_1 \mathbf{h}_i \;+\; W_2 \cdot \text{mean}_{j \in \mathcal{N}(i)} \mathbf{h}_j \Big)
    $$

    **Intuition:** keep **separate weights for "me" and "my neighbourhood"** instead of mixing yourself in
    as just another neighbour. The original paper also introduced **neighbour sampling**: use only a fixed number
    of random neighbours per node and layer, so cost no longer depends on hub degrees.

    - ✅ **inductive**: learns a function of features + neighbourhood, so it works on nodes/graphs never seen in training
    - ✅ scales to graphs with billions of edges (Pinterest's PinSage) via mini-batches of sampled subgraphs
    - ❌ on small homophilous graphs like Cora its separate self-weight makes it overfit a bit more than GCN

    Other aggregators (`aggr="max"`, `"lstm"`, ...) are one argument away in `SAGEConv`.
    """
)
_mlp = mo.md(
    r"""
    #### MLP: the "no graph" baseline

    $$
    \mathbf{h}_i' = \sigma(W \mathbf{h}_i)
    $$

    A plain 2-layer neural network on each paper's words **alone**: edges are ignored. It's the honest baseline:
    if a GNN can't beat it, the graph isn't helping. On Cora expect around **55–60%** for the MLP versus
    around **80%** for GNNs: 140 labelled bag-of-words vectors just aren't enough, but citations spread the label
    information to the whole graph.
    """
)
_cmp = mo.md(
    r"""
    | | **MLP** | **GCN** | **GAT** | **GraphSAGE** |
    |---|---|---|---|---|
    | uses edges | ✗ | ✓ | ✓ | ✓ |
    | neighbour weights | — | fixed: $1/\sqrt{d_i d_j}$ | **learned** attention $\alpha_{ij}$ | mean (+ separate self weight) |
    | PyG layer | `torch.nn.Linear` | `GCNConv` | `GATConv` | `SAGEConv` |
    | inductive | ✓ | partly | ✓ | ✓ (designed for it) |
    | cost per layer | $O(N)$ | $O(E)$ | $O(E \cdot K)$ heads | $O(E)$, or sampled |
    | typical Cora test acc. | ~57% | ~81% | ~82–83% | ~79–80% |

    Test these claims yourself in §4 with **Compare all models**.
    """
)
mo.ui.tabs(
    {
        "GCN": _gcn,
        "GAT": _gat,
        "GraphSAGE": _sage,
        "MLP baseline": _mlp,
        "Side by side": _cmp,
    }
)

mo.vstack(
    [
        mo.md("### Training techniques used in this notebook"),
        mo.accordion(
            {
                "Cross-entropy loss on the training mask only": mo.md(
                    r"""
                    The model outputs one score (*logit*) per class for **every** node. `F.cross_entropy` applies a softmax
                    and penalises $-\log p(\text{true class})$, but we index with `train_mask` so only the 140 labelled
                    papers contribute. The rest still pass messages, so the gradients flow through their features too.
                    """
                ),
                "Full-batch training": mo.md(
                    r"""
                    Cora is small enough that the **entire graph fits on the GPU**, so each epoch = one forward/backward pass
                    over all 2,708 nodes. That is why one "epoch" is a single optimizer step here.
                    """
                ),
                "Dropout (on features and on attention)": mo.md(
                    r"""
                    During training, randomly zero a fraction $p$ of the inputs and hidden units (scaling the rest by $1/(1-p)$).
                    This forces redundancy and is the main defence against overfitting with only 140 labels. GAT also drops
                    **attention coefficients**, i.e. randomly ignores some edges each step. `model.eval()` switches dropout off.
                    """
                ),
                "Weight decay (L2 regularization) + Adam": mo.md(
                    r"""
                    **Adam** adapts the learning rate per parameter using running averages of the gradient and its square, so
                    it's robust to the choice of `lr`. **Weight decay** $\lambda$ adds $\lambda \|W\|^2$ to the loss,
                    pulling weights towards zero. The Kipf & Welling defaults, `lr=0.01, weight_decay=5e-4`, are the
                    defaults below.
                    """
                ),
                "Model selection on the validation set (early-stopping style)": mo.md(
                    r"""
                    Every epoch we measure validation accuracy and **keep a copy of the weights from the best epoch**.
                    After training those weights are restored. The test set is never used to make choices, so the test
                    accuracy stays an honest estimate.
                    """
                ),
                "Feature normalization": mo.md(
                    r"""
                    `NormalizeFeatures()` makes each paper's word vector sum to 1, so papers with many words don't
                    produce larger activations than short ones.
                    """
                ),
                "Scaling up: mini-batches & neighbour sampling": mo.md(
                    r"""
                    For graphs that don't fit in memory, PyG's `NeighborLoader` samples e.g. 10 neighbours at hop 1 and 5
                    at hop 2 for a batch of seed nodes, and builds a small subgraph per batch (GraphSAGE-style):

                    ```python
                    from torch_geometric.loader import NeighborLoader
                    loader = NeighborLoader(data, num_neighbors=[10, 5], batch_size=512, input_nodes=data.train_mask)
                    for batch in loader:
                        out = model(batch.x, batch.edge_index)[: batch.batch_size]   # seed nodes come first
                    ```
                    *Note:* fast sampling uses the optional `pyg-lib` / `torch-sparse` CPU extensions; the sampled
                    batches can still be trained on the XPU.
                    """
                ),
            },
            multiple=True,
        ),
    ]
)

model_kind = mo.ui.dropdown(
    ["MLP", "GCN", "GAT", "GraphSAGE"], value="GCN", label="model"
)
hidden_dim = mo.ui.slider(
    16, 256, step=16, value=64, label="hidden size", show_value=True
)
gat_heads = mo.ui.dropdown(
    {"1": 1, "2": 2, "4": 4, "8": 8}, value="8", label="GAT heads"
)
learning_rate = mo.ui.dropdown(
    ["0.001", "0.005", "0.01", "0.05"], value="0.01", label="learning rate"
)
weight_decay = mo.ui.dropdown(
    ["0", "5e-4", "5e-3"], value="5e-4", label="weight decay"
)
dropout_rate = mo.ui.slider(
    0.0, 0.8, step=0.1, value=0.5, label="dropout", show_value=True
)
num_epochs = mo.ui.slider(
    50, 500, step=50, value=200, label="epochs", show_value=True
)
train_button = mo.ui.run_button(label="Train model", kind="success")

mo.vstack(
    [
        mo.md("### Training controls"),
        mo.hstack(
            [
                model_kind,
                gat_heads,
                learning_rate,
                weight_decay,
                train_button,
            ],
            justify="start",
            gap=1.5,
            wrap=True,
        ),
        mo.hstack(
            [hidden_dim, dropout_rate, num_epochs],
            justify="start",
            gap=1.5,
            wrap=True,
        ),
    ]
)

def evaluate(m: torch.nn.Module) -> dict[str, float]:
    """Accuracy on the train / val / test masks (dropout off, no gradients)."""
    m.eval()
    with torch.no_grad():
        pred = m(graph.x, graph.edge_index)[0].argmax(dim=1)
    return {
        split: (pred[mask] == graph.y[mask]).float().mean().item()
        for split, mask in (
            ("train", graph.train_mask),
            ("val", graph.val_mask),
            ("test", graph.test_mask),
        )
    }

def fit(
    m: torch.nn.Module,
    epochs: int,
    lr: float,
    wd: float,
    title: str | None = None,
) -> pd.DataFrame:
    """Full-batch training on graph.train_mask. Restores the best-validation weights; returns per-epoch history."""
    optimizer = torch.optim.Adam(m.parameters(), lr=lr, weight_decay=wd)
    best_val, best_state, rows = -1.0, None, []
    epoch_iter = range(1, epochs + 1)
    if title is not None:
        epoch_iter = mo.status.progress_bar(epoch_iter, title=title)
    for epoch in epoch_iter:
        m.train()
        optimizer.zero_grad()
        logits, _ = m(graph.x, graph.edge_index)
        loss = F.cross_entropy(
            logits[graph.train_mask], graph.y[graph.train_mask]
        )
        loss.backward()
        optimizer.step()

        acc = evaluate(m)
        rows.append({"epoch": epoch, "loss": loss.item(), **acc})
        if acc["val"] > best_val:
            best_val = acc["val"]
            best_state = {
                k: v.detach().clone() for k, v in m.state_dict().items()
            }
    m.load_state_dict(best_state)
    return pd.DataFrame(rows)

mo.stop(
    not train_button.value,
    mo.md("*Pick settings above and click **Train model**.*"),
)

torch.manual_seed(0)
model = GNN(
    model_kind.value,
    dataset.num_features,
    hidden_dim.value,
    dataset.num_classes,
    dropout=dropout_rate.value,
    heads=gat_heads.value,
).to(device)

_start = time.perf_counter()
history = fit(
    model,
    num_epochs.value,
    float(learning_rate.value),
    float(weight_decay.value),
    title=f"Training {model_kind.value} on {device}",
)
if device.type == "xpu":
    torch.xpu.synchronize()
train_seconds = time.perf_counter() - _start
final_acc = evaluate(model)

mo.hstack(
    [
        mo.stat(
            f"{final_acc['test']:.1%}",
            label=f"{model_kind.value} test accuracy",
            bordered=True,
        ),
        mo.stat(
            f"{final_acc['val']:.1%}",
            label="validation accuracy (best epoch)",
            bordered=True,
        ),
        mo.stat(
            f"{int(history['val'].idxmax()) + 1}",
            label="best epoch",
            bordered=True,
        ),
        mo.stat(
            f"{train_seconds:.2f}s",
            label=f"{num_epochs.value} epochs on {device}",
            bordered=True,
        ),
        mo.stat(
            f"{sum(p.numel() for p in model.parameters()):,}",
            label="parameters",
            bordered=True,
        ),
    ],
    justify="start",
    wrap=True,
)

_loss_chart = (
    alt.Chart(history)
    .mark_line()
    .encode(x="epoch:Q", y=alt.Y("loss:Q", title="training loss"))
    .properties(width=320, height=220, title="Loss")
)

_acc_long = history.melt(
    id_vars="epoch",
    value_vars=["train", "val", "test"],
    var_name="split",
    value_name="accuracy",
)
_acc_chart = (
    alt.Chart(_acc_long)
    .mark_line()
    .encode(
        x="epoch:Q",
        y=alt.Y(
            "accuracy:Q",
            scale=alt.Scale(domain=[0, 1]),
            axis=alt.Axis(format="%"),
        ),
        color=alt.Color("split:N", sort=["train", "val", "test"]),
        tooltip=[
            "epoch",
            "split",
            alt.Tooltip("accuracy:Q", format=".1%"),
        ],
    )
    .properties(width=320, height=220, title="Accuracy")
)

mo.vstack(
    [
        _loss_chart | _acc_chart,
        mo.md(
            "Train accuracy shoots to ~100% because only 140 nodes are labelled; the **val/test** curves are what "
            "matter. The gap between them is overfitting, which dropout and weight decay keep in check."
        ),
    ]
)

compare_seeds = mo.ui.slider(
    1, 5, value=3, label="seeds per model", show_value=True
)
compare_button = mo.ui.run_button(
    label="Compare all models", kind="neutral"
)
mo.hstack([compare_seeds, compare_button], justify="start", gap=2)

mo.stop(
    not compare_button.value,
    mo.md("*Click **Compare all models** to run the experiment.*"),
)

_kinds = ["MLP", "GCN", "GAT", "GraphSAGE"]
_runs = []
for _kind, _seed in mo.status.progress_bar(
    [(k, s) for k in _kinds for s in range(compare_seeds.value)],
    title="Training every model × seed",
):
    torch.manual_seed(_seed)
    _m = GNN(
        _kind,
        dataset.num_features,
        hidden_dim.value,
        dataset.num_classes,
        dropout=dropout_rate.value,
        heads=gat_heads.value,
    ).to(device)
    _t0 = time.perf_counter()
    fit(
        _m,
        num_epochs.value,
        float(learning_rate.value),
        float(weight_decay.value),
    )
    _runs.append(
        {
            "model": _kind,
            "seed": _seed,
            "test accuracy": evaluate(_m)["test"],
            "seconds": time.perf_counter() - _t0,
        }
    )
comparison_df = pd.DataFrame(_runs)

_summary = (
    comparison_df.groupby("model", sort=False)
    .agg(
        mean=("test accuracy", "mean"),
        std=("test accuracy", "std"),
        seconds=("seconds", "mean"),
    )
    .reset_index()
    .fillna(0)
)
_bars = (
    alt.Chart(_summary)
    .mark_bar(opacity=0.8)
    .encode(
        x=alt.X("model:N", sort=_kinds, title=None),
        y=alt.Y(
            "mean:Q",
            title="test accuracy",
            scale=alt.Scale(domain=[0, 1]),
            axis=alt.Axis(format="%"),
        ),
        color=alt.Color("model:N", sort=_kinds, legend=None),
        tooltip=[
            "model",
            alt.Tooltip("mean:Q", format=".1%"),
            alt.Tooltip("std:Q", format=".1%"),
            alt.Tooltip("seconds:Q", format=".1f", title="s / run"),
        ],
    )
)
_err = (
    alt.Chart(_summary)
    .mark_errorbar(ticks=True)
    .encode(
        x=alt.X("model:N", sort=_kinds),
        y=alt.Y("lo:Q", title=""),
        y2="hi:Q",
    )
    .transform_calculate(
        lo="datum.mean - datum.std", hi="datum.mean + datum.std"
    )
)
_pts = (
    alt.Chart(comparison_df)
    .mark_circle(color="black", size=30)
    .encode(
        x=alt.X("model:N", sort=_kinds),
        y="test accuracy:Q",
        tooltip=[
            "model",
            "seed",
            alt.Tooltip("test accuracy:Q", format=".1%"),
        ],
    )
)
_best = _summary.loc[_summary["mean"].idxmax()]
mo.vstack(
    [
        (_bars + _err + _pts).properties(
            width=420,
            height=260,
            title="Test accuracy (bar = mean, line = ±1 std, dots = seeds)",
        ),
        mo.md(
            f"Best: **{_best['model']}** at {_best['mean']:.1%} ± {_best['std']:.1%}. Graph models beat the MLP by "
            f"**{_summary.set_index('model').loc[['GCN', 'GAT', 'GraphSAGE'], 'mean'].mean() - _summary.set_index('model').loc['MLP', 'mean']:.0%} points** "
            "on average: that is the value of the citation graph."
        ),
    ]
)

inspect_node = mo.ui.number(
    start=0, stop=graph.num_nodes - 1, value=1701, label="paper id"
)
inspect_node

model.eval()
with torch.no_grad():
    _logits, _ = model(graph.x, graph.edge_index)
_probs = _logits.softmax(dim=1).cpu()
_pred_all = _probs.argmax(1)
_y = graph.y.cpu()
_i = int(inspect_node.value)

_wrong_test = (
    torch.nonzero((_pred_all != _y) & graph.test_mask.cpu())
    .flatten()[:12]
    .tolist()
)
_split = (
    "train"
    if graph.train_mask[_i]
    else "val"
    if graph.val_mask[_i]
    else "test"
    if graph.test_mask[_i]
    else "unlabelled"
)

_ei = graph.edge_index.cpu()
_nbrs = _ei[0][_ei[1] == _i]
_nbr_df = pd.DataFrame(
    {
        "neighbour": _nbrs.tolist(),
        "true topic": [CLASS_NAMES[int(_y[j])] for j in _nbrs],
        "predicted": [CLASS_NAMES[int(_pred_all[j])] for j in _nbrs],
        "labelled (train)": [bool(graph.train_mask[j]) for j in _nbrs],
    }
)
_vote = _nbr_df["true topic"].mode().iloc[0] if len(_nbr_df) else "n/a"

_prob_df = pd.DataFrame(
    {
        "topic": CLASS_NAMES,
        "probability": _probs[_i].numpy(),
        "is": [
            "true" if c == int(_y[_i]) else ""
            for c in range(len(CLASS_NAMES))
        ],
    }
)
_prob_chart = (
    alt.Chart(_prob_df)
    .mark_bar()
    .encode(
        x=alt.X(
            "probability:Q",
            scale=alt.Scale(domain=[0, 1]),
            axis=alt.Axis(format="%"),
        ),
        y=alt.Y("topic:N", sort=CLASS_NAMES, title=None),
        color=alt.condition(
            "datum.is == 'true'",
            alt.value("#54a24b"),
            alt.value("#9ecae9"),
        ),
        tooltip=["topic", alt.Tooltip("probability:Q", format=".1%")],
    )
    .properties(
        width=300,
        height=190,
        title="Model's probabilities (green = true topic)",
    )
)

_ok = int(_pred_all[_i]) == int(_y[_i])
mo.vstack(
    [
        mo.hstack(
            [
                mo.stat(
                    CLASS_NAMES[int(_y[_i])],
                    label="true topic",
                    bordered=True,
                ),
                mo.stat(
                    CLASS_NAMES[int(_pred_all[_i])],
                    label=f"predicted ({_probs[_i].max():.0%} confident)",
                    bordered=True,
                    direction="increase" if _ok else "decrease",
                    caption="correct" if _ok else "wrong",
                ),
                mo.stat(_split, label="split", bordered=True),
                mo.stat(
                    str(len(_nbrs)),
                    label=f"neighbours · majority topic: {_vote}",
                    bordered=True,
                ),
            ],
            justify="start",
            wrap=True,
        ),
        mo.hstack(
            [
                _prob_chart,
                mo.ui.table(
                    _nbr_df,
                    selection=None,
                    show_column_summaries=False,
                    page_size=8,
                ),
            ],
            align="start",
        ),
        mo.md(
            f"Some **misclassified test papers** to try: `{_wrong_test}`"
        ),
    ]
)

if model.kind != "GAT":
    _out = mo.callout(
        mo.md(
            "Select **GAT** in the training controls (§4) and click **Train model** to see attention weights."
        ),
        kind="info",
    )
else:
    model.eval()
    with torch.no_grad():
        _, (_att_ei, _alpha) = model.conv1(
            graph.x, graph.edge_index, return_attention_weights=True
        )
    _att_ei, _alpha = (
        _att_ei.cpu(),
        _alpha.cpu(),
    )  # _alpha: [num_edges + self-loops, heads]
    _i = int(inspect_node.value)
    _mask = _att_ei[1] == _i
    _src = _att_ei[0][_mask]
    _a = _alpha[_mask]
    _y = graph.y.cpu()
    _labels = [("self" if int(j) == _i else str(int(j))) for j in _src]
    _mean_df = pd.DataFrame(
        {
            "neighbour": _labels,
            "attention": _a.mean(1).numpy(),
            "topic": [CLASS_NAMES[int(_y[j])] for j in _src],
        }
    )
    _uniform = 1 / len(_src)
    _bars = (
        alt.Chart(_mean_df)
        .mark_bar()
        .encode(
            x=alt.X("neighbour:N", sort="-y", title="source node"),
            y=alt.Y("attention:Q", title="attention (mean over heads)"),
            color=alt.Color(
                "topic:N", scale=alt.Scale(domain=CLASS_NAMES)
            ),
            tooltip=[
                "neighbour",
                "topic",
                alt.Tooltip("attention:Q", format=".3f"),
            ],
        )
    )
    _rule = (
        alt.Chart(pd.DataFrame({"u": [_uniform]}))
        .mark_rule(strokeDash=[4, 4], color="black")
        .encode(y="u:Q")
    )
    _head_df = pd.DataFrame(
        [
            (lab, f"head {h}", float(_a[k, h]))
            for k, lab in enumerate(_labels)
            for h in range(_a.size(1))
        ],
        columns=["neighbour", "head", "attention"],
    )
    _heat = (
        alt.Chart(_head_df)
        .mark_rect()
        .encode(
            x=alt.X("neighbour:N", title="source node"),
            y=alt.Y("head:N", title=None),
            color=alt.Color(
                "attention:Q", scale=alt.Scale(scheme="oranges")
            ),
            tooltip=[
                "neighbour",
                "head",
                alt.Tooltip("attention:Q", format=".3f"),
            ],
        )
    )
    _out = mo.vstack(
        [
            mo.hstack(
                [
                    (_bars + _rule).properties(
                        width=380,
                        height=220,
                        title=f"Attention into paper #{_i} ({CLASS_NAMES[int(_y[_i])]})",
                    ),
                    _heat.properties(
                        width=320,
                        height=220,
                        title="Per head: heads specialise differently",
                    ),
                ],
                align="start",
            ),
            mo.md(
                f"Each column sums to 1 per head (softmax over the {len(_src)} incoming edges incl. self-loop)."
            ),
        ]
    )
_out

def node_embeddings(m: torch.nn.Module) -> tuple[np.ndarray, np.ndarray]:
    """Hidden-layer vectors and predicted classes for every node, as NumPy arrays."""
    m.eval()
    with torch.no_grad():
        logits, hidden = m(graph.x, graph.edge_index)
    return hidden.cpu().numpy(), logits.argmax(dim=1).cpu().numpy()

_hidden, _pred = node_embeddings(model)
_xy = TSNE(n_components=2, init="pca", random_state=0).fit_transform(
    _hidden
)
_true = graph.y.cpu().numpy()
embedding_df = pd.DataFrame(
    {
        "x": _xy[:, 0],
        "y": _xy[:, 1],
        "topic": [CLASS_NAMES[i] for i in _true],
        "predicted": [CLASS_NAMES[i] for i in _pred],
        "correct": _pred == _true,
    }
)

alt.Chart(embedding_df).mark_circle(size=18, opacity=0.75).encode(
    x=alt.X("x:Q", axis=None),
    y=alt.Y("y:Q", axis=None),
    color=alt.Color("topic:N", legend=alt.Legend(title="true topic")),
    shape=alt.Shape(
        "correct:N", legend=alt.Legend(title="prediction correct")
    ),
    tooltip=["topic", "predicted", "correct"],
).properties(
    width=560,
    height=420,
    title=f"t-SNE of {model_kind.value} hidden embeddings",
)