"""
=============================================================================
第 1 讲 · 工具模块：一个"可控的"玩具语料库 + 分词与词表工具
=============================================================================

对应讲义：lectures/lecture01_nlp_history_word_vectors.md 第 8.3 节（语料设计）、第 12 节（结果）
被 03_cooccurrence_svd.py 和 05_word2vec_training.py 使用。

-----------------------------------------------------------------------------
【为什么需要一个"玩具"语料库？】
-----------------------------------------------------------------------------
真正的 word2vec 是在几十亿个词的新闻、网页上训练的（Google News 约 1000 亿词）。
在笔记本上，用 numpy 手写的代码几秒钟内处理不了那么多文本。
更重要的是：在真实语料上，即使结果很好，你也很难说清"它为什么学会了这个"。

所以我们自己**按规则生成**一份小语料。好处是：
    (1) 小：几千个句子、一百来个不同的词，几秒钟就能训练完；
    (2) 可控：我们**事先知道**哪些词应该相似（同一类别的词出现在相似的上下文里），
        于是可以定量检查"学到的词向量是不是真的把相似的词放到了一起"；
    (3) 可解释：每个词的上下文完全由我们写的句子模板决定，
        你可以直接看模板，理解"分布假说"是怎么起作用的。

-----------------------------------------------------------------------------
【语料里有什么？】
-----------------------------------------------------------------------------
    1) 三类普通名词，每类 6 个词：
         动物 animal ：cat dog horse cow sheep pig
         水果 fruit  ：apple banana orange grape pear peach
         交通 vehicle：car bus train truck bike taxi
       同一类的词出现在同一组句子模板里（例如 "i ate the ___"、"the ___ is sweet"）。
       为了不让类别分得"太干净"，还有所有名词共用的通用模板（"i saw the ___"）
       和两类共用的模板（动物和交通工具都可以 "is fast"）。
       同一类里的词频率不同（每类第一个词大约是最后一个词的 3 倍），模仿真实语言。

    2) 8 个"人物"词，每个词由三个"特征"组合而成（2 × 2 × 2 = 8）：
         性别：male / female        身份：royal（王室）/ common（平民）
         年龄：adult（成年）/ young（年幼）
             king     = male   + royal  + adult      queen    = female + royal  + adult
             prince   = male   + royal  + young      princess = female + royal  + young
             man      = male   + common + adult      woman    = female + common + adult
             boy      = male   + common + young      girl     = female + common + young
       每生成一个人物句子，先随机挑一个人物，再随机挑一个特征，
       用这个特征对应的模板造句（比如性别是 female，就可能造出 "the queen said she was tired"）。
       （80% 的人物句这样生成，另外 20% 用通用模板，比如 "i saw the queen"。）
       于是"queen 的上下文" = "女性的上下文" + "王室的上下文" + "成年人的上下文" 的混合。
       这正是 king - man + woman ≈ queen 这类"类比"能成立的原因（讲义第 12 节）。

    3) 6 组"国家 - 首都 - 语言"：
         france-paris-french   japan-tokyo-japanese   italy-rome-italian
         china-beijing-chinese germany-berlin-german  spain-madrid-spanish
       国家和它的首都共享"people in ___ speak <语言>"这样的上下文，
       国家另有"国家的上下文"，城市另有"城市的上下文"。

    4) 约 8% 的"填充句"（"it was a nice day"），不含任何内容词，模拟与主题无关的文本。

-----------------------------------------------------------------------------
【边界：玩具语料能说明什么，不能说明什么】
-----------------------------------------------------------------------------
    - 能说明：分布假说的机制（相似上下文 → 相似向量）、训练算法是否写对、
              各种方法（计数 + SVD vs word2vec）在"理想条件"下能否学到结构。
    - 不能说明：哪种方法在真实语料上更好；真实语言比模板丰富得多、噪声也大得多。
      讲义里所有基于这份语料的结论都只针对这份语料。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    from corpus import generate_corpus, build_vocab, CATEGORIES
    sentences = generate_corpus()                 # list[list[str]]，每个句子是一个词列表
    word2id, id2word, counts = build_vocab(sentences)

直接运行本文件会打印语料的统计信息和几个样例句子：
    python corpus.py
"""

