"""
rec_world.py —— 第 17 讲：花园集市（Garden Market）——一个"真实偏好已知"的序列推荐世界
=====================================================================

【这个文件做什么】
    推荐系统的研究通常用真实的点击 / 购买日志，但我们永远不知道"用户真正想要什么"。这里造一个集市：
        - 144 件商品：花园世界的 43 个名词 × 几种形态（plum seeds、plum sapling、plum basket、plum jam ……），
          每件商品有一段文字描述（形容词 + 名词 + 形态，例如 ripe plum jam）。
        - 6 种用户（果园主、厨师、花匠、农场主、宠物主人、送礼的人），每种对品类和形态有不同的偏好；
          用户的类型是**隐藏**的，只能从他买过的东西里推断。
        - 购买序列：每一步以概率 ρ = 0.4 买"上一件商品的配套"（同一个名词、另一种形态：买了 plum seeds
          接着买 plum sapling），否则按自己类型的偏好重新挑一件。
    因为生成过程完全已知，任何一段购买历史之后"下一件商品"的**真实分布**都能精确算出来（对 6 种类型做贝叶斯
    后验）——它给出任何推荐模型都不可能超过的上限（贝叶斯最优），就像第 1 讲的真实困惑度下界。

【怎样使用】
    import rec_world as R
    users = R.make_users(6000, seed=1)                  # [(类型, [商品下标, ...]), ...]
    p = R.true_next_dist(users[0][1][:-1])              # 真实的下一件商品分布（144 维）
    model = R.train_sasrec("sasrec", users)             # 训练（或读缓存）一个 SASRec 式模型

【LLM 联系】
    SASRec（Kang & McAuley 2018）把"用户的购买序列"当作"句子"，"下一件商品"当作"下一个 token"，
    用因果 Transformer 建模——推荐和语言建模在形式上是同一个问题。
"""

import os
import random
import sys

import numpy as np
import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "lecture02"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

from garden_lm import load_or_train, pad_batch, train  # noqa: E402
from toy_world import ADJECTIVES, NOUNS, zipf_weights  # noqa: E402

CKPT_DIR = os.path.join(HERE, "checkpoints")

# =====================================================================
# 1. 商品
# =====================================================================

FORMS = {"fruit": ["seeds", "sapling", "basket", "jam"],
         "veg": ["seeds", "basket", "soup"],
         "flower": ["seeds", "bulb", "bouquet"],
         "animal": ["feed", "toy", "house"]}
CATS = list(FORMS)


def _make_items(seed=0):
    rng = random.Random(seed)
    items = []
    for cat in CATS:
        for noun in NOUNS[cat]:
            for form in FORMS[cat]:
                items.append({"cat": cat, "noun": noun, "form": form, "adj": rng.choice(ADJECTIVES[cat])})
    for i, it in enumerate(items):
        it["id"] = i
        it["text"] = [it["adj"], it["noun"], it["form"]]
    return items


ITEMS = _make_items()
N_ITEMS = len(ITEMS)                                         # 144
NOUN_P = {cat: dict(zip(NOUNS[cat], zipf_weights(len(NOUNS[cat]), 1.0))) for cat in CATS}   # 品类内的流行度（Zipf）
SIBLINGS = [[j for j in range(N_ITEMS) if ITEMS[j]["noun"] == it["noun"] and j != i] for i, it in enumerate(ITEMS)]

# =====================================================================
# 2. 用户类型与生成过程
# =====================================================================

TYPES = {
    "orchard": ({"fruit": 0.8, "veg": 0.1, "flower": 0.1}, {"seeds": 3, "sapling": 4, "bulb": 2}),
    "kitchen": ({"fruit": 0.45, "veg": 0.55}, {"basket": 3, "jam": 3, "soup": 3}),
    "florist": ({"flower": 0.85, "fruit": 0.15}, {"bulb": 3, "bouquet": 4, "seeds": 1, "sapling": 1}),
    "farm": ({"animal": 0.6, "veg": 0.4}, {"feed": 4, "house": 2, "seeds": 3}),
    "pets": ({"animal": 1.0}, {"toy": 3, "house": 2, "feed": 2}),
    "gifts": ({"fruit": 0.3, "flower": 0.35, "animal": 0.35}, {"jam": 3, "bouquet": 3, "toy": 3, "basket": 2}),
}
TYPE_NAMES = list(TYPES)
TYPE_PRIOR = np.full(len(TYPES), 1 / len(TYPES))
RHO = 0.4                                                     # "买配套"的概率


