"""
toy_corpus.py —— 第 1 讲的"农场玩具语言"（之后几讲也会用）
=====================================================================

【这个文件要回答的问题】
    要比较 n-gram 语言模型和 RNN 语言模型，我们需要一份语料：
    (1) 有足够多"局部规律"（n-gram 能学会的东西），
    (2) 又有"长程依赖"（只看最近几个词就一定会错的东西），
    (3) 并且我们**确切知道**它的真实分布，好把模型和"理论上最好的模型"作比较。
    真实的英语语料做不到第 (3) 点，所以我们自己造一门小语言。

【这门语言长什么样】
    每句话是一行，结构固定：

        the  <主语>  [<介词> the <干扰名词>] × k   <动词>  <宾语>  <时间/地点>  .

    例如：
        the bat makes noise at night .
        the cows near the dog behind the hens eat corn in the field .

    - 主语：8 种动物，单数或复数（bat / bats, cow / cows, ...），各占一半。
    - 中间可以插入 k 个介词短语（near / behind / beside the ...），里面的名词叫
      "干扰名词"（attractor）：它的单复数是随机的，和主语无关。
    - 动词必须和**主语**的单复数一致（the bat ... makes；the bats ... make），
      而不是和离动词最近的干扰名词一致。这就是英语的主谓一致（subject-verb agreement）。
    - 每种动物有自己的"习性"：吃什么、做什么、什么时候出现（蝙蝠多在夜里、
      奶牛多吃玉米和草……）。所以宾语和时间短语也依赖于很远之前的主语。

    这两条长程依赖（主谓一致、主语决定宾语与时间）跨越的距离是 3k+1 个词，
    k 越大越远。n-gram 只能看到最近 n-1 个词，所以 k 一大它就只能瞎猜；
    这正是第 1 讲要演示的"n-gram 的根本缺陷"，也是 RNN 想要解决的问题。

【真实概率】
    生成每句话时，我们把每一次随机选择的概率都乘起来，得到这句话在真实分布下的
    概率 p_true(句子)。因为这门语言的每个句子只有唯一的生成方式（介词、动词、
    宾语、时间短语用的词互不重叠），这个乘积就是精确的真实概率。
    用它可以算出"真实困惑度"——任何模型在期望意义下都不可能比它更好。

【怎样使用】
    from toy_corpus import make_splits, make_agreement_set, Vocab, check_sentence
    splits = make_splits(seed=0)          # {"train": [...], "valid": [...], "test": [...]}
    splits["train"][0]                    # 一个 Example：.tokens 是词列表
    直接运行本文件会打印几句样例和语料统计：python3 toy_corpus.py

【GenAI 联系】
    真实世界里，研究者也常常用"人造语言"来做受控实验，例如 Linzen et al. (2016)
    就是用主谓一致来测试 LSTM 能否学到句法结构。人造语言的好处是：我们知道正确答案，
    能把"模型学到了什么"和"数据里有什么"分开。
"""

import math
import random

# =====================================================================
# 1. 词表：名词、动词、宾语、时间短语
# =====================================================================

# 名词：单数 -> 复数。
# 【Python 小知识】字典（dict）用花括号写，"键: 值"成对出现。
#   NOUNS["bat"] 得到 "bats"。Python 3.7 起字典会记住插入顺序，
#   所以遍历 NOUNS 的顺序就是下面书写的顺序（这对可复现性很重要）。
NOUNS = {
    "bat": "bats",
    "cow": "cows",
    "dog": "dogs",
    "cat": "cats",
    "owl": "owls",
    "hen": "hens",
    "fox": "foxes",
    "horse": "horses",
}
ANIMALS = list(NOUNS)                     # ["bat", "cow", ...]，list(字典) 得到全部键

# 动词原形 -> (单数第三人称形式, 复数形式)
VERBS = {
    "make": ("makes", "make"),
    "eat": ("eats", "eat"),
    "see": ("sees", "see"),
    "chase": ("chases", "chase"),
    "hear": ("hears", "hear"),
}

# 这几个动词的宾语是"the + 某种动物"（单复数随机）；其余动词的宾语是一个普通名词。
ANIMAL_OBJECT_VERBS = ("see", "chase", "hear")