# 【Python 小知识】
#   random      ：Python 标准库的随机数模块。random.Random(seed) 创建一个"独立的随机数生成器"，
#                 只要 seed（种子）相同，每次运行生成的随机序列就完全相同 —— 这叫"可复现"。
#                 讲义里引用的数字都依赖这一点：你运行出来的结果应该和讲义一模一样。
#   collections.Counter：一个专门用来"数数"的字典。Counter(["a","b","a"]) → {"a": 2, "b": 1}。
import random
from collections import Counter

# =============================================================================
# 第 0 部分：语料的"设计图"——类别、人物特征、句子模板
# =============================================================================

# 三类普通名词。字典：类别名 → 该类的词列表。
CATEGORIES = {
    "animal": ["cat", "dog", "horse", "cow", "sheep", "pig"],
    "fruit": ["apple", "banana", "orange", "grape", "pear", "peach"],
    "vehicle": ["car", "bus", "train", "truck", "bike", "taxi"],
}

# 每个类别的专属模板。"{w}" 是占位符，造句时会被替换成该类别的某个词。
# 注意模板里离 {w} 很近（前后 1~2 个位置）的词最重要：
# 我们的窗口大小默认是 2，只有这些词会被算作 {w} 的"上下文"。
CATEGORY_TEMPLATES = {
    "animal": [
        "the {w} is sleeping",
        "i fed the {w}",
        "the hungry {w} ran away",
        "the {w} made a loud noise",
        "a baby {w} was born",
        "the farmer keeps a {w}",
    ],
    "fruit": [
        "i ate the {w}",
        "the {w} is sweet",
        "a ripe {w} fell down",
        "she cut the {w} in half",
        "fresh {w} juice is good",
        "the {w} tastes sour",
    ],
    "vehicle": [
        "i drove the {w}",
        "the {w} is fast",
        "the old {w} broke down",
        "we took the {w} to town",
        "the {w} stopped at the station",
        "a red {w} passed by",
    ],
}

# 同一类别里的词，出现频率并不相同（真实语言里 dog 比 sheep 常见得多）。
# 权重越大越常见：每类第一个词的频率大约是最后一个词的 3 倍。
CATEGORY_WEIGHTS = [1.0, 0.8, 0.6, 0.5, 0.4, 0.35]

# "通用"模板：任何名词都能放进去（动物、水果、交通工具、人物都可以）。
# 它们让不同类别的词也共享一部分上下文 —— 真实语言正是这样，语料不会把类别分得干干净净。
GENERIC_TEMPLATES = [
    "i saw the {w}",
    "look at that {w}",
    "there is a {w} here",
    "where is the {w}",
    "the {w} is big",
    "the {w} is small",
    "i like the {w}",
]

# 两个类别共享的模板：动物和交通工具都可以"快"，水果和交通工具都可以"贵"。
PAIR_TEMPLATES = [
    (("animal", "vehicle"), ["the {w} is fast", "the {w} went past us"]),
    (("fruit", "vehicle"), ["the {w} costs a lot", "we need a new {w}"]),
    (("animal", "fruit"), ["the {w} smells bad", "a fat {w}"]),
]

# 不含任何"内容词"的填充句：只由虚词和常见词组成，模拟真实文本里大量与主题无关的句子。
FILLER_SENTENCES = [
    "it was a nice day",
    "we went home early",
    "that is a good idea",
    "she said it is late",
    "he is not here",
    "they are at the door",
    "this is the end",
    "i do not know",
]

