"""
=============================================================================
CS230 · 第 16 讲共享模块：一个"知道真相"的合成小语料（word world）
=============================================================================

对应讲义：lectures/lecture16_word_embeddings_seq2seq_attention.md §1–§6

-----------------------------------------------------------------------------
【为什么要合成语料】
-----------------------------------------------------------------------------
真实的词向量（word2vec、GloVe）是在几十亿个词上训练的，下载要几百 MB，训练要几小时。
本讲义要求离线、几秒内跑完，所以我们用**模板语法**造一个小语料。它的好处是：**语义结构是我们亲手写进去的**，
所以可以检查词向量到底学到了什么：

    · 人物：性别（男 / 女）× 身份（平民 / 王室）× 年龄（大人 / 小孩），外加亲属称谓
            → 类比 man : woman :: king : queen、man : king :: woman : queen、man : boy :: woman : girl
    · 国家 × 首都 × 语言（每个国家还有一道"招牌菜"把三者绑在一起）
            → 类比 france : paris :: japan : tokyo、france : french :: japan : japanese
    · 形容词的原级 × 比较级（big / bigger ……，每个形容词有自己"典型"的动物）
            → 类比 big : bigger :: small : smaller
    · 水果（orange、apple、durian ……）：讲义 §1 "orange juice → apple juice" 的例子
    · 职业：**故意**让每个职业和性别有一个统计上的关联（OCCUPATION_P_MALE），用来做第 03 个脚本的去偏实验
    · 情感形容词：好评词出现在"loved / praised / smiled"附近，差评词出现在"hated / criticized / complained"附近，
            但我们**从不**告诉模型哪个词是好、哪个是坏。第 02 个脚本的情感分类要靠它们学出的向量。

**所有数字都只说明机制。** 这个语料只有几十万个词、三百来个不同的词，句子都是模板拼出来的，
和真实语言差得很远；真实词向量里的现象（比如类比的准确率）会弱很多、也乱很多。

-----------------------------------------------------------------------------
【这个模块提供什么】
-----------------------------------------------------------------------------
    generate_corpus(n_sentences, rng)   造 n 句话，每句是一个"词的列表"
    analogy_questions()                 根据写进语料的结构，列出全部类比题 (a, b, c, d)："a 之于 b 如同 c 之于 d"
    OCCUPATION_P_MALE                   每个职业在语料里和"男性"词一起出现的概率（去偏实验的"真相"）
    GENDER_PAIRS、DEFINITIONAL_PAIRS    性别词对（定义性的词，去偏时不能动它们）
    POS_ADJ_SEEN / POS_ADJ_UNSEEN ……   情感词表，以及在情感分类实验里"训练时见过 / 没见过"的划分

直接运行本文件：打印语料的统计量和几句样例。
    python3 word_world.py          # < 1 秒

【注释标签】【思路】【参数】【返回】【边界】【DL 联系】【Python 小知识】【numpy 小知识】
"""

import numpy as np

# =============================================================================
# 1. 人物：性别 × 身份 × 年龄
# =============================================================================
# (男性词, 女性词, 身份, 年龄, 这一对词专属的"身份标记"词)
#   身份 royal / common / family；年龄 adult / child / None（不写年龄语境）
#   最后一个词让每一对有自己的语境（king 和 queen 都"坐在 throne 上"），
#   否则 father 和 uncle 的语境会一模一样，类比就分不清它们。
PEOPLE = [
    ("man", "woman", "common", "adult", "street"),
    ("boy", "girl", "common", "child", "toy"),
    ("king", "queen", "royal", "adult", "throne"),
    ("prince", "princess", "royal", "child", "castle"),
    ("father", "mother", "family", "adult", "parent"),
    ("son", "daughter", "family", "child", "baby"),
    ("brother", "sister", "family", None, "sibling"),
    ("uncle", "aunt", "family", "adult", "cousin"),
    ("husband", "wife", "family", "adult", "wedding"),
]
GENDER_PAIRS = [(m, f) for m, f, *_ in PEOPLE]
# 去偏时用来"定义"性别方向、也不能被中和的词对（加上代词）
DEFINITIONAL_PAIRS = [("he", "she"), ("his", "her")] + GENDER_PAIRS

