"""
=============================================================================
第 1 讲 · 代码 3：基于计数的词向量 —— 共现矩阵、PPMI 加权、SVD 降维
=============================================================================

对应讲义：lectures/lecture01_nlp_history_word_vectors.md 第 8 节
对应作业：CS224N 作业 1 第一部分（Count-Based Word Vectors）

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
代码 02 的 tezguino 例子说明：一个词的意思可以由"它周围出现过哪些词"来刻画。
把这个想法推广到整个词表，就得到了最朴素的词向量：

    共现矩阵 M（co-occurrence matrix）：
        M[i, j] = 词 j 出现在词 i 的窗口内的次数。
        第 i 行就是词 i 的"上下文画像"，可以直接当作词 i 的向量。

但直接用它有三个问题，本文件逐一演示并解决：
    问题 1：太大、太稀疏。行和列的长度都是词表大小 V，真实词表 V 动辄几十万，
           矩阵有 V² 个格子，绝大部分是 0。
    问题 2：高频虚词（the、a、is）出现在几乎所有词旁边，主导了相似度计算。
           → 解决：PPMI 加权（只保留"比随机情况更常一起出现"的信息）。
    问题 3：维度太高，而且同一个"意思"被分散在很多列里（比如 drinks 和 milk 都暗示"年幼"）。
           → 解决：SVD 降维，把 V 维压缩到 k 维（k 远小于 V），合并相关的维度。

-----------------------------------------------------------------------------
【整体流程】
-----------------------------------------------------------------------------
    语料 → compute_co_occurrence_matrix → 共现矩阵 M（V × V）
         → ppmi（可选）                  → 加权后的矩阵（V × V）
         → reduce_to_k_dim               → 词向量矩阵（V × k）
         → 最近邻 / 类别纯度 / 类比 / 二维可视化（vector_tools.py）

函数名 distinct_words、compute_co_occurrence_matrix、reduce_to_k_dim、plot_embeddings
与 CS224N 作业 1 保持一致，方便你对照作业。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 03_cooccurrence_svd.py

需要 numpy；装了 matplotlib 会额外把二维图保存到 figures/ 文件夹。运行时间约 1 秒。
"""

from pathlib import Path

import numpy as np

# 【Python 小知识】corpus.py 和 vector_tools.py 与本文件在同一个文件夹里，可以直接 import。
from corpus import analogy_questions, build_vocab, classic_tiny_corpus, generate_corpus, word_categories
from vector_tools import (analogy, ascii_scatter, category_purity, evaluate_analogies, nearest_neighbors,
                          normalize_rows, pca_2d, save_scatter_png, similarity_gap)

FIGURE_DIR = Path(__file__).parent / "figures"


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


# =============================================================================
# 第 1 部分：构建共现矩阵
# =============================================================================
def distinct_words(sentences):
    """
    找出语料中所有不同的词，按**第一次出现的顺序**排列。

    【参数】
        sentences : list[list[str]]

    【返回】
        (words, n)：words 是 list[str]，n = len(words)

    【边界】
        - CS224N 作业 1 要求按字母排序；这里按第一次出现排序，是为了和课件里
          "I like deep learning" 例子的矩阵行列顺序一致。顺序不影响任何相似度结果。

    【Python 小知识】
        dict.fromkeys(序列) 生成一个"键按插入顺序排列、自动去重"的字典，
        list(...) 再取出它的键，就得到了"保持首次出现顺序的去重列表"。
    """
    words = list(dict.fromkeys(w for sent in sentences for w in sent))
    return words, len(words)


def compute_co_occurrence_matrix(sentences, window_size=2, word2id=None):
    """
    计算共现矩阵：M[i, j] = 词 j 出现在词 i 左右 window_size 个位置以内的总次数。

    【思路】
        对每个句子的每个位置 t（中心词 w_t），看它左右各 window_size 个位置的词 w_j，
        把 M[id(w_t), id(w_j)] 加 1。
        这样得到的 M 是**对称**的：i 在 j 的窗口里，j 也一定在 i 的窗口里。

    【参数】
        sentences   : list[list[str]]
        window_size : int >= 1，窗口半径
        word2id     : dict[str, int] 或 None。None 时用 distinct_words 自动建立。

    【返回】
        (M, word2id)
            M       : np.ndarray，形状 (V, V)，元素是非负整数（用 float 存储，方便后续计算）
            word2id : dict[str, int]

    【边界】
        - 窗口不跨句子。
        - 不在 word2id 中的词直接跳过（既不当中心词，也不当上下文词）。
        - 对角线 M[i, i] 只有在"同一个词在窗口内出现两次"时才非零（比如 "the cat and the dog"）。

    【Python 小知识】
        np.zeros((V, V))：注意参数是一个元组 (V, V)，表示二维形状。
    """
    if word2id is None:
        words, _ = distinct_words(sentences)
        word2id = {w: i for i, w in enumerate(words)}
    V = len(word2id)
    M = np.zeros((V, V))
    for sent in sentences:
        ids = [word2id.get(w) for w in sent]
        for t, center in enumerate(ids):
            if center is None:
                continue
            for j in range(max(0, t - window_size), min(len(ids), t + window_size + 1)):
                if j != t and ids[j] is not None:
                    M[center, ids[j]] += 1
    return M, word2id