# 8 个人物词，每个由 (性别, 身份, 年龄) 三个特征组成。
PERSONS = {
    "king": ("male", "royal", "adult"),
    "queen": ("female", "royal", "adult"),
    "prince": ("male", "royal", "young"),
    "princess": ("female", "royal", "young"),
    "man": ("male", "common", "adult"),
    "woman": ("female", "common", "adult"),
    "boy": ("male", "common", "young"),
    "girl": ("female", "common", "young"),
}

# 每个特征取值对应的模板。{w} 处放人物词。
# 例如性别为 female 时可能生成 "the queen said she was tired"。
FEATURE_TEMPLATES = {
    # 性别：he / his vs she / her
    "male": ["the {w} said he was tired", "the {w} lost his hat", "the {w} and his friend"],
    "female": ["the {w} said she was tired", "the {w} lost her hat", "the {w} and her friend"],
    # 身份：王室 vs 平民
    "royal": ["the royal {w} smiled", "the {w} rules the land", "the {w} wears a crown"],
    "common": ["the poor {w} smiled", "the {w} works on a farm", "the {w} sells bread"],
    # 年龄：成年 vs 年幼
    "adult": ["the {w} drinks coffee", "the {w} pays the bills", "the busy {w} laughed"],
    "young": ["the {w} drinks milk", "the {w} plays with toys", "the little {w} laughed"],
}

# 人物与物品"互动"的模板：让句子更像真实文本（同一句里同时出现人物和物品）。
# {p} 放人物，{x} 放某类物品。
INTERACTION_TEMPLATES = {
    "animal": ["the {p} fed the {x}", "the {p} saw a {x}"],
    "fruit": ["the {p} ate the {x}", "the {p} bought a {x}"],
    "vehicle": ["the {p} drove the {x}", "the {p} took the {x}"],
}

# 国家 - 首都 - 语言。
GEOGRAPHY = [
    ("france", "paris", "french"),
    ("japan", "tokyo", "japanese"),
    ("italy", "rome", "italian"),
    ("china", "beijing", "chinese"),
    ("germany", "berlin", "german"),
    ("spain", "madrid", "spanish"),
]

COUNTRY_TEMPLATES = ["{c} is a large country", "the government of {c}", "{c} has a long history"]
CITY_TEMPLATES = ["{t} is a busy city", "i flew to {t}", "the streets of {t}"]
# 国家和首都共享的模板：{x} 放国家或首都，{l} 放语言。
SHARED_GEO_TEMPLATES = ["people in {x} speak {l}", "{x} loves {l} food", "{x} is full of {l} people"]


