def walks_to_pairs(walks, window):
    """All (centre, context) pairs within ±window steps (plain Python, fine for small graphs)."""
    centres, contexts = [], []
    for walk in walks:
        for i, u in enumerate(walk):
            for v in walk[max(0, i - window) : i + window + 1]:
                if v != u:
                    centres.append(u)
                    contexts.append(v)
    return np.array(centres), np.array(contexts)

def train_sgns(
    centres,
    contexts,
    n_nodes,
    dim=16,
    n_neg=5,
    epochs=3,
    lr=0.02,
    batch=1024,
    seed=SEED,
):
    """Skip-gram with negative sampling on (centre, context) pairs, ONE table Z (decoder z_u . z_v).
    Returns (Z as numpy, per-epoch mean losses)."""
    torch.manual_seed(seed)
    centres, contexts = torch.as_tensor(centres), torch.as_tensor(contexts)
    noise = (
        torch.bincount(centres, minlength=n_nodes).float().pow(0.75)
    )  # P_n ∝ freq^(3/4)
    noise /= noise.sum()
    Z = torch.nn.Embedding(n_nodes, dim)
    torch.nn.init.normal_(Z.weight, std=0.1)
    opt = torch.optim.Adam(Z.parameters(), lr=lr)
    losses = []
    for _ in range(epochs):
        perm, total = torch.randperm(len(centres)), 0.0
        for k in range(0, len(perm), batch):
            idx = perm[k : k + batch]
            zu, zv = Z(centres[idx]), Z(contexts[idx])
            zn = Z(
                torch.multinomial(
                    noise, len(idx) * n_neg, replacement=True
                )
            ).view(len(idx), n_neg, dim)
            pos = F.logsigmoid((zu * zv).sum(-1))  # log σ(z_u·z_v)
            neg = F.logsigmoid(-(zn @ zu.unsqueeze(-1)).squeeze(-1)).sum(
                -1
            )  # Σ log σ(-z_u·z_n)
            loss = -(pos + neg).mean()
            opt.zero_grad()
            loss.backward()
            opt.step()
            total += loss.item() * len(idx)
        losses.append(total / len(perm))
    return Z.weight.detach().numpy(), losses

emb_controls = mo.ui.dictionary(
    {
        "p": mo.ui.slider(
            steps=[0.25, 0.5, 1, 2, 4], value=1, label="p", show_value=True
        ),
        "q": mo.ui.slider(
            steps=[0.25, 0.5, 1, 2, 4], value=1, label="q", show_value=True
        ),
        "window": mo.ui.slider(
            1, 6, value=3, label="window", show_value=True
        ),
        "dim": mo.ui.slider(
            steps=[2, 4, 8, 16, 32],
            value=16,
            label="dim d",
            show_value=True,
        ),
    }
)
mo.vstack(
    [
        mo.md(
            "**Try it:** train node2vec on Karate (≈ 1 s). Compare $q=4$ (local) with $q=0.25$ (exploratory), and $d=2$ with $d=16$."
        ),
        mo.hstack(list(emb_controls.elements.values()), justify="start"),
    ]
)

_t0 = time.time()
_r = np.random.default_rng(SEED)
_sets = {v: set(karate[v]) for v in karate}
_walks = [
    node2vec_walk(
        _sets, v, 10, emb_controls.value["p"], emb_controls.value["q"], _r
    )
    for v in karate
    for _ in range(20)
]
karate_Z, _losses = train_sgns(
    *walks_to_pairs(_walks, emb_controls.value["window"]),
    34,
    dim=emb_controls.value["dim"],
    epochs=5,
)
_secs = time.time() - _t0

_acc = cross_val_score(
    LogisticRegression(max_iter=2000),
    karate_Z,
    club,
    cv=StratifiedKFold(5, shuffle=True, random_state=SEED),
).mean()
_xy = (
    karate_Z
    if karate_Z.shape[1] == 2
    else PCA(2, random_state=SEED).fit_transform(karate_Z)
)
_fig, _axs = plt.subplots(
    1, 2, figsize=(8, 3.2), gridspec_kw={"width_ratios": [1, 1.6]}
)
_axs[0].plot(range(1, len(_losses) + 1), _losses, "o-")
_axs[0].set_xlabel("epoch")
_axs[0].set_title("SGNS loss", fontsize=9)
_axs[1].scatter(
    _xy[:, 0],
    _xy[:, 1],
    c=["#4C78A8" if c == 0 else "#F58518" for c in club],
    s=60,
    edgecolors="k",
    linewidths=0.4,
)
for _v in karate:
    _axs[1].annotate(
        str(_v), _xy[_v], fontsize=6, ha="center", va="center"
    )
