"""
=============================================================================
CMU 17-630 · 第 1 讲共享模块：pretrain_world.py —— 一个虚构世界和它的"预训练语料"
=============================================================================

【这个模块是做什么的】
    真实的大模型是在几万亿 token 的网页文本上预训练出来的。网页里有百科词条、问答论坛、
    考试题目列表、商品评论、小说……**每一种文档都有自己的格式**。提示工程里最重要的一条直觉
    （GPT-2、GPT-3 论文都在讲它）是：

        一个 base 模型只是在"续写文档"。你的提示有没有用，取决于它看起来像哪一种文档的开头，
        以及在那种文档里，你想要的回答是不是最自然的续写。

    为了让这条直觉能在你的笔记本电脑上"亲眼看到"，这个模块造了一个很小的虚构世界，
    再把它"写"成 5 种格式的文档，拼成一份迷你预训练语料（约 25 万个词）：

        article  百科词条  "<doc> article : varnia is a country in the north . the capital of varnia is lorn . ..."
        faq      问答页面  "<doc> faq : q : what is the capital of varnia ? a : lorn . q : ..."
        quiz     题目列表  "<doc> quiz : what is the capital of varnia ? what is dulmar known for ? ..."
        reviews  评论页面  "<doc> reviews : review : the soup was cold . rating : negative . ..."
        story    小故事    "<doc> story : the queen lives in the palace . she wears a crown . ..."

    世界里的"知识"（每个国家的首都、语言、出口商品）是随机生成的虚构名字，所以模型
    **不可能**从别的地方知道它们——它知道的一切都来自这份语料。这让我们可以做一些在真实模型上
    做不了的干净实验，例如：

        - 有 1/3 的国家，它们的知识**只**出现在百科词条里，从来没有以问答的形式出现过。
          用问答格式去问这些国家，模型还能答对吗？（知识能不能跨格式迁移，第 1 讲 §8）
        - 光秃秃的一个问题 "what is the capital of varnia ?"，在语料里既可能是问答页里的问题
          （后面跟着答案），也可能是题目列表里的一道题（后面跟着下一道题）。模型会怎样续写？

【设计上的几个选择（以及为什么）】
    1. **词级（word-level）分词**：每个词、每个标点都是一个 token，用空格分开。真实模型用的是
       子词（BPE，前置讲 P1）；这里用词级是为了让输出好读、模型好训练，不影响要讲的道理。
    2. 全部小写、标点前后有空格：这样 split() 就是分词器。
    3. 所有随机性都来自 random.Random(seed)：同一个 seed 永远生成同一份语料。
       训练集用 seed=0，测试集用 seed=1（同一个世界、不同的句子），所以测试集里的句子
       大多数是模型没见过的新组合。
    4. 评论的标签有三种难度：同一极性的简单句、"A but B"（以 but 后面为准）、
       带 "not" 的否定。后两种需要看更远的上下文，用来对比 n-gram 和 Transformer。
    5. 故事里的人物词（king/queen/man/woman……）带着"性别""王室""年轻"三种特征，
       "bank" 一半出现在金融语境、一半出现在河岸语境——这是给第 02 号脚本（词向量的类比和多义词）准备的。

【运行】python3 pretrain_world.py    # 自检：打印世界、每种文档的例子、语料规模
【依赖】只用 Python 标准库。第 1 讲的所有脚本都 import 这个文件；后面几讲会复用它的评论和国家。
"""

import random
from dataclasses import dataclass

# ---------------------------------------------------------------------------
# 1. 世界里的"词表"：商品、方位、评论用词、故事用词
# ---------------------------------------------------------------------------

EXPORTS = ["copper", "wheat", "salt", "tea", "wool", "iron", "silk", "rice", "coal", "timber", "olives", "glass"]
REGIONS = ["north", "south", "east", "west"]