# 介词短语用的介词
PREPOSITIONS = ("near", "behind", "beside")

# 时间 / 地点短语：每个短语是一个词列表
TIMES = {
    "night": ["at", "night"],
    "morning": ["in", "the", "morning"],
    "day": ["every", "day"],
    "barn": ["in", "the", "barn"],
    "field": ["in", "the", "field"],
}

# 每种动物的"习性"（权重不必加起来等于 1，抽样时会自动归一化）：
#   verbs：做各个动作的相对频率
#   objects：对于宾语是普通名词的动词，宾语的分布
#   targets：对于 see/chase/hear，宾语是哪种动物
#   times：句末时间/地点短语的分布
# 权重是随手设定的，只要求"不同动物差别明显"，好让主语携带的信息有用。
PROFILES = {
    "bat": {
        "verbs": {"make": 4, "eat": 4, "hear": 2, "see": 1},
        "objects": {"make": {"noise": 8, "tracks": 1}, "eat": {"flies": 6, "fruit": 3}},
        "targets": {"owl": 3, "cat": 2, "bat": 1},
        "times": {"night": 8, "barn": 2, "morning": 1},
    },
    "cow": {
        "verbs": {"eat": 6, "make": 3, "see": 1, "hear": 1},
        "objects": {"make": {"milk": 6, "noise": 3}, "eat": {"corn": 4, "grass": 3, "hay": 2, "seeds": 1}},
        "targets": {"horse": 3, "dog": 2, "hen": 1},
        "times": {"field": 5, "day": 3, "morning": 2},
    },
    "dog": {
        "verbs": {"chase": 5, "eat": 3, "hear": 2, "make": 2},
        "objects": {"make": {"noise": 5, "tracks": 3}, "eat": {"bones": 6, "corn": 1}},
        "targets": {"cat": 4, "fox": 3, "hen": 2},
        "times": {"day": 4, "morning": 3, "field": 2},
    },
    "cat": {
        "verbs": {"chase": 4, "eat": 3, "see": 2, "make": 1},
        "objects": {"make": {"noise": 4, "tracks": 1}, "eat": {"fish": 5, "mice": 4}},
        "targets": {"bat": 3, "hen": 2, "dog": 1},
        "times": {"night": 4, "barn": 3, "day": 1},
    },
    "owl": {
        "verbs": {"hear": 3, "see": 3, "eat": 3, "make": 2},
        "objects": {"make": {"noise": 4, "nests": 3}, "eat": {"mice": 7, "fish": 1}},
        "targets": {"bat": 3, "fox": 2, "cat": 1},
        "times": {"night": 9, "barn": 1},
    },
    "hen": {
        "verbs": {"make": 4, "eat": 4, "hear": 1, "see": 1},
        "objects": {"make": {"eggs": 7, "noise": 3}, "eat": {"seeds": 5, "corn": 4}},
        "targets": {"fox": 3, "dog": 2, "cow": 1},
        "times": {"morning": 5, "barn": 4, "day": 1},
    },
    "fox": {
        "verbs": {"chase": 5, "eat": 3, "see": 2},
        "objects": {"eat": {"eggs": 4, "mice": 4, "fruit": 1}},
        "targets": {"hen": 5, "cat": 1, "bat": 1},
        "times": {"night": 6, "field": 2},
    },
    "horse": {
        "verbs": {"eat": 5, "see": 2, "make": 2, "hear": 1},
        "objects": {"make": {"tracks": 5, "noise": 2}, "eat": {"hay": 5, "grass": 4}},
        "targets": {"cow": 3, "dog": 2, "fox": 1},
        "times": {"field": 6, "day": 2, "morning": 1},
    },
}

# 训练集里"插入几个介词短语"的分布：k = 0, 1, 2, 3 的相对权重。
# k 越大越少见；测试时我们还会专门造 k = 4..6 的句子（训练时从没见过这么长的）。
DEFAULT_K_WEIGHTS = (0.45, 0.30, 0.15, 0.10)

# 特殊记号（special tokens）
PAD = "<pad>"   # 填充：把不等长的句子补成等长，第 3 讲细讲
BOS = "<s>"     # 句子开始（START）
EOS = "</s>"    # 句子结束（END）