# =============================================================================
# 第 1 部分：生成语料
# =============================================================================
def generate_corpus(num_sentences=6000, seed=0):
    """
    按上面的"设计图"随机生成一份语料。

    【思路】
        每次随机决定这个句子属于哪一种（类别句 / 人物句 / 互动句 / 地理句），
        再随机挑模板和词，替换占位符，最后按空格切成词列表。
        各种句子的比例大约是：类别句 35%、人物句 30%、互动句 10%、地理句 17%、填充句 8%。

    【参数】
        num_sentences : int，要生成多少个句子（> 0）。默认 6000，约 2.8 万个词。
        seed          : int，随机种子。种子相同 → 语料完全相同（讲义里的数字依赖默认值 0）。

    【返回】
        list[list[str]]：句子列表，每个句子是一个"词的列表"，例如
            [["the", "cat", "is", "sleeping"], ["i", "flew", "to", "paris"], ...]

    【边界】
        - num_sentences <= 0 会抛出 ValueError。
        - 所有词都是小写英文，没有标点（我们只关心词与词的共现）。
        - 句子之间是独立的：后面数"上下文"时，窗口不会跨越句子边界。
          （真实的 word2vec 通常把整篇文本当作一个长序列，窗口可以跨句子，影响不大。）

    【Python 小知识】
        rng.random()        ：返回 [0, 1) 之间的随机小数，用来"按概率"选分支。
        rng.choice(列表)    ：从列表里等概率随机挑一个元素。
        rng.choices(列表, weights=权重)[0]：按权重随机挑一个（choices 返回列表，[0] 取出唯一的元素）。
        "abc {w}".format(w="cat") 或 str.replace：把占位符替换成具体的词。这里用 format。
        s.split()           ：按空白字符把字符串切开，返回词列表。
    """
    if num_sentences <= 0:
        raise ValueError("num_sentences 必须是正整数")

    rng = random.Random(seed)
    sentences = []
    for _ in range(num_sentences):           # 【Python 小知识】变量名 _ 表示"这个值我用不到"
        r = rng.random()
        if r < 0.35:
            # 类别句：先挑类别，再按频率权重挑词，最后挑模板（专属 / 通用 / 两类共享）
            cat = rng.choice(list(CATEGORIES))
            word = rng.choices(CATEGORIES[cat], weights=CATEGORY_WEIGHTS)[0]
            kind = rng.random()
            if kind < 0.55:
                template = rng.choice(CATEGORY_TEMPLATES[cat])
            elif kind < 0.85:
                template = rng.choice(GENERIC_TEMPLATES)
            else:
                shared = [t for cats, ts in PAIR_TEMPLATES if cat in cats for t in ts]
                template = rng.choice(shared)
            text = template.format(w=word)
        elif r < 0.65:
            # 人物句：先挑人物；80% 用它某一个特征的模板，20% 用通用模板
            person = rng.choice(list(PERSONS))
            if rng.random() < 0.8:
                feature = rng.choice(PERSONS[person])      # 三个特征里随机挑一个
                text = rng.choice(FEATURE_TEMPLATES[feature]).format(w=person)
            else:
                text = rng.choice(GENERIC_TEMPLATES).format(w=person)
        elif r < 0.75:
            # 互动句：人物 + 某类物品
            cat = rng.choice(list(CATEGORIES))
            text = rng.choice(INTERACTION_TEMPLATES[cat]).format(
                p=rng.choice(list(PERSONS)), x=rng.choices(CATEGORIES[cat], weights=CATEGORY_WEIGHTS)[0])
        elif r < 0.92:
            # 地理句：挑一组 (国家, 首都, 语言)，再决定写国家、写城市，还是写共享模板
            country, city, lang = rng.choice(GEOGRAPHY)
            kind = rng.random()
            if kind < 1 / 3:
                text = rng.choice(COUNTRY_TEMPLATES).format(c=country)
            elif kind < 2 / 3:
                text = rng.choice(CITY_TEMPLATES).format(t=city)
            else:
                # 共享模板里，国家和首都各占一半
                x = country if rng.random() < 0.5 else city
                text = rng.choice(SHARED_GEO_TEMPLATES).format(x=x, l=lang)
        else:
            # 填充句：没有内容词
            text = rng.choice(FILLER_SENTENCES)
        sentences.append(text.split())
    return sentences


def classic_tiny_corpus():
    """
    CS224N 课件里的经典三句话小语料，用来手算共现矩阵（讲义第 8.2 节）。

    【返回】
        list[list[str]]：
            [["I", "like", "deep", "learning", "."],
             ["I", "like", "NLP", "."],
             ["I", "enjoy", "flying", "."]]

    【边界】
        这里保留了大小写和句号 "."，和课件保持一致（句号也被当作一个"词"）。
    """
    return [s.split() for s in ["I like deep learning .", "I like NLP .", "I enjoy flying ."]]