# =============================================================================
# 第 2 部分：加权 —— 让高频虚词不再主导
# =============================================================================
def log_counts(M):
    """
    对数压缩：把每个计数 x 换成 log(1 + x)。

    【思路】
        "the" 和 cat 共现 300 次、"sleeping" 和 cat 共现 30 次，原始比值 10 倍；
        取对数后 log(301) ≈ 5.7、log(31) ≈ 3.4，差距缩小到不到 2 倍。
        这是最简单的"压低大数"的办法（GloVe 也用了对数，见第 2 讲）。

    【参数】M : np.ndarray，非负
    【返回】np.ndarray，同形状
    【Python 小知识】np.log1p(x) = log(1 + x)，在 x 很小时比 np.log(1 + x) 更精确。
    """
    return np.log1p(M)


def ppmi(M, eps=1e-12):
    """
    正点互信息（Positive Pointwise Mutual Information）加权。

    【思路】（讲义第 8.4 节）
        点互信息 PMI(w, c) = log [ P(w, c) / (P(w) · P(c)) ]
            分子 P(w, c)    ：w 和 c 实际一起出现的概率；
            分母 P(w) P(c)  ：如果 w 和 c 毫无关系（相互独立），它们"碰巧"一起出现的概率。
        比值 > 1（PMI > 0）：比碰巧更常一起出现 → 有关联；
        比值 ≈ 1（PMI ≈ 0）：和碰巧差不多 → 没信息（"the" 和几乎所有词都是这样）；
        比值 < 1（PMI < 0）：比碰巧还少 → 负值估计很不可靠（语料不够大），通常直接截成 0。
        所以 PPMI = max(PMI, 0)。

        概率都从共现矩阵里估计：
            P(w, c) = M[w, c] / 总和；  P(w) = 第 w 行的和 / 总和；  P(c) = 第 c 列的和 / 总和。

    【参数】
        M   : np.ndarray (V, V)，共现计数矩阵
        eps : 防止 log(0) 和除以 0 的小数

    【返回】
        np.ndarray (V, V)，所有元素 >= 0。M 中为 0 的格子，PPMI 也是 0。

    【边界】
        - PMI 对**罕见词**有偏：一个只出现过 1 次的词，和它唯一的邻居 PMI 会非常高。
          实践中常用平滑（例如把 P(c) 的计数开 0.75 次方）缓解，这里不做。

    【NLP 联系】
        第 2 讲会看到，带负采样的 word2vec 在某种意义上等价于分解一个"PMI - log k"的矩阵；
        "计数开 0.75 次方"也正是 word2vec 负采样分布的做法。

    【Python 小知识】
        M.sum(axis=1, keepdims=True)：每行求和，形状 (V, 1)；axis=0 则是每列求和，形状 (1, V)。
        (V,1) 的数组乘 (1,V) 的数组，numpy 会"广播"成 (V,V)：第 i 行第 j 列 = 行和_i × 列和_j。
        np.where(条件, A, B)：条件为真取 A，否则取 B（逐元素）。
    """
    total = M.sum()
    p_wc = M / total
    p_w = M.sum(axis=1, keepdims=True) / total
    p_c = M.sum(axis=0, keepdims=True) / total
    with np.errstate(divide="ignore"):                 # 暂时关掉 log(0) 的警告
        pmi = np.log(np.maximum(p_wc, eps) / np.maximum(p_w * p_c, eps))
    return np.where(M > 0, np.maximum(pmi, 0.0), 0.0)


