import warnings

warnings.filterwarnings("ignore", category=FutureWarning)
# harmless Intel-GPU probe message printed by torch on this machine
warnings.filterwarnings("ignore", message=".*Level Zero.*")

import hashlib
import json
import time
from collections import Counter

import marimo as mo
from gl.nb import json_cache, mean_std, reveal
import numpy as np
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import (
    GCNConv,
    GINConv,
    global_add_pool,
    global_mean_pool,
)
from torch_geometric.utils import from_networkx
from sklearn.svm import SVC
from sklearn.model_selection import StratifiedKFold, cross_val_score

SEED = 0
rng = np.random.default_rng(SEED)
torch.manual_seed(SEED)
# several open notebooks share the CPU; more threads = contention
torch.set_num_threads(4)
# shared datasets (Cora etc. are cached here) and this folder's disk cache
DATA_DIR = mo.notebook_dir().parents[1] / "data"
CACHE_DIR = mo.notebook_dir() / "cache"
CACHE_DIR.mkdir(exist_ok=True)
plt.rcParams.update(
    {
        "figure.dpi": 100,
        "axes.spines.top": False,
        "axes.spines.right": False,
    }
)

_rev = reveal(
    r"""
| pair | mean | max | sum |
|---|---|---|---|
| {a,b} vs {a,a,b,b} | ✗ same proportions | ✗ same set | ✓ |
| {a,b} vs {a,a,b} | ✓ | ✗ same set | ✓ |
| {a} vs {a,a} | ✗ | ✗ | ✓ |
| {a,a,b} vs {a,b,b} | ✓ | ✗ | ✓ |

**Mean** sees only proportions, **max** sees only *which* colours occur (the underlying set), **sum** of one-hot vectors is
the full **count vector**, so it separates every pair. With continuous features you need an MLP before the sum, which GIN
provides: $\text{MLP}\big(\sum_u f(\mathbf h_u)\big)$ can approximate any injective multiset function (Xu et al., 2019).
""",
    predict="Did your paper answers match the dropdown? The table below is the full answer, with the reason for each cell.",
    env="exact: no randomness involved",
)
_next = mo.md(r"""
## 2 · The Weisfeiler–Lehman (1-WL) test

Colour refinement: every node starts with the same colour; each round, a node's new colour is a hash of
(its colour, the multiset of its neighbours' colours). If the two graphs' colour histograms ever differ, they are **not**
isomorphic; if the histograms stay equal until the colouring stabilises, WL says "possibly isomorphic". A GNN with injective
aggregation computes exactly these colours as vectors, so **no MPNN can distinguish what WL can't**.

> **Predict first.** For each pair in the next dropdown, guess: does WL separate them, and at which round?
""")
mo.vstack([_rev, _next])

def wl_refine(graphs, max_iter=10):
    """Run 1-WL jointly on several graphs (shared colour dictionary).
    Returns a list of per-iteration colour dicts for each graph, and per-iteration histograms."""
    colours = [{v: 0 for v in g} for g in graphs]
    history = [[dict(c) for c in colours]]
    for _ in range(max_iter):
        palette = {}
        signatures = [
            {v: (c[v], tuple(sorted(c[u] for u in g[v]))) for v in g}
            for g, c in zip(graphs, colours)
        ]
        for sig in signatures:
            for s in sorted(set(sig.values())):
                palette.setdefault(s, len(palette))
        new = [
            {v: palette[s[v]] for v in g}
            for g, s in zip(graphs, signatures)
        ]
        stable = all(
            len(set(n.values())) == len(set(c.values()))
            for n, c in zip(new, colours)
        )
        colours = new
        history.append([dict(c) for c in colours])
        if stable:
            break
    return history

def wl_histograms(history):
    return [
        [tuple(sorted(Counter(c.values()).items())) for c in step]
        for step in history
    ]

def decalin():
    return nx.Graph(
        [
            (0, 1),
            (1, 2),
            (2, 3),
            (3, 4),
            (4, 5),
            (5, 0),
            (0, 6),
            (6, 7),
            (7, 8),
            (8, 9),
            (9, 5),
        ]
    )

def bicyclopentyl():
    return nx.Graph(
        [
            (0, 1),
            (1, 2),
            (2, 3),
            (3, 4),
            (4, 0),
            (0, 5),
            (5, 6),
            (6, 7),
            (7, 8),
            (8, 9),
            (9, 5),
        ]
    )

