M_LOOKBACK = 8

def grid_dataset(n=200, seed=SEED):
    r = np.random.default_rng(seed)
    return [
        nx.convert_node_labels_to_integers(
            nx.grid_2d_graph(int(r.integers(2, 5)), int(r.integers(3, 7)))
        )
        for _ in range(n)
    ]

def bfs_sequence(g, r):
    order = list(nx.bfs_tree(g, int(r.integers(g.number_of_nodes()))))
    S = np.zeros((len(order), M_LOOKBACK))
    for i in range(1, len(order)):
        for j in range(max(0, i - M_LOOKBACK), i):
            if g.has_edge(order[i], order[j]):
                S[i, i - 1 - j] = 1.0  # column k = edge to node i-1-k
    return S

class GraphRNN(nn.Module):
    def __init__(self, hidden=64):
        super().__init__()
        self.inp, self.gru, self.out = (
            nn.Linear(M_LOOKBACK, hidden),
            nn.GRU(hidden, hidden, batch_first=True),
            nn.Linear(hidden, M_LOOKBACK + 1),
        )

    def forward(self, x, h=None):
        o, h = self.gru(self.inp(x), h)
        return self.out(o), h  # M edge logits + 1 stop logit

def batch_from(seqs, max_n):
    B = len(seqs)
    X, Y, W = (
        np.zeros((B, max_n, M_LOOKBACK)),
        np.zeros((B, max_n, M_LOOKBACK + 1)),
        np.zeros((B, max_n, M_LOOKBACK + 1)),
    )
    for b, S in enumerate(seqs):
        n = len(S)
        X[b, 0] = 1.0
        X[b, 1:n] = S[
            1:n
        ]  # start-of-sequence row of ones, then teacher forcing
        Y[b, : n - 1, :M_LOOKBACK] = S[1:n]
        Y[b, n - 1, M_LOOKBACK] = (
            1.0  # next row's edges; stop after the last node
        )
        for i in range(n):
            W[b, i, : min(i + 1, M_LOOKBACK)] = (
                1.0 if i < n - 1 else 0.0
            )  # only edges to nodes that exist
            W[b, i, M_LOOKBACK] = 1.0
    t = lambda a: torch.tensor(a, dtype=torch.float)
    return t(X), t(Y), t(W)

def train_graphrnn(graphs, steps=3000, seed=SEED):
    torch.manual_seed(seed)
    r = np.random.default_rng(seed)
    model = GraphRNN()
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    max_n = max(g.number_of_nodes() for g in graphs) + 1
    losses = []
    for step in range(steps):
        X, Y, W = batch_from(
            [
                bfs_sequence(graphs[i], r)
                for i in r.choice(len(graphs), 32)
            ],
            max_n,
        )
        loss = (
            F.binary_cross_entropy_with_logits(
                model(X)[0], Y, reduction="none"
            )
            * W
        ).sum() / W.sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
        losses.append(loss.item())
    return model, losses

@torch.no_grad()
def sample_graph(model, max_nodes=40, generator=None):
    g = nx.Graph()
    g.add_node(0)
    x, h = torch.ones(1, 1, M_LOOKBACK), None
    for i in range(1, max_nodes):
        logits, h = model(x, h)
        p = torch.sigmoid(logits[0, 0])
        if torch.rand(1, generator=generator) < p[M_LOOKBACK]:
            break  # stop
        for _ in range(
            20
        ):  # reject empty rows (GraphRNN-S simplification)
            row = (
                torch.rand(M_LOOKBACK, generator=generator)
                < p[:M_LOOKBACK]
            ).float()
            row[i:] = 0
            if row.sum() > 0:
                break
        g.add_node(i)
        g.add_edges_from(
            (i, i - 1 - k) for k in range(M_LOOKBACK) if row[k]
        )
        x = row.view(1, 1, M_LOOKBACK)
    return g

grid_graphs = grid_dataset()
graphrnn_run = mo.ui.run_button(
    label="Train GraphRNN-lite (≈ 15 s) and sample 100 graphs"
)
graphrnn_run

mo.stop(
    not graphrnn_run.value, mo.md("Press the button to train and sample.")
)
_t = time.time()
graphrnn_model, _losses = train_graphrnn(grid_graphs)
_gen = torch.Generator().manual_seed(SEED)
generated = [
    sample_graph(graphrnn_model, generator=_gen) for _ in range(100)
]
_er = [
    nx.gnm_random_graph(g.number_of_nodes(), g.number_of_edges(), seed=i)
    for i, g in enumerate(generated)
]  # same sizes, random wiring