def _type_dists():
    """
    每种类型"重新挑一件"时的分布 g_k(商品) = P(品类 | k) · P(形态 | 品类, k) · P(名词 | 品类)。
    【边界】某类型偏好表里没提到的形态权重为 0.3（偶尔也会买），所以每种类型对每件商品的概率都 > 0。
    """
    G = np.zeros((len(TYPES), N_ITEMS))
    for k, (cat_w, form_w) in enumerate(TYPES.values()):
        for cat, pc in cat_w.items():
            fw = np.array([form_w.get(f, 0.3) for f in FORMS[cat]])
            fw = fw / fw.sum()
            for i, it in enumerate(ITEMS):
                if it["cat"] == cat:
                    G[k, i] = pc * fw[FORMS[cat].index(it["form"])] * NOUN_P[cat][it["noun"]]
    return G / G.sum(1, keepdims=True)


G = _type_dists()
# 配套转移矩阵：F[i, j] = 买了 i 之后、"买配套"时选 j 的概率（同名词的其他形态，均匀）
F_MAT = np.zeros((N_ITEMS, N_ITEMS))
for _i, _sib in enumerate(SIBLINGS):
    F_MAT[_i, _sib] = 1 / len(_sib)


def sample_user(rng, min_len=5, max_len=20, exclude=frozenset()):
    """
    采样一个用户：(类型下标, 购买序列)。第一件商品按类型的偏好挑。

    【参数】exclude：不存在的商品（"还没上架"）：从偏好分布和配套里去掉；如果上一件商品的配套全被去掉，
            这一步改为按偏好重新挑。02 号脚本用它模拟冷启动——训练数据里从没出现过的新商品。
    """
    w = G if not exclude else G * np.array([i not in exclude for i in range(N_ITEMS)])
    k = rng.randrange(len(TYPES))
    L = rng.randint(min_len, max_len)
    seq = [rng.choices(range(N_ITEMS), weights=w[k])[0]]
    while len(seq) < L:
        sib = [j for j in SIBLINGS[seq[-1]] if j not in exclude]
        if sib and rng.random() < RHO:
            seq.append(rng.choice(sib))
        else:
            seq.append(rng.choices(range(N_ITEMS), weights=w[k])[0])
    return k, seq


def make_users(n, seed, exclude=frozenset()):
    rng = random.Random(seed)
    return [sample_user(rng, exclude=exclude) for _ in range(n)]


def type_posterior(hist):
    """
    看过购买历史 hist 之后，对用户类型的后验 P(k | hist)。

    【思路】p(x_1 | k) = g_k(x_1)；之后每一步 p(x_t | x_{t-1}, k) = ρ F[x_{t-1}, x_t] + (1 − ρ) g_k(x_t)。
            在 log 空间里累加，最后归一化。
    """
    with np.errstate(divide="ignore"):             # 某些类型从不买某些品类：概率 0，log 是 −inf，后验正好是 0
        lp = np.log(TYPE_PRIOR) + np.log(G[:, hist[0]])
        for a, b in zip(hist[:-1], hist[1:]):
            lp += np.log(RHO * F_MAT[a, b] + (1 - RHO) * G[:, b])
    lp -= lp.max()
    p = np.exp(lp)
    return p / p.sum()


def true_next_dist(hist):
    """下一件商品的真实分布：Σ_k P(k | hist) [ρ F[最后一件, ·] + (1 − ρ) g_k(·)]。【返回】144 维数组。"""
    post = type_posterior(hist)
    return RHO * F_MAT[hist[-1]] + (1 - RHO) * post @ G


# =====================================================================
# 3. 评测：留一法（每个用户的最后一件商品作为答案）
# =====================================================================


def rank_metrics(scores, targets, k=10, exclude=None):
    """
    【参数】scores：(n, 144) 的分数（越大越推荐）；targets：(n,) 真实的下一件商品；
            exclude：可选，每个用户要排除的商品集合（例如"已经买过的"——这里默认不排除，配套商品可能和以前买过的相同）。
    【返回】dict：HR@k（答案在前 k 名的比例）、NDCG@k（答案在第 r 名时得 1/log2(r+1)）、MRR。
    """
    scores = np.array(scores, dtype=float)
    if exclude is not None:
        for r, ex in enumerate(exclude):
            scores[r, list(ex)] = -np.inf
    tgt = scores[np.arange(len(targets)), targets]
    rank = (scores > tgt[:, None]).sum(1) + 1
    hit = rank <= k
    return {"HR@10": float(hit.mean()), "NDCG@10": float(np.where(hit, 1 / np.log2(rank + 1), 0).mean()),
            "MRR": float((1 / rank).mean())}


