"""
wordvec.py —— 从共现到词向量：PPMI + SVD，以及从零实现的 word2vec（SGNS）
=====================================================================

【这个文件要回答的问题】
    "意思相近的词，向量也相近"——这样的向量从哪里来？
    答案来自**分布假设**（distributional hypothesis）：一个词的意思由它周围经常出现的词决定。
    这里实现两种从语料学词向量的方法：
      1. 计数法：统计"词 × 上下文词"的共现次数 → 转成 PPMI（正的点互信息）→ 用 SVD 降到 d 维。
      2. 预测法：word2vec 的 skip-gram + 负采样（SGNS）。让模型区分"真实出现在一起的词对"
         和"随机拼凑的词对"，用梯度下降学出向量。
    Levy & Goldberg (2014) 证明了 SGNS 隐式地在分解一个移位的 PMI 矩阵，所以两种方法本质上很接近。

【怎样使用】
    from wordvec import Vocabulary, cooccurrence, ppmi, svd_vectors, train_sgns, nearest
    voc = Vocabulary(tokens, min_count=5)
    C = cooccurrence(voc.encode(tokens), len(voc), window=2)
    vecs = svd_vectors(ppmi(C), dim=16)
    nearest(vecs, voc, "apple", k=5)

【LLM 联系】
    第 1 讲的神经语言模型、以及所有 LLM 的嵌入层，学到的也是这种"分布式表示"，
    只不过是作为"预测下一个词"的副产品学出来的。word2vec（2013）是"预训练表示"思想的起点：
    先在大量无标注文本上学表示，再拿去做别的任务——LLM 把这个思路推到了极致。
"""

from collections import Counter

import numpy as np


class Vocabulary:
    """
    词表：只保留出现次数 >= min_count 的词，按频率从高到低编号。

    【边界】低频词被丢弃：encode 时它们变成 -1，cooccurrence / train_sgns 会跳过 -1。
    """

    def __init__(self, tokens, min_count=1, max_size=None):
        counts = Counter(tokens)
        words = sorted((w for w, c in counts.items() if c >= min_count), key=lambda w: (-counts[w], w))
        if max_size is not None:
            words = words[:max_size]
        self.itos = words
        self.stoi = {w: i for i, w in enumerate(words)}
        self.counts = np.array([counts[w] for w in words], dtype=np.float64)

    def __len__(self):
        return len(self.itos)

    def encode(self, tokens):
        return np.array([self.stoi.get(t, -1) for t in tokens], dtype=np.int64)


def cooccurrence(ids, V, window=2):
    """
    共现矩阵：C[i, j] = 词 j 出现在词 i 左右 window 个位置以内的次数（对称）。

    【参数】ids：词编号数组（-1 表示被丢弃的词，跳过）；V：词表大小；window：窗口半径。
    【返回】V × V 的 float64 数组。
    【思路】对每个偏移 k = 1..window，把序列和"自己平移 k 位"配对，一次性用 np.add.at 计数。
    """
    C = np.zeros((V, V))
    for k in range(1, window + 1):
        a, b = ids[:-k], ids[k:]
        keep = (a >= 0) & (b >= 0)
        np.add.at(C, (a[keep], b[keep]), 1.0)
        np.add.at(C, (b[keep], a[keep]), 1.0)
    return C


def ppmi(C):
    """
    正的点互信息：PPMI[i, j] = max(0, log P(i, j) / (P(i) P(j)))。

    【直觉】PMI 衡量"i 和 j 一起出现的次数，比它们各自随机出现时的预期多多少倍"。
            the 和谁都经常一起出现，但这只是因为 the 本身很常见——PMI 把这种"常见"除掉了。
            负的 PMI（比预期还少）估计得很不可靠，所以截断成 0。
    【边界】从未共现的词对 log 0 = -inf，截断后是 0。
    """
    total = C.sum()
    row = C.sum(axis=1, keepdims=True)
    col = C.sum(axis=0, keepdims=True)
    with np.errstate(divide="ignore", invalid="ignore"):
        pmi = np.log(C * total / (row * col))
    pmi[~np.isfinite(pmi)] = 0.0
    return np.maximum(pmi, 0.0)


def svd_vectors(M, dim):
    """
    用截断 SVD 把 V × V 的矩阵压缩成 V × dim 的词向量：M ≈ U S Vᵀ，取 U[:, :dim] · sqrt(S[:dim])。

    【为什么要降维】原始的 PPMI 行向量有 V 维、大部分是 0（稀疏）；SVD 找到最能解释
            矩阵结构的 dim 个方向，把"经常出现在相似上下文里"的词压到相近的位置。
    【边界】SVD 的每个奇异向量的正负号是任意的；我们把每一列调整成"第一个非零元素为正"，保证结果可复现。
    """
    U, S, _ = np.linalg.svd(M, full_matrices=False)
    U = U[:, :dim] * np.sqrt(S[:dim])
    signs = np.sign(U[np.argmax(np.abs(U) > 1e-12, axis=0), np.arange(dim)])
    return U * np.where(signs == 0, 1.0, signs)


