def make_database(n_users=1500, n_products=300, days=365, seed=SEED):
    r = np.random.default_rng(seed)
    categories = np.array(
        ["electronics", "books", "grocery", "fashion", "toys"]
    )
    products = pd.DataFrame(
        {
            "product_id": np.arange(n_products),
            "category": r.choice(categories, n_products),
            "price": np.round(np.exp(r.normal(3, 0.8, n_products)), 2),
        }
    )
    bad = r.random(n_products) < 0.2  # hidden quality — NOT stored
    users = pd.DataFrame(
        {
            "user_id": np.arange(n_users),
            "signup_day": r.integers(0, 200, n_users),
            "age": r.integers(18, 70, n_users),
        }
    )
    rate = r.gamma(2.0, 0.06, n_users)  # purchases per day
    pref = r.dirichlet(np.ones(5) * 0.5, n_users)  # category preferences
    by_cat = {
        c: products.index[products.category == c].to_numpy()
        for c in categories
    }
    rows = []
    for u in range(n_users):
        hazard = 0.002  # daily churn hazard
        for day in range(users.signup_day[u], days):
            if r.random() < hazard:
                break  # churned: no more purchases
            for _ in range(r.poisson(rate[u])):
                p = int(
                    r.choice(by_cat[categories[r.choice(5, p=pref[u])]])
                )
                rows.append(
                    (
                        len(rows),
                        u,
                        p,
                        day,
                        round(products.price[p] * r.uniform(0.8, 1.2), 2),
                    )
                )
                hazard += (
                    0.03 if bad[p] else -0.0005
                )  # bad products drive churn
                hazard = max(hazard, 0.0005)
    transactions = pd.DataFrame(
        rows, columns=["txn_id", "user_id", "product_id", "day", "amount"]
    )
    return users, products, transactions, bad

users, products, transactions, hidden_bad = (
    make_database()
)  # hidden_bad: ground truth, never shown to models
mo.vstack(
    [
        mo.md(
            f"**users** {len(users):,} rows · **products** {len(products):,} rows · **transactions** {len(transactions):,} rows"
        ),
        mo.hstack(
            [users.head(3), products.head(3), transactions.head(3)],
            justify="start",
        ),
    ]
)

def sql_features(table, leak=False):
    """One row of hand-made features per (user, cutoff). leak=True deliberately ignores the cutoff (see §6)."""
    time_filter = "TRUE" if leak else "t.day < tt.cutoff"
    return duckdb.sql(f"""
        SELECT tt.user_id, tt.cutoff,
               COUNT(t.txn_id) FILTER (WHERE t.day >= tt.cutoff - 30)                AS n_30d,
               COUNT(t.txn_id) FILTER (WHERE t.day >= tt.cutoff - 90)                AS n_90d,
               COUNT(t.txn_id)                                                      AS n_all,
               COALESCE(SUM(t.amount) FILTER (WHERE t.day >= tt.cutoff - 30), 0)    AS spend_30d,
               COALESCE(SUM(t.amount) FILTER (WHERE t.day >= tt.cutoff - 90), 0)    AS spend_90d,
               tt.cutoff - MAX(t.day)                                               AS days_since_last,
               COUNT(DISTINCT p.category)                                           AS n_categories,
               tt.cutoff - ANY_VALUE(u.signup_day)                                  AS tenure,
               ANY_VALUE(u.age)                                                     AS age
        FROM train_table tt
        JOIN users u ON u.user_id = tt.user_id
        LEFT JOIN transactions t ON t.user_id = tt.user_id AND {time_filter}
        LEFT JOIN products p ON p.product_id = t.product_id
        GROUP BY tt.user_id, tt.cutoff
    """).df()

def fit_gbm(feature_frame):
    df = train_table.merge(feature_frame, on=["user_id", "cutoff"])
    cols = [
        c for c in feature_frame.columns if c not in ("user_id", "cutoff")
    ]
    tr, te = df[df.split == "train"], df[df.split == "test"]
    clf = HistGradientBoostingClassifier(
        max_iter=200, learning_rate=0.05, random_state=SEED
    ).fit(tr[cols], tr.label)
    return (
        roc_auc_score(te.label, clf.predict_proba(te[cols])[:, 1]),
        clf,
        cols,
    )