def expected_hr(dists, k=10):
    """贝叶斯最优排序的**期望** HR@k：每个用户真实分布里前 k 件的概率之和的平均（不受抽样噪声影响）。"""
    return float(np.mean([np.sort(d)[-k:].sum() for d in dists]))


# =====================================================================
# 4. 基线
# =====================================================================


def popularity_scores(train_users):
    cnt = np.zeros(N_ITEMS)
    for _, s in train_users:
        for i in s:
            cnt[i] += 1
    return cnt


def markov_scores(train_users, alpha=0.1):
    """一阶马尔可夫：P(下一件 | 上一件)，加 alpha 平滑并混入一点流行度（打破并列）。【返回】144 × 144 矩阵。"""
    T = np.full((N_ITEMS, N_ITEMS), alpha)
    for _, s in train_users:
        for a, b in zip(s[:-1], s[1:]):
            T[a, b] += 1
    T = T / T.sum(1, keepdims=True)
    pop = popularity_scores(train_users)
    return T + 1e-6 * pop / pop.sum()


# =====================================================================
# 5. SASRec：商品序列上的因果 Transformer
# =====================================================================

PAD, BOS = N_ITEMS, N_ITEMS + 1                               # 商品 0..143，之后是 <pad>、<bos>
SEQ_VOCAB = N_ITEMS + 2


def new_sasrec(seed=0, d_model=64, n_heads=4, n_layers=2, d_ff=256, pos="learned", vocab=SEQ_VOCAB):
    """SASRec ≈ 第 2 讲的 TransformerLM，词表换成商品 ID（输入嵌入和输出层共享，和 SASRec 原文一样）。"""
    from transformer import TransformerLM
    torch.manual_seed(seed)
    return TransformerLM(vocab, d_model=d_model, n_heads=n_heads, n_layers=n_layers, d_ff=d_ff, max_len=64,
                         attn="causal", pos=pos)


def seq_batches(seqs, batch_size=64, bos=BOS, pad=PAD):
    """每个序列变成 [<bos>, x_1, ..., x_n]，按长度排序后切批。【返回】list[(ids, pad_mask)]。"""
    ss = sorted(([bos] + list(s) for s in seqs), key=len)
    return [pad_batch(ss[i:i + batch_size], pad) for i in range(0, len(ss), batch_size)]


def ce_loss(model, batch, n_out=N_ITEMS):
    """全 softmax 的交叉熵：每个位置预测下一件商品（只在商品上做 softmax，不含 <pad>/<bos>）。"""
    ids, pad = batch
    logits = model(ids[:, :-1], pad=pad[:, :-1])[..., :n_out]
    tgt, keep = ids[:, 1:], ~pad[:, 1:]
    nll = F.cross_entropy(logits.reshape(-1, n_out), tgt.clamp(max=n_out - 1).reshape(-1), reduction="none")
    return (nll * keep.reshape(-1)).sum(), keep.sum()


def make_bce_loss(n_neg=1, seed=0):
    """
    SASRec 原文的损失：每个位置一个正例（真实的下一件）和 n_neg 个随机负例，二元交叉熵。
    【LLM 联系】Klenitskiy & Vasilev (2023)：把它换成全 softmax 的交叉熵，SASRec 会好很多——
                "BERT4Rec 比 SASRec 好"的很多结论其实来自损失函数的差别。
    """
    g = torch.Generator().manual_seed(seed)

    def loss(model, batch):
        ids, pad = batch
        logits = model(ids[:, :-1], pad=pad[:, :-1])[..., :N_ITEMS]
        tgt, keep = ids[:, 1:].clamp(max=N_ITEMS - 1), (~pad[:, 1:]).float()
        pos = logits.gather(2, tgt[..., None])[..., 0]
        neg_ids = torch.randint(0, N_ITEMS, (*tgt.shape, n_neg), generator=g)
        neg = logits.gather(2, neg_ids)
        l = -F.logsigmoid(pos) - F.logsigmoid(-neg).sum(-1)
        return (l * keep).sum(), keep.sum()
    return loss


VALID = None


def valid_users():
    """验证集：1000 个用户（种子 2），用来挑最好的一遍（提前停止）。"""
    global VALID
    if VALID is None:
        VALID = make_users(1000, seed=2)
    return VALID


def valid_nll(model, users=None):
    """验证集上每个位置预测下一件商品的平均交叉熵（越小越好）。"""
    users = users or valid_users()
    s_tot, n_tot = 0.0, 0
    with torch.no_grad():
        for b in seq_batches([s for _, s in users], 256):
            s, n = ce_loss(model, b)
            s_tot += float(s)
            n_tot += int(n)
    return s_tot / n_tot