# =====================================================================
# 2. 抽样小工具
# =====================================================================

def weighted_choice(rng, weights):
    """
    按权重从一个字典里抽一个键，同时返回它的概率。

    【思路】把权重归一化成概率，再用 rng.random() 落在哪个区间来决定。
        我们自己写而不用 random.choices，是因为还需要"抽中项的概率"
        来累计真实对数概率。
    【参数】
        rng     : random.Random 对象（自带种子的随机数发生器）
        weights : 字典 {选项: 非负权重}，至少有一个正权重
    【返回】(选项, 概率)，概率 = 该选项权重 / 总权重
    【边界】
        - 权重为 0 的选项永远不会被抽中。
        - 遍历顺序 = 字典插入顺序，所以同一个种子总得到同样的结果。
        - 总权重为 0 会抛出 ValueError（这是配置错误，不应悄悄继续）。
    """
    total = sum(weights.values())
    if total <= 0:
        raise ValueError("权重之和必须为正数")
    r = rng.random() * total
    acc = 0.0
    last = None
    for key, w in weights.items():
        if w <= 0:
            continue
        acc += w
        last = key
        if r < acc:
            return key, w / total
    # 浮点误差可能让 r 恰好等于 total，这时返回最后一个正权重的选项
    return last, weights[last] / total


def noun_form(animal, plural):
    """单复数形式：noun_form("fox", True) -> "foxes"。"""
    return NOUNS[animal] if plural else animal


# =====================================================================
# 3. 句子与生成
# =====================================================================

class Example:
    """
    一句话和它的"生成档案"。

    【属性】
        tokens       : 词列表，例如 ["the", "bat", "makes", "noise", "at", "night", "."]
        subject      : 主语动物（单数原形），例如 "bat"
        plural       : 主语是否复数（True / False）
        k            : 插入的介词短语个数
        attractors   : 每个干扰名词的 (动物, 是否复数)，按出现顺序
        verb         : 动词原形，例如 "make"
        verb_index   : 动词在 tokens 里的下标（从 0 开始）
        logprob_true : 这句话在真实分布下的自然对数概率 log p_true(句子)，
                       已经包括"句子在这里结束"（即 </s>）那一步（结束是确定的，概率 1）
    【Python 小知识】
        __slots__ 列出这个类允许有的属性，可以省内存、也能防止拼错属性名。
        __repr__ 决定 print(对象) 时显示什么。
    """

    __slots__ = ("tokens", "subject", "plural", "k", "attractors",
                 "verb", "verb_index", "logprob_true")

    def __init__(self, tokens, subject, plural, k, attractors, verb, verb_index, logprob_true):
        self.tokens = tokens
        self.subject = subject
        self.plural = plural
        self.k = k
        self.attractors = attractors
        self.verb = verb
        self.verb_index = verb_index
        self.logprob_true = logprob_true

    def text(self):
        return " ".join(self.tokens)

    def __repr__(self):
        return f"Example({self.text()!r}, k={self.k}, logp={self.logprob_true:.3f})"