# =============================================================================
# 第 3 部分：SVD 降维
# =============================================================================
def reduce_to_k_dim(M, k=2):
    """
    用奇异值分解（SVD）把 V 维的行向量压缩到 k 维。

    【思路】（讲义第 8.5 节）
        任何矩阵都可以分解成 M = U · diag(S) · Vᵀ：
            U  的列：一组正交的"新坐标轴"（在行空间里），
            S       ：每个坐标轴的"重要程度"（奇异值），从大到小排列，
            Vᵀ      ：另一组正交坐标轴（在列空间里）。
        只保留最大的 k 个奇异值及其坐标轴，U[:, :k] · diag(S[:k]) 就是每个词的 k 维坐标。
        这是在"平方误差最小"意义下对 M 最好的秩 k 近似（Eckart–Young 定理）。
        直觉：大的奇异值对应语料中"最主要的规律"（例如"这是动物还是水果"），
              小的奇异值往往是噪声或偶然的细节，丢掉它们反而让向量更"干净"。

    【参数】
        M : np.ndarray，形状 (V, n)
        k : int，1 <= k <= min(V, n)

    【返回】
        np.ndarray，形状 (V, k)

    【边界】
        - k 超过矩阵的秩时，多出来的维度全是 0（不会报错）。
        - numpy 的完整 SVD 复杂度约 O(V³)，V = 几十万时完全算不动；
          作业 1 用 sklearn 的 TruncatedSVD，只算前 k 个奇异值，快得多。
          我们的 V 只有一百多，用 numpy 完整分解即可。
        - 奇异向量的正负号是任意的（整体乘 -1 也是合法分解），不影响余弦相似度。

    【NLP 联系】
        "用低秩矩阵近似一个大矩阵"的思想在后面还会出现：
        第 11 讲的 LoRA 就假设微调时权重的变化量是低秩的，只训练两个很"瘦"的矩阵。

    【Python 小知识】
        np.linalg.svd(M, full_matrices=False) 返回 (U, S, Vt)，S 是一维数组（已从大到小排序）。
        U[:, :k] 取前 k 列；U[:, :k] * S[:k] 利用广播把第 j 列乘以 S[j]。
    """
    if not 1 <= k <= min(M.shape):
        raise ValueError(f"k 必须在 1 到 {min(M.shape)} 之间")
    U, S, Vt = np.linalg.svd(M, full_matrices=False)
    return U[:, :k] * S[:k]


def svd_word_vectors(sentences=None, window_size=2, k=10, weighting="ppmi"):
    """
    一站式函数：语料 → 共现矩阵 → 加权 → SVD → 词向量。供 05 号文件对比使用。

    【参数】
        sentences   : list[list[str]] 或 None（None 时使用默认的玩具语料）
        window_size : 窗口半径
        k           : 降到多少维
        weighting   : "raw"（原始计数）、"log"（对数）或 "ppmi"

    【返回】
        (vectors, word2id, id2word)，vectors 形状 (V, k)

    【边界】
        - weighting 取其他值会抛出 ValueError。
    """
    if sentences is None:
        sentences = generate_corpus()
    word2id, id2word, _ = build_vocab(sentences)
    M, _ = compute_co_occurrence_matrix(sentences, window_size, word2id)
    weights = {"raw": lambda x: x, "log": log_counts, "ppmi": ppmi}
    if weighting not in weights:
        raise ValueError(f"weighting 必须是 {list(weights)} 之一")
    return reduce_to_k_dim(weights[weighting](M), k), word2id, id2word


def plot_embeddings(vectors, word2id, words, title, filename):
    """
    把指定词的向量投影到二维并画出来：终端打印字符画；装了 matplotlib 再存一张 PNG。

    【参数】
        vectors  : np.ndarray (V, d)，d >= 2
        word2id  : dict[str, int]
        words    : list[str]，要画的词
        title    : str，图标题
        filename : str，PNG 文件名（保存在 figures/ 下）

    【边界】
        - 先把每个向量归一化成长度 1（我们关心方向，不关心长度；作业 1 画图前也做了归一化），
          d > 2 时再用 PCA 投影到二维（会丢失信息，图只用来建立直觉）。
    """
    pts = pca_2d(normalize_rows(vectors[[word2id[w] for w in words]]))
    print(ascii_scatter(pts, words))
    FIGURE_DIR.mkdir(exist_ok=True)
    if save_scatter_png(pts, words, FIGURE_DIR / filename, title):
        print(f"  （已保存图片：figures/{filename}）")