_t = time.time()
hand_features = sql_features(train_table)
gbm_auc, _, _ = fit_gbm(hand_features)
mo.vstack(
    [
        hand_features.head(4),
        mo.md(
            f"Gradient boosting on 9 SQL features: **test AUC = {gbm_auc:.3f}** (features + training: {time.time() - _t:.1f} s)"
        ),
    ]
)

def temporal_neighbourhoods(table):
    """For every example (user, cutoff): the user's transactions strictly before the cutoff (temporal sampling)."""
    by_user = {u: g for u, g in transactions.groupby("user_id")}
    ex_idx, txn_idx = [], []
    for i, (u, t) in enumerate(zip(table.user_id, table.cutoff)):
        g = by_user.get(u)
        if g is not None:
            keep = g.txn_id.values[
                g.day.values < t
            ]  # no peeking past the cutoff
            ex_idx.append(np.full(len(keep), i))
            txn_idx.append(keep)
    return torch.from_numpy(np.concatenate(ex_idx)), torch.from_numpy(
        np.concatenate(txn_idx)
    )

nbr_example, nbr_txn = temporal_neighbourhoods(train_table)
_txn_day = torch.tensor(transactions.day.values)
_delta = (
    torch.tensor(train_table.cutoff.values)[nbr_example]
    - _txn_day[nbr_txn]
).float()  # days before cutoff
txn_features = torch.stack(
    [
        db_graph["transaction"].x[nbr_txn, 0],
        _delta / 90,
        torch.exp(-_delta / 30),
    ],
    1,
)
txn_product = torch.tensor(transactions.product_id.values)[nbr_txn]
user_features = torch.tensor(
    np.c_[
        users.age.values[train_table.user_id] / 70,
        (
            train_table.cutoff.values
            - users.signup_day.values[train_table.user_id]
        )
        / 365,
    ],
    dtype=torch.float,
)
labels = torch.tensor(train_table.label.values, dtype=torch.float)
split_mask = {
    s: torch.tensor((train_table.split == s).values)
    for s in ("train", "val", "test")
}

class RelationalGNN(nn.Module):
    def __init__(self, d=32, product_embedding=True):
        super().__init__()
        self.prod_emb = (
            nn.Embedding(len(products), 16) if product_embedding else None
        )
        d_msg_in = (
            3
            + db_graph["product"].x.size(1)
            + (16 if product_embedding else 0)
        )
        self.msg = nn.Sequential(
            nn.Linear(d_msg_in, d), nn.ReLU(), nn.Linear(d, d)
        )  # learnable "JOIN" message
        self.head = nn.Sequential(
            nn.Linear(2 * d + 2 + 1, d), nn.ReLU(), nn.Linear(d, 1)
        )

    def forward(self):
        parts = [txn_features, db_graph["product"].x[txn_product]]
        if self.prod_emb is not None:
            parts.append(self.prod_emb(txn_product))
        m = self.msg(torch.cat(parts, 1))
        n_ex = len(labels)
        s = torch.zeros(n_ex, m.size(1)).index_add_(
            0, nbr_example, m
        )  # "GROUP BY user": SUM
        cnt = (
            torch.bincount(nbr_example, minlength=n_ex)
            .float()
            .unsqueeze(1)
        )
        h = torch.cat(
            [s, s / cnt.clamp(min=1), user_features, torch.log1p(cnt)], 1
        )  # + MEAN, own row, count
        return self.head(h).squeeze(1)