# 评论：方面词 + 褒义 / 贬义形容词
ASPECTS = ["food", "soup", "bread", "coffee", "staff", "service", "room", "bed", "price", "music", "view", "waiter"]
POSITIVE = ["great", "fresh", "warm", "friendly", "clean", "lovely", "quick", "tasty", "cheap", "quiet"]
NEGATIVE = ["cold", "stale", "rude", "slow", "dirty", "awful", "noisy", "bland", "pricey", "soggy"]
LABELS = {+1: "positive", -1: "negative"}

# 故事：成对的人物词（男性, 女性），以及每对人物带的特征
GENDER_PAIRS = [
    ("king", "queen"), ("prince", "princess"), ("lord", "lady"),
    ("man", "woman"), ("boy", "girl"), ("father", "mother"), ("son", "daughter"),
    ("brother", "sister"), ("uncle", "aunt"), ("husband", "wife"),
]
ROYAL = {"king", "queen", "prince", "princess", "lord", "lady"}
YOUNG = {"prince", "princess", "boy", "girl", "son", "daughter"}
FAMILY = {"father", "mother", "son", "daughter", "brother", "sister", "uncle", "aunt", "husband", "wife"}

# 每种特征对应的"语境短语"：一个人物词出现在什么短语旁边，取决于它带哪些特征
CONTEXTS = {
    "royal": ["wears a crown", "lives in the palace", "rules the land", "sits on the throne"],
    "young": ["plays in the yard", "goes to school", "loves toys", "is very small"],
    "adult": ["works in the field", "pays the taxes", "reads the news", "drives a cart"],
    "family": ["cooks for the family", "visits the cousins", "cares for the house"],
    "any": ["walks to the market", "is very kind", "sings a song", "eats an apple"],
}

# "bank" 的两种意思，以及只有一种意思的对照词（vault 只在金融语境，shore 只在河岸语境）
BANK_FINANCE = ["she put money in the {w} .", "he got a loan from the {w} .",
                "they keep cash at the {w} .", "the {w} counts the coins ."]
BANK_RIVER = ["we sat on the river {w} .", "fish swim near the {w} .",
              "boats are tied to the {w} .", "the {w} was muddy after the rain ."]

# ---------------------------------------------------------------------------
# 2. 百科、问答、题目列表的句子模板
# ---------------------------------------------------------------------------

# 四种"关系"：问答里会被问到的知识。值是 Country 的属性名。
RELATIONS = ["capital", "language", "export", "leader"]

# 问答格式里的问题（faq 和 quiz 用的是同一套问法）
QUESTIONS = {
    "capital": "what is the capital of {c} ?",
    "language": "what language do people speak in {c} ?",
    "export": "what is {c} known for ?",
    "leader": "who leads {c} ?",
}

# 百科词条里的说法：每种关系有好几种措辞（paraphrase）。
# 【为什么要多种措辞】同一个事实以不同的方式出现，模型才更可能把它存成"关于 varnia 的知识"，
#   而不是死记一个固定的句子。（Allen-Zhu & Li 2023 "Physics of Language Models 3.1" 系统地研究了这一点。）
ARTICLE_TEMPLATES = {
    "intro": ["{c} is a country in the {region} .", "{c} lies in the {region} ."],
    "capital": ["the capital of {c} is {capital} .", "{capital} is the capital of {c} .",
                "{c} 's capital city is {capital} .", "the government of {c} sits in {capital} ."],
    "language": ["people in {c} speak {language} .", "the main language of {c} is {language} .",
                 "{language} is spoken in {c} ."],
    "export": ["{c} is known for its {export} .", "the main export of {c} is {export} .",
               "{c} sells {export} to its neighbours ."],
    # 注意：领导人的三种说法**都是"国家在前、人名在后"**；而首都既有"首都 of 国家"也有"首都 is the capital of 国家"。
    # 这个不对称是故意的：05 号脚本用它演示"反转诅咒"（reversal curse）——
    # 只以"A 的领导人是 B"的方向见过的知识，反过来问"B 是谁的领导人"时，模型答不出来。
    "leader": ["the leader of {c} is {leader} .", "{c} is led by {leader} .", "{c} is ruled by {leader} ."],
    "filler": ["{capital} lies on a wide river .", "many visitors come to {c} in summer .",
               "the roads of {c} are old .", "{c} has a long history ."],
}