def generate_sentence(rng, k=None, k_weights=DEFAULT_K_WEIGHTS, subject=None, plural=None,
                      attractor_plurals=None):
    """
    按上面的语法生成一句话。

    【思路】依次做 6 个随机选择：主语 → 单复数 → k 个介词短语 → 动词 → 宾语 → 时间短语，
        每一步都把"抽中项的概率"取对数累加到 logp 里。
    【参数】
        rng       : random.Random 对象
        k         : 介词短语个数。None 表示按 k_weights 随机抽（这时 k 的概率也计入 logp）；
                    给定整数时 k 是"条件"，不计入 logp
        k_weights : k = 0, 1, 2, ... 的相对权重
        subject   : 指定主语动物（None 表示随机）。指定时同样不计入 logp
        plural    : 指定主语单复数（None 表示随机）
        attractor_plurals : 指定每个干扰名词的单复数（长度必须等于 k）；None 表示随机。
                    用来专门造"干扰名词和主语单复数相反"的困难句子
    【返回】Example 对象
    【边界】
        - 只有当 k、subject、plural、attractor_plurals 都是 None 时，logprob_true 才是
          训练分布下完整的 log p(句子)；否则它是给定这些条件下的条件对数概率。
        - k 可以超过 k_weights 的长度（用于造训练中没见过的长句）。
    """
    logp = 0.0
    if subject is None:
        subject = rng.choice(ANIMALS)                 # 8 种动物等概率
        logp += math.log(1.0 / len(ANIMALS))
    if plural is None:
        plural = rng.random() < 0.5
        logp += math.log(0.5)
    if k is None:
        k_dict = {i: w for i, w in enumerate(k_weights)}
        k, pk = weighted_choice(rng, k_dict)
        logp += math.log(pk)

    tokens = ["the", noun_form(subject, plural)]
    attractors = []
    for j in range(k):
        prep = rng.choice(PREPOSITIONS)
        logp += math.log(1.0 / len(PREPOSITIONS))
        animal = rng.choice(ANIMALS)
        logp += math.log(1.0 / len(ANIMALS))
        if attractor_plurals is None:
            a_plural = rng.random() < 0.5
            logp += math.log(0.5)
        else:
            a_plural = attractor_plurals[j]
        attractors.append((animal, a_plural))
        tokens += [prep, "the", noun_form(animal, a_plural)]

    profile = PROFILES[subject]
    verb, pv = weighted_choice(rng, profile["verbs"])
    logp += math.log(pv)
    verb_index = len(tokens)
    singular_form, plural_form = VERBS[verb]
    tokens.append(plural_form if plural else singular_form)

    if verb in ANIMAL_OBJECT_VERBS:
        target, pt = weighted_choice(rng, profile["targets"])
        logp += math.log(pt)
        t_plural = rng.random() < 0.5
        logp += math.log(0.5)
        tokens += ["the", noun_form(target, t_plural)]
    else:
        obj, po = weighted_choice(rng, profile["objects"][verb])
        logp += math.log(po)
        tokens.append(obj)

    time_key, ptime = weighted_choice(rng, profile["times"])
    logp += math.log(ptime)
    tokens += TIMES[time_key]
    tokens.append(".")
    return Example(tokens, subject, plural, k, attractors, verb, verb_index, logp)


def make_splits(seed=0, n_train=20000, n_valid=2000, n_test=2000, k_weights=DEFAULT_K_WEIGHTS):
    """
    生成训练 / 验证 / 测试三份语料（同一个分布、互相独立的样本）。

    【参数】seed：随机种子；n_*：各份的句子数；k_weights：介词短语个数的分布
    【返回】字典 {"train": [Example...], "valid": [...], "test": [...]}
    【边界】
        - 同一个 seed 每次得到完全相同的语料（Python 3.9 和 3.13 也相同）。
        - 三份语料可能包含相同的句子（短句子组合有限，重复很正常），
          这和真实语料里"测试集里也会出现训练集见过的常用句"是一回事。
    """
    rng = random.Random(seed)
    out = {}
    for name, n in (("train", n_train), ("valid", n_valid), ("test", n_test)):
        out[name] = [generate_sentence(rng, k_weights=k_weights) for _ in range(n)]
    return out


def make_agreement_set(seed=123, k=1, n=400, incongruent=None):
    """
    造一组专门用来测"主谓一致"的句子：每句恰好有 k 个干扰名词。

    【思路】主语、单复数随机；干扰名词的单复数：
        incongruent=None  → 随机（自然分布）
        incongruent=True  → 全部和主语相反（最难：最近的名词会"误导"模型）
        incongruent=False → 全部和主语相同（最容易：就算看错了名词也能蒙对）
    【参数】seed、k（≥0）、n（句子数）、incongruent（见上）
    【返回】Example 列表
    【边界】k = 0 时没有干扰名词，incongruent 参数不起作用。
    """
    rng = random.Random(seed * 1000 + k)        # 不同 k 用不同的随机流，互不影响
    out = []
    for _ in range(n):
        plural = rng.random() < 0.5
        if incongruent is None or k == 0:
            att = None
        else:
            att = [(not plural) if incongruent else plural for _ in range(k)]
        out.append(generate_sentence(rng, k=k, plural=plural, attractor_plurals=att))
    return out