# =============================================================================
# 演示
# =============================================================================
def demo_classic_example():
    """演示 1：课件里的经典三句话例子，手算共现矩阵。"""
    print_header("演示 1：经典小例子 —— I like deep learning / I like NLP / I enjoy flying")
    sents = classic_tiny_corpus()
    M, word2id = compute_co_occurrence_matrix(sents, window_size=1)
    words = list(word2id)
    print("\n  窗口大小 1 的共现矩阵：")
    print(" " * 12 + "".join(f"{w:>9s}" for w in words))
    for w in words:
        print(f"    {w:<8s}" + "".join(f"{int(x):>9d}" for x in M[word2id[w]]))
    print(f"\n  矩阵是对称的吗？ {np.allclose(M, M.T)}")

    U, S, Vt = np.linalg.svd(M)
    print(f"  奇异值（从大到小）：{np.round(S, 2)}")
    pts = reduce_to_k_dim(M, 2)
    print("  降到 2 维后的坐标：")
    for w in words:
        x, y = pts[word2id[w]]
        print(f"    {w:<9s}({x:+.2f}, {y:+.2f})")
    print("""
  观察：
    - like 和 enjoy 的行很像（都和 I 共现），降维后坐标也接近；
      deep、NLP、flying 都出现在"动词之后"，也被放到了相近的位置。
    - 3 句话的语料太小，别对坐标过度解读。下面换成更大的玩具语料。""")


