A_karate = torch.tensor(
    nx.to_numpy_array(karate, nodelist=range(34), weight=None),
    dtype=torch.float,
)
X_karate = torch.randn(
    34, 8, generator=torch.Generator().manual_seed(SEED)
)

class MeanGNN(nn.Module):
    """K layers of h' = ReLU(W · mean_neighbours(h) + B · h)."""

    def __init__(self, d_in, d_hid, n_layers=2):
        super().__init__()
        dims = [d_in] + [d_hid] * n_layers
        self.W = nn.ModuleList(
            nn.Linear(a, b, bias=False) for a, b in zip(dims, dims[1:])
        )
        self.B = nn.ModuleList(
            nn.Linear(a, b) for a, b in zip(dims, dims[1:])
        )

    def forward(self, A, H):
        P = A / A.sum(1, keepdim=True).clamp(min=1)  # D^-1 A
        for W, B in zip(self.W, self.B):
            H = F.relu(W(P @ H) + B(H))
        return H

torch.manual_seed(SEED)
_gnn = MeanGNN(8, 16)
_row_mlp = nn.Sequential(
    nn.Linear(34 + 8, 16), nn.ReLU(), nn.Linear(16, 16)
)  # input: [row of A, features]
_perm = torch.randperm(34, generator=torch.Generator().manual_seed(1))
_P = torch.eye(34)[_perm]
with torch.no_grad():
    _gnn_gap = (
        (
            _gnn(_P @ A_karate @ _P.T, _P @ X_karate)
            - _P @ _gnn(A_karate, X_karate)
        )
        .abs()
        .max()
        .item()
    )
    _mlp_in = lambda A, X: torch.cat([A, X], 1)
    _mlp_gap = (
        (
            _row_mlp(_mlp_in(_P @ A_karate @ _P.T, _P @ X_karate))
            - _P @ _row_mlp(_mlp_in(A_karate, X_karate))
        )
        .abs()
        .max()
        .item()
    )
    _readout_gap = abs(
        _gnn(_P @ A_karate @ _P.T, _P @ X_karate).sum().item()
        - _gnn(A_karate, X_karate).sum().item()
    )
pd.DataFrame(
    {
        "test": [
            "GNN node outputs: f(PAPᵀ, PX) − P f(A, X)",
            "GNN + sum readout (graph level): |f(PAPᵀ,PX) − f(A,X)|",
            "MLP on [row of A, x]: f(PAPᵀ, PX) − P f(A, X)",
        ],
        "max |gap|": [_gnn_gap, _readout_gap, _mlp_gap],
    }
)

class GCN(nn.Module):
    """2-layer GCN. Dropout only on the hidden layer: dropping out the 2708×1433 input every epoch costs
    more CPU time than the whole model, for little accuracy."""

    def __init__(self, d_in, d_hid, n_cls, dropout=0.5):
        super().__init__()
        self.c1, self.c2, self.p = (
            GCNLayerScratch(d_in, d_hid),
            GCNLayerScratch(d_hid, n_cls),
            dropout,
        )

    def forward(self, x, A_hat):
        x = F.relu(self.c1(x, A_hat))  # layer 1: 1-hop
        x = F.dropout(x, self.p, self.training)
        return self.c2(x, A_hat)  # layer 2: 2-hop

class MLP(nn.Module):
    def __init__(self, d_in, d_hid, n_cls, dropout=0.5):
        super().__init__()
        self.l1, self.l2, self.p = (
            nn.Linear(d_in, d_hid),
            nn.Linear(d_hid, n_cls),
            dropout,
        )

    def forward(self, x, A_hat=None):  # ignores the graph
        x = F.relu(self.l1(x))
        x = F.dropout(x, self.p, self.training)
        return self.l2(x)

def train_node_classifier(
    model, data, A_hat, epochs=200, lr=0.01, wd=5e-4, train_mask=None
):
    """Full-batch training; returns (history DataFrame, test accuracy at the best-validation epoch)."""
    train_mask = data.train_mask if train_mask is None else train_mask
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=wd)
    hist, best_val, test_at_best = [], 0, 0
    for ep in range(epochs):
        model.train()
        opt.zero_grad()
        out = model(data.x, A_hat)
        loss = F.cross_entropy(out[train_mask], data.y[train_mask])
        loss.backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            pred = model(data.x, A_hat).argmax(1)
        accs = [
            (pred[m] == data.y[m]).float().mean().item()
            for m in (train_mask, data.val_mask, data.test_mask)
        ]
        hist.append([loss.item(), *accs])
        if accs[1] >= best_val:
            best_val, test_at_best = accs[1], accs[2]
    return pd.DataFrame(
        hist, columns=["loss", "train", "val", "test"]
    ), test_at_best

