"""
=============================================================================
CS525 · 第 6 讲模块：迷你语言模型、评测、质量分类器与去重（textlab.py）
=============================================================================

对应讲义：lectures/lecture06_pretraining_data_dclm.md（第 7 讲也会只读地导入它）

-----------------------------------------------------------------------------
【这个模块提供什么】
-----------------------------------------------------------------------------
    Vocab / encode              词 ↔ 整数编号
    NGramStats                  一个数据集的三元 / 二元 / 一元计数（"训练"只需要数一遍）
    ExposureLM                  "慢慢学"的三元语言模型：看过某个上下文越多次，越相信看到过的下一个词
    train_lm                    在固定的训练预算（看多少个词）下"训练"：预算比数据少就随机抽一部分文档，多就重复看
    evaluate                    评测：知识题准确率、居中准确率、答案损失
    FastTextLike                fastText 式的质量分类器（词袋 + 二元组 → 线性 → sigmoid）
    exact_dedup / bff_dedup     完全去重；DCLM 式的"段落 + 文档"级去重（可分片）

-----------------------------------------------------------------------------
【为什么用一个"慢慢学"的 n-gram 模型，而不是一个神经网络】
-----------------------------------------------------------------------------
真正的预训练是用梯度下降训练 Transformer：模型每看一次某个上下文，就把对下一个词的预测
往看到的词挪一小步。我们用最简单的方式保留这个性质：

    对每个上下文 c（前两个词），模型的预测分布 p_c 从"先验"出发；
    每看到一次 (c, w)，就做一次  p_c ← (1 − η) · p_c + η · onehot(w)。

看过 n 次之后，取期望（假设这 n 次出现的顺序是随机的），
    p_c = (1 − (1 − η)^n) · q̂_c  +  (1 − η)^n · 先验
其中 q̂_c 是数据里 c 后面各个词出现的频率。先验用"更短的上下文"的同样的模型（二元 → 一元），
这就是一个带回退（backoff）的三元模型，只不过"相信多少"由看过的次数决定。

这个模型有两个对数据研究很重要的性质，而且计算只需要数数（几秒钟）：
    (1) 训练预算 T（看多少个词）有意义：数据集有 |S| 个词。T 小于 |S| 时只看了随机的一部分文档（每篇看一遍）；
        T 大于 |S| 时整个数据集被看了 T / |S| 遍，每个上下文被看的次数 = 它的计数 × T / |S|；
    (2) 重复有用，但有上限：重复让 n 变大、更接近 q̂_c，却不能带来数据里没有的事实——
        第 7 讲用它研究"数据不够时，重复几遍还有用"。
它当然不是一个真正的语言模型：没有泛化（没见过的三元组只能靠回退），没有容量限制。
讲义会说明哪些结论依赖于这些简化。
"""

import zlib

import numpy as np

from webtext_world import FOREIGNW, JUNK, RELATIONS

BOS, EOS = "<s>", "</s>"


class Vocab:
    """
    词表：把世界里所有可能出现的词编上号。

    【参数】world：WebWorld
    【属性】itos（编号 → 词）、stoi（词 → 编号）、V（词表大小）
    """

    def __init__(self, world):
        words = [BOS, EOS] + sorted(world.english) + list(world.entities) + list(JUNK) + list(FOREIGNW)
        self.itos = list(dict.fromkeys(words))            # 去重并保持顺序
        self.stoi = {w: i for i, w in enumerate(self.itos)}
        self.V = len(self.itos)


def encode(docs_lines, vocab):
    """
    把一组文档编码成一个整数数组：每个文档前面加两个 <s>（保证三元上下文不跨文档），后面加 </s>。

    【参数】docs_lines：文档列表，每个文档是"行的列表"（每行是词列表）；vocab
    【返回】int64 数组
    【边界】空文档也会贡献 <s> <s> </s> 三个词——调用方通常先把空文档去掉。
    """
    out = []
    s = vocab.stoi
    for lines in docs_lines:
        out.extend((0, 0))
        for l in lines:
            out.extend(s[t] for t in l)
        out.append(1)
    return np.array(out, dtype=np.int64)