wl_pairs = {
    "hexagon C6  vs  two triangles": (
        nx.cycle_graph(6),
        nx.disjoint_union(nx.cycle_graph(3), nx.cycle_graph(3)),
    ),
    "triangular prism  vs  K3,3 (both 3-regular)": (
        nx.circular_ladder_graph(3),
        nx.complete_bipartite_graph(3, 3),
    ),
    "decalin  vs  bicyclopentyl (molecules)": (decalin(), bicyclopentyl()),
    "path P6  vs  star S5": (nx.path_graph(6), nx.star_graph(5)),
    "two random trees, 8 nodes": (
        nx.random_labeled_tree(8, seed=1),
        nx.random_labeled_tree(8, seed=2),
    ),
}
wl_picker = mo.ui.dropdown(
    list(wl_pairs),
    value="hexagon C6  vs  two triangles",
    label="graph pair",
)
wl_picker

_rev = reveal(
    r"""
The lecture's worked example — hexagon $C_6$ vs. two triangles — keeps **identical** histograms; colouring is
stable after one round. WL (and so GCN, GAT, GraphSAGE, GIN, …) says "possibly isomorphic", yet one graph is connected and the
other isn't. The triangular prism vs. $K_{3,3}$ (both 3-regular) and **decalin vs. bicyclopentyl** (two real molecules) also fool
WL. Path vs. star and the two random trees are separated at the first round (different degrees). WL succeeds on almost
all real graphs; it fails on highly regular structures. (Exercise 4 below asks you to prove the regular case.)
""",
    predict="Which of the five pairs fooled WL? What do the fooling pairs have in common?",
    env="exact: WL is deterministic",
)
_next = mo.md(r"""
## 3 · GIN: the most expressive message-passing GNN

$$\mathbf h_v^{(k)}=\text{MLP}^{(k)}\Big((1+\epsilon^{(k)})\,\mathbf h_v^{(k-1)}+\sum_{u\in N(v)}\mathbf h_u^{(k-1)}\Big)$$
Sum (injective on multisets) + MLP (an injective map of the result) + a learnable weight on the node itself.

**Counting task.** Random graphs with **no features** (every node gets 1). Predict each node's **degree** (0–7, capped).

> **Predict first.** Three 2-layer models: mean aggregation, GCN (degree-normalised sum), GIN (sum + MLP). Which can learn
> to output a node's degree from constant inputs, and which are stuck near the majority class (≈ 0.18)?
""")
mo.vstack([_rev, _next])

def mlp(d_in, d_out):
    return nn.Sequential(
        nn.Linear(d_in, d_out),
        nn.BatchNorm1d(d_out),
        nn.ReLU(),
        nn.Linear(d_out, d_out),
    )

class MeanConv(nn.Module):
    """h_v' = MLP(h_v + mean_{u∈N(v)} h_u): the GIN layer with mean instead of sum."""

    def __init__(self, d_in, d_out):
        super().__init__()
        self.mlp = mlp(d_in, d_out)

    def forward(self, x, edge_index):
        src, dst = edge_index
        agg = torch.zeros_like(x).index_add_(0, dst, x[src])
        deg = (
            torch.bincount(dst, minlength=x.size(0))
            .clamp(min=1)
            .unsqueeze(1)
        )
        return self.mlp(x + agg / deg)

class NodeNet(nn.Module):
    def __init__(self, kind, d_hid=32, n_out=8, n_layers=2):
        super().__init__()
        make = {
            "mean aggregation": lambda a: MeanConv(a, d_hid),
            "GCN (degree-normalised sum)": lambda a: GCNConv(a, d_hid),
            "GIN (sum + MLP)": lambda a: GINConv(
                mlp(a, d_hid), train_eps=True
            ),
        }[kind]
        self.convs = nn.ModuleList(
            make(1 if i == 0 else d_hid) for i in range(n_layers)
        )
        self.head = nn.Linear(d_hid, n_out)

    def forward(self, x, edge_index):
        for conv in self.convs:
            x = F.relu(conv(x, edge_index))
        return self.head(x)

def degree_task(n_graphs=60, n=30, seed=SEED):
    r = np.random.default_rng(seed)
    graphs = [
        nx.gnp_random_graph(
            n, r.uniform(0.05, 0.2), seed=int(r.integers(1e9))
        )
        for _ in range(n_graphs)
    ]
    big = nx.disjoint_union_all(graphs)
    data = from_networkx(big)
    data.x = torch.ones(data.num_nodes, 1)
    data.y = torch.tensor([min(d, 7) for _, d in big.degree()])
    data.train_mask = (
        torch.arange(data.num_nodes) < data.num_nodes * 2 // 3
    )  # first 40 graphs train, rest test
    return data