def graph_stats(gs):
    degs = Counter(d for g in gs for _, d in g.degree())
    tot = sum(degs.values())
    return {
        "nodes (mean)": np.mean([g.number_of_nodes() for g in gs]),
        "mean degree": np.mean(
            [2 * g.number_of_edges() / g.number_of_nodes() for g in gs]
        ),
        "clustering": np.mean([nx.average_clustering(g) for g in gs]),
        "connected": np.mean([nx.is_connected(g) for g in gs]),
        "share of degree-1 nodes": degs[1] / tot,
        "share of degree ≥ 5 nodes": sum(
            v for k, v in degs.items() if k >= 5
        )
        / tot,
    }

graphrnn_eval = pd.DataFrame(
    {
        "data (grids)": graph_stats(grid_graphs),
        "GraphRNN-lite samples": graph_stats(generated),
        "Erdős–Rényi, same sizes": graph_stats(_er),
    }
).T
_fig, _axs = plt.subplots(1, 5, figsize=(11, 2.3))
for _ax, _g in zip(_axs, generated[:5]):
    nx.draw(
        _g, nx.kamada_kawai_layout(_g), ax=_ax, node_size=25, width=0.8
    )
    _ax.set_title(f"{_g.number_of_nodes()} nodes", fontsize=8)
mo.vstack(
    [
        mo.md(
            f"trained in {time.time() - _t:.0f} s; final BCE {np.mean(_losses[-100:]):.3f}"
        ),
        graphrnn_eval.round(3),
        _fig,
    ]
)

_rev = reveal(
    r"""
Measured (100 samples): the samples look like **grids and ladders** (square faces, degrees 2–4) and match the data's mean degree
(2.88 vs 2.83); all are connected. They are not perfect: a few triangles (clustering 0.08 vs 0 in grids) and a few degree-1 / degree-5
nodes slip in, the price of sampling the $M$ edge bits independently (the full GraphRNN's edge-level RNN fixes exactly this).
Samples are a bit smaller (12 vs 14 nodes): the stop head fires slightly early. An **Erdős–Rényi** graph with the same sizes is much further
off (clustering 0.23, 11% leaves, 13% hubs, 22% disconnected): the model has learned *structure*, not just density.

""",
    predict="Compared with the training grids, what will GraphRNN-lite samples get right (mean degree, connectivity) and wrong (triangles, size)?",
)
_post = mo.md(r"""
## 4 · GCPN-lite: goal-directed generation with RL

GCPN builds molecules atom by atom with a **GNN policy**, checks **chemical validity** at each step (step reward), and optimises a
**property** score at the end with policy gradients. Our miniature: 10 nodes, 12 "add an edge" actions; **valence rule** = max degree
3 (an invalid action is penalised −0.2 and ignored); **property** = number of triangles. The policy is a 2-layer GNN scoring every
non-edge; training is REINFORCE with a moving-average baseline:
$\nabla J\approx\sum_t (R-b)\nabla\log\pi(a_t\mid G_t)$.
""")
mo.vstack([_rev, _post])

RL_N, RL_STEPS, RL_MAXDEG = 10, 12, 3
RL_PAIRS = torch.tensor(list(itertools.combinations(range(RL_N), 2)))

class EdgePolicy(nn.Module):
    def __init__(self, d=32):
        super().__init__()
        self.inp = nn.Linear(RL_MAXDEG + 2, d)
        self.layers = nn.ModuleList(nn.Linear(2 * d, d) for _ in range(2))
        self.edge = nn.Sequential(
            nn.Linear(2 * d, d), nn.ReLU(), nn.Linear(d, 1)
        )

    def forward(self, A):
        deg = A.sum(1).long().clamp(max=RL_MAXDEG + 1)
        h = F.relu(self.inp(F.one_hot(deg, RL_MAXDEG + 2).float()))
        for lin in self.layers:
            h = F.relu(
                lin(torch.cat([h, A @ h], 1))
            )  # message passing on the current graph
        u, v = RL_PAIRS[:, 0], RL_PAIRS[:, 1]
        return self.edge(torch.cat([h[u] + h[v], h[u] * h[v]], 1)).squeeze(
            1
        )  # one logit per candidate edge

def count_triangles(A):
    return int(torch.trace(A @ A @ A).item() / 6)