class NGramStats:
    """
    一个数据集的计数：三元 (a, b, w)、二元 (b, w)、一元 w，以及每个上下文出现的总次数。

    【参数】docs_lines：文档（行的列表）的列表；vocab
    【思路】把 (a, b, w) 编码成一个整数 key = (a·V + b)·V + w，用 np.unique 一次数完——比 Python 字典快得多。
            查询时用 np.searchsorted 在排好序的 key 里二分查找。
    【属性】n_tokens：数据集的词数（不含 <s>，含 </s>）——训练预算按它换算成"看了几遍"
    """

    def __init__(self, docs_lines, vocab):
        V = vocab.V
        self.V = V
        x = encode([d for d in docs_lines if any(d)], vocab)
        a, b, w = x[:-2], x[1:-1], x[2:]
        valid = w != 0                                    # 预测目标不能是 <s>
        a, b, w = a[valid], b[valid], w[valid]
        self.n_tokens = int(valid.sum())
        self.k3, self.c3 = np.unique((a * V + b) * V + w, return_counts=True)
        self.kc3, self.n3 = np.unique(a * V + b, return_counts=True)
        self.k2, self.c2 = np.unique(b * V + w, return_counts=True)
        self.kc2, self.n2 = np.unique(b, return_counts=True)
        self.c1 = np.bincount(w, minlength=V).astype(np.float64)
        # 每个上下文后面出现过几种不同的词（Witten–Bell 平滑要用）。
        # k3 是排好序的，k3 // V 就是它的上下文，也是排好序的，和 kc3 一一对应。
        _, self.t3 = np.unique(self.k3 // V, return_counts=True)
        _, self.t2 = np.unique(self.k2 // V, return_counts=True)

    @staticmethod
    def _lookup(keys, counts, query):
        """在排好序的 keys 里查 query，返回对应的计数（查不到为 0）。"""
        if len(keys) == 0:
            return np.zeros(len(query))
        i = np.searchsorted(keys, query)
        i = np.clip(i, 0, len(keys) - 1)
        return np.where(keys[i] == query, counts[i], 0).astype(np.float64)


class ExposureLM:
    """
    "慢慢学"的三元模型（见模块说明）。

    【参数】stats：模型看过的数据的 NGramStats；epochs：这些数据被看了几遍（可以是 2.5 这样的小数）；
            eta：每看一次往观测挪多少（学习率的类比）。默认 0.5：一条事实看一次，模型就有一半相信它；
            看 4 次约 94%。第一版用 0.1，结果"看一次几乎学不到"，热门事实的频率偏差压倒了一切（见 README）。
    【方法】prob(a, b, w)：P(w | a, b)，三个参数都是同样长度的整数数组
    【边界】这里用"期望"代替真正的随机顺序：每个上下文被看的次数 = 它的计数 × epochs。
            通常不直接构造它，而是用 train_lm（它会处理"预算比数据少"的情况）。
    """

    def __init__(self, stats, epochs, eta=0.5):
        self.s = stats
        self.f = epochs                                    # 平均每个词被看了几遍
        self.eta = eta

    def _w(self, n):
        # 看了 f·n 次之后，"相信观测"的权重。
        # 【边界】看了很多次时 (1 − η)^(f·n) 会小到 0，模型就会给没见过的词概率 0、损失变成无穷大；
        #         所以永远给先验留至少千分之一（真实的模型也不会给任何词严格为 0 的概率）。
        return 1.0 - np.maximum((1.0 - self.eta) ** (self.f * n), 1e-3)

    def prob(self, a, b, w):
        """
        P(w | a, b)。

        【思路】三层回退，每一层都是  "相信观测的权重" × 平滑后的观测频率 + 其余 × 下一层：
            P1(w)       = 一元频率（加一平滑）
            P2(w | b)   = w2 · q̃2 + (1 − w2) · P1,   q̃2 = (f·c(b,w) + t(b) · P1) / (f·n(b) + t(b))
            P3(w | a,b) = w3 · q̃3 + (1 − w3) · P2,   q̃3 = (f·c(a,b,w) + t(a,b) · P2) / (f·n(a,b) + t(a,b))
        n 是上下文在数据里出现的次数，t 是它后面出现过几种不同的词，f 是数据被看了几遍。q̃ 是 Witten–Bell 平滑：
        一个上下文后面的词越五花八门，越要给"没见过的词"留余地；同一条数据看得越多遍，留的余地越小。
        """
        s, V = self.s, self.s.V
        n3 = s._lookup(s.kc3, s.n3, a * V + b)
        t3 = s._lookup(s.kc3, s.t3, a * V + b)
        c3 = s._lookup(s.k3, s.c3, (a * V + b) * V + w)
        n2 = s._lookup(s.kc2, s.n2, b)
        t2 = s._lookup(s.kc2, s.t2, b)
        c2 = s._lookup(s.k2, s.c2, b * V + w)
        f = self.f
        p1 = (f * s.c1[w] + 1.0) / (f * s.c1.sum() + V)                # 一元：加一平滑
        # 平滑用的计数也乘以"看了几遍"f：同一条数据看得越多，模型越相信它（越接近死记硬背）。
        # f = 1（每篇看一遍）时就是普通的 Witten–Bell；第 7 讲重复训练时 f > 1。
        q2 = np.divide(f * c2 + t2 * p1, f * n2 + t2, out=p1.copy(), where=n2 > 0)
        p2 = self._w(n2) * q2 + (1 - self._w(n2)) * p1
        q3 = np.divide(f * c3 + t3 * p2, f * n3 + t3, out=p2.copy(), where=n3 > 0)
        return self._w(n3) * q3 + (1 - self._w(n3)) * p2


def make_eval(world, vocab):
    """
    准备评测：3000 道知识题，每个 (国家, 关系) 一道。

    【参数】world；vocab
    【返回】字典：a, b（题目的上下文：国家、关系的编号）、cand（每题的候选值编号，按关系补齐成同样长度）、
            n_cand（每题真实的候选个数）、truth（正确答案在候选里的位置）、rank（国家的热度排名）
    【数据联系】DCLM 用 53 个下游任务；我们只有一种任务——"这个模型从数据里学到了多少条事实"。
            这让结论容易解释，也让它们只适用于"知识"这一种能力（讲义会讨论这个局限）。
    """
    qs = world.questions()
    K = max(len(world.values[r]) for r in RELATIONS)
    a = np.array([vocab.stoi[world.entities[e]] for e, _, _ in qs])
    b = np.array([vocab.stoi[r] for _, r, _ in qs])
    cand = np.zeros((len(qs), K), dtype=np.int64)
    n_cand = np.zeros(len(qs), dtype=np.int64)
    for i, (_, r, _) in enumerate(qs):
        ids = [vocab.stoi[v] for v in world.values[r]]
        cand[i, :len(ids)] = ids
        cand[i, len(ids):] = ids[0]                       # 补齐用的占位，评测时会被屏蔽
        n_cand[i] = len(ids)
    truth = np.array([t for _, _, t in qs])
    rank = np.array([e for e, _, _ in qs])
    return {"a": a, "b": b, "cand": cand, "n_cand": n_cand, "truth": truth, "rank": rank}


def train_lm(docs_lines, budget, vocab, seed=0, eta=0.5):
    """
    在固定的训练预算下"训练"一个模型（DCLM 的做法：模型和训练 token 数固定，只换数据）。

    【参数】docs_lines：数据集（文档列表）；budget：训练时看多少个词；vocab；seed：抽文档用的随机种子；eta
    【返回】(ExposureLM, 看了几遍)
    【思路】预算 < 数据集大小：把文档随机打乱，从头取，直到凑够预算——每篇被选中的文档看一遍，其余没看到。
            预算 ≥ 数据集大小：整个数据集都看，看 budget / |S| 遍（第 7 讲研究的"重复"）。
    【边界】空数据集无法训练：返回一个只有 <s> 的模型（什么都不知道），看了 0 遍。
    """
    docs = [d for d in docs_lines if any(d)]
    total = sum(doc_tokens(d) + 1 for d in docs)          # +1：每篇文档末尾的 </s>
    if total == 0:
        return ExposureLM(NGramStats([[["."]]], vocab), 0.0, eta), 0.0
    if budget < total:
        order = np.random.default_rng(seed).permutation(len(docs))
        lens = np.array([doc_tokens(docs[i]) + 1 for i in order])
        k = int(np.searchsorted(np.cumsum(lens), budget)) + 1   # 凑够预算需要的文档数
        stats = NGramStats([docs[i] for i in order[:k]], vocab)
        return ExposureLM(stats, 1.0, eta), budget / total
    stats = NGramStats(docs, vocab)
    return ExposureLM(stats, budget / stats.n_tokens, eta), budget / stats.n_tokens


def evaluate(lm, ev):
    """
    评测一个"训练好"的模型。

    【参数】lm：ExposureLM（通常来自 train_lm）；ev：make_eval 的结果
    【返回】字典：
        acc       知识题准确率（3000 道题，每个国家一样重要——长尾国家和热门国家同等计分）
        acc_head  前 50 个热门国家的准确率；acc_tail：排名 200 以后的国家的准确率
        centered  居中准确率：(准确率 − 随机猜) / (1 − 随机猜)，0 = 随机，1 = 全对（DCLM 用的报告方式）
        ans_loss  答案损失：−log P(正确答案)，只在候选里归一化——比准确率更"连续"，重复训练的效果会体现在这里
    【边界】候选的概率完全相同（模型什么都没学到）时，按"随机打破平局"的期望计分：和 m 个候选并列第一，记 1/m 分。
    """
    N, K = ev["cand"].shape
    p = lm.prob(np.repeat(ev["a"], K), np.repeat(ev["b"], K), ev["cand"].ravel()).reshape(N, K)
    mask = np.arange(K)[None, :] < ev["n_cand"][:, None]
    p = np.where(mask, p, 0.0)
    pt = p[np.arange(N), ev["truth"]]
    better = (p > pt[:, None]).sum(1)
    ties = (p == pt[:, None]).sum(1)                      # 包括正确答案自己
    correct = np.where(better == 0, 1.0 / ties, 0.0)
    chance = 1.0 / ev["n_cand"]
    centered = (correct - chance) / (1 - chance)
    ans_loss = -np.log(pt / p.sum(1))
    return {"acc": float(correct.mean()), "acc_head": float(correct[ev["rank"] < 50].mean()),
            "acc_tail": float(correct[ev["rank"] >= 200].mean()), "centered": float(centered.mean()),
            "ans_loss": float(ans_loss.mean())}


# =============================================================================
# 质量分类器：fastText 的玩具版
# =============================================================================
class FastTextLike:
    """
    fastText 式的线性分类器：一个文档 = 它的词和相邻二元组的"平均"，打一个分数 → sigmoid → P(高质量)。

    【参数】dim：特征哈希的桶数（词和二元组被哈希到这么多个桶里）；l2：权重的 L2 正则；
            steps / lr：全批量梯度下降的步数和学习率
    【思路】真正的 fastText 先把每个 n-gram 映射成一个 100 维向量再平均、再接线性层；
            把 100 维去掉、直接给每个 n-gram 一个权重，就是这里的"平均词袋 + 逻辑回归"——同样又快又简单。
    【边界】特征哈希会把不同的 n-gram 撞到同一个桶里；dim 足够大时影响很小。
    """

    def __init__(self, dim=1 << 15, l2=1e-4, steps=300, lr=2.0):
        self.dim, self.l2, self.steps, self.lr = dim, l2, steps, lr
        self.w = np.zeros(dim)
        self.bias = 0.0

    def _feats(self, docs_lines):
        """把文档列表变成 (文档编号, 特征桶, 权重) 三个数组：每个文档的特征权重加起来等于 1（"平均"）。"""
        rows, cols, vals = [], [], []
        for i, lines in enumerate(docs_lines):
            toks = [t for l in lines for t in l]
            if not toks:
                continue
            grams = toks + [a + " " + b for a, b in zip(toks, toks[1:])]
            # 【Python 小知识】zlib.crc32 是一个稳定的哈希：每次运行、每台机器结果都一样（内置的 hash() 不是）
            ids = [zlib.crc32(g.encode("utf-8")) % self.dim for g in grams]
            rows.extend([i] * len(ids))
            cols.extend(ids)
            vals.extend([1.0 / len(ids)] * len(ids))
        return np.array(rows, dtype=np.int64), np.array(cols, dtype=np.int64), np.array(vals)

    def _scores(self, n, rows, cols, vals):
        # 【numpy 小知识】np.bincount(rows, weights=...) 把同一个文档的特征贡献加起来——相当于稀疏矩阵乘向量
        return np.bincount(rows, weights=vals * self.w[cols], minlength=n) + self.bias

    def fit(self, pos, neg):
        """
        训练：pos 是正例文档（"高质量参照"），neg 是负例文档（通常是候选池里随机抽的）。

        【参数】pos, neg：文档（行的列表）的列表
        【返回】self
        """
        docs = list(pos) + list(neg)
        y = np.r_[np.ones(len(pos)), np.zeros(len(neg))]
        rows, cols, vals = self._feats(docs)
        n = len(docs)
        for _ in range(self.steps):
            p = 1.0 / (1.0 + np.exp(-self._scores(n, rows, cols, vals)))
            err = p - y                                   # 交叉熵对分数的梯度
            g = np.bincount(cols, weights=vals * err[rows], minlength=self.dim) / n + self.l2 * self.w
            self.w -= self.lr * g
            self.bias -= self.lr * err.mean()
        return self

    def score(self, docs_lines):
        """每个文档的分数（logit，越大越像正例）。"""
        rows, cols, vals = self._feats(docs_lines)
        return self._scores(len(docs_lines), rows, cols, vals)


# =============================================================================
# 去重
# =============================================================================
def exact_dedup(docs_lines):
    """完全去重：正文完全相同的文档只留第一个。返回保留的下标列表。"""
    seen, keep = set(), []
    for i, lines in enumerate(docs_lines):
        key = zlib.crc32(" ".join(" ".join(l) for l in lines).encode("utf-8"))
        if key not in seen:
            seen.add(key)
            keep.append(i)
    return keep


def bff_dedup(docs_lines, n=5, min_len=5, threshold=0.8, shards=1, rng=None):
    """
    DCLM 式的"段落 + 文档"级去重（BFF 的玩具版）。

    【参数】docs_lines：文档列表；n：n-gram 的长度；min_len：一行至少多少个词才参与去重
            （DCLM 的 min_ngram_size：太短的段落——列表项、选项——直接保留）；
            threshold：一行（或整个文档）里有多少比例的 n-gram 已经见过，就删掉它；
            shards：把文档随机分成几片、每片独立去重（并行的代价：跨片的重复留了下来）；rng：分片用
    【返回】新的文档列表（被删的行去掉了；整篇被删的文档变成空列表，下标不变）
    【思路】按顺序扫描。对每一行：算它的 n-gram，查"见过的集合"；见过的比例 > threshold 就删掉这一行，
            否则把它的 n-gram 加进集合。整篇文档的 n-gram 里见过的比例 > threshold，整篇删掉。
            真实的 BFF 用 Bloom filter 代替这里的 Python 集合（几十亿个 n-gram 放不进集合，第 4 讲）。
    """
    out = [None] * len(docs_lines)
    shard_of = np.zeros(len(docs_lines), dtype=int) if shards == 1 else rng.integers(0, shards, size=len(docs_lines))
    for s in range(shards):
        seen = set()
        for i in np.flatnonzero(shard_of == s):
            kept, n_seen, n_all = [], 0, 0
            for l in docs_lines[i]:
                if len(l) < max(min_len, n):
                    kept.append(l)
                    continue
                grams = {zlib.crc32(" ".join(l[j:j + n]).encode("utf-8")) for j in range(len(l) - n + 1)}
                hit = sum(g in seen for g in grams)
                n_seen += hit
                n_all += len(grams)
                if hit / len(grams) > threshold:
                    continue
                seen.update(grams)
                kept.append(l)
            out[i] = [] if (n_all > 0 and n_seen / n_all > threshold) else kept
    return out


def doc_tokens(lines):
    """一个文档（行的列表）的词数。"""
    return sum(len(l) for l in lines)


if __name__ == "__main__":
    from webtext_world import WebWorld
    W = WebWorld()
    voc = Vocab(W)
    rng = np.random.default_rng(1)
    ev = make_eval(W, voc)
    docs = [W.make_doc("WIKI", rng, wrap=False).lines for _ in range(2000)]
    print("词表大小", voc.V, "；2000 个百科网页的词数", NGramStats(docs, voc).n_tokens)
    for T in [1e4, 1e5, 1e6]:
        lm, ep = train_lm(docs, T, voc)
        print(f"预算 {T:>9,.0f}（看了 {ep:.2f} 遍）：", {k: round(v, 3) for k, v in evaluate(lm, ev).items()})