def train_sasrec(name, users, epochs=20, lr=1e-3, loss="ce", seed=0, verbose=False, **arch):
    """
    训练（或读缓存）SASRec。

    【参数】loss："ce"（全 softmax）或 "bce"（每个位置 1 个随机负例）。
    【边界】每遍结束在验证集上算一次交叉熵，最后恢复最好的一遍（提前停止）——推荐数据很小，模型很容易过拟合。
    """
    seqs = [s for _, s in users]
    batches = seq_batches(seqs)
    loss_fn = ce_loss if loss == "ce" else make_bce_loss()

    def fit(m):
        return train(m, batches, loss_fn, epochs=epochs, lr=lr, warmup=100, seed=seed, verbose=verbose, tag=name,
                     valid_fn=valid_nll)

    cfg = {"n": len(seqs), "tok": sum(map(len, seqs)), "epochs": epochs, "lr": lr, "loss": loss, "seed": seed,
           "arch": arch, "v": 2}
    m, _ = load_or_train(name, cfg, lambda: new_sasrec(seed, **arch), fit, verbose=verbose, ckpt_dir=CKPT_DIR)
    return m


@torch.no_grad()
def sasrec_scores(model, hists, n_out=N_ITEMS, bos=BOS):
    """每段历史之后对所有商品的打分（最后一个位置的 logits）。【返回】(n, n_out) numpy 数组。"""
    out = np.zeros((len(hists), n_out))
    groups = {}
    for i, h in enumerate(hists):
        groups.setdefault(len(h), []).append(i)
    for L, idx in groups.items():
        ids = torch.tensor([[bos] + list(hists[i]) for i in idx])
        out[idx] = model(ids)[:, -1, :n_out].numpy()
    return out


# =====================================================================
# 6. 商品怎样变成 token（02 号脚本）：原子 ID vs 语义 ID
# =====================================================================


def content_vectors(w_cat=2.0, w_form=1.5, w_noun=1.0, w_adj=0.3):
    """
    每件商品的"内容向量"：品类、形态、名词、形容词的 one-hot，按权重拼起来（真实系统里是商品标题 / 描述的
    文本嵌入，例如 TIGER 用 Sentence-T5）。【返回】(144, d) numpy 数组。
    """
    nouns = sorted({it["noun"] for it in ITEMS})
    forms = sorted({it["form"] for it in ITEMS})
    adjs = sorted({it["adj"] for it in ITEMS})
    X = np.zeros((N_ITEMS, len(CATS) + len(forms) + len(nouns) + len(adjs)))
    for i, it in enumerate(ITEMS):
        X[i, CATS.index(it["cat"])] = w_cat
        X[i, len(CATS) + forms.index(it["form"])] = w_form
        X[i, len(CATS) + len(forms) + nouns.index(it["noun"])] = w_noun
        X[i, len(CATS) + len(forms) + len(nouns) + adjs.index(it["adj"])] = w_adj
    return X


def kmeans(X, k, seed=0, iters=100, restarts=5):
    """朴素的 k-means（多次随机初始化取误差最小的一次）。【返回】(每个点的簇号, 簇中心)。"""
    best = None
    for r in range(restarts):
        rng = np.random.default_rng(seed * 100 + r)
        C = X[rng.choice(len(X), k, replace=False)].copy()
        for _ in range(iters):
            a = ((X[:, None] - C[None]) ** 2).sum(-1).argmin(1)
            for j in range(k):
                if (a == j).any():
                    C[j] = X[a == j].mean(0)
        sse = ((X - C[a]) ** 2).sum()
        if best is None or sse < best[0]:
            best = (sse, a, C)
    return best[1], best[2]


def residual_codes(X, ks=(13, 8)):
    """
    残差 k-means 的语义 ID（TIGER 用 RQ-VAE；这是它"不训练编码器"的最简版本）：
        第 1 层对向量聚类得到 c1；减去第 1 层的簇中心，对残差再聚类得到 c2；……
        最后加一位"消歧"编号 c3，区分前缀完全相同的商品。
    【返回】list[tuple]：每件商品一个 (c1, c2, c3)。
    【边界】前缀相同的商品越多，c3 的取值越多；内容向量越没有结构，后几层的码越接近随机分组。
    """
    codes, R_ = [], X.copy()
    for lvl, k in enumerate(ks):
        a, C = kmeans(R_, k, seed=lvl)
        codes.append(a)
        R_ = R_ - C[a]
    seen, out = {}, []
    for i in range(len(X)):
        key = tuple(int(c[i]) for c in codes)
        seen[key] = seen.get(key, 0) + 1
        out.append(key + (seen[key] - 1,))
    return out


