"""
=============================================================================
CMU 17-630 · 第 1 讲共享模块：embeddings.py —— 从"数共现"到 word2vec 的词向量
=============================================================================

【这个模块是做什么的】
    "看一个词和谁在一起，就知道它是什么意思"（Firth 1957）。这句话叫**分布式假设**
    （distributional hypothesis），是所有词向量方法的出发点。这个模块实现两条路线：

    1. 基于计数（count-based）：数出"词 × 上下文词"的共现矩阵 → 换成 PPMI（正的点互信息）
       → 用 SVD 压缩成低维稠密向量。这是 LSA（1990）一脉的做法。
    2. 基于预测（prediction-based）：skip-gram + 负采样（SGNS，Mikolov et al. 2013，即 word2vec）。
       训练一个"给定中心词，判断某个词是不是它的真实邻居"的小分类器，分类器的参数就是词向量。

    Levy & Goldberg（2014）证明了 SGNS 在隐式地分解一个（平移过的）PMI 矩阵，所以这两条路线
    其实是一回事的两种算法——02 号脚本会看到它们给出的邻居非常相似。

【模块里有什么】
    cooccurrence(ids, V, window)        共现计数矩阵
    ppmi(C, alpha)                      正的点互信息（带上下文分布平滑）
    svd_vectors(M, k)                   SVD 压缩成 k 维
    train_sgns(ids, V, ...)             numpy 写的 skip-gram 负采样
    normalize / cosine / nearest        余弦相似度与最近邻
    analogy / analogy_accuracy          "a 之于 a* 如同 b 之于 ?"

【运行】python3 embeddings.py    # 自检：在训练语料上算两种词向量，打印几个邻居
【依赖】numpy。02、03 号脚本会 import 它。
"""

import numpy as np


def cooccurrence(ids, V, window=2, skip=()):
    """
    共现矩阵：C[w, c] = 词 c 出现在词 w 左右 window 个位置之内的次数。

    【参数】ids：token id 流；V：词表大小；window：窗口半径；skip：不参与计数的 id（例如 "<doc>"）
    【返回】V × V 的 float64 矩阵（对称）
    【思路】对每个偏移量 o = 1…window，把 (ids[i], ids[i+o]) 和 (ids[i+o], ids[i]) 都加一次。
        用 numpy 的 np.add.at 一次性完成"按下标累加"，比 Python 双重循环快几百倍。
    【边界】跨文档边界的共现也会被数进去；把 "<doc>" 放进 skip 可以去掉"<doc>"本身，但它两边的词
        仍可能被算作邻居——影响很小，不再处理。
    【Python 小知识】np.add.at(A, (rows, cols), 1)：对每一对 (rows[i], cols[i]) 执行 A[r, c] += 1，
        重复的下标会被**累加**（普通的 A[rows, cols] += 1 遇到重复下标只加一次）。
    """
    a = np.asarray(ids)
    C = np.zeros((V, V))
    skip = np.isin(a, list(skip)) if skip else np.zeros(len(a), bool)
    for o in range(1, window + 1):
        left, right = a[:-o], a[o:]
        ok = ~(skip[:-o] | skip[o:])
        np.add.at(C, (left[ok], right[ok]), 1)
        np.add.at(C, (right[ok], left[ok]), 1)
    return C


def ppmi(C, alpha=0.75):
    """
    正的点互信息：PPMI(w, c) = max(0, log P(w, c) / (P(w) · P_α(c)))。

    【思路】原始计数被高频词（"the"、"."）主导：几乎每个词都常和 "the" 一起出现，这说明不了什么。
        PMI 问的是"w 和 c 一起出现的次数，比它们各自独立出现时预期的多多少倍"。
        负的 PMI（比预期少）在小语料上很不可靠，所以截成 0。
        α = 0.75：把上下文的分布 P(c) 开 0.75 次方再归一化（Levy et al. 2015 发现这能让罕见的上下文
        不至于得到夸张的高分；word2vec 的负采样里也有同样的 0.75）。
    【参数】C：共现矩阵；alpha：上下文分布平滑指数
    【返回】同样形状的矩阵，元素 ≥ 0
    【边界】某行或某列全为 0 时（从没出现过的词），对应的 PMI 记为 0。
    【Python 小知识】with np.errstate(divide="ignore", invalid="ignore")：在这个代码块里，
        除以 0、log(0) 不报警告；之后用 np.nan_to_num / np.maximum 把 inf、nan 处理掉。
    """
    total = C.sum()
    pw = C.sum(axis=1) / total
    pc = C.sum(axis=0) ** alpha
    pc = pc / pc.sum()
    with np.errstate(divide="ignore", invalid="ignore"):
        pmi = np.log((C / total) / np.outer(pw, pc))
    pmi = np.nan_to_num(pmi, nan=0.0, posinf=0.0, neginf=0.0)
    return np.maximum(pmi, 0.0)