# 生成虚构名字用的音节。名字由"声母 + 元音 + 韵尾"拼成，读起来像地名，但不是真实的词。
_ONSETS = ["b", "d", "f", "g", "k", "l", "m", "n", "p", "r", "s", "t", "v", "z",
           "br", "dr", "gr", "kr", "tr", "st", "sh", "th"]
_VOWELS = ["a", "e", "i", "o", "u"]
_CODAS = ["l", "n", "r", "s", "m", "nd", "rk", "st", "rn", "lt"]
_COUNTRY_ENDINGS = ["ia", "a", "or", "an", "el", "is", "ova", "ania"]
_LANGUAGE_ENDINGS = ["ish", "ese", "ic", "i"]
_PERSON_ENDINGS = ["o", "a", "u", "ek", "ia"]
# 首都也加一个结尾，变成两个音节：第一版首都是单音节的，结果拼出了 "last" 这样的真实英文单词。
_CITY_ENDINGS = ["a", "o", "en", "ar", "um", "is"]


def _template_words():
    """
    所有模板、语境、评论里用到的英文单词。虚构名字不许和它们重名（否则 "the capital of x is last" 这种句子会把模型和读者都搞糊涂）。

    【Python 小知识】set().union(*列表) 把很多个集合合并成一个。
    """
    texts = [t for ts in ARTICLE_TEMPLATES.values() for t in ts] + list(QUESTIONS.values())
    texts += [t for ts in CONTEXTS.values() for t in ts] + BANK_FINANCE + BANK_RIVER
    words = set().union(*(set(t.replace("{", " ").replace("}", " ").split()) for t in texts))
    words |= set(EXPORTS + REGIONS + ASPECTS + POSITIVE + NEGATIVE + list(LABELS.values()))
    words |= {w for pair in GENDER_PAIRS for w in pair}
    words |= {"article", "faq", "quiz", "reviews", "review", "rating", "story", "bank", "vault", "shore",
              "he", "she", "but", "not", "and", "q", "a"}
    return words


@dataclass(frozen=True)
class Country:
    """
    世界里的一个国家。

    【字段】name 国名；capital 首都；language 语言；export 主要出口商品；leader 领导人；region 方位；
        qa_seen：True 表示它的知识会以问答（faq）的形式出现在训练语料里；
                 False 表示它的知识**只**出现在百科词条里（用来测"跨格式迁移"）。
    【Python 小知识】@dataclass 自动生成 __init__ 等方法；frozen=True 表示创建后不能修改字段，
        这样一个 Country 可以放心地被很多地方共享，不会被谁偷偷改掉。
    """
    name: str
    capital: str
    language: str
    export: str
    leader: str
    region: str
    qa_seen: bool

    def attr(self, relation):
        """按关系名取答案，例如 attr("capital") → 首都名。getattr(对象, "字段名") 就是"按名字取字段"。"""
        return getattr(self, relation)


@dataclass
class World:
    """整个虚构世界：一组国家 + 生成它用的 seed。"""
    countries: list
    seed: int

    @property
    def seen(self):
        """知识以问答形式出现过的国家。【Python 小知识】@property 让方法可以像字段一样用：world.seen。"""
        return [c for c in self.countries if c.qa_seen]

    @property
    def held(self):
        """知识只出现在百科词条里的国家（问答格式从没见过）。"""
        return [c for c in self.countries if not c.qa_seen]


def _make_name(rng, endings=None, used=None):
    """
    随机拼一个虚构名字，保证不和已经用过的名字重复。

    【参数】rng：random.Random 对象；endings：可选的结尾列表（国名、语言名会加结尾）；
        used：已经用过的名字集合（会被就地更新）
    【返回】一个新名字（字符串）
    【边界】音节组合有几千种，24 个国家 × 4 个名字远远用不完，所以 while 循环很快就会退出。
    """
    used = used if used is not None else set()
    while True:
        name = rng.choice(_ONSETS) + rng.choice(_VOWELS) + rng.choice(_CODAS)
        if endings:
            name += rng.choice(endings)
        if name not in used and len(name) >= 4:
            used.add(name)
            return name