def train_sgns(ids, V, counts, dim=16, window=2, negatives=5, epochs=3, lr=0.05, batch=1024, seed=0,
               verbose=False):
    """
    从零实现 word2vec：skip-gram + 负采样（SGNS），纯 NumPy。

    【目标】对每一对真实的 (中心词 c, 上下文词 o)，以及 k 个随机抽的"负样本"词 n_1..n_k，最大化
                log σ(u_c · v_o) + Σ_j log σ(-u_c · v_{n_j})
            即：让真实词对的点积大、随机词对的点积小。σ 是 sigmoid 函数。
            u 是"中心词向量"，v 是"上下文向量"，最后用 u 作为词向量。
    【负样本分布】按 count^0.75 抽（word2vec 原论文的经验做法：比按频率抽更照顾罕见词）。
    【参数】ids：词编号序列（-1 跳过）；V：词表大小；counts：每个词的出现次数（长度 V）；
            dim：向量维度；window：窗口半径；negatives：每个正样本配几个负样本；
            epochs：遍历语料几遍；lr：学习率（线性衰减到 0）；batch：每批多少个正样本。
    【返回】(中心词向量 U, 每个 epoch 的平均损失列表)。
    【边界】为了简单和可复现，这里用小批量 SGD。原版 word2vec 是一次更新一对词；
            小批量里高频词会出现很多次，所以对每个词的梯度取"批内平均"再更新（否则会发散，见代码 README）。
    """
    rng = np.random.default_rng(seed)
    pairs = []
    for k in range(1, window + 1):
        a, b = ids[:-k], ids[k:]
        keep = (a >= 0) & (b >= 0)
        pairs.append(np.stack([a[keep], b[keep]], axis=1))
        pairs.append(np.stack([b[keep], a[keep]], axis=1))
    pairs = np.concatenate(pairs)
    noise = counts ** 0.75
    noise = noise / noise.sum()
    U = (rng.random((V, dim)) - 0.5) / dim          # word2vec 的初始化：中心向量小随机数
    W = np.zeros((V, dim))                           # 上下文向量初始化为 0
    total_steps = epochs * int(np.ceil(len(pairs) / batch))
    step = 0
    history = []

    def log_sigmoid(x):
        return -np.logaddexp(0.0, -x)               # 数值稳定的 log σ(x)

    for ep in range(epochs):
        order = rng.permutation(len(pairs))
        losses = []
        for start in range(0, len(pairs), batch):
            idx = order[start:start + batch]
            c, o = pairs[idx, 0], pairs[idx, 1]
            neg = rng.choice(V, size=(len(idx), negatives), p=noise)
            uc = U[c]                                # (B, d)
            vo = W[o]                                # (B, d)
            vn = W[neg]                              # (B, k, d)
            pos_score = np.sum(uc * vo, axis=1)                  # (B,)
            neg_score = np.einsum("bd,bkd->bk", uc, vn)          # (B, k)
            losses.append(float(-(log_sigmoid(pos_score) + log_sigmoid(-neg_score).sum(axis=1)).mean()))
            # 梯度：d/dx [-log σ(x)] = σ(x) - 1；d/dx [-log σ(-x)] = σ(x)
            g_pos = 1.0 / (1.0 + np.exp(-pos_score)) - 1.0       # (B,)
            g_neg = 1.0 / (1.0 + np.exp(-neg_score))             # (B, k)
            grad_uc = g_pos[:, None] * vo + np.einsum("bk,bkd->bd", g_neg, vn)
            grad_vo = g_pos[:, None] * uc
            grad_vn = g_neg[:, :, None] * uc[:, None, :]
            cur_lr = lr * (1.0 - step / total_steps)             # 学习率线性衰减
            # 同一个词在一批里可能出现几百次（例如句号）。把它的梯度直接加起来，步子会大几百倍而发散；
            # 所以先累加、再除以它在这一批里出现的次数（每个词的梯度取平均）。
            gU = np.zeros_like(U)
            np.add.at(gU, c, grad_uc)
            nU = np.bincount(c, minlength=V)
            gW = np.zeros_like(W)
            ctx_ids = np.concatenate([o, neg.reshape(-1)])
            np.add.at(gW, ctx_ids, np.concatenate([grad_vo, grad_vn.reshape(-1, dim)]))
            nW = np.bincount(ctx_ids, minlength=V)
            U -= cur_lr * gU / np.maximum(nU, 1)[:, None]
            W -= cur_lr * gW / np.maximum(nW, 1)[:, None]
            step += 1
        history.append(float(np.mean(losses)))
        if verbose:
            print(f"    SGNS 第 {ep + 1} 遍：平均损失 {history[-1]:.4f}")
    return U, history


def normalize(vecs):
    """每一行除以自己的长度（零向量保持为 0）。"""
    norms = np.linalg.norm(vecs, axis=1, keepdims=True)
    return vecs / np.where(norms == 0, 1.0, norms)


def nearest(vecs, voc, word, k=5, candidates=None):
    """
    余弦相似度最高的 k 个词。

    【参数】candidates：可选，只在这些词里找（例如只在名词里找）。
    【返回】[(词, 相似度), ...]。
    """
    X = normalize(vecs)
    i = voc.stoi[word]
    sims = X @ X[i]
    pool = range(len(voc)) if candidates is None else [voc.stoi[w] for w in candidates if w in voc.stoi]
    ranked = sorted((j for j in pool if j != i), key=lambda j: -sims[j])
    return [(voc.itos[j], float(sims[j])) for j in ranked[:k]]


def category_nn_accuracy(vecs, voc, groups):
    """
    "最近邻是否同类"的比例：对 groups 里的每个词，在所有分组词中找最近的一个，看是否同组。

    【参数】groups：{组名: 词列表}。
    【返回】(准确率, 随机猜的期望准确率)。
    """
    X = normalize(vecs)
    members = [(w, g) for g, ws in groups.items() for w in ws if w in voc.stoi]
    idx = np.array([voc.stoi[w] for w, _ in members])
    labels = [g for _, g in members]
    S = X[idx] @ X[idx].T
    np.fill_diagonal(S, -np.inf)
    hits = sum(labels[int(np.argmax(S[r]))] == labels[r] for r in range(len(members)))
    sizes = Counter(labels)
    n = len(members)
    chance = sum(s * (s - 1) for s in sizes.values()) / (n * (n - 1))
    return hits / n, chance