cora_hidden = mo.ui.slider(
    steps=[4, 8, 16, 32, 64],
    value=16,
    label="hidden units",
    show_value=True,
)
cora_seeds = mo.ui.slider(1, 5, value=3, label="seeds", show_value=True)
mo.vstack(
    [
        mo.md(
            "**Try it:** shrink the hidden layer to 4 units, or add seeds to see the spread."
        ),
        mo.hstack([cora_hidden, cora_seeds], justify="start"),
    ]
)

HOMOPHILY_SETS = ["cora", "lastfm-asia", "facebook-pages", "twitch-en", "deezer-europe", "penn94"]

class SparseGCN(nn.Module):
    def __init__(self, d_in, d_hid, n_cls):
        super().__init__()
        self.c1, self.c2 = GCNConv(d_in, d_hid), GCNConv(d_hid, n_cls)

    def forward(self, x, ei):
        return self.c2(F.dropout(F.relu(self.c1(x, ei)), 0.5, self.training), ei)

class SparseMLP(nn.Module):
    def __init__(self, d_in, d_hid, n_cls):
        super().__init__()
        self.l1, self.l2 = nn.Linear(d_in, d_hid), nn.Linear(d_hid, n_cls)

    def forward(self, x, ei=None):
        return self.l2(F.dropout(F.relu(self.l1(x)), 0.5, self.training))