def build_world(seed=0, n_countries=24, n_held=8):
    """
    生成虚构世界。

    【参数】seed：随机种子；n_countries：国家数；n_held：其中"只在百科里出现"的国家数
    【返回】World
    【边界】n_held 必须小于 n_countries（至少要有一些国家出现在问答里，模型才能学会问答格式）。
    【Python 小知识】random.Random(seed) 创建一个**独立的**随机数生成器：只要 seed 相同，
        它产生的随机数序列就完全相同。用它而不用全局的 random.xxx()，可以避免不同的代码
        互相"偷走"随机数、让结果变得不可复现。
    """
    assert 0 < n_held < n_countries
    rng = random.Random(seed)
    # 先把语料里所有的英文单词放进 used，防止虚构名字恰好和它们重名
    used = _template_words()
    countries = []
    for i in range(n_countries):
        countries.append(Country(
            name=_make_name(rng, _COUNTRY_ENDINGS, used),
            capital=_make_name(rng, _CITY_ENDINGS, used),
            language=_make_name(rng, _LANGUAGE_ENDINGS, used),
            export=rng.choice(EXPORTS),
            leader=_make_name(rng, _PERSON_ENDINGS, used),
            region=rng.choice(REGIONS),
            qa_seen=(i >= n_held),   # 前 n_held 个国家留作"问答里没见过"
        ))
    return World(countries=countries, seed=seed)


# ---------------------------------------------------------------------------
# 3. 把世界"写"成 5 种文档
# ---------------------------------------------------------------------------

def _fill(template, c):
    """把模板里的 {c} {capital} {language} {export} {region} 换成国家 c 的值。"""
    return template.format(c=c.name, capital=c.capital, language=c.language, export=c.export,
                           leader=c.leader, region=c.region)


def make_article(rng, c):
    """
    一篇百科词条：介绍句 + 首都 / 语言 / 出口各一句（措辞随机）+ 0–2 句填充（这几句顺序打乱），
    最后**固定**以领导人那一句结尾。

    【返回】文档正文（不含 "<doc>"）
    【为什么领导人那句必须放在最后】反转诅咒的实验要求训练语料里**领导人的名字从来不出现在国名前面**。
        第一版把领导人那句和别的句子一起打乱，结果 "... the leader of x is kiro . x is known for ..."
        这样的文本里，kiro 后面紧跟着 x——模型于是学会了"kiro 后面接 x"，反着问也能答对 21/24，
        诅咒没有出现。放在最后之后，kiro 后面只会是下一篇无关文档的 "<doc>"。
    """
    body = [_fill(rng.choice(ARTICLE_TEMPLATES[k]), c) for k in RELATIONS if k != "leader"]
    body += [_fill(t, c) for t in rng.sample(ARTICLE_TEMPLATES["filler"], rng.randint(0, 2))]
    rng.shuffle(body)
    intro = _fill(rng.choice(ARTICLE_TEMPLATES["intro"]), c)
    last = _fill(rng.choice(ARTICLE_TEMPLATES["leader"]), c)
    return "article : " + " ".join([intro] + body + [last])


def make_faq(rng, countries):
    """
    一个问答页面：2–5 个"q : 问题 ? a : 答案 ."。只从 qa_seen 的国家里出题。

    【参数】countries：可以出题的国家列表（调用方传 world.seen）
    """
    pairs = []
    for _ in range(rng.randint(2, 5)):
        c = rng.choice(countries)
        rel = rng.choice(RELATIONS)
        pairs.append(f"q : {_fill(QUESTIONS[rel], c)} a : {c.attr(rel)} .")
    return "faq : " + " ".join(pairs)