def fit_rdl(product_embedding=True, epochs=300, seed=SEED):
    torch.manual_seed(seed)
    model = RelationalGNN(product_embedding=product_embedding)
    opt = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=1e-4)
    best_val, best_test = 0, 0
    for _ in range(epochs):
        model.train()
        opt.zero_grad()
        out = model()
        F.binary_cross_entropy_with_logits(
            out[split_mask["train"]], labels[split_mask["train"]]
        ).backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            p = model()
        val = roc_auc_score(
            labels[split_mask["val"]], p[split_mask["val"]]
        )
        if val > best_val:
            best_val, best_test = (
                val,
                roc_auc_score(
                    labels[split_mask["test"]], p[split_mask["test"]]
                ),
            )
    return best_test, model

_t = time.time()
rdl_auc, rdl_model = fit_rdl(True)
rdl_auc_noemb, _ = fit_rdl(False)
pd.DataFrame(
    [
        {
            "model": "gradient boosting on 9 SQL features",
            "test AUC": gbm_auc,
        },
        {
            "model": "relational GNN, product features only",
            "test AUC": rdl_auc_noemb,
        },
        {
            "model": "relational GNN + learned product embeddings",
            "test AUC": rdl_auc,
        },
    ]
).round(3).assign(
    note=["", "", f"both GNNs trained in {time.time() - _t:.0f} s"]
)

_rev = reveal(
    r"""
Measured (one seed): **gradient boosting 0.840 · GNN with product features 0.851 · GNN + learned product embeddings 0.873.**
The extra 2 points come from the product embeddings, which can only be learned from *other users' outcomes*: a multi-hop
signal (user → transaction → product ← other transactions ← other users' churn) that none of the nine SQL features express.
""",
    predict="Can a GNN beat gradient boosting on nine hand-written SQL features? What signal could it use that the SQL features cannot express?",
)
_post = mo.md(r"""
Did the embeddings really discover the hidden "bad" products? We know the truth, so we can check:
""")
mo.vstack([_rev, _post])

ex1_visible = None
ex2_hops = None

def distinct_products_60d():
    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_sql():
    try:
        got = distinct_products_60d()
        if got is None:
            return _verdict(False, False)
        ref = []
        for u, t in zip(train_table.user_id, train_table.cutoff):
            g = transactions[
                (transactions.user_id == u)
                & (transactions.day < t)
                & (transactions.day >= t - 60)
            ]
            ref.append((u, t, g.product_id.nunique()))
        ref = pd.DataFrame(
            ref, columns=["user_id", "cutoff", "n_products_60d"]
        ).drop_duplicates()
        m = ref.merge(
            got,
            on=["user_id", "cutoff"],
            how="left",
            suffixes=("", "_got"),
        ).fillna(-1)
        return _verdict(
            bool((m.n_products_60d == m.n_products_60d_got).all()), True
        )
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

_SQL_SOLUTION = """```python
def distinct_products_60d():
    return duckdb.sql('''
        SELECT tt.user_id, tt.cutoff, COUNT(DISTINCT t.product_id) AS n_products_60d
        FROM train_table tt
        LEFT JOIN transactions t
          ON t.user_id = tt.user_id AND t.day < tt.cutoff AND t.day >= tt.cutoff - 60
        GROUP BY tt.user_id, tt.cutoff
    ''').df()
```
The time condition belongs in the `JOIN … ON`, not in a `WHERE`: a `WHERE` would drop users with no recent purchase instead of giving them 0."""

_visible = int(
    ((transactions.user_id == 5) & (transactions.day < 200)).sum()
)
mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(ex1_visible == _visible, ex1_visible is not None),
        mo.md("**Exercise 2**"),
        _verdict(ex2_hops == 4, ex2_hops is not None),
        mo.md("**Exercise 3**"),
        _check_sql(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    "`int(((transactions.user_id == 5) & (transactions.day < 200)).sum())`: strictly before the cutoff."
                ),
                "Solution 2": mo.md(
                    "user → transaction → product → (other) transaction → (other) user: **4 hops**. A 2-hop GNN cannot see other users' labels directly; the learned product embedding in §5 is a shortcut, since it absorbs that signal during training."
                ),
                "Solution 3": mo.md(_SQL_SOLUTION),
            }
        ),
    ]
)