def sbm_task(homophilous, n=600, k=3, seed=SEED):
    r = np.random.default_rng(seed)
    sizes = [n // k] * k
    p_in, p_out = (0.05, 0.005) if homophilous else (0.005, 0.03)
    g = nx.stochastic_block_model(
        sizes,
        [[p_in if i == j else p_out for j in range(k)] for i in range(k)],
        seed=seed,
    )
    y = torch.tensor(np.repeat(np.arange(k), sizes))
    x = torch.tensor(
        np.eye(k)[y.numpy()] * 0.6 + r.normal(0, 1, (n, k)),
        dtype=torch.float,
    )  # weak features
    x = torch.cat(
        [x, torch.tensor(r.normal(0, 1, (n, 13)), dtype=torch.float)], 1
    )
    data = from_networkx(g)
    data.x, data.y = x, y
    perm = torch.randperm(n, generator=torch.Generator().manual_seed(seed))
    data.train_mask = torch.zeros(n, dtype=torch.bool)
    data.train_mask[perm[:60]] = True
    data.val_mask = torch.zeros(n, dtype=torch.bool)
    data.val_mask[perm[60:200]] = True
    data.test_mask = torch.zeros(n, dtype=torch.bool)
    data.test_mask[perm[200:]] = True
    return data

TASKS = {
    "Cora": planetoid("Cora")[0],
    "CiteSeer": planetoid("CiteSeer")[0],
    "SBM-homophilous": sbm_task(True),
    "SBM-heterophilous": sbm_task(False),
}
DESIGN_SPACE = {
    "layer": ["GCN", "SAGE", "GAT"],
    "n_layers": [1, 2, 3],
    "batchnorm": [False, True],
    "skip": [False, True],
    "hidden": [32, 64],
}

def edge_homophily(d):
    return (
        (d.y[d.edge_index[0]] == d.y[d.edge_index[1]])
        .float()
        .mean()
        .item()
    )

pd.DataFrame(
    [
        {
            "task": n,
            "nodes": d.num_nodes,
            "edges": d.edge_index.size(1) // 2,
            "classes": int(d.y.max()) + 1,
            "train labels": int(d.train_mask.sum()),
            "edge homophily": edge_homophily(d),
        }
        for n, d in TASKS.items()
    ]
).round(3)

class DesignGNN(nn.Module):
    def __init__(self, design, d_in, n_cls):
        super().__init__()
        h, self.design = design["hidden"], design
        make = {
            "GCN": lambda: GCNConv(h, h),
            "SAGE": lambda: SAGEConv(h, h),
            "GAT": lambda: GATConv(h, h // 4, heads=4),
        }[design["layer"]]
        self.pre = nn.Linear(d_in, h)
        self.convs = nn.ModuleList(
            make() for _ in range(design["n_layers"])
        )
        self.norms = nn.ModuleList(
            nn.BatchNorm1d(h) if design["batchnorm"] else nn.Identity()
            for _ in range(design["n_layers"])
        )
        self.head = nn.Linear(h, n_cls)

    def forward(self, x, ei):
        h = F.relu(self.pre(x))
        for conv, norm in zip(self.convs, self.norms):
            out = F.dropout(F.relu(norm(conv(h, ei))), 0.5, self.training)
            h = h + out if self.design["skip"] else out
        return self.head(h)

RESULTS_PATH = CACHE_DIR / "L19_design_results.json"

def evaluate_design(design, task, seed=SEED, epochs=100):
    """Test accuracy at the best-validation epoch; cached on disk by (design, task, seed)."""
    key = json.dumps(
        {"design": design, "task": task, "seed": seed, "epochs": epochs},
        sort_keys=True,
    )
    cache = (
        json.loads(RESULTS_PATH.read_text())
        if RESULTS_PATH.exists()
        else {}
    )
    if key in cache:
        return cache[key]
    d = TASKS[task]
    torch.manual_seed(seed)
    model = DesignGNN(design, d.num_features, int(d.y.max()) + 1)
    opt = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
    best_val, best_test = -1.0, 0.0
    for _ in range(epochs):
        model.train()
        opt.zero_grad()
        F.cross_entropy(
            model(d.x, d.edge_index)[d.train_mask], d.y[d.train_mask]
        ).backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            pred = model(d.x, d.edge_index).argmax(1)
        val = (pred[d.val_mask] == d.y[d.val_mask]).float().mean().item()
        if val > best_val:
            best_val, best_test = (
                val,
                (pred[d.test_mask] == d.y[d.test_mask])
                .float()
                .mean()
                .item(),
            )
    cache = (
        json.loads(RESULTS_PATH.read_text())
        if RESULTS_PATH.exists()
        else {}
    )
    cache[key] = best_test
    RESULTS_PATH.write_text(json.dumps(cache))
    return best_test

def evaluate_mean(design, task, seeds=SEEDS):
    """Mean test accuracy over `seeds` (each seed cached separately)."""
    return float(np.mean([evaluate_design(design, task, seed=s) for s in seeds]))

def random_design(r):
    return {
        k: (
            v[int(r.integers(len(v)))]
            if not isinstance(v[0], bool)
            else bool(r.integers(2))
        )
        for k, v in DESIGN_SPACE.items()
    }

_t = time.time()
_demo = evaluate_design(
    {
        "layer": "GCN",
        "n_layers": 2,
        "batchnorm": False,
        "skip": False,
        "hidden": 64,
    },
    "Cora",
)
mo.md(
    f"One design on Cora (2-layer GCN, 64 hidden): test acc **{_demo:.3f}** in {time.time() - _t:.1f} s (cached afterwards)."
)

    _lines = []
    for _dim, _df in crs_results.items():
        _r = _df[f"rank of {_dim}=True"]
        _wins, _losses = int((_r == 1).sum()), int((_r == 2).sum())
        _per = crs_summary[crs_summary.dimension == _dim].sort_values("mean accuracy change when switched on")
        _span = ", ".join(
            f"{t}: {g * 100:+.1f} pts ({n} configs)"
            for t, g, n in zip(_per.task, _per["mean accuracy change when switched on"], _per.configs)
        )
        _lines.append(f"* **{_dim}** wins {_wins} and loses {_losses} of {len(_r)} paired comparisons (mean rank {_r.mean():.2f}). By task: {_span}.")
    mo.md(
        "**This run** (20 random configurations per dimension, each accuracy the mean of "
        + f"{len(SEEDS)} seeds):\n\n"
        + "\n".join(_lines)
        + """

How to read it: a mean rank near 1.5 with large per-task effects of opposite sign means the choice is **task-dependent**,
not good or bad. The design-space paper found BatchNorm *generally* helpful across its 32 tasks; small, label-scarce citation
tasks trained full-batch can disagree. That is exactly why the method reports *distributions over tasks and seeds*, not one
number. (With a single seed, a 1–2 point difference between two runs of the same design is common; averaging 3 seeds cuts that
noise by about 40%, a factor of $1/\\sqrt3$.)

## 4 · Task similarity and transfer

Characterise each task by how a fixed set of **anchor designs** rank on it; tasks with similar rankings (high Kendall τ) should share
good designs. For a **new task**: run only the anchors, find the most similar known task, and reuse *its* best design, without searching
the whole space. We evaluate 16 random designs on all 4 tasks (3 seeds each); the first 8 are the anchors.

> **Predict first.** Which pair of tasks will have the highest τ? Which task has no good neighbour?
"""
    )

mo.stop(
    not transfer_run.value,
    mo.md("Press the button (cached per design and task)."),
)
pool_acc = pd.DataFrame(
    {
        task: [evaluate_mean(d, task) for d in POOL]
        for task in mo.status.progress_bar(list(TASKS), title="tasks")
    }
)
_anchor = pool_acc.iloc[:N_ANCHORS]
task_tau = pd.DataFrame(
    {
        a: {b: kendalltau(_anchor[a], _anchor[b]).statistic for b in TASKS}
        for a in TASKS
    }
)
_rows = []
for new in TASKS:
    nearest = task_tau[new].drop(new).idxmax()
    chosen = int(
        pool_acc[nearest].idxmax()
    )  # best design on the most similar known task
    _rows.append(
        {
            "new task": new,
            "most similar known task": nearest,
            "τ": task_tau.loc[nearest, new],
            "transferred design acc": pool_acc.loc[chosen, new],
            "best design in pool": pool_acc[new].max(),
            "average design": pool_acc[new].mean(),
        }
    )
_fig, _ax = plt.subplots(figsize=(4.2, 3.4))
_im = _ax.imshow(task_tau.values, cmap="RdBu", vmin=-1, vmax=1)
_ax.set_xticks(range(4), list(TASKS), rotation=30, fontsize=7)
_ax.set_yticks(range(4), list(TASKS), fontsize=7)
for (_i, _j), _v in np.ndenumerate(task_tau.values):
    _ax.text(_j, _i, f"{_v:.2f}", ha="center", va="center", fontsize=8)
_ax.set_title("Kendall τ of anchor rankings", fontsize=9)
transfer_results = pd.DataFrame(_rows)
mo.vstack(
    [
        mo.hstack(
            [
                _fig,
                pool_acc.round(3)
                .assign(
                    design=[
                        f"{d['layer']}-{d['n_layers']}L{'-BN' if d['batchnorm'] else ''}{'-skip' if d['skip'] else ''}-{d['hidden']}"
                        for d in POOL
                    ]
                )
                .set_index("design"),
            ],
            widths=[2, 3],
        ),
        transfer_results.round(3),
    ]
)

    _pairs = [(a_, b_, task_tau.loc[a_, b_]) for i, a_ in enumerate(task_tau.index) for b_ in task_tau.index[i + 1 :]]
    _hi = max(_pairs, key=lambda t: t[2])
    _lonely = task_tau.where(~np.eye(len(task_tau), dtype=bool)).max().idxmin()
    _rows = "\n".join(
        f"| {r['new task']} | {r['most similar known task']} | {r['τ']:.2f} | {r['transferred design acc']:.3f} | "
        f"{r['best design in pool']:.3f} | {r['average design']:.3f} |"
        for _, r in transfer_results.iterrows()
    )
    mo.md(
        f"""**This run.** Most similar pair: **{_hi[0]} ↔ {_hi[1]} (τ = {_hi[2]:.2f})**. Task with the weakest best neighbour: **{_lonely}**.

| new task | nearest known task | τ | transferred design | best in pool | average design |
|---|---|---|---|---|---|
{_rows}

How to read it: when τ is high the transferred design should sit near the best in the pool and well above the average.
When the nearest neighbour has τ near or below 0, transfer is a guess: that is the signal to *search* instead. An easy task can
get a good design from a dissimilar one by luck, because many designs solve it. GraphGym (`torch_geometric.graphgym`) runs
this at scale from YAML configs."""
    )

ex1_tau = None
ex2_designs = None

def nearest_task(anchor_accs):
    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_nearest():
    try:
        if "pool_acc" not in globals():
            return mo.callout(
                mo.md("Run §4 first (press its button)."), kind="neutral"
            )
        noisy = pool_acc["Cora"].iloc[
            :N_ANCHORS
        ].values + np.random.default_rng(0).normal(0, 0.003, N_ANCHORS)
        got = nearest_task(list(noisy))
        if got is None:
            return _verdict(False, False)
        return _verdict(got in ("Cora", "CiteSeer"), True)
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_tau is not None and abs(ex1_tau - 1 / 3) < 2e-3,
            ex1_tau is not None,
        ),
        mo.md("**Exercise 2**"),
        _verdict(ex2_designs == 72 * 9, ex2_designs is not None),
        mo.md("**Exercise 3**"),
        _check_nearest(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"6 pairs; (M1,M2) and (M3,M4) are flipped → 4 concordant, 2 discordant → $\tau=(4-2)/6\approx 0.333$."
                ),
                "Solution 2": mo.md(
                    "72 × 3 × 3 = **648**. The paper's space has 315,000 designs, which is why it samples instead of enumerating."
                ),
                "Solution 3": mo.md("""```python
def nearest_task(anchor_accs):
    taus = {t: kendalltau(anchor_accs, pool_acc[t].iloc[:N_ANCHORS]).statistic for t in TASKS}
    return max(taus, key=taus.get)
```"""),
            }
        ),
    ]
)