deg_data = degree_task()
_rows = []
for _kind in [
    "mean aggregation",
    "GCN (degree-normalised sum)",
    "GIN (sum + MLP)",
]:
    torch.manual_seed(SEED)
    _m = NodeNet(_kind)
    _opt = torch.optim.Adam(_m.parameters(), lr=0.01)
    for _ in range(200):
        _m.train()
        _opt.zero_grad()
        F.cross_entropy(
            _m(deg_data.x, deg_data.edge_index)[deg_data.train_mask],
            deg_data.y[deg_data.train_mask],
        ).backward()
        _opt.step()
    _m.eval()
    with torch.no_grad():
        _pred = _m(deg_data.x, deg_data.edge_index).argmax(1)
    _rows.append(
        {
            "model": _kind,
            "test acc (predict degree)": (
                _pred[~deg_data.train_mask]
                == deg_data.y[~deg_data.train_mask]
            )
            .float()
            .mean()
            .item(),
        }
    )
_rows.append(
    {
        "model": "majority class",
        "test acc (predict degree)": deg_data.y[~deg_data.train_mask]
        .bincount()
        .max()
        .item()
        / int((~deg_data.train_mask).sum()),
    }
)
deg_table = pd.DataFrame(_rows).round(3)
deg_table

_acc = dict(zip(deg_table["model"], deg_table["test acc (predict degree)"]))
mo.md(rf"""
**This run:** GIN **{_acc["GIN (sum + MLP)"]:.2f}**, GCN {_acc["GCN (degree-normalised sum)"]:.2f}, mean aggregation
{_acc["mean aggregation"]:.2f} (majority class: {_acc["majority class"]:.2f}).

Why: with constant inputs every neighbour "looks like" the node itself, so the **mean** of the neighbours is the same vector
whatever their number, and the mean layer can do no better than guess the most common degree. **GCN** gets a little through its
$1/\sqrt{{d_ud_v}}$ normalisation, which leaks degree information into the sum. **Sum** aggregation counts neighbours exactly, so
GIN solves the task (the reference run gave 1.00 / 0.29 / 0.20).

## 4 · Real benchmarks: GIN vs. GCN, and the WL kernel

Graph classification (TU datasets, downloaded to `data/TUDataset/`):

* **MUTAG** — 188 molecules, mutagenic or not; atom-type features.
* **PROTEINS** — 1,113 proteins, enzyme or not; 3 node features.
* **IMDB-BINARY** — 1,000 actor ego-networks, Action vs Romance; **no node features** (we use constant 1s), so structure is everything.

**A controlled comparison.** "GIN vs. GCN" usually changes three things at once: the aggregator (sum vs. normalised sum), the
**readout** (sum vs. mean over nodes) and the update (MLP vs. one linear layer). To see *which* one matters we cross the two
layers with the two readouts, a $2\times2$ ablation, and run every cell with **3 seeds × 5 folds** (15 test scores).

Classical baseline: the **WL subtree kernel** (Shervashidze et al., 2011): count WL colours over 3 rounds (the same depth as the
3-layer GNNs) and train a linear SVM. On IMDB it starts from the **same constant colour** the GNNs get; on the molecules from the atom type.

This is the protocol of Errica et al., *A Fair Comparison of Graph Neural Networks for Graph Classification* (ICLR 2020), scaled
down: they found that on these very datasets a structure-agnostic baseline often matches published GNNs once model selection and
splits are fixed.

> **Predict first.** On featureless IMDB-BINARY, which of the four GNN cells collapses to chance (0.50), and is it the
> *aggregator* or the *readout* that causes it? Does any GNN beat the WL kernel?
""")

def load_tu(name):
    ds = TUDataset(str(DATA_DIR / "TUDataset"), name)
    if ds.num_features == 0:  # IMDB-BINARY: constant features
        graphs = []
        for g in ds:
            g = g.clone()
            g.x = torch.ones(g.num_nodes, 1)
            graphs.append(g)
        return graphs
    return list(ds)

class GraphNet(nn.Module):
    """`conv` in {GIN, GCN}; `readout` in {sum, mean}; reads out every layer (jumping knowledge)."""

    def __init__(self, conv, d_in, readout=None, d_hid=64, n_layers=3, n_cls=2):
        super().__init__()
        readout = readout or ("sum" if conv == "GIN" else "mean")
        self.pool = global_add_pool if readout == "sum" else global_mean_pool
        self.convs = nn.ModuleList(
            GINConv(mlp(d_in if i == 0 else d_hid, d_hid), train_eps=True)
            if conv == "GIN"
            else GCNConv(d_in if i == 0 else d_hid, d_hid)
            for i in range(n_layers)
        )
        self.head = nn.Sequential(
            nn.Linear(d_hid * n_layers, d_hid),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(d_hid, n_cls),
        )

    def forward(self, x, edge_index, batch):
        pooled = []
        for conv in self.convs:
            x = F.relu(conv(x, edge_index))
            pooled.append(self.pool(x, batch))
        return self.head(torch.cat(pooled, 1))

