def scatter_sum(src, index, n):
    return torch.zeros(n, src.size(1)).index_add_(0, index, src)

class StandardHeteroMP(nn.Module):
    def __init__(self, n_layers, d=32):
        super().__init__()
        self.enc_c, self.enc_t, self.enc_p = (
            nn.Linear(1, d),
            nn.Linear(1, d),
            nn.Linear(6, d),
        )
        self.layers = nn.ModuleList(
            nn.ModuleDict(
                {
                    k: nn.Linear(d, d)
                    for k in [
                        "self_c",
                        "self_t",
                        "self_p",
                        "t2c",
                        "c2t",
                        "p2t",
                        "t2p",
                    ]
                }
            )
            for _ in range(n_layers)
        )
        self.out = nn.Linear(d, 1)

    def forward(self):
        c, t, p = self.enc_c(x_cust), self.enc_t(x_txn), self.enc_p(x_prod)
        for W in self.layers:
            c, t, p = (
                F.relu(
                    W["self_c"](c)
                    + scatter_sum(W["t2c"](t), txn_cust, n_cust)
                ),
                F.relu(
                    W["self_t"](t)
                    + W["c2t"](c)[txn_cust]
                    + W["p2t"](p)[txn_prod]
                ),
                F.relu(
                    W["self_p"](p)
                    + scatter_sum(W["t2p"](t), txn_prod, n_prod)
                ),
            )
        return self.out(c).squeeze(1)

class CompositeMP(nn.Module):
    """RelGNN-style: ONE step along the atomic route customer ← (transaction ⊕ product)."""

    def __init__(self, d=32):
        super().__init__()
        self.enc_c, self.enc_t, self.enc_p = (
            nn.Linear(1, d),
            nn.Linear(1, d),
            nn.Linear(6, d),
        )
        self.fuse = nn.Sequential(
            nn.Linear(2 * d, d), nn.ReLU(), nn.Linear(d, d)
        )  # FUSE(T, P) keeps the pairing
        self.self_c, self.out = nn.Linear(d, d), nn.Linear(d, 1)

    def forward(self):
        msg = self.fuse(
            torch.cat([self.enc_t(x_txn), self.enc_p(x_prod)[txn_prod]], 1)
        )
        return self.out(
            F.relu(
                self.self_c(self.enc_c(x_cust))
                + scatter_sum(msg, txn_cust, n_cust)
            )
        ).squeeze(1)

def fit_customer_model(model, epochs=300, seed=SEED):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    t0 = time.time()
    for _ in range(epochs):
        opt.zero_grad()
        F.binary_cross_entropy_with_logits(
            model()[train_mask], y_cust[train_mask]
        ).backward()
        opt.step()
    with torch.no_grad():
        acc = (
            ((model() > 0).float()[~train_mask] == y_cust[~train_mask])
            .float()
            .mean()
            .item()
        )
    return {
        "test acc": acc,
        "parameters": sum(p.numel() for p in model.parameters()),
        "seconds": time.time() - t0,
    }

mp_results = pd.DataFrame(
    [
        {
            "model": "standard MP, 1 layer",
            **fit_customer_model(
                (torch.manual_seed(SEED), StandardHeteroMP(1))[1]
            ),
        },
        {
            "model": "standard MP, 2 layers",
            **fit_customer_model(
                (torch.manual_seed(SEED), StandardHeteroMP(2))[1]
            ),
        },
        {
            "model": "standard MP, 3 layers",
            **fit_customer_model(
                (torch.manual_seed(SEED), StandardHeteroMP(3))[1]
            ),
        },
        {
            "model": "composite MP (RelGNN-style), 1 step",
            **fit_customer_model(
                (torch.manual_seed(SEED), CompositeMP())[1]
            ),
        },
    ]
)
mp_results.round(3)