def homophily_benchmark(name):
    d = node_data(name)
    x = (d.x - d.x.mean(0)) / (d.x.std(0) + 1e-6)
    y = d.y
    lab = torch.nonzero(y >= 0).squeeze(1)  # Penn94 has unknown genders (-1)
    n_cls = int(y.max()) + 1
    out = {"nodes": d.num_nodes, "edges": d.edge_index.size(1) // 2, "classes": n_cls,
           "edge homophily": edge_homophily(d.edge_index, y), "adjusted homophily": adjusted_homophily(d.edge_index, y),
           "majority": float(torch.bincount(y[lab]).max() / len(lab)), "MLP": [], "GCN": []}
    for seed in range(3):
        g = torch.Generator().manual_seed(seed)
        perm = lab[torch.randperm(len(lab), generator=g)]
        tr, va, te = perm[: len(perm) // 2], perm[len(perm) // 2 : 3 * len(perm) // 4], perm[3 * len(perm) // 4 :]
        for mname, cls in [("MLP", SparseMLP), ("GCN", SparseGCN)]:
            torch.manual_seed(seed)
            model = cls(x.size(1), 64, n_cls)
            opt = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
            best_val, test_at_best = 0.0, 0.0
            for _ in range(200):
                model.train()
                opt.zero_grad()
                F.cross_entropy(model(x, d.edge_index)[tr], y[tr]).backward()
                opt.step()
                model.eval()
                with torch.no_grad():
                    pred = model(x, d.edge_index).argmax(1)
                val = (pred[va] == y[va]).float().mean().item()
                if val >= best_val:
                    best_val, test_at_best = val, (pred[te] == y[te]).float().mean().item()
            out[mname].append(test_at_best)
    return out

homophily_run = mo.ui.run_button(label="Run on six graphs (≈ 10 min the first time, downloads ≈ 30 MB; then cached)")
homophily_run

_t = homophily_table
_lose = [n for n in _t.index if _t.loc[n, "GCN − MLP"] < 0]
_high = [n for n in _t.index if _t.loc[n, "adjusted homophily"] > 0.5]
_low = [n for n in _t.index if n not in _high]
_odd = [n for n in _low if _t.loc[n, "GCN − MLP"] > 0.03]
mo.md(rf"""
**Reading the run.** On the homophilous graphs ({", ".join(_high)}; adjusted homophily
{_t.loc[_high, "adjusted homophily"].min():.2f}–{_t.loc[_high, "adjusted homophily"].max():.2f}) the GCN gains
{_t.loc[_high, "GCN − MLP"].min():+.3f} to {_t.loc[_high, "GCN − MLP"].max():+.3f} over the MLP: neighbours share the
label, and averaging them is extra evidence. On the low-homophily graphs ({", ".join(_low)}) the picture splits:
{", ".join(f"{n} {_t.loc[n, 'GCN − MLP']:+.3f}" for n in _low)}{" — the GCN *loses* on " + ", ".join(_lose) if _lose else ""}.
Raw edge homophily would have misled here: Twitch and Deezer score {_t.loc["twitch-en", "edge homophily"]:.2f} and
{_t.loc["deezer-europe", "edge homophily"]:.2f}, which looks moderate until you remember that a binary label gives about
0.5 by chance.

{("**Low homophily is not the same as useless neighbours.** " + ", ".join(_odd) + " has almost no label homophily, yet the GCN still gains. Homophily asks whether neighbours share *my label*; a GNN only needs the neighbourhood to be *informative* about it. A plausible reading, not tested here: in a university friendship graph who your friends are (their dorm, year, major) can say something about you even when their gender does not match yours. Adjusted homophily is a warning light, not a verdict: compare against an MLP before concluding the graph does not help.") if _odd else ""}

For a research question on social data this is the first check to run: *is my label homophilous on my graph, and
does a graph model beat a feature-only one?* Political camps on a retweet graph are usually homophilous; "is this
account a bot" or a demographic attribute often is not, and then a mean-aggregating GCN can average the signal away.
L04's GraphSAGE (separate self and neighbour weights) and the heterophily-aware designs of later lectures exist for
exactly that case.
""")

def inductive_experiment(hide_frac=0.3, seed=SEED):
    g = torch.Generator().manual_seed(seed)
    unseen = torch.zeros(cora.num_nodes, dtype=torch.bool)
    unseen[
        torch.randperm(cora.num_nodes, generator=g)[
            : int(hide_frac * cora.num_nodes)
        ]
    ] = True
    seen_idx = (~unseen).nonzero().squeeze()
    sub_ei, _ = subgraph(
        seen_idx,
        cora.edge_index,
        relabel_nodes=True,
        num_nodes=cora.num_nodes,
    )
    sub = cora.subgraph(
        seen_idx
    )  # features, labels, masks of seen nodes only
    torch.manual_seed(seed)
    model = GCN(cora.num_features, 16, 7)
    train_node_classifier(
        model, sub, gcn_norm_dense(sub_ei, sub.num_nodes).to_sparse()
    )
    model.eval()
    with torch.no_grad():
        pred = model(cora.x, cora_A_hat_sparse).argmax(
            1
        )  # full graph, new nodes included
    return {
        "train labels used (seen nodes only)": int(sub.train_mask.sum()),
        "accuracy on the UNSEEN 30% of nodes": round(
            (pred[unseen] == cora.y[unseen]).float().mean().item(), 3
        ),
        "accuracy on seen test nodes": round(
            (
                pred[~unseen & cora.test_mask]
                == cora.y[~unseen & cora.test_mask]
            )
            .float()
            .mean()
            .item(),
            3,
        ),
    }

cora_A_hat_sparse = cora_A_hat.to_sparse()
inductive_experiment()

ex1_h3 = None

def lazy_layer(A, H):
    return None  # replace

def bfs_layer(A, h):
    return None

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_lazy():
    try:
        out = lazy_layer(A_karate, X_karate)
        if out is None:
            return _verdict(False, False)
        ref = (
            0.5 * X_karate
            + 0.5 * (A_karate / A_karate.sum(1, keepdim=True)) @ X_karate
        )
        return _verdict(torch.allclose(out, ref, atol=1e-5), True)
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

def _check_bfs():
    try:
        h = torch.zeros(34, 1)
        h[0] = 1.0
        dist = nx.single_source_shortest_path_length(karate, 0)
        for k in range(1, 4):
            h = bfs_layer(A_karate, h)
            if h is None:
                return _verdict(False, False)
            want = torch.tensor(
                [[1.0 if dist[v] <= k else 0.0] for v in range(34)]
            )
            if not torch.allclose(h.float(), want):
                return _verdict(False, True)
        return _verdict(True, True)
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_h3 is not None and abs(ex1_h3 - 14 / 3) < 1e-2,
            ex1_h3 is not None,
        ),
        mo.md("**Exercise 2**"),
        _check_lazy(),
        mo.md("**Exercise 3**"),
        _check_bfs(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"$W\cdot\text{mean}(1,2,4)+0=2\cdot 7/3\approx 4.667$."
                ),
                "Solution 2": mo.md("""```python
def lazy_layer(A, H):
    P = A / A.sum(1, keepdim=True)        # D^-1 A (row-normalised)
    return 0.5 * H + 0.5 * P @ H
```
Without the skip, repeated application converges to the stationary distribution (over-smoothing); the skip keeps
half of each node's own signal per step, which slows convergence."""),
                "Hint 3": mo.md(
                    "Message = neighbour's flag; AGGREGATE = **max** over neighbours; UPDATE = max(own flag, aggregate)."
                ),
                "Solution 3": mo.md("""```python
def bfs_layer(A, h):
    neighbour_max = (A * h.T).max(dim=1, keepdim=True).values   # max over neighbours' flags
    return torch.maximum(h, neighbour_max)
```
Sum or mean would also give non-zeros, but not exact 0/1 flags; **max** is the aggregator that matches the algorithm
(an example of *algorithmic alignment*)."""),
            }
        ),
    ]
)