# =============================================================================
# 第 2 部分：词表
# =============================================================================
def build_vocab(sentences, min_count=1):
    """
    统计词频，给每个词分配一个整数编号（id）。

    【思路】
        计算机不能直接处理字符串 "cat"，我们需要把每个词映射到一个整数，
        之后再用这个整数去"查"它的向量（矩阵的第几行）。
        编号规则：**按词频从高到低**排序，词频相同按字母顺序 —— 这样编号是确定的。
        （word2vec 的原始实现也是按词频排序的；高频词排在前面，便于观察。）

    【参数】
        sentences : list[list[str]]，句子列表。
        min_count : int，出现次数少于 min_count 的词丢弃（不进词表）。默认 1 = 全部保留。
                    真实的 word2vec 默认 min_count=5，用来丢掉错别字、极罕见词。

    【返回】
        (word2id, id2word, counts)
            word2id : dict[str, int]，词 → 编号，例如 {"the": 0, "a": 1, ...}
            id2word : list[str]，编号 → 词，id2word[word2id[w]] == w
            counts  : Counter，词 → 出现次数（包含被丢弃的词）

    【边界】
        - 空语料 → 返回空词表（word2id == {}）。
        - 被 min_count 丢掉的词，在后续处理中要么跳过，要么当作未知词——调用方自己决定。

    【Python 小知识】
        Counter.update(列表)：把列表里的元素逐个计数并累加。
        sorted(序列, key=函数)：按 key 排序。key=lambda w: (-counts[w], w)
            表示"先按次数从大到小（取负号），次数相同再按字母顺序"。
        {w: i for i, w in enumerate(列表)}：字典推导式；enumerate 同时给出下标 i 和元素 w。

    【NLP 联系】
        今天的大语言模型也有"词表"，只不过里面是子词 token（见 CS336 第 1 讲的 BPE），
        token id 的作用完全一样：去嵌入矩阵里"查"对应的那一行向量。
    """
    counts = Counter()
    for sent in sentences:
        counts.update(sent)
    words = sorted((w for w in counts if counts[w] >= min_count), key=lambda w: (-counts[w], w))
    word2id = {w: i for i, w in enumerate(words)}
    return word2id, words, counts


def sentences_to_ids(sentences, word2id):
    """
    把"词的句子"变成"编号的句子"。不在词表里的词直接跳过。

    【参数】
        sentences : list[list[str]]
        word2id   : dict[str, int]，来自 build_vocab

    【返回】
        list[list[int]]，例如 [[0, 5, 2, 17], ...]

    【边界】
        - 跳过未知词会让它两边的词在序列里"变相邻"。玩具语料里没有未知词，不受影响。
    """
    return [[word2id[w] for w in sent if w in word2id] for sent in sentences]


def skipgram_pairs(id_sentences, window=2):
    """
    从编号句子里抽出所有 (中心词, 上下文词) 对 —— skip-gram 的训练样本。

    【思路】
        对句子中的每个位置 t，中心词是 w_t，
        它左右各 window 个位置以内的词都是"上下文词"（不包括它自己）。
        例如句子 [the, cat, is, sleeping]、window=2、中心词 cat：
            上下文是 the（左 1）、is（右 1）、sleeping（右 2） → 3 个样本对。

    【参数】
        id_sentences : list[list[int]]
        window       : int >= 1，窗口半径 m。讲义公式里的 m。

    【返回】
        list[tuple[int, int]]：所有 (center_id, outside_id) 对。
        出现在不同位置的同一对会重复出现 —— 这正是"按频率加权"：常一起出现的词对，训练时被看到的次数也多。

    【边界】
        - 句首、句尾的词上下文不足 window 个，只取实际存在的那些（不补零、不跨句）。
        - 只有 1 个词的句子不产生任何样本对。

    【Python 小知识】
        max(0, t - window) 和 min(len(sent), t + window + 1) 保证下标不越界。
        range(a, b) 生成 a, a+1, ..., b-1（不含 b）。

    【NLP 联系】
        原始 word2vec 用的是"动态窗口"：每个位置的窗口半径在 1 ~ window 之间随机取，
        效果相当于离中心词越近的上下文权重越大。第 2 讲会讨论这类技巧。
    """
    if window < 1:
        raise ValueError("window 至少为 1")
    pairs = []
    for sent in id_sentences:
        for t, center in enumerate(sent):
            for j in range(max(0, t - window), min(len(sent), t + window + 1)):
                if j != t:
                    pairs.append((center, sent[j]))
    return pairs