IDENTITY_VERBS = {           # 每一对的身份标记词用在哪个动词短语里
    "street": "walked down the", "toy": "played with the", "throne": "sat on the",
    "castle": "ran around the", "parent": "was a proud", "baby": "was once a tiny",
    "sibling": "argued like a", "cousin": "visited the", "wedding": "remembered the",
}

# =============================================================================
# 2. 国家 × 首都 × 语言 × 招牌菜
# =============================================================================
COUNTRIES = [
    ("france", "paris", "french", "baguette"),
    ("italy", "rome", "italian", "risotto"),
    ("spain", "madrid", "spanish", "paella"),
    ("germany", "berlin", "german", "pretzel"),
    ("japan", "tokyo", "japanese", "sushi"),
    ("china", "beijing", "chinese", "dumpling"),
    ("russia", "moscow", "russian", "borscht"),
    ("greece", "athens", "greek", "souvlaki"),
    ("sweden", "stockholm", "swedish", "meatball"),
    ("poland", "warsaw", "polish", "pierogi"),
]

# =============================================================================
# 3. 形容词：原级 × 比较级，每个形容词有两种"典型"的动物
# =============================================================================
ADJECTIVES = [
    ("big", "bigger", ("elephant", "whale")),
    ("small", "smaller", ("mouse", "ant")),
    ("fast", "faster", ("cheetah", "falcon")),
    ("slow", "slower", ("turtle", "snail")),
    ("tall", "taller", ("giraffe", "camel")),
    ("strong", "stronger", ("bear", "ox")),
    ("loud", "louder", ("lion", "parrot")),
    ("quiet", "quieter", ("owl", "rabbit")),
]

# =============================================================================
# 4. 水果
# =============================================================================
FRUITS = ["orange", "apple", "mango", "grape", "lemon", "banana", "peach", "durian"]

# =============================================================================
# 5. 职业：和性别的统计关联（去偏实验的"真相"）
# =============================================================================
# p_male = 这个职业在"代词 / 亲属称谓"语境里配男性词的概率。0.5 表示没有偏向。
# 数字是为了做实验随手设的，**不代表任何真实的统计**。
OCCUPATION_P_MALE = {
    "engineer": 0.9, "mechanic": 0.9, "carpenter": 0.9, "pilot": 0.85, "programmer": 0.85,
    "surgeon": 0.8, "scientist": 0.75, "lawyer": 0.65,
    "doctor": 0.6, "chef": 0.55, "writer": 0.5, "singer": 0.5, "artist": 0.45, "teacher": 0.4,
    "librarian": 0.25, "dancer": 0.2, "hairdresser": 0.2, "secretary": 0.15,
    "receptionist": 0.15, "nurse": 0.1, "babysitter": 0.1, "homemaker": 0.1,
}
# 每个职业专属的"工作对象"（中性的语境：和性别无关，只让职业彼此区分开）
OCCUPATION_OBJECT = {
    "engineer": "bridge", "mechanic": "engine", "carpenter": "table", "pilot": "plane",
    "programmer": "computer", "surgeon": "operation", "scientist": "experiment", "lawyer": "court",
    "doctor": "patient", "chef": "kitchen", "writer": "novel", "singer": "song", "artist": "painting",
    "teacher": "class", "librarian": "library", "dancer": "stage", "hairdresser": "haircut",
    "secretary": "letter", "receptionist": "phone", "nurse": "patient", "babysitter": "baby",
    "homemaker": "house",
}