def make_quiz(rng, countries):
    """
    一份题目列表：3–6 道**不带答案**的题，所有国家都可能出现。

    【为什么要有它】真实网页里有大量"只有问题、没有答案"的文本（考试题、讨论帖标题……）。
        GPT-3 论文之前人们就发现，给 base 模型一个光秃秃的问题，它常常接着写更多的问题。
        这份语料里也有这种文档，所以我们能在迷你模型上复现这个现象。
    """
    qs = []
    for _ in range(rng.randint(3, 6)):
        c = rng.choice(countries)
        qs.append(_fill(QUESTIONS[rng.choice(RELATIONS)], c))
    return "quiz : " + " ".join(qs)


def _clause(rng, polarity, negated):
    """
    评论里的一个分句，例如 "the soup was cold" 或 "the staff was not friendly"。

    【参数】polarity：+1 褒 / −1 贬（指整句**最终**的意思）；negated：是否用 "not" 表达
    【返回】(分句文字, 极性)
    【思路】否定句要"反着取词"：想表达褒义却用 not，就得说 "was not 贬义词"。
    """
    aspect = rng.choice(ASPECTS)
    if negated:
        adj = rng.choice(NEGATIVE if polarity > 0 else POSITIVE)
        return f"the {aspect} was not {adj}", polarity
    adj = rng.choice(POSITIVE if polarity > 0 else NEGATIVE)
    return f"the {aspect} was {adj}", polarity


def make_review(rng):
    """
    生成一条评论和它的标签。

    【返回】(评论正文, 标签 "positive"/"negative", 难度类型 "simple"/"but"/"not")
    【思路】三种类型的比例 60% / 20% / 20%：
        simple：1–2 个同极性分句，用 and 连接——看任何一个形容词都能判对；
        but   ："A but B"，A、B 极性相反，以 B 为准——要知道 but 的作用；
        not   ：一个带 not 的分句（可能再接一个同极性的 and 分句）——要知道 not 会反转意思。
    """
    polarity = rng.choice([+1, -1])
    r = rng.random()
    if r < 0.6:
        kind = "simple"
        parts = [_clause(rng, polarity, False)[0] for _ in range(rng.randint(1, 2))]
        text = " and ".join(parts)
    elif r < 0.8:
        kind = "but"
        text = _clause(rng, -polarity, False)[0] + " but " + _clause(rng, polarity, False)[0]
    else:
        kind = "not"
        parts = [_clause(rng, polarity, True)[0]]
        if rng.random() < 0.5:
            parts.append(_clause(rng, polarity, False)[0])
            rng.shuffle(parts)
        text = " and ".join(parts)
    return text, LABELS[polarity], kind


def make_reviews_page(rng):
    """一个评论页面：2–4 条 "review : 正文 . rating : 标签 ."。"""
    items = []
    for _ in range(rng.randint(2, 4)):
        text, label, _ = make_review(rng)
        items.append(f"review : {text} . rating : {label} .")
    return "reviews : " + " ".join(items)


def _person_sentence(rng, word, pronoun):
    """
    关于一个人物词的两句话："the {word} {语境} . {代词} {语境} ."

    【思路】语境从这个词带的特征里抽：王室词会"戴王冠"，年轻的词会"去上学"，成年的词会"去田里干活"……
        所以 king 的上下文 ≈ 王室 + 成年 + he，queen ≈ 王室 + 成年 + she。
        词向量的类比（king − man + woman ≈ queen）就是从这种**可组合**的上下文结构里来的。
    """
    feats = ["any"]
    if word in ROYAL:
        feats.append("royal")
    feats.append("young" if word in YOUNG else "adult")
    if word in FAMILY:
        feats.append("family")
    a = rng.choice(CONTEXTS[rng.choice(feats)])
    b = rng.choice(CONTEXTS[rng.choice(feats)])
    return f"the {word} {a} . {pronoun} {b} ."