def cv_splits(y, seeds, folds=5):
    for seed in seeds:
        for tr, te in StratifiedKFold(folds, shuffle=True, random_state=seed).split(y, y):
            yield seed, tr, te

def cv_graph_model(graphs, conv, readout, seeds=(0, 1, 2), epochs=100):
    """Test accuracy on every (seed, fold): 15 numbers for 3 seeds × 5 folds. Fixed epochs, no model selection."""
    y = np.array([int(g.y) for g in graphs])
    accs = []
    for seed, tr, te in cv_splits(y, seeds):
        torch.manual_seed(seed)
        model = GraphNet(conv, graphs[0].num_features, readout)
        opt = torch.optim.Adam(model.parameters(), lr=0.01)
        loader = DataLoader([graphs[i] for i in tr], batch_size=64, shuffle=True)
        for _ in range(epochs):
            model.train()
            for b in loader:
                opt.zero_grad()
                F.cross_entropy(model(b.x, b.edge_index, b.batch), b.y).backward()
                opt.step()
        model.eval()
        b = next(iter(DataLoader([graphs[i] for i in te], batch_size=len(te))))
        with torch.no_grad():
            accs.append((model(b.x, b.edge_index, b.batch).argmax(1) == b.y).float().mean().item())
    return accs

def wl_kernel_features(graphs, h=3):
    """WL subtree features: counts of every colour seen in rounds 0..h, palette shared by all graphs.
    Initial colour = node label (argmax of one-hot features) or, for featureless graphs, one constant colour,
    exactly the information the GNNs receive."""
    adj, colours = [], []
    for g in graphs:
        nbrs = [[] for _ in range(g.num_nodes)]
        for s_, d_ in g.edge_index.t().tolist():
            nbrs[d_].append(s_)
        adj.append(nbrs)
        init = g.x.argmax(1).tolist() if g.num_features > 1 else [0] * g.num_nodes
        colours.append([("init", c) for c in init])
    palette, feats = {}, [Counter() for _ in graphs]
    for _ in range(h + 1):
        for i, cols in enumerate(colours):
            for c in cols:
                feats[i][palette.setdefault(c, len(palette))] += 1
        colours = [
            [
                (palette[cols[v]], tuple(sorted(palette[cols[u]] for u in adj[i][v])))
                for v in range(len(cols))
            ]
            for i, cols in enumerate(colours)
        ]
    X = np.zeros((len(graphs), len(palette)))
    for i, f in enumerate(feats):
        for k, v in f.items():
            X[i, k] = v
    return X

def wl_svm(graphs, seeds=(0, 1, 2)):
    y = np.array([int(g.y) for g in graphs])
    X = wl_kernel_features(graphs)
    X = X / np.linalg.norm(X, axis=1, keepdims=True)
    accs = []
    for _, tr, te in cv_splits(y, seeds):
        clf = SVC(kernel="linear", C=1.0).fit(X[tr], y[tr])
        accs.append(float((clf.predict(X[te]) == y[te]).mean()))
    return accs

ABLATION = [("GCN", "mean"), ("GCN", "sum"), ("GIN", "mean"), ("GIN", "sum")]

def benchmark(name, graphs=None):
    """All models on one dataset; per-model lists of 15 test accuracies. Cached in cache/L06_benchmarks.json."""

    def compute():
        gs = graphs if graphs is not None else load_tu(name)
        out = {"graphs": len(gs), "WL kernel + SVM": wl_svm(gs)}
        for conv, readout in ABLATION:
            out[f"{conv} + {readout} readout"] = cv_graph_model(gs, conv, readout)
        return out

    if graphs is not None:  # ad-hoc data (tests): never cache
        return compute()
    return json_cache(
        CACHE_DIR / "L06_benchmarks.json",
        {"dataset": name, "seeds": [0, 1, 2], "folds": 5, "epochs": 100, "v": 2},
        compute,
    )

bench_run = mo.ui.run_button(
    label="Run benchmarks: 3 datasets × 5 models × 15 runs (≈ 15 min first time, then cached)"
)
bench_run

def _get(d, m):
    r = bench_table[(bench_table.dataset == d) & (bench_table.model == m)].iloc[0]
    return f"{r['mean']:.3f} ± {r['std']:.3f}"