# =============================================================================
# 6. 情感：形容词、方面、程度词
# =============================================================================
POS_ADJ_SEEN = ["good", "great", "nice", "tasty"]                # 情感分类的训练集里出现
POS_ADJ_UNSEEN = ["excellent", "wonderful", "superb", "lovely"]  # 只在无标签语料里出现
NEG_ADJ_SEEN = ["bad", "awful", "poor", "bland"]
NEG_ADJ_UNSEEN = ["terrible", "horrible", "dreadful", "disgusting"]
POS_ADJ = POS_ADJ_SEEN + POS_ADJ_UNSEEN
NEG_ADJ = NEG_ADJ_SEEN + NEG_ADJ_UNSEEN
ASPECTS = ["food", "service", "soup", "dessert", "coffee", "staff", "music", "room"]
# 每个情感词"最常形容"的方面（一半的时候用它），让同一极性的词之间也有点差别，不至于完全可以互换
FAVORITE_ASPECT = dict(zip(POS_ADJ, ASPECTS))
FAVORITE_ASPECT.update(zip(NEG_ADJ, ASPECTS))
INTENSIFIERS = ["very", "really", "extremely"]

STATES = ["tired", "happy", "busy", "late", "hungry", "ready"]
MONTHS = ["january", "february", "march", "april", "may", "june", "july", "august",
          "september", "october", "november", "december"]


# =============================================================================
# 句子模板
# =============================================================================
def _pick(rng, items):
    """
    从列表里等概率挑一个元素。

    【Python 小知识】rng.integers(n) 返回 0..n-1 里的一个随机整数；列表[下标] 取出元素。
            np.random.default_rng(seed) 造一个"自己的"随机数生成器：每个演示用自己的生成器，
            改一个演示不会影响别的演示的随机数（全课的约定）。
    """
    return items[int(rng.integers(len(items)))]


def _gendered(rng, p_male, male_word, female_word):
    """以概率 p_male 返回男性词，否则返回女性词。"""
    return male_word if rng.random() < p_male else female_word


def _person_sentence(rng):
    """
    一句关于人物的话。随机选一对 (男, 女) 里的一个人，再随机选一种语境：
        性别语境：the king said that he was tired / the queen lost her crown
        身份语境：royal → ruled the kingdom、lived in the palace；common → took the bus；family → called home
        年龄语境：adult → went to work、drove the car；child → went to school
        专属语境：the king sat on the throne（king 和 queen 共享，让每一对有自己的"身份"）
    """
    male, female, status, age, ident = PEOPLE[int(rng.integers(len(PEOPLE)))]
    is_male = rng.random() < 0.5
    p = male if is_male else female
    he, his = ("he", "his") if is_male else ("she", "her")
    kind = int(rng.integers(5))
    if kind == 0:
        return ["the", p, "said", "that", he, "was", _pick(rng, STATES)]
    if kind == 1:
        return ["the", p, "lost", his, _pick(rng, ["keys", "bag", "phone", "hat"])]
    if kind == 2:
        if status == "royal":
            return ["the", p] + _pick(rng, [["ruled", "the", "kingdom"], ["lived", "in", "the", "palace"],
                                             ["wore", "a", "golden", "crown"]])
        if status == "common":
            return ["the", p] + _pick(rng, [["took", "the", "bus"], ["lived", "in", "a", "small", "house"],
                                             ["bought", "some", "bread"]])
        return ["my", p] + _pick(rng, [["called", "home"], ["cooked", "dinner", "for", "the", "family"],
                                        ["joined", "the", "family", "trip"]])
    if kind == 3:
        if age == "adult":
            return ["the", p] + _pick(rng, [["went", "to", "work"], ["drove", "the", "car"], ["paid", "the", "bills"]])
        if age == "child":
            return ["the", p] + _pick(rng, [["went", "to", "school"], ["played", "in", "the", "garden"],
                                             ["drank", "some", "milk"]])
        # 没有年龄语境的词（brother / sister）：改用性别语境
        return ["the", p, "said", "that", he, "was", _pick(rng, STATES)]
    return ["the", p] + IDENTITY_VERBS[ident].split() + [ident]