def svd_vectors(M, k=20):
    """
    用截断 SVD 把 V × V 的矩阵压成 V × k 的稠密向量：M ≈ U S Vᵀ，取 U_k · sqrt(S_k)。

    【为什么要压缩】PPMI 的每一行有 V 维、大部分是 0；压到 k 维后，"出现在相似上下文里的词"
        会被迫共用少数几个方向，相似性才能"传递"（两个词即使没有共同的邻居，也可能有相近的向量）。
    【边界】k 不能超过 V。sqrt(S) 是 Levy et al. 2015 推荐的权重（比直接用 U·S 更好）。
    """
    U, S, _ = np.linalg.svd(M, full_matrices=False)
    return U[:, :k] * np.sqrt(S[:k])


def train_sgns(ids, V, dim=32, window=2, negatives=5, epochs=3, lr=0.05, batch=1024,
               subsample=1e-3, skip=(), seed=0, verbose=False):
    """
    skip-gram + 负采样（SGNS），numpy 实现。

    【思路】正样本：(中心词 w, 窗口里的真实邻居 c)；负样本：(w, 按 unigram^0.75 随机抽的 k 个词)。
        对每一对，计算 x = v_w · u_c（v 是"输入"向量，u 是"输出"向量），
        损失 = −log σ(x_正) − Σ log σ(−x_负)：让真实邻居的点积变大，随机词的点积变小。
        梯度：∂L/∂x_正 = σ(x) − 1，∂L/∂x_负 = σ(x)；再乘上另一边的向量就得到 v、u 的梯度。
        "高频词下采样"：每个 token 以 1 − sqrt(t / f) 的概率被丢掉（f 是它的频率），
        这样 "the"、"." 不会占满训练时间（Mikolov et al. 2013）。
    【参数】ids：token 流；V：词表大小；dim：向量维度；window：窗口半径；negatives：每个正样本配几个负样本；
        epochs：遍数；lr：学习率（线性衰减到 0）；batch：每批多少对；subsample：下采样阈值 t；
        skip：不参与训练的 id；seed：随机种子
    【返回】(输入向量 V × dim, 每个 epoch 的平均损失列表)
    【边界】np.add.at 让同一批里重复出现的词的梯度**全部累加**：一个高频词（"."、"the"）在一批 1024 对里
        可能出现上百次，相当于一步走了上百倍的学习率。所以关闭下采样（subsample=1.0）时，lr=0.05 会发散，
        要把 lr 调到 0.005 左右；发散时函数会抛出 RuntimeError，而不是悄悄返回一堆 nan。
        【下采样的副作用】被丢掉的高频词是在**构造窗口之前**删掉的，所以窗口实际上"撑大"了：
        "cold . rating : negative" 里的 "." 和 ":" 被删掉后，negative 就进了 cold 的 ±2 窗口
        （Levy, Goldberg & Dagan 2015 讨论过这种"脏下采样"）。02 号脚本演示 5 会看到它的后果。
    """
    rng = np.random.default_rng(seed)
    a = np.asarray(ids)
    freq = np.bincount(a, minlength=V).astype(float)
    freq /= freq.sum()
    keep_prob = np.minimum(1.0, np.sqrt(subsample / np.maximum(freq, 1e-12)))
    neg_dist = np.bincount(a, minlength=V).astype(float) ** 0.75
    for s in skip:
        neg_dist[s] = 0.0
    neg_dist /= neg_dist.sum()
    W_in = (rng.random((V, dim)) - 0.5) / dim
    W_out = np.zeros((V, dim))
    losses = []
    skip_mask = np.isin(a, list(skip)) if skip else np.zeros(len(a), bool)
    for ep in range(epochs):
        kept = a[(rng.random(len(a)) < keep_prob[a]) & ~skip_mask]   # 下采样
        centers, contexts = [], []
        for o in range(1, window + 1):
            centers += [kept[:-o], kept[o:]]
            contexts += [kept[o:], kept[:-o]]
        centers, contexts = np.concatenate(centers), np.concatenate(contexts)
        order = rng.permutation(len(centers))
        centers, contexts = centers[order], contexts[order]
        total, n_batches = 0.0, 0
        for b0 in range(0, len(centers), batch):
            step_lr = lr * max(0.02, 1.0 - (ep + b0 / len(centers)) / epochs)
            c, o = centers[b0:b0 + batch], contexts[b0:b0 + batch]
            neg = rng.choice(V, size=(len(c), negatives), p=neg_dist)
            v = W_in[c]                                   # (B, d)
            u_pos, u_neg = W_out[o], W_out[neg]           # (B, d), (B, k, d)
            x_pos = np.clip(np.sum(v * u_pos, axis=1), -10, 10)           # 截断到 ±10，防止 exp 溢出
            x_neg = np.clip(np.einsum("bd,bkd->bk", v, u_neg), -10, 10)
            s_pos = 1 / (1 + np.exp(-x_pos))                               # σ(x_正)
            s_neg = 1 / (1 + np.exp(-x_neg))                               # σ(x_负)
            total += float(np.mean(-np.log(s_pos + 1e-10) - np.sum(np.log(1 - s_neg + 1e-10), axis=1)))
            n_batches += 1
            g_pos = (s_pos - 1)[:, None]                  # ∂L/∂x_正
            g_neg = s_neg[:, :, None]                     # ∂L/∂x_负
            grad_v = g_pos * u_pos + np.sum(g_neg * u_neg, axis=1)
            np.add.at(W_out, o, -step_lr * g_pos * v)
            np.add.at(W_out, neg, -step_lr * g_neg * v[:, None, :])
            np.add.at(W_in, c, -step_lr * grad_v)
        losses.append(total / n_batches)
        if not (np.isfinite(W_in).all() and np.isfinite(losses[-1])):
            raise RuntimeError("SGNS 训练发散了（出现了 inf / nan）：请调小 lr，或打开高频词下采样。")
        if verbose:
            print(f"    SGNS 第 {ep + 1} 遍：平均损失 {losses[-1]:.3f}（{len(centers):,} 对）")
    return W_in, losses