_kind = (
    "raw 2-D"
    if karate_Z.shape[1] == 2
    else f"PCA of {karate_Z.shape[1]}-D"
)
_axs[1].set_title(
    f"embeddings ({_kind}); faction accuracy {_acc:.2f}", fontsize=9
)
plt.tight_layout()
mo.vstack(
    [_fig, mo.md(f"trained in {_secs:.1f} s on {len(_walks)} walks")]
)

cora = planetoid("Cora")[0]
print(cora)

def node2vec_walks_fast(
    edge_index, n, walks_per_node, length, p, q, seed=SEED
):
    """Vectorised node2vec walks (rejection sampling). Returns an int array [n * walks_per_node, length]."""
    r = np.random.default_rng(seed)
    src, dst = edge_index.numpy()
    order = np.argsort(src, kind="stable")
    src, dst = src[order], dst[order]
    rowptr = np.searchsorted(src, np.arange(n + 1))
    deg = np.diff(rowptr)
    edge_keys = np.sort(
        src.astype(np.int64) * n + dst
    )  # for "is x adjacent to prev?" lookups
    is_edge = lambda a, b: np.isin(
        a.astype(np.int64) * n + b, edge_keys, assume_unique=False
    )
    starts = np.repeat(np.arange(n), walks_per_node)
    starts = starts[deg[starts] > 0]
    walks = np.empty((len(starts), length), dtype=np.int64)
    walks[:, 0] = starts
    walks[:, 1] = dst[
        rowptr[starts] + (r.random(len(starts)) * deg[starts]).astype(int)
    ]
    max_w = max(1 / p, 1.0, 1 / q)
    for t in range(2, length):
        prev, cur = walks[:, t - 2], walks[:, t - 1]
        nxt = np.empty_like(cur)
        todo = np.arange(len(cur))
        while len(todo):
            c = cur[todo]
            x = dst[
                rowptr[c] + (r.random(len(todo)) * deg[c]).astype(int)
            ]  # uniform proposal
            alpha = np.where(
                x == prev[todo],
                1 / p,
                np.where(is_edge(prev[todo], x), 1.0, 1 / q),
            )
            accept = r.random(len(todo)) < alpha / max_w
            nxt[todo[accept]] = x[accept]
            todo = todo[~accept]
        walks[:, t] = nxt
    return walks

def walks_to_pairs_array(walks, window):
    """Vectorised (centre, context) pairs for a [num_walks, length] walk array."""
    c, o = [], []
    for k in range(1, window + 1):
        a, b = walks[:, :-k].ravel(), walks[:, k:].ravel()
        c += [a, b]
        o += [b, a]
    return np.concatenate(c), np.concatenate(o)

_t = time.time()
_demo_walks = node2vec_walks_fast(
    cora.edge_index, cora.num_nodes, 10, 20, p=1, q=0.5
)
mo.md(
    f"Generated **{_demo_walks.shape[0]:,} walks** of length {_demo_walks.shape[1]} in **{time.time() - _t:.1f} s**. "
    f"First walk: `{_demo_walks[0].tolist()}`"
)

def cached_array(tag, params, compute):
    """Load `tag_<hash(params)>.npy` from CACHE_DIR, or compute + save. Returns (array, was_cached)."""
    key = hashlib.sha1(
        json.dumps(params, sort_keys=True).encode()
    ).hexdigest()[:12]
    path = CACHE_DIR / f"{tag}_{key}.npy"
    if path.exists():
        return np.load(path), True
    arr = compute()
    np.save(path, arr)
    return arr, False