def _country_sentence(rng):
    """
    一句关于国家 / 首都 / 语言的话。三类词各有自己的"类型语境"，又通过招牌菜和互相提及绑在一起：
        类型：the government of france / the streets of paris / a book written in french
        互相提及：paris is the capital of france；people in france speak french
        招牌菜：people in paris eat baguette；baguette is popular in france；the french menu offered baguette
    """
    country, capital, lang, dish = COUNTRIES[int(rng.integers(len(COUNTRIES)))]
    kind = int(rng.integers(9))
    return [
        [capital, "is", "the", "capital", "of", country],
        ["people", "in", country, "speak", lang],
        ["the", "government", "of", country, "met", "today"],
        ["we", "walked", "through", "the", "streets", "of", capital],
        ["she", "read", "a", "book", "written", "in", lang],
        ["people", "in", capital, "eat", dish],
        [dish, "is", "popular", "in", country],
        ["the", "menu", "was", "in", lang, "and", "offered", dish],
        ["we", "flew", "to", capital, "in", _pick(rng, MONTHS)],
    ][kind]


def _adjective_sentence(rng):
    """
    一句关于形容词的话：原级出现在 "is X" / "a very X"，比较级出现在 "is Xer than"；
    同一个形容词的原级和比较级共享两种"典型"动物（big、bigger 都常和 elephant、whale 一起出现）。
    """
    base, comp, animals = ADJECTIVES[int(rng.integers(len(ADJECTIVES)))]
    a = _pick(rng, list(animals))
    other = _pick(rng, ["dog", "cat", "goat", "duck"])
    kind = int(rng.integers(4))
    if kind == 0:
        return ["the", a, "is", base]
    if kind == 1:
        return ["a", "very", base, a]
    if kind == 2:
        return ["the", a, "is", comp, "than", "the", other]
    return ["the", a, "became", comp, "than", "before"]


def _fruit_sentence(rng):
    """一句关于水果的话：i want a glass of orange juice / the boy ate a ripe mango / the apple tree grew tall。"""
    f = _pick(rng, FRUITS)
    kind = int(rng.integers(4))
    return [
        ["i", "want", "a", "glass", "of", f, "juice"],
        ["the", _pick(rng, ["boy", "girl", "man", "woman"]), "ate", "a", "ripe", f],
        ["the", f, "tree", "grew", "in", "the", "garden"],
        ["she", "bought", "a", "bag", "of", f, "at", "the", "market"],
    ][kind]


def _occupation_sentence(rng):
    """
    一句关于职业的话。一半的语境带性别（按 OCCUPATION_P_MALE 选男 / 女），一半是中性的：
        带性别：the engineer said that he was busy / my uncle works as an engineer / the nurse finished her shift
        中性：  the engineer repaired the bridge（每个职业有自己的工作对象）
    """
    occ = _pick(rng, list(OCCUPATION_P_MALE))
    pm = OCCUPATION_P_MALE[occ]
    kind = int(rng.integers(4))
    if kind == 0:
        return ["the", occ, "said", "that", _gendered(rng, pm, "he", "she"), "was", _pick(rng, STATES)]
    if kind == 1:
        kin = _pick(rng, GENDER_PAIRS[4:])                       # 亲属称谓：father/mother ……
        return ["my", _gendered(rng, pm, kin[0], kin[1]), "works", "as", "a", occ]
    if kind == 2:
        return ["the", occ, "finished", _gendered(rng, pm, "his", "her"), "shift"]
    return ["the", occ, "worked", "on", "the", OCCUPATION_OBJECT[occ]]