def agreement_pair(ex):
    """
    给定一个 Example，返回 (前缀词列表, 正确的动词形式, 错误的动词形式)。

    前缀 = 动词之前的所有词。评测时问模型：在这个前缀之后，
    p(正确形式) 是否大于 p(错误形式)？（这是 Linzen et al. 2016 的做法：
    只比较同一个动词的两种形式，所以和"模型想用哪个动词"无关。）
    """
    sg, pl = VERBS[ex.verb]
    correct, wrong = (pl, sg) if ex.plural else (sg, pl)
    return ex.tokens[:ex.verb_index], correct, wrong


def is_congruent(ex):
    """最近的那个干扰名词是否和主语单复数相同（k=0 时返回 None）。"""
    if ex.k == 0:
        return None
    return ex.attractors[-1][1] == ex.plural


# =====================================================================
# 4. 语法检查：判断一句（模型生成的）话是否合法、主谓是否一致
# =====================================================================

_FORM_TO_ANIMAL = {}
for _a in ANIMALS:
    _FORM_TO_ANIMAL[_a] = (_a, False)
    _FORM_TO_ANIMAL[NOUNS[_a]] = (_a, True)
_VERB_FORMS = {}
for _v, (_sg, _pl) in VERBS.items():
    _VERB_FORMS[_sg] = (_v, False)
    _VERB_FORMS[_pl] = (_v, True)
_TIME_PHRASES = [tuple(p) for p in TIMES.values()]


def check_sentence(tokens):
    """
    检查一个词列表是不是这门语言的合法句子。

    【思路】按语法从左到右"走"一遍：the 主语 (介词 the 名词)* 动词 宾语 时间 .
        任何一步对不上就判为不合法，并给出原因。
    【参数】tokens：词列表（不含 <s>、</s>）
    【返回】字典：
        well_formed : 结构是否合法（忽略单复数是否一致）
        agree       : 主谓是否一致（只有结构合法时才有意义，否则为 None）
        reason      : 不合法时的原因（字符串），合法时为 ""
    【边界】
        - 这个检查只看语法，不看"习性"：the fox makes milk at night . 是合法的
          （只是在真实分布里概率为 0——狐狸的 profile 里没有 make）。
        - 空列表、缺句号、多余的词都会判为不合法。
    """
    def bad(reason):
        return {"well_formed": False, "agree": None, "reason": reason}

    i = 0
    n = len(tokens)
    if n < 2 or tokens[0] != "the" or tokens[1] not in _FORM_TO_ANIMAL:
        return bad("开头不是 the + 动物")
    subj_plural = _FORM_TO_ANIMAL[tokens[1]][1]
    i = 2
    while i < n and tokens[i] in PREPOSITIONS:
        if i + 2 >= n or tokens[i + 1] != "the" or tokens[i + 2] not in _FORM_TO_ANIMAL:
            return bad("介词短语不完整")
        i += 3
    if i >= n or tokens[i] not in _VERB_FORMS:
        return bad("该出现动词的位置不是动词")
    verb, verb_plural = _VERB_FORMS[tokens[i]]
    i += 1
    if verb in ANIMAL_OBJECT_VERBS:
        if i + 1 >= n or tokens[i] != "the" or tokens[i + 1] not in _FORM_TO_ANIMAL:
            return bad("宾语应是 the + 动物")
        i += 2
    else:
        all_objects = set()
        for prof in PROFILES.values():
            for v, objs in prof["objects"].items():
                if v == verb:
                    all_objects.update(objs)
        if i >= n or tokens[i] not in all_objects:
            return bad("宾语不是这个动词能带的名词")
        i += 1
    rest = tuple(tokens[i:])
    if rest[-1:] != (".",) or rest[:-1] not in _TIME_PHRASES:
        return bad("时间/地点短语或句号不对")
    return {"well_formed": True, "agree": verb_plural == subj_plural, "reason": ""}


# =====================================================================
# 5. 词表：词 <-> 整数编号
# =====================================================================

