def _ensure_movielens():
    """Download and unzip ml-latest-small into data/ if it is not there yet (≈ 1 MB)."""
    folder = DATA_DIR / "ml-latest-small"
    if (folder / "ratings.csv").exists():
        return True
    import io
    import urllib.request
    import zipfile

    url = "https://files.grouplens.org/datasets/movielens/ml-latest-small.zip"
    try:
        with urllib.request.urlopen(url, timeout=60) as resp:
            zipfile.ZipFile(io.BytesIO(resp.read())).extractall(DATA_DIR)
        return (folder / "ratings.csv").exists()
    except Exception:
        return False

mo.stop(
    not _ensure_movielens(),
    mo.callout(
        mo.md(
            "MovieLens `ml-latest-small` is missing and could not be downloaded. Get it from "
            "<https://grouplens.org/datasets/movielens/latest/> and unzip it into `data/` "
            "(you should have `data/ml-latest-small/ratings.csv`). See `data/README.md`."
        ),
        kind="warn",
    ),
)
_ratings = pd.read_csv(DATA_DIR / "ml-latest-small" / "ratings.csv")
movies = pd.read_csv(
    DATA_DIR / "ml-latest-small" / "movies.csv"
).set_index("movieId")
_pos = _ratings[_ratings.rating >= 4].sort_values(["userId", "timestamp"])
_pos = _pos[
    _pos.groupby("userId").movieId.transform("size") >= 10
]  # users with ≥ 10 positives
_rank = _pos.groupby("userId").cumcount()
_n = _pos.groupby("userId").movieId.transform("size")
_pos = _pos.assign(is_test=_rank >= (0.8 * _n).astype(int))
_train_items = _pos[~_pos.is_test].movieId.unique()
_pos = _pos[
    _pos.movieId.isin(_train_items)
]  # can only recommend items seen in training
user_ids = {u: i for i, u in enumerate(sorted(_pos.userId.unique()))}
item_ids = {m: i for i, m in enumerate(sorted(_train_items))}
item_movie = np.array(sorted(_train_items))
n_users, n_items = len(user_ids), len(item_ids)
train_pairs = np.stack(
    [
        _pos[~_pos.is_test].userId.map(user_ids),
        _pos[~_pos.is_test].movieId.map(item_ids),
    ],
    1,
)
test_pairs = np.stack(
    [
        _pos[_pos.is_test].userId.map(user_ids),
        _pos[_pos.is_test].movieId.map(item_ids),
    ],
    1,
)
train_by_user = (
    pd.Series(train_pairs[:, 1])
    .groupby(train_pairs[:, 0])
    .apply(set)
    .to_dict()
)
test_by_user = (
    pd.Series(test_pairs[:, 1])
    .groupby(test_pairs[:, 0])
    .apply(set)
    .to_dict()
)
mo.md(
    f"**{n_users} users × {n_items:,} items**; train edges {len(train_pairs):,}, test edges {len(test_pairs):,}; "
    f"density {len(train_pairs) / (n_users * n_items):.4f}."
)

def bipartite_norm_adj():
    """Sparse D^-1/2 A D^-1/2 over n_users + n_items nodes (items offset by n_users)."""
    u, i = (
        torch.from_numpy(train_pairs[:, 0]),
        torch.from_numpy(train_pairs[:, 1]) + n_users,
    )
    row, col = torch.cat([u, i]), torch.cat([i, u])
    deg = torch.bincount(row, minlength=n_users + n_items).float()
    val = deg[row].pow(-0.5) * deg[col].pow(-0.5)
    return torch.sparse_coo_tensor(
        torch.stack([row, col]), val, (n_users + n_items,) * 2
    ).coalesce()

A_norm = bipartite_norm_adj()

class Recommender(nn.Module):
    def __init__(self, n_layers=0, dim=64):
        super().__init__()
        self.emb = nn.Embedding(n_users + n_items, dim)
        nn.init.normal_(self.emb.weight, std=0.1)
        self.n_layers = n_layers  # 0 layers = plain matrix factorisation

    def final_embeddings(self):
        e = self.emb.weight
        layers = [e]
        for _ in range(self.n_layers):
            e = torch.sparse.mm(A_norm, e)
            layers.append(e)
        return torch.stack(layers).mean(0)

    def all_scores(self):
        e = self.final_embeddings()
        return e[:n_users] @ e[n_users:].T

def sample_bpr_batch(r, size):
    idx = r.integers(len(train_pairs), size=size)
    users, pos = train_pairs[idx, 0], train_pairs[idx, 1]
    neg = r.integers(
        n_items, size=size
    )  # uniform negatives (rarely a positive by accident)
    return (
        torch.from_numpy(users),
        torch.from_numpy(pos),
        torch.from_numpy(neg),
    )