def _best(d):
    sub = bench_table[bench_table.dataset == d]
    return sub.loc[sub["mean"].idxmax(), "model"]

_imdb = {m: _get("IMDB-BINARY", m) for m in bench_table[bench_table.dataset == "IMDB-BINARY"].model}
mo.md(rf"""
**Reading the ablation (this run).**

* IMDB-BINARY, the featureless case: GCN + mean readout **{_imdb["GCN + mean readout"]}**, GCN + sum readout
  **{_imdb["GCN + sum readout"]}**, GIN + mean readout **{_imdb["GIN + mean readout"]}**, GIN + sum readout
  **{_imdb["GIN + sum readout"]}**. Compare along a row (same layer, different readout) and down a column (same readout,
  different layer) to see which choice carries the effect: here it is the **layer**, not the readout. A standard deviation of
  0.000 at 0.500 means the model predicted one class in every run. Why the GCN layer: on a $d$-regular graph the normalised
  sum of a constant signal is $\sum_{{j}}\frac{{c}}{{\sqrt{{(d+1)(d+1)}}}}=\frac{{(d+1)c}}{{d+1}}=c$, the same at every node
  *and for every degree*. IMDB ego-networks are dense and close to regular, so the layer sees almost nothing: a blind spot of
  the normalisation, which GIN's plain sum does not have (the sum gives $(d+1)c$, which counts).
* Best model per dataset: MUTAG → {_best("MUTAG")}; PROTEINS → {_best("PROTEINS")}; IMDB-BINARY → {_best("IMDB-BINARY")}.
* The WL kernel + linear SVM (2011) is a strong baseline. "As powerful as WL" bounds what a GNN *can* represent; it does not
  promise that training finds it, and small datasets favour simple models. Differences smaller than one standard deviation are
  not evidence of anything.

## 5 · The WL ceiling: what no MPNN can see

A GIN can only be as good as WL. Task: random **2-regular** graphs on $2n$ nodes — one big cycle $C_{{2n}}$ (connected) vs.
two cycles $C_n\cup C_n$ (disconnected). Every node in both graphs has the same infinite computation tree.

> **Predict first.** Trained on all 20 graphs (no test set at all), what *training* accuracy can GIN reach?
""")

ex1_left = None
ex1_right = None

def wl_rounds(g1, g2):
    return "todo"

def _verdict(ok, answered):
    if not answered:
        return mo.callout(mo.md("Not answered yet."), kind="neutral")
    return mo.callout(
        mo.md("Correct!" if ok else "Not yet — try again or open a hint."),
        kind="success" if ok else "warn",
    )

def _check_ex1():
    if ex1_left is None or ex1_right is None:
        return _verdict(False, False)
    l, r = np.array(ex1_left, float), np.array(ex1_right, float)
    return _verdict(
        np.isclose(l.mean(), r.mean())
        and np.isclose(l.max(), r.max())
        and not np.isclose(l.sum(), r.sum()),
        True,
    )

def _check_ex2():
    tests = [
        (
            wl_pairs[k],
            None
            if k.startswith(("hexagon", "triangular", "decalin"))
            else 1,
        )
        for k in wl_pairs
    ]
    tests.append(((nx.path_graph(5), nx.path_graph(5)), None))
    tests.append(
        (
            (
                nx.path_graph(7),
                nx.Graph([(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (2, 6)]),
            ),
            1,
        )
    )
    try:
        if wl_rounds(nx.path_graph(3), nx.path_graph(3)) == "todo":
            return _verdict(False, False)
        return _verdict(
            all(wl_rounds(*g) == want for g, want in tests), True
        )
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _check_ex1(),
        mo.md("**Exercise 2**"),
        _check_ex2(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    "`ex1_left = [1]`, `ex1_right = [1, 1]`: mean 1 = 1, max 1 = 1, sum 1 ≠ 2. (A node with one neighbour vs. two identical neighbours.)"
                ),
                "Solution 2": mo.md("""```python
def wl_rounds(g1, g2):
    hist = wl_histograms(wl_refine([g1, g2]))
    return next((i for i, (a, b) in enumerate(hist) if a != b), None)
```"""),
                "Solution 3": mo.md(r"""Induction on the layer $k$: if two nodes have the same WL colour after $k$ rounds, they have the same colour and the same
multiset of neighbour colours at round $k-1$; by the induction hypothesis the neighbours' GNN embeddings form the same multiset, so any
AGGREGATE/COMBINE (functions of the multiset) gives the same output. So equal WL colours ⇒ equal embeddings; equal histograms ⇒ equal
READOUT. Contrapositive: if the GNN separates two graphs, WL does too."""),
            }
        ),
    ]
)