def demo_toy_corpus():
    """演示 2~5：在玩具语料上，比较原始计数 / 对数 / PPMI，以及 SVD 降维。"""
    sents = generate_corpus()
    word2id, id2word, counts = build_vocab(sents)
    cats = word_categories()
    content = list(cats)                                # 有类别的"内容词"，用来找邻居
    M, _ = compute_co_occurrence_matrix(sents, 2, word2id)
    V = len(word2id)

    # ---------------- 演示 2：矩阵的规模与稀疏性 ----------------
    print_header("演示 2：玩具语料的共现矩阵")
    print(f"\n  语料：{len(sents)} 句，{sum(map(len, sents))} 个词；词表大小 V = {V}")
    print(f"  共现矩阵形状：{M.shape}，非零元素占比：{np.mean(M > 0):.1%}")
    for real_V in [100_000, 400_000]:
        print(f"  如果 V = {real_V:,}：矩阵有 {real_V**2:.1e} 个格子，按 float32 存需要 {real_V**2 * 4 / 1e9:,.0f} GB")
    row = M[word2id["cat"]]
    top = np.argsort(-row)[:8]
    print("\n  'cat' 这一行里次数最多的 8 个上下文词：")
    print("    " + ", ".join(f"{id2word[j]}:{int(row[j])}" for j in top))
    print("  → 排在最前面的是 the、a 这类虚词：它们几乎出现在每个名词旁边，说明不了 cat 是什么。")

    # ---------------- 演示 3：三种加权的对比（不降维） ----------------
    print_header("演示 3：原始计数 vs 对数 vs PPMI（不降维，直接用 V 维行向量）")
    variants = {"原始计数": M, "log(1+计数)": log_counts(M), "PPMI": ppmi(M)}
    print(f"\n  共 {len(content)} 个内容词（7 类）。两个指标：")
    print("    类别纯度：每个内容词的 5 个最近邻中同类词的比例（只看排名）")
    print("    区分度  ：同类词平均余弦 - 异类词平均余弦（看绝对值，越大越好）\n")
    for name, mat in variants.items():
        purity, chance = category_purity(mat, word2id, cats, k=5)
        within, between, gap = similarity_gap(mat, word2id, cats)
        print(f"    {name:<12s}纯度 {purity:.3f}（随机 ≈ {chance:.3f}）  "
              f"同类余弦 {within:.3f}  异类余弦 {between:.3f}  区分度 {gap:.3f}")
    print("""
  观察：三种做法的纯度都是满分 —— 我们的玩具语料里，类别本身很容易分开。
    真正的区别在"异类余弦"：原始计数下，cat 和 car 这种毫不相干的词平均也有 0.4 的相似度，
    因为它们共享了 the、a、is 这些高频虚词；PPMI 把这部分"碰巧的共现"去掉，异类相似度降到 0.1 左右。""")
    print("\n  最近邻（只在内容词中找）：")
    for name, mat in variants.items():
        print(f"    [{name}]")
        for q in ["cat", "apple", "king"]:
            nbrs = nearest_neighbors(q, mat, word2id, id2word, k=4, restrict_to=content)
            print(f"      {q:<6s}→ " + ", ".join(f"{w}({s:.2f})" for w, s in nbrs))

    # ---------------- 演示 4：SVD 降维 ----------------
    print_header("演示 4：对 PPMI 矩阵做 SVD 降维")
    P = ppmi(M)
    S = np.linalg.svd(P, compute_uv=False)               # 只要奇异值
    energy = np.cumsum(S**2) / np.sum(S**2)
    print("\n  前 k 个奇异值保留的'能量'（奇异值平方和的占比）：")
    for k in [2, 5, 10, 20, 50]:
        print(f"    k = {k:>2d}：{energy[k - 1]:.1%}")
    print("\n  不同 k 下的类别纯度、区分度与类比准确率（96 道类比题）：")
    questions = analogy_questions()
    for k in [2, 5, 10, 20, 50, V]:
        vec = reduce_to_k_dim(P, k)
        purity, _ = category_purity(vec, word2id, cats, k=5)
        _, _, gap = similarity_gap(vec, word2id, cats)
        acc, _ = evaluate_analogies(questions, vec, word2id, id2word, restrict_to=content)
        tag = "（= 不降维）" if k == V else ""
        print(f"    k = {k:>3d}：纯度 {purity:.3f}，区分度 {gap:.3f}，类比准确率 {acc:.1%} {tag}")
    print("""
  观察：
    - k 太小（2）时，7 个类别挤在 2 维里，纯度明显下降；k = 5 起类别就分开了。
    - 适度降维（k = 10、20）的区分度**高于**不降维：丢掉小奇异值 ≈ 丢掉噪声，同类词被"压"得更紧。
    - 类比需要的是**类别内部**的细微差别（男/女、王室/平民），它们藏在较小的奇异值里，
      k 小时首先被丢掉，所以类比准确率整体上随 k 增大而提高（中间有波动）。
    - 所以"最好的 k"取决于你要做什么：分大类只要很少几维，捕捉细微关系需要更多维。""")

    # ---------------- 演示 5：看看 k=10 的词向量 ----------------
    print_header("演示 5：PPMI + SVD（k = 10）得到的词向量")
    vec = reduce_to_k_dim(P, 10)
    print("\n  最近邻：")
    for q in ["cat", "apple", "car", "king", "girl", "paris", "french"]:
        nbrs = nearest_neighbors(q, vec, word2id, id2word, k=4, restrict_to=content)
        print(f"    {q:<7s}→ " + ", ".join(f"{w}({s:.2f})" for w, s in nbrs))
    acc, per = evaluate_analogies(questions, vec, word2id, id2word, restrict_to=content)
    print(f"\n  类比总准确率：{acc:.1%}，按关系：" +
          "，".join(f"{name} {r}/{t}" for name, (r, t) in per.items()))
    for a, b, c in [("man", "woman", "king"), ("france", "paris", "japan"), ("boy", "man", "princess")]:
        top = analogy(a, b, c, vec, word2id, id2word, k=3, restrict_to=content)
        print(f"    {a} : {b} :: {c} : ?   → " + ", ".join(f"{w}({s:.2f})" for w, s in top))

    print("\n  二维投影（PCA），只画一部分内容词：")
    show = ["cat", "dog", "horse", "apple", "banana", "pear", "car", "bus", "train",
            "king", "queen", "man", "woman", "boy", "girl", "paris", "france", "tokyo", "japan"]
    plot_embeddings(vec, word2id, show, "PPMI + SVD (k=10), PCA to 2D", "lecture01_svd_2d.png")


def self_check():
    """自我检查：经典例子的矩阵必须和课件一致；PPMI 非负；SVD 维度正确。"""
    M, w2i = compute_co_occurrence_matrix(classic_tiny_corpus(), window_size=1)
    assert M[w2i["I"], w2i["like"]] == 2 and M[w2i["I"], w2i["enjoy"]] == 1
    assert M[w2i["like"], w2i["deep"]] == 1 and M[w2i["deep"], w2i["learning"]] == 1
    assert np.allclose(M, M.T) and M.sum() == 20          # 4 + 3 + 3 = 10 对相邻关系 × 2 个方向
    P = ppmi(M)
    assert (P >= 0).all()
    assert reduce_to_k_dim(M, 3).shape == (len(w2i), 3)


if __name__ == "__main__":
    self_check()
    demo_classic_example()
    demo_toy_corpus()