def id_scheme(name):
    """
    【返回】(每件商品的 token 列表, token 词表)：
        "atomic"：每件商品一个独立的 token（传统的 ID 嵌入表）；
        "title" ：商品名里的 (名词, 形态) 两个 token——"plum jam" 和 "plum seeds" 共享 plum；
        "rq"    ：残差 k-means 的 (c1, c2, c3)。
    """
    if name == "atomic":
        toks = [[f"i{i}"] for i in range(N_ITEMS)]
    elif name == "title":
        toks = [[it["noun"], it["form"]] for it in ITEMS]
    elif name == "rq":
        toks = [[f"a{c1}", f"b{c2}", f"c{c3}"] for c1, c2, c3 in residual_codes(content_vectors())]
    else:
        raise ValueError(name)
    vocab = sorted({t for ts in toks for t in ts}) + ["<pad>", "<bos>"]
    return toks, vocab


class CodeRec:
    """
    "生成式检索"的推荐模型：把每件商品写成若干个 token，购买序列就是这些 token 的拼接，用因果 Transformer
    预测下一个 token（TIGER，Rajput et al. 2023）。atomic 方案就是普通的 SASRec。

    【打分】一件商品的分数 = 它的 token 序列在历史之后的 log 概率之和（精确计算全部 144 件，而不是束搜索）。
            prefix=True 时只用除最后一个 token 以外的前缀（TIGER 处理冷启动的办法：生成前缀，再把前缀相同的
            新商品也放进候选）。
    """

    def __init__(self, scheme, users, epochs=30, seed=0, tag="", exclude=frozenset()):
        self.scheme = scheme
        self.toks, self.vocab = id_scheme(scheme)
        self.stoi = {t: i for i, t in enumerate(self.vocab)}
        self.ids = [[self.stoi[t] for t in ts] for ts in self.toks]
        self.pad, self.bos = self.stoi["<pad>"], self.stoi["<bos>"]
        seqs = [[t for i in s for t in self.ids[i]] for _, s in users]
        batches = seq_batches(seqs, 64, bos=self.bos, pad=self.pad)
        V = len(self.vocab)
        # 验证集和训练数据来自同一个世界（冷启动实验里同样没有新商品），否则提前停止会被新商品带偏
        vu = valid_users() if not exclude else make_users(1000, seed=2, exclude=exclude)
        valid = [[t for i in s for t in self.ids[i]] for _, s in vu]
        vb = seq_batches(valid, 256, bos=self.bos, pad=self.pad)
        loss = lambda m, b: ce_loss(m, b, n_out=V - 2)

        def vfn(m):
            with torch.no_grad():
                tot = [loss(m, b) for b in vb]
            return float(sum(float(a) for a, _ in tot) / sum(int(n) for _, n in tot))

        def fit(m):
            return train(m, batches, loss, epochs=epochs, lr=1e-3, warmup=100, seed=seed, verbose=False,
                         valid_fn=vfn)

        cfg = {"scheme": scheme, "n": len(seqs), "tok": sum(map(len, seqs)), "epochs": epochs, "seed": seed,
               "tag": tag, "excl": sorted(exclude), "v": 2}
        self.model, _ = load_or_train(f"coderec_{scheme}", cfg,
                                      lambda: new_sasrec(seed, vocab=V, pos="sinusoidal"), fit, verbose=False,
                                      ckpt_dir=CKPT_DIR)

    def backoff_scores(self, hists, cold):
        """
        老商品用完整的 log 概率；新商品（训练里没出现过）用"前缀的 log 概率 + log(1/同前缀的商品数)"——
        相当于假设最后一个 token 在同前缀的商品里均匀分布。两者放在同一个排序里。
        """
        full, pre = self.scores(hists), self.scores(hists, prefix=True)
        L = len(self.ids[0])
        if L == 1:
            return full
        groups = {}
        for i, ids in enumerate(self.ids):
            groups.setdefault(tuple(ids[:-1]), []).append(i)
        n_same = np.array([len(groups[tuple(ids[:-1])]) for ids in self.ids])
        out = full.copy()
        cl = sorted(cold)
        out[:, cl] = pre[:, cl] - np.log(n_same[cl])
        return out

    @torch.no_grad()
    def scores(self, hists, prefix=False):
        """【返回】(n, 144)：每段历史之后每件商品的 log 概率（prefix=True 时是前缀的 log 概率）。"""
        out = np.zeros((len(hists), N_ITEMS))
        L = len(self.ids[0]) - (1 if prefix and len(self.ids[0]) > 1 else 0)
        cand = torch.tensor([ids[:L] for ids in self.ids])                       # 144 × L
        for r, h in enumerate(hists):
            ctx = torch.tensor([self.bos] + [t for i in h for t in self.ids[i]])
            seq = torch.cat([ctx.expand(N_ITEMS, -1), cand], 1)
            lp = torch.log_softmax(self.model(seq[:, :-1])[:, -L:, :len(self.vocab) - 2], -1)
            out[r] = lp.gather(2, cand[..., None])[..., 0].sum(1).numpy()
        return out


