def ring_of_cycles(n_comp=40, seed=SEED):
    """Triangles and squares (random mix) in a ring; component i's first node links to component i+1's first node."""
    r = np.random.default_rng(seed)
    g, y, firsts = nx.Graph(), [], []
    for _ in range(n_comp):
        size, base = int(r.choice([3, 4])), g.number_of_nodes()
        nx.add_cycle(g, range(base, base + size))
        y += [int(size == 3)] * size
        firsts.append(base)
    g.add_edges_from(zip(firsts, firsts[1:] + firsts[:1]))
    return g, torch.tensor(y)

def encode(g, kind, k=8):
    n = g.number_of_nodes()
    if kind == "constant":
        return torch.ones(n, 1)
    if kind == "closed walks log(1+diag(A^k))":
        return torch.tensor(
            np.log1p(closed_walks(g, 6, nodelist=range(n))),
            dtype=torch.float,
        )
    data = from_networkx(g)
    if kind == "random-walk PE diag((D⁻¹A)^k)":
        return T.AddRandomWalkPE(walk_length=6, attr_name="pe")(data).pe
    if kind == "Laplacian eigenvector PE (k=8)":
        return T.AddLaplacianEigenvectorPE(k=k, attr_name="pe")(data).pe
    raise ValueError(kind)

class NodeGIN(nn.Module):
    def __init__(self, d_in, d=32, n_layers=3, n_out=2):
        super().__init__()
        self.convs = nn.ModuleList(
            GINConv(
                nn.Sequential(
                    nn.Linear(d_in if i == 0 else d, d),
                    nn.ReLU(),
                    nn.Linear(d, d),
                )
            )
            for i in range(n_layers)
        )
        self.head = nn.Linear(d, 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 fit_nodes(x, edge_index, y, epochs=300, seed=SEED, augment=None):
    torch.manual_seed(seed)
    model = NodeGIN(x.size(1))
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    for _ in range(epochs):
        opt.zero_grad()
        F.cross_entropy(
            model(augment(x) if augment else x, edge_index), y
        ).backward()
        opt.step()
    return model

def node_acc(model, x, edge_index, y):
    with torch.no_grad():
        return (model(x, edge_index).argmax(1) == y).float().mean().item()

ring_a, y_a = ring_of_cycles(seed=0)
ring_b, y_b = ring_of_cycles(seed=1)
ei_a, ei_b = (
    from_networkx(ring_a).edge_index,
    from_networkx(ring_b).edge_index,
)
encodings = [
    "constant",
    "closed walks log(1+diag(A^k))",
    "random-walk PE diag((D⁻¹A)^k)",
    "Laplacian eigenvector PE (k=8)",
]
_rows = []
for _kind in encodings:
    _xa, _xb = encode(ring_a, _kind), encode(ring_b, _kind)
    _m = fit_nodes(_xa, ei_a, y_a)
    _rows.append(
        {
            "encoding": _kind,
            "train acc (ring A)": node_acc(_m, _xa, ei_a, y_a),
            "test acc (new ring B)": node_acc(_m, _xb, ei_b, y_b),
        }
    )
pd.DataFrame(_rows).round(3)

_rev = reveal(
    r"""
Measured: **constant** inputs reach 0.84 / 0.87. They are not hopeless here, because the ring's linking edges create degree-3
nodes that break some symmetry, but they are far from perfect. **Closed walks: 1.00 on the new ring**; random-walk PE 0.99.
These features are computed the same way on any graph, so they transfer. The **Laplacian PE** fits ring A perfectly (1.00)
but gets only ≈ 0.58 on ring B: eigenvectors are a *global* coordinate system specific to one graph, and their signs are arbitrary.

""",
    predict="Trained on ring A and tested on a new ring B, rank constant inputs, closed-walk counts, random-walk PE and Laplacian PE. Which transfers worst, and why?",
)
_post = mo.md(r"""
## 3 · Fix A: eigenvector positional encodings — and their sign ambiguity

If $v$ is an eigenvector, so is $-v$; with repeated eigenvalues any rotation inside the eigenspace works too. A network fed raw
eigenvectors can give **different answers for the same graph** depending on which signs the solver returned.
Test it on ring A itself: flip the sign of each PE column at random (same graph, equally valid eigenvectors).
""")
mo.vstack([_rev, _post])

_rev = reveal(
    r"""
Measured: **constant** and **closed-walk** features are at or below chance (0.16–0.19 test, chance = 0.25): the four
quadrants are mirror images, so structure cannot tell them apart. **Anchor-set distances reach ≈ 0.91**. Across settings of $c$
the test accuracy ranged 0.81–0.94 (random anchor sets; P-GNN re-samples them every epoch to stay inductive).
""",
    predict="On the mirrored grid, which features can tell the four quadrants apart: constant, closed walks, or anchor distances? Can anything score *below* chance (0.25)?",
)
_post = mo.md(r"""
**Try it:** change the number of anchor sets per scale.

| failure | example | fix |
|---|---|---|
| different structure, same embedding | triangle vs. square node | structural features: closed walks, cycle counts, RWPE (§2) |
| same structure, should differ | mirror positions on a grid | positional features: anchor distances (§4), Laplacian PE (§3, mind the signs) |

---

## 5 · Exercises

**Exercise 1 (by hand).** Closed walks $(A^k)_{vv}$ for a node of the 5-cycle $C_5$, $k=0..5$. Use $\mathrm{tr}(A^k)/5$ and
$\lambda_j=2\cos(2\pi j/5)$, or count. → `ex1_c5` (list of 6 integers).

**Exercise 2 (code).** Write `my_rwpe(g, K)` returning an $n\times K$ tensor whose column $k$ is $\mathrm{diag}\big((D^{-1}A)^{k}\big)$, $k=1..K$.
It must match `T.AddRandomWalkPE(walk_length=K)`.
""")
mo.vstack([_rev, _post])

ex1_c5 = None

def my_rwpe(g, K):
    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_rwpe():
    try:
        g = nx.karate_club_graph()
        got = my_rwpe(g, 5)
        if got is None:
            return _verdict(False, False)
        ref = T.AddRandomWalkPE(walk_length=5, attr_name="pe")(
            from_networkx(g)
        ).pe
        return _verdict(
            torch.allclose(
                torch.as_tensor(got, dtype=torch.float), ref, atol=1e-5
            ),
            True,
        )
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_c5 is not None and list(ex1_c5) == [1, 0, 2, 0, 6, 2],
            ex1_c5 is not None,
        ),
        mo.md("**Exercise 2**"),
        _check_rwpe(),
        mo.accordion(
            {
                "Solution 1": mo.md(r"""$[1, 0, 2, 0, 6, 2]$. Length 2: step out and back (2 ways). Length 3: 0 (no triangle). Length 4: 6.
Length 5: **2**, going once round the cycle either way. A 5-cycle node is the first to show a non-zero odd entry at $k=5$,
just as a triangle node does at $k=3$. Check: `closed_walks(nx.cycle_graph(5), 6)[0]`."""),
                "Solution 2": mo.md("""```python
def my_rwpe(g, K):
    A = torch.tensor(nx.to_numpy_array(g, weight=None), dtype=torch.float)
    P = A / A.sum(1, keepdim=True)            # D^-1 A
    M, cols = torch.eye(len(A)), []
    for _ in range(K):
        M = M @ P
        cols.append(torch.diagonal(M))        # return probability after k steps
    return torch.stack(cols, 1)
```"""),
            }
        ),
    ]
)