def make_story(rng):
    """
    一篇小故事：2–3 段人物描写，有 25% 的概率再加一句含 bank / vault / shore 的句子。

    【思路】bank 有两种意思：金融语境（money、loan）和河岸语境（river、fish）。
        vault 只出现在金融语境，shore 只出现在河岸语境，作为"单义词"的对照（02 号脚本用）。
    """
    parts = []
    for _ in range(rng.randint(2, 3)):
        male, female = rng.choice(GENDER_PAIRS)
        if rng.random() < 0.5:
            parts.append(_person_sentence(rng, male, "he"))
        else:
            parts.append(_person_sentence(rng, female, "she"))
    if rng.random() < 0.25:
        if rng.random() < 0.5:   # bank 的两种意思各占一半
            word = "bank"
            tpl = rng.choice(BANK_FINANCE if rng.random() < 0.5 else BANK_RIVER)
        else:                    # 对照词：vault 只用金融模板，shore 只用河岸模板
            if rng.random() < 0.5:
                word, tpl = "vault", rng.choice(BANK_FINANCE)
            else:
                word, tpl = "shore", rng.choice(BANK_RIVER)
        parts.append(tpl.format(w=word))
    return "story : " + " ".join(parts)


@dataclass(frozen=True)
class Document:
    """一篇文档：kind 是 "article"/"faq"/"quiz"/"reviews"/"story" 之一；text 是正文（不含 "<doc>"）。"""
    kind: str
    text: str


# 每种文档的数量（训练语料的"配方"）。缩放 scale 可以得到更小的测试集。
# 【为什么是这些数】第一版只有现在的 1/3（7.5 万词），迷你 GPT 训练 3000 步相当于把语料看了 80 遍，
#   验证损失从第 1000 步起就开始上升（死记硬背训练句子）。放大到 3 倍后验证损失不再上升。
CORPUS_MIX = {"article": 36, "faq": 900, "quiz": 450, "reviews": 1350, "story": 2100}


def build_corpus(world, seed=0, scale=1.0):
    """
    按配方生成一份语料（文档列表，顺序打乱）。

    【参数】world：build_world() 的结果；seed：生成句子用的随机种子（训练集 0，测试集 1）；
        scale：数量缩放（测试集通常用 0.2）
    【返回】Document 列表
    【边界】"article" 的配方是**每个国家**的篇数；其余是总篇数。scale 很小时每种至少 1 篇。
    【提示工程联系】这份配方决定了模型"见过什么样的文档"。§8 的实验会显示：
        提示的效果几乎完全由"它像不像训练语料里的某种文档"决定。
    """
    rng = random.Random(1000 + seed)
    docs = []
    n_art = max(1, round(CORPUS_MIX["article"] * scale))
    for c in world.countries:
        docs += [Document("article", make_article(rng, c)) for _ in range(n_art)]
    for kind, maker in [("faq", lambda: make_faq(rng, world.seen)),
                        ("quiz", lambda: make_quiz(rng, world.countries)),
                        ("reviews", lambda: make_reviews_page(rng)),
                        ("story", lambda: make_story(rng))]:
        n = max(1, round(CORPUS_MIX[kind] * scale))
        docs += [Document(kind, maker()) for _ in range(n)]
    rng.shuffle(docs)
    return docs


def corpus_text(docs):
    """
    把文档列表拼成一条长长的文本流，每篇文档前加一个 "<doc>" 标记。

    【返回】字符串，例如 "<doc> article : ... <doc> faq : ..."
    【提示工程联系】"<doc>" 相当于 GPT-2 的 <|endoftext|>：告诉模型"一篇新文档开始了"。
        用它开头的提示，模型会认为自己在写一篇全新文档的开头。
    """
    return " ".join("<doc> " + d.text for d in docs)


# ---------------------------------------------------------------------------
# 4. 分词器：词级，空格切分
# ---------------------------------------------------------------------------