# =====================================================================
# 7. 把解码器语言模型变成文本编码器（03 号脚本，LLM2Vec）
# =====================================================================

TEXT_EXTRA = ["bought", "then", "<mask>", "<pad>"] + sorted({f for fs in FORMS.values() for f in fs})


def text_vocab():
    from toy_world import Vocab, all_words
    return Vocab(["<bos>", "<eos>", "<unk>"] + all_words() + [w for w in TEXT_EXTRA if w not in all_words()])


TV = text_vocab()
T_BOS, T_PAD, T_MASK = TV.stoi["<bos>"], TV.stoi["<pad>"], TV.stoi["<mask>"]


def market_corpus(users, catalog_copies=5, seed=0):
    """
    集市的"文本语料"：商品目录（每件商品的描述，例如 ripe plum jam .，重复 catalog_copies 次）+ 用户的购物日记
    （I bought ripe plum seeds . then I bought sweet plum sapling . …）。
    【LLM 联系】LLM 预训练语料里的商品评论、购物清单、博客，让模型"知道"哪些东西常被一起买。
    """
    rng = random.Random(seed)
    docs = [ITEMS[i]["text"] + [".", "<eos>"] for _ in range(catalog_copies) for i in range(N_ITEMS)]
    for _, s in users:
        d = ["I", "bought"] + ITEMS[s[0]]["text"] + ["."]
        for i in s[1:]:
            d += ["then", "I", "bought"] + ITEMS[i]["text"] + ["."]
        docs.append(d + ["<eos>"])
    rng.shuffle(docs)
    return docs


def text_batches(docs, bs=32):
    ss = sorted(([T_BOS] + TV.encode(d) for d in docs), key=len)
    return [pad_batch(ss[i:i + bs], T_PAD) for i in range(0, len(ss), bs)]


def text_clm_loss(m, b):
    ids, pad = b
    lg = m(ids[:, :-1], pad=pad[:, :-1])
    keep = ~pad[:, 1:]
    nll = F.cross_entropy(lg.reshape(-1, lg.shape[-1]), ids[:, 1:].reshape(-1), reduction="none")
    return (nll * keep.reshape(-1)).sum(), keep.sum()


def train_text_lm(name, docs, epochs=4):
    """在集市语料上预训练一个因果语言模型（第 2 讲的 TransformerLM，d = 64、2 层）。"""
    from transformer import TransformerLM

    def build():
        torch.manual_seed(0)
        return TransformerLM(len(TV), d_model=64, n_heads=4, n_layers=2, d_ff=256, attn="causal")

    b = text_batches(docs)
    m, _ = load_or_train(name, {"n": len(docs), "tok": sum(map(len, docs)), "epochs": epochs, "v": 1}, build,
                         lambda m: train(m, b, text_clm_loss, epochs=epochs, lr=3e-3, warmup=100, verbose=False),
                         verbose=False, ckpt_dir=CKPT_DIR)
    return m


@torch.no_grad()
def embed_texts(m, texts, pool="mean", bidirectional=False):
    """
    文本嵌入：最后一层隐藏状态（去掉 <bos>）的平均（pool="mean"）或最后一个 token（pool="last"）。
    bidirectional=True 时把注意力临时换成双向（参数不变——LLM2Vec 的第一步）。【返回】L2 归一化的 numpy 数组。
    """
    old = m.attn_type
    m.attn_type = "bidirectional" if bidirectional else "causal"
    out = []
    for tx in texts:
        h = m.hidden_states(torch.tensor([[T_BOS] + TV.encode(tx)]))[0, 1:]
        out.append(h.mean(0) if pool == "mean" else h[-1])
    m.attn_type = old
    return F.normalize(torch.stack(out), dim=-1).numpy()