def _review_sentence(rng):
    """
    一句"无标签"的评论语料。好评词只出现在正面的语境里，差评词只出现在负面的语境里：
        正面：i loved the soup because it was superb / the critic praised the lovely music / he smiled because ...
        负面：i hated the soup because it was dreadful / the critic criticized the awful service / she complained ...
        中性：the coffee was very good（两种词都会出现；程度词 very / really / extremely 只出现在这里）
    "not" 和 "lacking" 只出现在和评论无关的中性句子里（the boy did not ...），所以它们的向量本身不带褒贬。
    【DL 联系】我们从来没有给任何词贴"好 / 坏"的标签。词向量能把好评词和差评词分开，靠的只是"和哪些词一起出现"。
            这就是 Firth 的那句话："You shall know a word by the company it keeps."
    """
    positive = rng.random() < 0.5
    adj = _pick(rng, POS_ADJ if positive else NEG_ADJ)
    asp = FAVORITE_ASPECT[adj] if rng.random() < 0.5 else _pick(rng, ASPECTS)
    kind = int(rng.integers(5))
    if kind == 0:
        return ["i", "loved" if positive else "hated", "the", asp, "because", "it", "was", adj]
    if kind == 1:
        return ["the", "critic", "praised" if positive else "criticized", "the", adj, asp]
    if kind == 2:
        return [_pick(rng, ["he", "she"]), "smiled" if positive else "complained", "because", "the", asp,
                "was", adj]
    if kind == 3:
        return ["we", "will" if positive else "never", "come", "back", "for", "the", adj, asp]
    return ["the", asp, "was", _pick(rng, INTENSIFIERS), adj]


def _filler_sentence(rng):
    """
    几句和上面的结构无关的中性句子，让 not、lacking、very 这类功能词也有自己的语境（向量本身不带褒贬）。
    """
    kind = int(rng.integers(4))
    who = _pick(rng, ["boy", "girl", "man", "woman"])
    return [
        ["the", who, "did", "not", _pick(rng, ["go", "sleep", "wait", "answer"])],
        ["the", "report", "was", "lacking", "in", _pick(rng, ["detail", "data", "pages"])],
        ["it", "was", _pick(rng, INTENSIFIERS), "cold", "in", _pick(rng, MONTHS)],
        ["the", "train", "was", "not", "late", "in", _pick(rng, MONTHS)],
    ][kind]


# 每类句子的比例（加起来不必是 1，会被归一化）
SENTENCE_MIX = [
    (_person_sentence, 0.26), (_country_sentence, 0.18), (_adjective_sentence, 0.12),
    (_fruit_sentence, 0.06), (_occupation_sentence, 0.20), (_review_sentence, 0.12),
    (_filler_sentence, 0.06),
]


def generate_corpus(n_sentences, rng):
    """
    造一个合成语料。

    【思路】每句话先按 SENTENCE_MIX 的比例选一种句子，再调用对应的模板函数。
    【参数】n_sentences：句子数；rng：np.random.Generator
    【返回】list[list[str]]：每句话是一个词的列表（全部小写，没有标点）
    【边界】同样的 n_sentences 和种子，结果逐字相同。
    【Python 小知识】zip(*列表) 把 [(a1, b1), (a2, b2)] 拆成 (a1, a2) 和 (b1, b2) 两组。
    """
    makers, weights = zip(*SENTENCE_MIX)
    w = np.array(weights) / np.sum(weights)
    choice = rng.choice(len(makers), size=n_sentences, p=w)
    return [makers[k](rng) for k in choice]


# =============================================================================
# 情感分类的有标签数据（第 02 个脚本）：合成的短评论
# =============================================================================
STARS = ["1★ 很差", "2★ 差", "4★ 好", "5★ 很好"]          # 4 个类别的名字：0, 1, 2, 3
REVIEW_KINDS = ["普通", "程度词", "否定", "否定+程度词", "lacking in", "并列"]