def make_link_split(edge_index, n, frac=0.1, seed=SEED):
    """Hide `frac` of the undirected edges; return (train edge_index, test positives, test negatives)."""
    g = torch.Generator().manual_seed(seed)
    und = edge_index[
        :, edge_index[0] < edge_index[1]
    ]  # each undirected edge once
    perm = torch.randperm(und.size(1), generator=g)
    n_test = int(frac * und.size(1))
    test_pos, train_und = und[:, perm[:n_test]], und[:, perm[n_test:]]
    torch.manual_seed(seed)
    test_neg = negative_sampling(
        edge_index, num_nodes=n, num_neg_samples=n_test
    )  # true non-edges
    return to_undirected(train_und), test_pos, test_neg

cora_train_ei, cora_test_pos, cora_test_neg = make_link_split(
    cora.edge_index, cora.num_nodes
)

def node2vec_embedding(
    edge_index, tag, p, q, dim=64, walks_per_node=10, length=20, window=5
):
    params = dict(
        tag=tag,
        p=p,
        q=q,
        dim=dim,
        walks_per_node=walks_per_node,
        length=length,
        window=window,
        seed=SEED,
    )

    def compute():
        walks = node2vec_walks_fast(
            edge_index, cora.num_nodes, walks_per_node, length, p, q
        )
        Z, _ = train_sgns(
            *walks_to_pairs_array(walks, window),
            cora.num_nodes,
            dim=dim,
            epochs=1,
            lr=0.01,
            batch=4096,
        )
        return Z

    return cached_array("cora_n2v", params, compute)

def evaluate_cora(Z_full, Z_train):
    tr, te, y = (
        cora.train_mask.numpy(),
        cora.test_mask.numpy(),
        cora.y.numpy(),
    )
    acc = (
        LogisticRegression(max_iter=3000)
        .fit(Z_full[tr], y[tr])
        .score(Z_full[te], y[te])
    )
    score = lambda Z, e: (Z[e[0].numpy()] * Z[e[1].numpy()]).sum(1)
    auc = roc_auc_score(
        np.r_[
            np.ones(cora_test_pos.size(1)), np.zeros(cora_test_neg.size(1))
        ],
        np.r_[
            score(Z_train, cora_test_pos), score(Z_train, cora_test_neg)
        ],
    )
    return acc, auc

cora_setting = mo.ui.dropdown(
    {
        "DeepWalk  (p=1, q=1)": (1, 1),
        "BFS-like  (p=1, q=4)": (1, 4),
        "DFS-like  (p=1, q=0.25)": (1, 0.25),
    },
    value="DeepWalk  (p=1, q=1)",
    label="walk setting",
)
cora_run = mo.ui.run_button(label="Train / load embeddings")
mo.hstack([cora_setting, cora_run], justify="start")

mo.stop(
    not cora_run.value,
    mo.md(
        "Press **Train / load embeddings** (≈ 1 min first time per setting, instant afterwards)."
    ),
)
_p, _q = cora_setting.value
_t = time.time()
cora_Z, _hit1 = node2vec_embedding(cora.edge_index, "full", _p, _q)
cora_Z_train, _hit2 = node2vec_embedding(
    cora_train_ei, "train-edges", _p, _q
)
_acc, _auc = evaluate_cora(cora_Z, cora_Z_train)
_tr, _te, _y = (
    cora.train_mask.numpy(),
    cora.test_mask.numpy(),
    cora.y.numpy(),
)
_raw_acc = (
    LogisticRegression(max_iter=3000)
    .fit(cora.x.numpy()[_tr], _y[_tr])
    .score(cora.x.numpy()[_te], _y[_te])
)
_xy = PCA(2, random_state=SEED).fit_transform(cora_Z)
_fig, _ax = plt.subplots(figsize=(5.5, 3.6))
_ax.scatter(_xy[:, 0], _xy[:, 1], c=_y, cmap="tab10", s=4, alpha=0.7)
_ax.set_title(
    f"Cora node2vec embeddings (PCA), colour = topic — {cora_setting.selected_key}",
    fontsize=8,
)
_ax.set_xticks([])
_ax.set_yticks([])
mo.hstack(
    [
        _fig,
        mo.md(f"""
| metric | value |
|---|---|
| node classification acc (embedding) | **{_acc:.3f}** |
| node classification acc (raw bag-of-words features) | {_raw_acc:.3f} |
| link prediction AUC (held-out edges) | **{_auc:.3f}** |
| cache | {"hit (loaded from disk)" if _hit1 and _hit2 else "computed and saved"} |
| time | {time.time() - _t:.1f} s |
"""),
    ],
    widths=[3, 2],
)