_rev = reveal(
    r"""
Measured (one seed): **standard MP with 1 layer = 0.724**, exactly "always negative": it never reaches the categories.
**2 layers: 0.964**, 3 layers: 0.944 (more layers, more redundancy and over-fitting). **Composite MP: 0.966 in one step with
4.6k parameters** instead of 15k, and roughly half the time. On RelBench the lecture reports up to 25% gains. Here the
benefit shows up as efficiency: one hop of the *right* shape does the work of two generic hops.

""",
    predict="The label needs customer → transaction → product. What does 1-layer standard message passing score? What do 2 layers, and one composite step, achieve?",
)
_post = mo.md(r"""
## 4 · RelGT: a relational graph transformer

RelGT samples a local context around each seed entity and represents **every token** by a 5-tuple:
**(features, node type, hop distance, relative time, positional encoding)**. Here: the seed customer, its ≤ 11 transactions and
their products → ≤ 23 tokens. Time is omitted (this toy database is static); as positional encoding we use a simple
**local-structure code**: a transaction and *its* product share a one-hot slot id, so attention can tell which product belongs
to which purchase.
""")
mo.vstack([_rev, _post])

MAX_TXN = 11
N_TOK = 1 + 2 * MAX_TXN
TYPE_NAMES = ["customer", "transaction", "product"]

def relgt_tokens():
    """Per customer: token features (7 = amount | category one-hot | log price), type, hop, slot PE, padding mask."""
    order = np.argsort(txn_table.customer_id.values, kind="stable")
    start = np.searchsorted(
        txn_table.customer_id.values[order], np.arange(n_cust + 1)
    )
    feat = torch.zeros(n_cust, N_TOK, 7)
    typ = torch.zeros(n_cust, N_TOK, dtype=torch.long)
    hop = torch.zeros(n_cust, N_TOK, dtype=torch.long)
    pe = torch.zeros(n_cust, N_TOK, MAX_TXN)
    mask = torch.zeros(n_cust, N_TOK, dtype=torch.bool)
    for c in range(n_cust):
        feat[c, 0, 0] = 1.0
        mask[c, 0] = True  # seed token
        for j, ti in enumerate(order[start[c] : start[c + 1]][:MAX_TXN]):
            a, b = 1 + 2 * j, 2 + 2 * j
            feat[c, a, 0] = x_txn[ti, 0]
            typ[c, a] = 1
            hop[c, a] = 1
            feat[c, b, 1:] = x_prod[txn_prod[ti]]
            typ[c, b] = 2
            hop[c, b] = 2
            pe[c, a, j] = pe[c, b, j] = 1.0  # shared slot id
            mask[c, a] = mask[c, b] = True
    return feat, typ, hop, pe, mask

tok_feat, tok_type, tok_hop, tok_pe, tok_mask = relgt_tokens()
_c = 0
pd.DataFrame(
    [
        {
            "token": i,
            "type": TYPE_NAMES[tok_type[_c, i]],
            "hop": int(tok_hop[_c, i]),
            "features (amount | category | log price)": [
                round(v, 2) for v in tok_feat[_c, i].tolist()
            ],
            "slot PE": int(tok_pe[_c, i].argmax())
            if tok_pe[_c, i].any()
            else "—",
        }
        for i in range(N_TOK)
        if tok_mask[_c, i]
    ]
)

class RelGTLite(nn.Module):
    def __init__(self, use_pe, d=32):
        super().__init__()
        self.use_pe = use_pe
        self.feat_enc = nn.Linear(7, d)
        self.type_emb = nn.Embedding(3, d)
        self.hop_emb = nn.Embedding(3, d)
        self.pe_enc = nn.Linear(MAX_TXN, d)
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d, 4, 64, dropout=0.0, batch_first=True
            ),
            2,
            enable_nested_tensor=False,
        )
        self.out = nn.Linear(d, 1)

    def forward(self, idx):
        h = (
            self.feat_enc(tok_feat[idx])
            + self.type_emb(tok_type[idx])
            + self.hop_emb(tok_hop[idx])
        )
        if self.use_pe:
            h = h + self.pe_enc(tok_pe[idx])
        h = self.encoder(h, src_key_padding_mask=~tok_mask[idx])
        return self.out(h[:, 0]).squeeze(1)  # read out the seed token