def mntp_adapt(m0, docs, epochs=2, p=0.2, seed=0):
    """
    LLM2Vec 的第二步：掩码下一个 token 预测（MNTP）。打开双向注意力，随机把 p 比例的 token 换成 <mask>，
    用**前一个位置**的输出预测被遮住的 token（保持"位置 i−1 预测 i"的习惯，和预训练对齐）。
    【返回】适配后的新模型（attn_type 恢复为 causal；用 embed_texts(..., bidirectional=True) 取嵌入）。
    """
    import copy
    m = copy.deepcopy(m0)
    m.attn_type = "bidirectional"
    g = torch.Generator().manual_seed(seed)

    def loss(m, b):
        ids, pad = b
        msk = (torch.rand(ids.shape, generator=g) < p) & ~pad
        msk[:, 0] = False
        lg = m(ids.masked_fill(msk, T_MASK), pad=pad)
        sel = msk[:, 1:]
        nll = F.cross_entropy(lg[:, :-1].reshape(-1, lg.shape[-1]), ids[:, 1:].reshape(-1), reduction="none")
        return (nll * sel.reshape(-1)).sum(), sel.sum().clamp(min=1)

    train(m, text_batches(docs), loss, epochs=epochs, lr=1e-3, warmup=50, seed=seed, verbose=False)
    m.attn_type = "causal"
    return m


def contrastive_adapt(m0, pairs, steps=300, tau=0.05, batch=64, lr=3e-4, seed=0):
    """
    对比学习（InfoNCE）：pairs 里每一对 (文本 a, 文本 b) 是正例，同一批里其他的 b 是负例。
        unsupervised SimCSE：a、b 是同一段文本的两个随机遮盖版本（"同一句话"就是正例）；
        有监督：a、b 是"先后被同一个用户买"的两件商品的描述（推荐系统真正关心的相似）。
    【参数】pairs：函数 rng -> (a 的词列表, b 的词列表)，或 (a, b) 列表。
    【边界】同一批里 b 相同的其他样例不算负例（避免把正例当成负例）。
    """
    import copy
    m = copy.deepcopy(m0)
    m.attn_type = "bidirectional"
    g = torch.Generator().manual_seed(seed)
    rng = random.Random(seed)
    opt = torch.optim.Adam(m.parameters(), lr=lr)

    def enc(texts, mask_p=0.0):
        ids, pad = pad_batch([[T_BOS] + TV.encode(t) for t in texts], T_PAD)
        if mask_p:
            msk = (torch.rand(ids.shape, generator=g) < mask_p) & ~pad
            msk[:, 0] = False
            ids = ids.masked_fill(msk, T_MASK)
        h = m.hidden_states(ids, pad=pad)[:, 1:]
        keep = (~pad[:, 1:]).float()[..., None]
        return F.normalize((h * keep).sum(1) / keep.sum(1), dim=-1)

    for _ in range(steps):
        if callable(pairs):
            ab = [pairs(rng) for _ in range(batch)]
            a, b = enc([x for x, _ in ab], 0.3), enc([y for _, y in ab], 0.3)
            keys = [tuple(y) for _, y in ab]
        else:
            ab = [pairs[rng.randrange(len(pairs))] for _ in range(batch)]
            a, b = enc([x for x, _ in ab]), enc([y for _, y in ab])
            keys = [tuple(y) for _, y in ab]
        lg = a @ b.T / tau
        same = torch.tensor([[ki == kj and i != j for j, kj in enumerate(keys)] for i, ki in enumerate(keys)])
        loss = F.cross_entropy(lg.masked_fill(same, -1e9), torch.arange(batch))
        opt.zero_grad()
        loss.backward()
        opt.step()
    m.eval()
    m.attn_type = "causal"
    return m


def true_affinity():
    """
    两件商品的"真实亲和度"：同一个用户（按类型的偏好重新挑时）买它们的概率，相对于各自独立被买的比值
    PMI(i, j) = log Σ_k π_k g_k(i) g_k(j) − log g(i) − log g(j)。【返回】144 × 144 矩阵。
    """
    g = TYPE_PRIOR @ G
    A = (G.T * TYPE_PRIOR) @ G
    return np.log(A + 1e-12) - np.log(np.outer(g, g))


def spearman(a, b):
    ra, rb = np.argsort(np.argsort(a)), np.argsort(np.argsort(b))
    return float(np.corrcoef(ra, rb)[0, 1])


# =====================================================================
# 8. 重排序（04 号脚本）：逐点、成对、列表
# =====================================================================