_rev = reveal(
    r"""
Measured on this machine (one run each, seed 0, 64-D, 1 epoch):

| walk setting | node classification acc | link-prediction AUC |
|---|---|---|
| DeepWalk (p=1, q=1) | 0.654 | 0.854 |
| BFS-like (p=1, q=4) | 0.594 | 0.853 |
| DFS-like (p=1, q=0.25) | 0.636 | 0.855 |

All beat logistic regression on the raw bag-of-words features (0.576) **without seeing a single label during training**.
Topic labels on Cora follow **communities** (homophily), so the local BFS-like walks are worst for classification, while
link prediction barely cares. In the PCA plot most topics (colours) occupy their own region. One topic (red) is clearly
separated; the rest overlap, which is why a linear classifier tops out around 0.65.

""",
    predict="Which walk setting (DeepWalk p=q=1, BFS-like q=4, DFS-like q=0.25) classifies Cora's topics best, and do any beat logistic regression on the raw words (≈ 0.58) without seeing a label? Hint: are Cora's topics communities or structural roles?",
)
_post = mo.md(r"""
## 6 · It was matrix factorisation all along (NetMF)

Qiu et al. (2018) showed that DeepWalk with window $T$ and $b$ negatives implicitly factorises
$$M=\log\!\Big(\tfrac{\text{vol}(G)}{bT}\Big(\textstyle\sum_{r=1}^{T}P^r\Big)D^{-1}\Big),\qquad P=D^{-1}A,\ \ \text{vol}(G)=\textstyle\sum_v d_v.$$
So we can skip walks and SGD: build $M$ (clip at 1 before the log, as in the paper), take a rank-$d$ SVD $M\approx U_d\Sigma_dV_d^\top$,
and use $Z=U_d\sqrt{\Sigma_d}$. Exact, deterministic, and on Cora only a few seconds.
""")
mo.vstack([_rev, _post])

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

_A = nx.to_numpy_array(karate, nodelist=range(34), weight=None)
_pi_true = _A.sum(1) / _A.sum()
_ex2_ok = (
    ex2_pi is not None
    and np.shape(ex2_pi) == (34,)
    and np.allclose(ex2_pi, _pi_true, atol=1e-4)
)
mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_back_prob is not None
            and abs(ex1_back_prob - 2 / 3.5) < 2e-3,
            ex1_back_prob is not None,
        ),
        mo.md("**Exercise 2**"),
        _verdict(_ex2_ok, ex2_pi is not None),
        mo.accordion(
            {
                "Hint 1": mo.md(
                    "Weights: $t: 1/p$, $s_1: 1$, $s_2, s_3: 1/q$. Normalise."
                ),
                "Solution 1": mo.md(
                    r"Weights $2, 1, 0.25, 0.25$, sum 3.5, so $P(\text{back})=2/3.5\approx 0.571$. Or call `node2vec_step_probs(toy_sets, 't', 'w', 0.5, 4)`."
                ),
                "Hint 2": mo.md(
                    "Try power iteration: start with uniform $\\pi$ and repeat `pi = pi @ P`. Then look at the result next to the degrees."
                ),
                "Solution 2": mo.md(r"""For an undirected connected graph $\pi_v = d_v / 2|E|$: detailed balance $\pi_u P_{uv} = \frac{d_u}{2|E|}\frac{1}{d_u} = \frac{1}{2|E|}$ is symmetric.
```python
A = nx.to_numpy_array(karate, nodelist=range(34), weight=None)
ex2_pi = A.sum(1) / A.sum()
```
(Power iteration also converges here because Karate is connected and non-bipartite.) Hubs are visited most, which is why
negative sampling uses $\text{freq}^{3/4}$ to flatten the distribution a little."""),
            }
        ),
    ]
)