def rollout(policy, greedy=False):
    A, logps, penalty = torch.zeros(RL_N, RL_N), [], 0.0
    for _ in range(RL_STEPS):
        logits = policy(A.clone()).masked_fill(
            A[RL_PAIRS[:, 0], RL_PAIRS[:, 1]] > 0, -1e9
        )  # can't re-add an edge
        dist = torch.distributions.Categorical(logits=logits)
        a = logits.argmax() if greedy else dist.sample()
        logps.append(dist.log_prob(a))
        u, v = RL_PAIRS[a]
        if A[u].sum() >= RL_MAXDEG or A[v].sum() >= RL_MAXDEG:
            penalty -= 0.2  # invalid: violates the valence rule
        else:
            A[u, v] = A[v, u] = 1.0
    return A, torch.stack(logps), count_triangles(A) + penalty

def train_gcpn(iters=250, batch=16, seed=SEED):
    torch.manual_seed(seed)
    policy = EdgePolicy()
    opt = torch.optim.Adam(policy.parameters(), lr=3e-3)
    baseline, curve = 0.0, []
    for _ in range(iters):
        eps = [rollout(policy) for _ in range(batch)]
        R = torch.tensor([e[2] for e in eps])
        baseline = 0.9 * baseline + 0.1 * R.mean().item()
        loss = (
            -sum(
                (r - baseline) * lp.sum() for (_, lp, _), r in zip(eps, R)
            )
            / batch
        )
        opt.zero_grad()
        loss.backward()
        opt.step()
        curve.append(R.mean().item())
    return policy, curve

gcpn_run = mo.ui.run_button(label="Train the RL generator (≈ 20 s)")
gcpn_run

_rev = reveal(
    r"""
Measured: the untrained policy averages **0.28** triangles; after ≈ 60 iterations REINFORCE reaches **8**, the optimum under the
rules: 12 edges and max degree 3 allow at most two disjoint $K_4$'s, 4 triangles each, and that is what the greedy sample builds,
leaving two nodes unused. The policy also learned to **never** violate the valence rule (no −0.2 penalties at the end). Step rewards
for validity plus a final property reward is GCPN's recipe; the real system adds an adversarial "looks like a real molecule" reward and
imitation of real molecules, so it doesn't drift into chemically odd but high-scoring structures.

""",
    predict="A small olympiad problem first: with 12 edges and maximum degree 3, what is the largest possible number of triangles? Will REINFORCE find that optimum?",
)
_post = mo.md(r"""
---

## 5 · Exercises

**Exercise 1 (likelihood).** A model generates a triangle in order 1, 2, 3: $p(2\text{–}1)=0.8$ (edge), $p(3\text{–}1)=0.5$ (edge),
$p(3\text{–}2)=0.4$ (edge), then stops with $p=0.9$. What is $p(S^\pi)$? → `ex1_p` (4 decimals).

**Exercise 2 (BFS look-back).** What maximum look-back $M$ does BFS need on an $a\times b$ grid started at a corner, with $a \le b$
(rows ≤ columns)? Check your formula numerically and put the value for a $4\times7$ grid in `ex2_M`.

**Exercise 3 (reward design).** Change the RL property to "number of **squares** (4-cycles)". Which structure should the optimal
12-edge, degree-≤3 graph on 10 nodes have? Write your guess, then open the solution.
""")
mo.vstack([_rev, _post])

ex1_p = None
ex2_M = 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",
    )

_g47 = nx.convert_node_labels_to_integers(
    nx.grid_2d_graph(4, 7), ordering="sorted"
)
_M47 = max_lookback(_g47, list(nx.bfs_tree(_g47, 0)))
mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_p is not None
            and abs(ex1_p - 0.8 * 0.5 * 0.4 * 0.9) < 1e-4,
            ex1_p is not None,
        ),
        mo.md("**Exercise 2**"),
        _verdict(ex2_M == _M47, ex2_M is not None),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"$0.8\cdot0.5\cdot0.4\cdot0.9=$ **0.144**."
                ),
                "Solution 2": mo.md(
                    f"Measured with `max_lookback(nx.grid_2d_graph(4, 7) …, bfs order)`: **{_M47}**. BFS from a corner sweeps diagonals, "
                    "so the look-back equals the length of the longest diagonal frontier: the shorter side, **M = a**."
                ),
                "Solution 3": mo.md(
                    "A **cube** $Q_3$ (8 nodes, 12 edges, every degree 3) has 6 square faces, the maximum here; so the optimum is a cube plus 2 "
                    "isolated nodes. Try it: replace `count_triangles` with a 4-cycle counter (e.g. via `(trace(A⁴) − 2|E| − 4·Σ C(d,2))/8`)."
                ),
            }
        ),
    ]
)