RR_SPECIALS = ["<pad>", "<bos>", "<sep>", "<yes>", "<no>", "<A>", "<B>", "<sep2>"]
RR = {s: N_ITEMS + i for i, s in enumerate(RR_SPECIALS)}
RR_VOCAB = N_ITEMS + len(RR_SPECIALS)


def rr_point(h, c):
    """逐点：<bos> 历史 <sep> 候选 → 下一个 token 是 <yes> 还是 <no>（"这件商品是用户接下来要买的吗？"）。"""
    return [RR["<bos>"]] + list(h) + [RR["<sep>"], int(c)]


def rr_pair(h, a, b):
    """成对：<bos> 历史 <sep> A <sep2> B → <A> 还是 <B>（"这两件里哪件更可能？"，Qin et al. 2023 的 PRP）。"""
    return [RR["<bos>"]] + list(h) + [RR["<sep>"], int(a), RR["<sep2>"], int(b)]


def rr_list(h, cands):
    """列表：<bos> 历史 <sep> 候选 1 … 候选 20 <sep2> → 直接说出最可能的那件商品（只在候选里取 argmax）。"""
    return [RR["<bos>"]] + list(h) + [RR["<sep>"]] + [int(x) for x in cands] + [RR["<sep2>"]]


def rr_batches(examples, bs=64):
    """examples：[(token 序列, 答案 token)]。【返回】list[(ids, pad, 最后位置, 答案)]。"""
    ex = sorted(examples, key=lambda e: len(e[0]))
    out = []
    for i in range(0, len(ex), bs):
        ch = ex[i:i + bs]
        ids, pad = pad_batch([e[0] for e in ch], RR["<pad>"])
        out.append((ids, pad, torch.tensor([len(e[0]) - 1 for e in ch]), torch.tensor([e[1] for e in ch])))
    return out


def rr_loss(m, b):
    """只在最后一个位置算交叉熵：预测"答案 token"（像让 LLM 回答 yes / no、A / B 或说出商品名）。"""
    ids, pad, last, y = b
    lg = m(ids, pad=pad)[torch.arange(len(last)), last]
    return F.cross_entropy(lg, y, reduction="sum"), torch.tensor(len(last))


def new_reranker(init):
    """
    重排序模型：从预训练的 SASRec（init）出发——像用一个预训练的 LLM 做重排序。商品的嵌入、注意力和前馈层
    全部复制过来；新增的 8 个特殊 token 的嵌入随机初始化。
    """
    m = new_sasrec(0, vocab=RR_VOCAB)
    own = m.state_dict()
    for k, v in init.state_dict().items():
        if k == "emb.weight":
            own[k][:N_ITEMS] = v[:N_ITEMS]
        else:
            own[k] = v
    m.load_state_dict(own)
    return m


def train_reranker(name, examples, init, epochs=5, lr=1e-3):
    b = rr_batches(examples)
    m, _ = load_or_train(name, {"n": len(examples), "tok": sum(len(e[0]) for e in examples), "epochs": epochs,
                                "lr": lr, "v": 1},
                         lambda: new_reranker(init),
                         lambda m: train(m, b, rr_loss, epochs=epochs, lr=lr, warmup=100, verbose=False),
                         verbose=False, ckpt_dir=CKPT_DIR)
    return m


@torch.no_grad()
def rr_last_logits(m, seqs):
    """一批序列在最后一个位置的 logits。【返回】(n, RR_VOCAB) 张量。"""
    ids, pad = pad_batch(seqs, RR["<pad>"])
    return m(ids, pad=pad)[torch.arange(len(seqs)), torch.tensor([len(s) - 1 for s in seqs])]


__all__ = ["ITEMS", "N_ITEMS", "FORMS", "CATS", "SIBLINGS", "TYPES", "TYPE_NAMES", "RHO", "G", "F_MAT", "make_users",
           "sample_user", "type_posterior", "true_next_dist", "rank_metrics", "expected_hr", "popularity_scores",
           "markov_scores", "PAD", "BOS", "SEQ_VOCAB", "new_sasrec", "seq_batches", "ce_loss", "make_bce_loss",
           "train_sasrec", "sasrec_scores", "valid_users", "valid_nll", "content_vectors", "kmeans",
           "residual_codes", "id_scheme", "CodeRec", "TV", "market_corpus", "train_text_lm", "embed_texts", "mntp_adapt",
           "contrastive_adapt", "true_affinity", "spearman", "RR", "RR_VOCAB", "rr_point", "rr_pair", "rr_list",
           "train_reranker", "rr_last_logits"]