def normalize(X):
    """把每一行缩放成长度 1（全零行保持为 0）。之后行与行的点积就是余弦相似度。"""
    n = np.linalg.norm(X, axis=1, keepdims=True)
    return X / np.where(n == 0, 1.0, n)


def cosine(X, tok, a, b):
    """词 a 和词 b 的余弦相似度。X 是（不必归一化的）词向量矩阵。"""
    Xn = normalize(X)
    return float(Xn[tok.stoi[a]] @ Xn[tok.stoi[b]])


def nearest(X, tok, word, k=5, exclude=()):
    """
    和 word 余弦相似度最高的 k 个词。

    【返回】[(词, 相似度), ...]
    【边界】自动排除 word 自己、"<unk>"、"<doc>" 和 exclude 里的词。
        相似度并列时按 id 排序，保证输出稳定。
    """
    Xn = normalize(X)
    sims = Xn @ Xn[tok.stoi[word]]
    bad = {word, "<unk>", "<doc>", *exclude}
    order = sorted(range(len(sims)), key=lambda i: (-sims[i], i))
    return [(tok.itos[i], float(sims[i])) for i in order if tok.itos[i] not in bad][:k]


def analogy(X, tok, a, a_star, b, candidates=None):
    """
    类比："a 之于 a* 如同 b 之于 ?"，答案 = argmax_x cos(x, a* − a + b)（排除 a、a*、b 本身）。

    【参数】candidates：只在这些词里找答案（None 表示整个词表）
    【返回】预测的词
    【思路】如果"性别"在向量空间里是一个固定的方向，queen − king ≈ woman − man，
        那么 king 加上这个方向就落在 queen 附近。这是 word2vec 论文最出名的例子。
    """
    Xn = normalize(X)
    target = Xn[tok.stoi[a_star]] - Xn[tok.stoi[a]] + Xn[tok.stoi[b]]
    pool = candidates if candidates is not None else tok.itos
    best, best_s = None, -np.inf
    for w in pool:
        if w in (a, a_star, b) or w in ("<unk>", "<doc>"):
            continue
        s = float(Xn[tok.stoi[w]] @ target)
        if s > best_s:
            best, best_s = w, s
    return best


def analogy_accuracy(X, tok, pairs, candidates=None):
    """
    在一组 (x, x*) 对上做所有"a:a* :: b:?"的类比（a、b 取不同的对），返回 (对了几个, 一共几个)。

    【例子】pairs = [("king","queen"), ("man","woman"), ...] → 测 king:queen :: man:? 等 90 道题。
    """
    ok = n = 0
    for a, a_star in pairs:
        for b, b_star in pairs:
            if a == b:
                continue
            n += 1
            ok += analogy(X, tok, a, a_star, b, candidates) == b_star
    return ok, n


if __name__ == "__main__":
    from pretrain_world import WordTokenizer, build_corpus, build_world, corpus_text

    world = build_world()
    text = corpus_text(build_corpus(world))
    tok = WordTokenizer(text)
    ids = tok.encode(text)
    skip = (tok.stoi["<doc>"],)
    X_svd = svd_vectors(ppmi(cooccurrence(ids, len(tok), 2, skip)), 20)
    X_sg, losses = train_sgns(ids, len(tok), epochs=2, skip=skip, verbose=True)
    for w in ["queen", "copper", "bank"]:
        print(f"  {w:8s} PPMI-SVD：{nearest(X_svd, tok, w, 4)}")
        print(f"  {w:8s} SGNS    ：{nearest(X_sg, tok, w, 4)}")