class WordTokenizer:
    """
    最简单的分词器：按空格切分，每个不同的词分配一个整数编号（id）。

    【思路】词表 = 训练语料里出现过的所有词（排序后编号，保证每次都一样）+ 两个特殊符号：
        "<unk>"（unknown，没见过的词都映射到它）和 "<doc>"（文档开始）。
    【Python 小知识】dict 的 .get(key, default) 在 key 不存在时返回 default 而不报错。
    """
    SPECIALS = ["<unk>", "<doc>"]

    def __init__(self, text):
        words = sorted(set(text.split()) - set(self.SPECIALS))
        self.itos = self.SPECIALS + words                     # id → 词（index to string）
        self.stoi = {w: i for i, w in enumerate(self.itos)}   # 词 → id（string to index）
        self.unk = self.stoi["<unk>"]

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

    def encode(self, text):
        """
        文字 → id 列表。

        【边界】词表外的词变成 <unk>（id 0）。大写字母不会自动转小写：
            "Q" 和 "q" 是两个不同的词，而训练语料里只有小写，所以 "Q" 会变成 <unk>。
            这是故意的——它模拟了"提示里出现了训练时很少见的写法"。
        """
        return [self.stoi.get(w, self.unk) for w in text.split()]

    def decode(self, ids):
        """id 列表 → 用空格连接的文字。"""
        return " ".join(self.itos[i] for i in ids)

    def unknown_words(self, text):
        """返回 text 里所有不在词表中的词（用来提醒"这个提示里有模型没见过的词"）。"""
        return [w for w in text.split() if w not in self.stoi]


# ---------------------------------------------------------------------------
# 5. 探针：用来"问"模型的测试题
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class FactProbe:
    """一道知识题：country 国家；relation 关系；answer 标准答案（单个词）；seen 问答格式里见过没有。"""
    country: str
    relation: str
    answer: str
    seen: bool


def fact_probes(world):
    """所有国家 × 四种关系的知识题（24 × 4 = 96 道）。"""
    return [FactProbe(c.name, rel, c.attr(rel), c.qa_seen) for c in world.countries for rel in RELATIONS]


def question_text(country_name, relation):
    """某国某关系的问题文字，例如 question_text("varnia", "capital") → "what is the capital of varnia ?"。"""
    return QUESTIONS[relation].format(c=country_name)


def review_probes(seed=7, n=300):
    """
    一批**新生成**的评论（和训练语料用不同的 seed），用来测情感分类。

    【返回】[(正文, 标签, 类型), ...]
    【边界】是"新组合"但不是"新词"：用的方面词和形容词都在训练语料里出现过。
    """
    rng = random.Random(5000 + seed)
    return [make_review(rng) for _ in range(n)]


if __name__ == "__main__":
    world = build_world()
    print(f"世界：{len(world.countries)} 个国家，其中 {len(world.held)} 个的知识只出现在百科词条里")
    for c in world.countries[:3] + world.countries[-2:]:
        print(f"  {c.name:10s} 首都 {c.capital:8s} 语言 {c.language:10s} 出口 {c.export:7s} 领导人 {c.leader:8s} "
              f"{'问答见过' if c.qa_seen else '只在百科'}")
    train = build_corpus(world, seed=0)
    test = build_corpus(world, seed=1, scale=0.2)
    text = corpus_text(train)
    tok = WordTokenizer(text)
    print(f"\n训练语料：{len(train)} 篇文档，{len(text.split()):,} 个词；词表 {len(tok)} 个词")
    print(f"测试语料：{len(test)} 篇文档，{len(corpus_text(test).split()):,} 个词")
    for kind in CORPUS_MIX:
        d = next(d for d in train if d.kind == kind)
        n = sum(len(x.text.split()) for x in train if x.kind == kind)
        print(f"\n[{kind}] 共 {n:,} 个词。例子：\n  <doc> {d.text[:220]}{' ...' if len(d.text) > 220 else ''}")
    probes = review_probes()
    kinds = {k: sum(1 for p in probes if p[2] == k) for k in ["simple", "but", "not"]}
    print(f"\n评论探针 {len(probes)} 条：{kinds}；例子：{probes[0]}")
    unseen = set(corpus_text(test).split()) - set(text.split())
    print(f"测试语料里训练时没见过的词：{sorted(unseen) if unseen else '无'}")