def make_review(rng, kind, adjectives_pos, adjectives_neg, n_items=None):
    """
    造一条有标签的短评论。

    【思路】标签规则（我们写下的"真相"）：
        普通        the food was good                → 好（2）；the food was bad → 差（1）
        程度词      the food was very good           → 很好（3）；the food was very bad → 很差（0）
        否定        the food was not good            → 差（1）；the food was not bad → 好（2）
        否定+程度词 the food was not very good       → 差（1）；not very bad → 好（2）（否定让程度变弱）
        lacking in  it was lacking in good food and good service → 差（1）：好评词再多，也是在说"缺"
        并列        the food was good and the music was nice     → 好（2）（同一极性）
    【参数】rng；kind：REVIEW_KINDS 之一；adjectives_pos / adjectives_neg：可用的好评 / 差评词；
            n_items：lacking in / 并列 里有几项（None 时随机 1–2 项）
    【返回】(tokens, label)：词的列表和 0–3 的整数标签
    【DL 联系】Ng 在 C5M2 "Sentiment Classification" 里的反例 "Completely lacking in good taste, good service,
            and good ambience" 就是 lacking in 这一类：把词向量平均起来，三个 good 会把整句拉向好评。
    """
    positive = rng.random() < 0.5
    adj = _pick(rng, adjectives_pos if positive else adjectives_neg)
    asp = _pick(rng, ASPECTS)
    if kind == "普通":
        return ["the", asp, "was", adj], 2 if positive else 1
    if kind == "程度词":
        return ["the", asp, "was", _pick(rng, INTENSIFIERS), adj], 3 if positive else 0
    if kind == "否定":
        return ["the", asp, "was", "not", adj], 1 if positive else 2
    if kind == "否定+程度词":
        return ["the", asp, "was", "not", _pick(rng, INTENSIFIERS), adj], 1 if positive else 2
    n = n_items if n_items is not None else 1 + int(rng.integers(2))
    aspects = [ASPECTS[i] for i in rng.permutation(len(ASPECTS))[:n]]
    if kind == "lacking in":
        toks = ["it", "was", "lacking", "in"]
        for i, a in enumerate(aspects):
            toks += (["and"] if i > 0 else []) + [_pick(rng, adjectives_pos), a]
        return toks, 1
    if kind == "并列":
        toks = []
        for i, a in enumerate(aspects):
            toks += (["and"] if i > 0 else []) + ["the", a, "was", _pick(rng, adjectives_pos if positive else adjectives_neg)]
        return toks, 2 if positive else 1
    raise ValueError(kind)


# =============================================================================
# 类比题：从写进语料的结构里枚举
# =============================================================================
def analogy_questions():
    """
    列出所有类比题，按"关系"分组。

    【思路】一个关系是一组词对 [(a1, b1), (a2, b2), ...]，任意两对 (a, b)、(c, d)（顺序不同算不同的题）
            组成一道题 "a : b :: c : ?"，标准答案是 d。例如性别关系里有 man : woman :: king : queen。
    【返回】dict：关系名 → list[(a, b, c, d)]
    【边界】只列出我们写进语料的关系；真实的类比测试集（Mikolov et al. 2013）有一万多道题。
    """
    relations = {
        "性别 man→woman": GENDER_PAIRS,
        "身份 平民→王室": [("man", "king"), ("woman", "queen"), ("boy", "prince"), ("girl", "princess")],
        "年龄 大人→小孩": [("man", "boy"), ("woman", "girl"), ("king", "prince"), ("queen", "princess"),
                         ("father", "son"), ("mother", "daughter")],
        "国家→首都": [(c, cap) for c, cap, _, _ in COUNTRIES],
        "国家→语言": [(c, lang) for c, _, lang, _ in COUNTRIES],
        "原级→比较级": [(b, c) for b, c, _ in ADJECTIVES],
    }
    questions = {}
    for name, pairs in relations.items():
        qs = []
        for i, (a, b) in enumerate(pairs):
            for j, (c, d) in enumerate(pairs):
                if i != j:
                    qs.append((a, b, c, d))
        questions[name] = qs
    return questions


def _selftest():
    rng = np.random.default_rng(0)
    corpus = generate_corpus(60000, rng)
    tokens = [w for s in corpus for w in s]
    vocab = sorted(set(tokens))
    print(f"句子数 {len(corpus)}，词数（token）{len(tokens)}，不同的词（词表大小 V）{len(vocab)}")
    print("样例：")
    for s in corpus[:8]:
        print("   ", " ".join(s))
    qs = analogy_questions()
    print("类比题：" + "，".join(f"{k} {len(v)} 道" for k, v in qs.items()))


if __name__ == "__main__":
    _selftest()