def train_recommender(
    n_layers,
    epochs=40,
    batch=2048,
    lr=0.005,
    reg=1e-4,
    seed=SEED,
    neg_sampler=None,
):
    torch.manual_seed(seed)
    r = np.random.default_rng(seed)
    model = Recommender(n_layers)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    history = []
    for ep in range(epochs):
        for _ in range(len(train_pairs) // batch):
            u, p, n = (neg_sampler or sample_bpr_batch)(r, batch)
            e = model.final_embeddings()
            eu, ep_, en = e[u], e[n_users + p], e[n_users + n]
            loss = -F.logsigmoid(
                (eu * ep_).sum(1) - (eu * en).sum(1)
            ).mean()  # BPR
            loss = (
                loss
                + reg
                * (
                    model.emb.weight[u].pow(2).sum()
                    + model.emb.weight[n_users + p].pow(2).sum()
                    + model.emb.weight[n_users + n].pow(2).sum()
                )
                / batch
            )  # L2 on used rows
            opt.zero_grad()
            loss.backward()
            opt.step()
        if (ep + 1) % 10 == 0:
            with torch.no_grad():
                history.append(
                    {
                        "epoch": ep + 1,
                        "Recall@20": recall_at_k(model.all_scores()),
                    }
                )
    return model, pd.DataFrame(history)

rec_run = mo.ui.run_button(
    label="Train MF and LightGCN (K = 1, 2, 3) (≈ 1–2 min)"
)
rec_run

_rev = reveal(
    r"""
Measured (Recall@20 after 40 epochs, one seed): **MF 0.110 · LightGCN K=1 0.124 · K=2 0.128 · K=3 0.122** vs. popularity 0.082.
Propagating over the graph adds ≈ 15% relative recall over plain MF, with **zero extra parameters**: the only parameters
are the embedding table in both. Two layers are the sweet spot here; a third starts to over-smooth.

""",
    predict="Will LightGCN beat plain matrix factorisation on Recall@20 with zero extra parameters? Is K = 3 better than K = 2?",
)
_post = mo.md(r"""
**Check against PyG.** PyG's `LightGCN` implements the same propagation. Copy our trained table into it and compare final embeddings:
""")
mo.vstack([_rev, _post])

ex1_recall = None
ex2_bpr = None

def lightgcn_propagate(E, A, 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_prop():
    try:
        m = rec_models["LightGCN, K=2"]
        with torch.no_grad():
            got = lightgcn_propagate(m.emb.weight, A_norm, 2)
            if got is None:
                return _verdict(False, False)
            return _verdict(
                torch.allclose(got, m.final_embeddings(), 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_recall is not None and abs(ex1_recall - 0.5) < 1e-6,
            ex1_recall is not None,
        ),
        mo.md("**Exercise 2**"),
        _verdict(
            ex2_bpr is not None and abs(ex2_bpr - np.log1p(np.e)) < 2e-3,
            ex2_bpr is not None,
        ),
        mo.md("**Exercise 3**"),
        _check_prop(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    "Hits $\\{a, c\\}$ out of 4 future positives → **0.5**. (Precision@3 would be 2/3.)"
                ),
                "Solution 2": mo.md(
                    r"$-\log\sigma(0.5-1.5)=-\log\sigma(-1)=\log(1+e)\approx 1.313$. The ranking is wrong, so the loss is large."
                ),
                "Solution 3": mo.md("""```python
def lightgcn_propagate(E, A, K):
    layers = [E]
    for _ in range(K):
        layers.append(torch.sparse.mm(A, layers[-1]))
    return torch.stack(layers).mean(0)
```"""),
            }
        ),
    ]
)

with torch.no_grad():
    _scores = rec_models["LightGCN, K=2"].all_scores()
_top = _scores.topk(
    20, dim=1
).indices.numpy()  # NO masking of training items
_unmasked = float(
    np.mean(
        [
            len(set(_top[u]) & test_by_user[u]) / len(test_by_user[u])
            for u in test_by_user
        ]
    )
)
_already_seen = float(
    np.mean(
        [len(set(_top[u]) & train_by_user[u]) / 20 for u in test_by_user]
    )
)
mo.callout(
    mo.md(
        f"Without masking: Recall@20 = **{_unmasked:.4f}** (vs **{recall_at_k(_scores):.4f}** masked); "
        f"on average **{_already_seen:.0%}** of each top-20 list are movies the user already liked. "
        "**Fix:** set the scores of training interactions to −∞ before `topk`, as `recall_at_k` does. "
        "PyG's `LightGCN.recommend()` does *not* do this for you."
    ),
    kind="warn",
)