def fit_relgt(use_pe, epochs=60, seed=SEED):
    torch.manual_seed(seed)
    model = RelGTLite(use_pe)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    tr_idx, te_idx = (
        train_mask.nonzero().squeeze(),
        (~train_mask).nonzero().squeeze(),
    )
    t0 = time.time()
    for _ in range(epochs):
        model.train()
        for b in tr_idx[torch.randperm(len(tr_idx))].split(100):
            opt.zero_grad()
            F.binary_cross_entropy_with_logits(
                model(b), y_cust[b]
            ).backward()
            opt.step()
    model.eval()
    with torch.no_grad():
        acc = (
            ((model(te_idx) > 0).float() == y_cust[te_idx])
            .float()
            .mean()
            .item()
        )
    return {
        "test acc": acc,
        "parameters": sum(p.numel() for p in model.parameters()),
        "seconds": time.time() - t0,
    }

relgt_run = mo.ui.run_button(
    label="Train RelGT-lite with / without the slot PE (≈ 10 s)"
)
relgt_run

_rev = reveal(
    r"""
Measured: both RelGT-lite variants reach **0.884**, better than "always negative" (0.724) and than 1-layer MP, but below
2-layer and composite MP (≈ 0.965), with 4× the parameters of the composite model. The slot PE made no difference at this data
size (1,000 training customers, 60 epochs). The transformer could in principle use it, but has too little data to learn to.
That matches the lecture's framing: RelGT's strength is **scale and transfer**. One token format for every table, every
schema and every time point makes it the backbone of a pretrained relational foundation model (L14). On a single small task,
a well-shaped message-passing route is hard to beat.

""",
    predict="On this small task, will RelGT-lite beat composite message passing? Will the slot positional encoding matter?",
)
_post = mo.md(r"""
| | standard hetero MP | composite MP (RelGNN) | RelGT |
|---|---|---|---|
| unit of computation | one PK–FK hop | an atomic route (e.g. txn ⊕ product) | a token in a sampled context |
| handles junction tables | needs 2 hops, redundancy | **1 step**, pairing preserved | via attention + PE |
| best when | simple schemas | many-to-many / junction-heavy schemas | large, multi-schema, pretraining |

---

## 5 · Exercises

**Exercise 1 (tokens).** How many RelGT tokens does a customer with 4 transactions get (seed + 1-hop + 2-hop)? → `ex1_tokens`.

**Exercise 2 (hops).** In the PK–FK graph, how many hops separate a customer from a product they bought (standard MP), and how
many composite steps (RelGNN)? → `ex2_hops = (standard, composite)`.

**Exercise 3 (code: the composite message as SQL).** Write `spend_by_category()` returning an `n_cust × 5` tensor with each
customer's total `amount` per product category, using only `txn_table`, `category`, `txn_cust`, `txn_prod` and tensor ops
(no Python loop over transactions). This is the information a composite message must preserve.
""")
mo.vstack([_rev, _post])

ex1_tokens = None
ex2_hops = None

def spend_by_category():
    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_spend():
    try:
        got = spend_by_category()
        if got is None:
            return _verdict(False, False)
        ref = (
            txn_table.assign(cat=category[txn_table.product_id])
            .pivot_table(
                index="customer_id",
                columns="cat",
                values="amount",
                aggfunc="sum",
            )
            .reindex(index=range(n_cust), columns=range(5))
            .fillna(0)
            .values
        )
        return _verdict(
            np.allclose(np.asarray(got, dtype=float), ref, atol=1e-3), True
        )
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(ex1_tokens == 9, ex1_tokens is not None),
        mo.md("**Exercise 2**"),
        _verdict(
            tuple(ex2_hops) == (2, 1) if ex2_hops is not None else False,
            ex2_hops is not None,
        ),
        mo.md("**Exercise 3**"),
        _check_spend(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    "1 seed + 4 transactions + 4 products = **9** tokens."
                ),
                "Solution 2": mo.md(
                    "Standard: customer → transaction → product = **2** hops. Composite: **1** step along the route customer ← (transaction ⊕ product)."
                ),
                "Solution 3": mo.md("""```python
def spend_by_category():
    amount = torch.tensor(txn_table.amount.values, dtype=torch.float)
    onehot = torch.eye(5)[torch.tensor(category)[txn_prod]]          # category of each transaction's product
    return torch.zeros(n_cust, 5).index_add_(0, txn_cust, amount[:, None] * onehot)
```
`SUM(amount) GROUP BY customer, category` as one scatter: exactly a composite message (amount ⊗ category) summed per customer."""),
            }
        ),
    ]
)