class Vocab:
    """
    词表：把词映射成整数编号（神经网络只认数字），以及反过来。

    【思路】编号 0、1、2 固定留给 <pad>、<s>、</s>，其余的词按"第一次出现的顺序"编号。
        用出现顺序而不是 set 排序，是为了在任何 Python 版本、任何哈希种子下都一样。
    【用法】
        vocab = Vocab.from_examples(splits["train"])
        ids = vocab.encode(["the", "bat"])    # -> [3, 4]
        vocab.decode(ids)                     # -> ["the", "bat"]
        len(vocab)                            # 词表大小
    【边界】
        - 遇到词表里没有的词会抛出 KeyError：这门语言的词是封闭的，
          出现生词说明程序有 bug，宁可报错也不要悄悄映射成 <unk>。
    """

    def __init__(self, words):
        self.itos = [PAD, BOS, EOS]            # index to string：编号 -> 词
        for w in words:
            if w not in self.itos:
                self.itos.append(w)
        # 【Python 小知识】字典推导式 {键: 值 for ...}，enumerate 同时给出下标和元素
        self.stoi = {w: i for i, w in enumerate(self.itos)}   # string to index

    @classmethod
    def from_examples(cls, examples):
        """
        从语料里收集所有出现过的词，再补上语法里所有可能的词（保证词表封闭）。
        【Python 小知识】@classmethod 表示这是"类方法"：用 Vocab.from_examples(...) 调用，
            第一个参数 cls 就是 Vocab 这个类本身，return cls(...) 等于 return Vocab(...)。
        """
        words = []
        for ex in examples:
            words.extend(ex.tokens)
        words.extend(all_words())
        return cls(words)

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

    def encode(self, tokens):
        return [self.stoi[t] for t in tokens]

    def decode(self, ids):
        return [self.itos[i] for i in ids]

    @property
    def pad_id(self):
        return self.stoi[PAD]

    @property
    def bos_id(self):
        return self.stoi[BOS]

    @property
    def eos_id(self):
        return self.stoi[EOS]


def all_words():
    """这门语言里所有可能出现的词（按固定顺序，不含特殊记号）。"""
    words = ["the"]
    for a in ANIMALS:
        words += [a, NOUNS[a]]
    words += list(PREPOSITIONS)
    for sg, pl in VERBS.values():
        words += [sg, pl]
    for prof in PROFILES.values():
        for objs in prof["objects"].values():
            for o in objs:
                if o not in words:
                    words.append(o)
    for phrase in TIMES.values():
        for w in phrase:
            if w not in words:
                words.append(w)
    words.append(".")
    return words


def true_perplexity(examples):
    """
    用生成时记录下来的真实概率算"真实困惑度"：
        exp( - Σ log p_true(句子) / Σ 每句的预测次数 )
    每句的预测次数 = 词数 + 1（最后还要预测一次 </s>）。
    这是任何模型在这份数据上期望能达到的最好成绩（理论下界）。
    【边界】只对 make_splits 生成的句子有意义（它们的 logprob_true 是完整的联合概率）。
    """
    total_logp = math.fsum(ex.logprob_true for ex in examples)
    total_tokens = sum(len(ex.tokens) + 1 for ex in examples)
    return math.exp(-total_logp / total_tokens)


# =====================================================================
# 6. 直接运行：看看这门语言
# =====================================================================

if __name__ == "__main__":
    # 【Python 小知识】只有直接运行 python3 toy_corpus.py 时这一段才会执行；
    #   被别的文件 import 时不会执行。
    splits = make_splits(seed=0)
    print("前 8 句训练语料：")
    for ex in splits["train"][:8]:
        print(f"  k={ex.k}  {ex.text()}")
    vocab = Vocab.from_examples(splits["train"])
    print(f"\n词表大小（含 3 个特殊记号）：{len(vocab)}")
    lengths = [len(ex.tokens) for ex in splits["train"]]
    print(f"训练句子数：{len(lengths)}，平均长度 {sum(lengths) / len(lengths):.2f} 个词，"
          f"最短 {min(lengths)}，最长 {max(lengths)}")
    for name in ("train", "valid", "test"):
        print(f"{name:5s} 的真实困惑度（理论下界）：{true_perplexity(splits[name]):.3f}")
    bad = [ex for ex in splits["train"] if not check_sentence(ex.tokens)["agree"]]
    print(f"语法检查器认为不合法 / 不一致的训练句子数：{len(bad)}（应为 0）")