# =============================================================================
# 第 3 部分：评估用的"标准答案"
# =============================================================================
def word_categories():
    """
    返回"词 → 类别"的字典，用来检查学到的向量是不是把同类词放在了一起。

    【返回】
        dict[str, str]，例如 {"cat": "animal", ..., "king": "person", ..., "paris": "city", ...}
        类别一共 7 种：animal、fruit、vehicle、person、country、city、language。

    【边界】
        - 模板里的功能词（the、is、a……）不属于任何类别，不在这个字典里。
    """
    cats = {w: c for c, words in CATEGORIES.items() for w in words}
    for p in PERSONS:
        cats[p] = "person"
    for country, city, lang in GEOGRAPHY:
        cats[country], cats[city], cats[lang] = "country", "city", "language"
    return cats


def analogy_questions():
    """
    生成类比题 a : b :: c : d（"a 之于 b，相当于 c 之于 ?"，答案是 d）。

    【思路】
        一种"关系"由若干词对组成，例如"男 → 女"关系：(king, queen)、(man, woman)……
        同一关系里任取两个不同的词对 (a, b)、(c, d)，就得到一道题：a : b :: c : ?，答案 d。
        n 个词对可以组成 n × (n - 1) 道题。

    【返回】
        list[tuple[str, str, str, str, str]]：(关系名, a, b, c, d)。
        共 5 种关系：
            gender（男→女）4 对 → 12 题；royal（平民→王室）4 对 → 12 题；
            age（年幼→成年）4 对 → 12 题；capital（国家→首都）6 对 → 30 题；
            language（国家→语言）6 对 → 30 题。合计 96 题。

    【边界】
        - 这些关系都是我们**设计进语料里**的，答得好说明方法能捕捉到"可加的结构"，
          不说明它在真实语料上的类比表现（真实语料要复杂得多，见讲义第 12、13 节）。
    """
    relations = {
        "gender": [("king", "queen"), ("prince", "princess"), ("man", "woman"), ("boy", "girl")],
        "royal": [("man", "king"), ("woman", "queen"), ("boy", "prince"), ("girl", "princess")],
        "age": [("prince", "king"), ("princess", "queen"), ("boy", "man"), ("girl", "woman")],
        "capital": [(c, t) for c, t, _ in GEOGRAPHY],
        "language": [(c, l) for c, _, l in GEOGRAPHY],
    }
    questions = []
    for name, pairs in relations.items():
        for a, b in pairs:
            for c, d in pairs:
                if (a, b) != (c, d):
                    questions.append((name, a, b, c, d))
    return questions


# =============================================================================
# 直接运行本文件：看看语料长什么样
# =============================================================================
if __name__ == "__main__":
    # 【Python 小知识】
    #   if __name__ == "__main__": 下面的代码只在"直接运行本文件"时执行；
    #   被别的文件 import 时不会执行。这样一个文件既能当模块用，也能单独运行做演示。
    sents = generate_corpus()
    word2id, id2word, counts = build_vocab(sents)
    n_tokens = sum(len(s) for s in sents)
    print(f"句子数：{len(sents)}，总词数（token）：{n_tokens}，不同的词（词表大小）：{len(word2id)}")
    print("\n前 8 个句子：")
    for s in sents[:8]:
        print("   ", " ".join(s))
    print("\n最常见的 12 个词：")
    for w, c in counts.most_common(12):
        print(f"    {w:<10s}{c:>6d}")
    print("\n每类词出现的次数（动物、水果、交通工具按权重有高有低；人物和地理词大致均匀）：")
    cats = word_categories()
    for cat in ["animal", "fruit", "vehicle", "person", "country", "city", "language"]:
        ws = [w for w in cats if cats[w] == cat]
        print(f"    {cat:<9s}" + "  ".join(f"{w}:{counts[w]}" for w in ws))
    print(f"\n窗口 2 下的 (中心词, 上下文词) 样本对数：{len(skipgram_pairs(sentences_to_ids(sents, word2id), 2))}")
    print(f"类比题数量：{len(analogy_questions())}")