"""
=============================================================================
CS230 · 第 15 讲共享模块：三个合成的序列世界（都是在代码里现场生成的，不需要下载）
=============================================================================

对应讲义：lectures/lecture15_recurrent_neural_networks.md §6–7（恐龙名）、§8–10（主谓一致）、§11（Teddy）

-----------------------------------------------------------------------------
【为什么用合成数据】
-----------------------------------------------------------------------------
Coursera 的 Dinosaur Island 作业用的是一份真实的恐龙名单（约 1500 个名字）。这里不下载任何东西，
而是用**手写的音节规则**生成"恐龙名风格"的名字。好处和前几讲一样：**我们知道真相**。
因为规则是我们写的，可以精确算出任何一个名字在"真实分布"下的概率，于是能回答两个平常回答不了的问题：
    · 困惑度（perplexity）的下限是多少？—— 任何模型都不可能比"真实分布本身"更好（类似第 10 讲的 Bayes 误差）；
    · 模型采样出来的新名字，是"背下来的训练名字"、"符合规则的新名字"，还是"违反规则的乱码"？

三个世界：

    1. 恐龙名（§6–7，代码 03）
         名字 = k 个音节 + 连接字母（可选的 "o"）+ 后缀 + 换行符
         k ∈ {1, 2, 3}，概率 0.20 / 0.55 / 0.25
         音节 = 声母（23 种）+ 元音（6 种，有权重）+ 韵尾（"" 或 n / r / s / l）
         后缀：saurus（最常见）、raptor、don、ceratops、mimus、venator、titan、nykus、pteryx、suchus、dromeus、lophus
         例：ste·go + saurus = "stegosaurus"；bra + o + raptor = "braoraptor"。

    2. 主谓一致 "The cat(s), which ..., was / were"（§8–10，代码 05）
         "the" + 主语（cat / cats ……）+ "," + "which" + G 个随机词 + "," → 预测 was（单数）还是 were（复数）
         中间的 G 个词里也会出现 cat / cats 这类名词（"干扰项"），所以"看最近的名词"不管用，
         必须记住**第一个**名词的单复数，跨过 G 个词。

    3. Teddy Roosevelt vs Teddy bears（§11，代码 01）
         每个词标 1（人名的一部分）或 0。"teddy"、"april" 这类词是不是人名，**要看下一个词**：
         后面是姓（roosevelt、smith……）就是人名，后面是 bears、river……就不是。只看过去的单向 RNN 判断不了。

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

直接运行本文件会打印三个世界的例子和统计量（约 1 秒）：
    python3 rnn_data.py
"""

import math

import numpy as np

# =============================================================================
# 1. 恐龙名风格的名字
# =============================================================================
# 每个列表是 (字符串, 权重)。权重不必加起来等于 1，用的时候会归一化。
ONSETS = [(s, 1.0) for s in ("b", "br", "c", "ch", "d", "dr", "g", "k", "l", "m", "n", "p", "ph",
                             "pr", "r", "s", "sp", "st", "t", "th", "tr", "v", "z")]
VOWELS = [("a", 3.0), ("e", 2.0), ("i", 2.0), ("o", 3.0), ("u", 1.0), ("y", 1.0)]
CODAS = [("", 6.0), ("n", 1.0), ("r", 1.0), ("s", 1.0), ("l", 1.0)]
N_SYLLABLES = [(1, 0.20), (2, 0.55), (3, 0.25)]
LINKS = [("", 0.5), ("o", 0.5)]
SUFFIXES = [("saurus", 8.0), ("raptor", 2.0), ("don", 2.0), ("ceratops", 2.0), ("mimus", 1.0), ("venator", 1.0),
            ("titan", 1.0), ("nykus", 1.0), ("pteryx", 1.0), ("suchus", 1.0), ("dromeus", 1.0), ("lophus", 1.0)]

# 字符表：和 Coursera 作业一样，"\n"（名字结束，相当于 <EOS>）编号 0，a–z 编号 1–26
CHARS = ["\n"] + [chr(ord("a") + i) for i in range(26)]
CHAR_TO_IX = {ch: i for i, ch in enumerate(CHARS)}
IX_TO_CHAR = {i: ch for i, ch in enumerate(CHARS)}
VOCAB_SIZE = len(CHARS)                                     # 27


def _normalize(pairs):
    """把 [(值, 权重), ...] 变成 [(值, 概率), ...]，概率加起来等于 1。"""
    total = sum(w for _, w in pairs)
    return [(v, w / total) for v, w in pairs]


_ON, _VO, _CO = _normalize(ONSETS), _normalize(VOWELS), _normalize(CODAS)
_NS, _LI, _SU = _normalize(N_SYLLABLES), _normalize(LINKS), _normalize(SUFFIXES)


def _pick(pairs, rng):
    """按概率从 [(值, 概率), ...] 里抽一个值。"""
    probs = np.array([p for _, p in pairs])
    return pairs[int(rng.choice(len(pairs), p=probs))][0]


def make_dino_name(rng):
    """
    按音节规则生成一个名字（不含结尾的换行符）。

    【参数】rng：np.random.default_rng(...)
    【返回】字符串，例如 "stegosaurus"
    【Python 小知识】"".join(列表) 把一串字符串首尾相接；range(k) 产生 0, 1, ..., k−1。
    """
    k = _pick(_NS, rng)
    syllables = "".join(_pick(_ON, rng) + _pick(_VO, rng) + _pick(_CO, rng) for _ in range(k))
    return syllables + _pick(_LI, rng) + _pick(_SU, rng)


def make_dino_dataset(n_train, n_dev, rng):
    """
    生成训练集和 dev 集：每个名字都独立地从同一个规则分布里抽（允许偶尔重复，就像真实语料里的重复）。

    【参数】n_train、n_dev：名字个数；rng：随机数生成器
    【返回】(train, dev)：两个字符串列表
    【边界】没有去重：去重会改变分布（高概率的名字被"压扁"），那样"真实分布的困惑度"就不再是严格的下限了。
    """
    train = [make_dino_name(rng) for _ in range(n_train)]
    dev = [make_dino_name(rng) for _ in range(n_dev)]
    return train, dev


def _syllable_steps(s, i):
    """
    从位置 i 开始，s 能被读成哪些"一个音节"？返回 [(结束位置 j, 这个音节的概率), ...]。

    【思路】同一个字符串可能有多种切法（"tra" 可以是 t·ra 吗？不行，音节必须"声母 + 元音"开头；
            但 "stra" 可以是 st + r? 不行……）。最稳妥的办法是把所有切法都枚举出来，概率相加。
    """
    out = []
    for on, p_on in _ON:
        if not s.startswith(on, i):                         # str.startswith(前缀, 起点)：从 i 开始是否以 on 开头
            continue
        j1 = i + len(on)
        for vo, p_vo in _VO:
            if not s.startswith(vo, j1):
                continue
            j2 = j1 + len(vo)
            for co, p_co in _CO:
                if s.startswith(co, j2):
                    out.append((j2 + len(co), p_on * p_vo * p_co))
    return out


def dino_log_prob(name):
    """
    名字在"真实规则分布"下的对数概率 log P(name + "\\n")（自然对数）。

    【思路】一个名字可能有多种生成方式（音节怎么切、"o" 算连接字母还是音节的一部分），真实概率是所有方式的概率之和。
            用动态规划：f[n][j] = "前 j 个字符恰好是 n 个音节" 的概率之和，一步步往后推。
            然后对每个可能的音节结束位置 j，看剩下的 s[j:] 能不能读成"连接字母 + 后缀"。
    【参数】name：不含换行符的字符串
    【返回】float；不符合规则的名字返回 −inf
    【DL 联系】有了它就能算出"真实分布的困惑度"——任何语言模型在这份数据上都不可能比它更低（讲义 §6.4）。
    """
    L = len(name)
    f = [dict() for _ in range(4)]                          # f[n]：{位置 j: 概率}
    f[0][0] = 1.0
    for n in range(3):
        for i, p in f[n].items():                           # dict.items()：同时取出键和值
            for j, q in _syllable_steps(name, i):
                f[n + 1][j] = f[n + 1].get(j, 0.0) + p * q
    total = 0.0
    for k, p_k in _NS:
        for j, p in f[k].items():
            for li, p_li in _LI:
                if not name.startswith(li, j):
                    continue
                rest = name[j + len(li):]
                for su, p_su in _SU:
                    if rest == su:
                        total += p_k * p * p_li * p_su
    return math.log(total) if total > 0 else -math.inf


def is_valid_dino(name):
    """名字是否能由规则生成（概率 > 0）。"""
    return dino_log_prob(name) > -math.inf


def name_to_xy(name):
    """
    把一个名字变成语言模型的训练样本（Coursera 作业的写法，讲义 §6.2）。

    【返回】(X, Y)：X = [None, c1, ..., cn]，Y = [c1, ..., cn, "\\n"]（都是编号）。None 表示 x<1> = 全 0 向量。
    【Python 小知识】[CHAR_TO_IX[ch] for ch in name] 是列表推导式：对 name 的每个字符查编号，组成一个新列表。
    """
    ids = [CHAR_TO_IX[ch] for ch in name]
    return [None] + ids, ids + [CHAR_TO_IX["\n"]]


def indices_to_name(indices):
    """编号列表 → 字符串（去掉结尾的换行符）。"""
    return "".join(IX_TO_CHAR[i] for i in indices).rstrip("\n")


def classify_samples(samples, train_set):
    """
    把采样出来的名字分成三类：背下来的（和某个训练名字一模一样）/ 符合规则的新名字 / 违反规则的。

    【参数】samples：字符串列表；train_set：训练名字组成的 set
    【返回】dict：copied、novel_valid、invalid（三者是比例，加起来等于 1），distinct（不同名字的比例），mean_len（平均长度）
    【Python 小知识】set（集合）查找 "x in s" 是常数时间，比在 list 里逐个比较快得多。
    """
    n = len(samples)
    copied = sum(1 for s in samples if s in train_set)
    valid_new = sum(1 for s in samples if s not in train_set and is_valid_dino(s))
    return {"copied": copied / n, "novel_valid": valid_new / n, "invalid": (n - copied - valid_new) / n,
            "distinct": len(set(samples)) / n, "mean_len": float(np.mean([len(s) for s in samples]))}


# =============================================================================
# 2. 主谓一致：The cat(s), which ..., was / were
# =============================================================================
SUBJ_SING = ["cat", "dog", "bird", "child", "fox", "horse"]
SUBJ_PLUR = ["cats", "dogs", "birds", "children", "foxes", "horses"]
FILLERS = ["ate", "saw", "chased", "liked", "found", "the", "a", "some", "big", "small", "red", "old",
           "quickly", "already", "again", "today", "with", "near", "after", "fish", "apple", "apples",
           "mouse", "mice", "tree", "trees"]
AGREE_VOCAB = ["<pad>", "the", ",", "which"] + SUBJ_SING + SUBJ_PLUR + FILLERS
AGREE_IX = {w: i for i, w in enumerate(AGREE_VOCAB)}
CLAUSE_WORDS = FILLERS + SUBJ_SING + SUBJ_PLUR             # 从句里的词从这 38 个里均匀地抽（其中 12 个是干扰名词）
_SING_IDS = {AGREE_IX[w] for w in SUBJ_SING}
_PLUR_IDS = {AGREE_IX[w] for w in SUBJ_PLUR}


def make_agreement_batch(rng, n, gap):
    """
    生成 n 个"主谓一致"句子，每个句子的从句长度都是 gap。

    【思路】句子 = the + 主语 + , + which + gap 个随机词 + ,     → 标签：0 = was（单数），1 = were（复数）
            主语和标签之间隔着 gap + 3 个词。从句里约三分之一的词是 cat / cats 这样的名词，单复数随机，
            它们是"干扰项"：模型必须记住**第一个**名词，而不是最近的名词。
    【参数】rng：随机数生成器；n：句子个数；gap：从句长度 G（≥ 0）
    【返回】(X, y)：X 是 (n, gap + 5) 的 int64 数组（词的编号，PyTorch 的习惯：每行一个样本）；y 是 (n,) 的 0/1
    【边界】同一批里所有句子一样长，所以不需要 padding。不同的 batch 可以用不同的 gap。
    """
    plural = rng.integers(0, 2, n)
    noun = rng.integers(0, len(SUBJ_SING), n)
    clause = rng.integers(0, len(CLAUSE_WORDS), (n, gap))
    X = np.zeros((n, gap + 5), dtype=np.int64)
    X[:, 0] = AGREE_IX["the"]
    X[:, 1] = [AGREE_IX[(SUBJ_PLUR if p else SUBJ_SING)[k]] for p, k in zip(plural, noun)]
    X[:, 2] = AGREE_IX[","]
    X[:, 3] = AGREE_IX["which"]
    clause_ids = np.array([AGREE_IX[w] for w in CLAUSE_WORDS])
    X[:, 4:4 + gap] = clause_ids[clause]                    # 花式索引：用整数数组一次查好所有词的编号
    X[:, -1] = AGREE_IX[","]
    return X, plural.astype(np.int64)


def agreement_text(row):
    """把一行编号变回句子（打印例子用）。"""
    return " ".join(AGREE_VOCAB[i] for i in row)


def last_noun_baseline(X):
    """
    "看最近的名词"的基线：从句末往前找第一个 cat / cats 类的词，按它的单复数预测。

    【返回】(n,) 的 0/1 预测
    【DL 联系】如果一个模型只有"短期记忆"，它能做到的大概就是这个水平（讲义 §10.4）。
    """
    pred = np.zeros(len(X), dtype=np.int64)
    for r, row in enumerate(X):
        for tok in row[::-1]:                               # row[::-1]：倒序
            if tok in _SING_IDS or tok in _PLUR_IDS:
                pred[r] = 1 if tok in _PLUR_IDS else 0
                break
    return pred


# =============================================================================
# 3. Teddy Roosevelt vs Teddy bears：需要"看后文"的标注任务
# =============================================================================
AMBIG_FIRST = ["teddy", "april", "jordan", "paris", "georgia", "chelsea"]   # 可能是人名，也可能不是
PLAIN_FIRST = ["harry", "hermione", "alice", "bob"]                         # 一定是人名
SURNAMES = ["roosevelt", "smith", "lee", "brown", "garcia", "miller"]
OBJECTS = ["bears", "showers", "river", "hotel", "peaches", "streets"]
TAG_FILLER = ["he", "she", "said", "that", "the", "was", "on", "sale", "great", "we", "saw", "a",
              "today", "in", "love", "new", "visited", "big", "old", "with"]
TAG_VOCAB = ["<pad>"] + TAG_FILLER + AMBIG_FIRST + PLAIN_FIRST + SURNAMES + OBJECTS
TAG_IX = {w: i for i, w in enumerate(TAG_VOCAB)}
AMBIG_IDS = np.array([TAG_IX[w] for w in AMBIG_FIRST])


def make_tagging_batch(rng, n, length=10):
    """
    生成 n 个长度为 length 的句子，每个词标 1（人名的一部分）或 0。

    【思路】句子由随机的填充词组成，中间插入一个两词短语：
                35%：模糊名 + 姓      例 "teddy roosevelt"   → 1 1
                35%：模糊名 + 物品    例 "teddy bears"       → 0 0
                30%：普通名 + 姓      例 "harry smith"       → 1 1
            "teddy" 是不是人名只取决于**下一个词**。姓一定是人名，普通名（harry）也一定是人名。
    【参数】rng；n：句子数；length：句子长度（≥ 2）
    【返回】(X, Y, ambig_mask)：X (n, length) 词编号；Y (n, length) 0/1 标签；
            ambig_mask (n, length) 布尔，标出"模糊名"所在的位置（评估时单独统计这些位置）
    """
    X = rng.integers(1, 1 + len(TAG_FILLER), (n, length))   # 填充词的编号是 1..20
    Y = np.zeros((n, length), dtype=np.int64)
    kind = rng.choice(3, size=n, p=[0.35, 0.35, 0.30])
    pos = rng.integers(0, length - 1, n)                    # 短语占 pos、pos + 1 两个位置
    first_amb = rng.integers(0, len(AMBIG_FIRST), n)
    first_plain = rng.integers(0, len(PLAIN_FIRST), n)
    second_sur = rng.integers(0, len(SURNAMES), n)
    second_obj = rng.integers(0, len(OBJECTS), n)
    for r in range(n):
        p = pos[r]
        if kind[r] == 0:
            X[r, p], X[r, p + 1] = TAG_IX[AMBIG_FIRST[first_amb[r]]], TAG_IX[SURNAMES[second_sur[r]]]
            Y[r, p] = Y[r, p + 1] = 1
        elif kind[r] == 1:
            X[r, p], X[r, p + 1] = TAG_IX[AMBIG_FIRST[first_amb[r]]], TAG_IX[OBJECTS[second_obj[r]]]
        else:
            X[r, p], X[r, p + 1] = TAG_IX[PLAIN_FIRST[first_plain[r]]], TAG_IX[SURNAMES[second_sur[r]]]
            Y[r, p] = Y[r, p + 1] = 1
    ambig_mask = np.isin(X, AMBIG_IDS)                      # np.isin(a, 集合)：a 的每个元素是否在集合里
    return X, Y, ambig_mask


def tagging_text(row, labels=None):
    """把一行编号变回句子；给了 labels 时在人名的词后面加 /1。"""
    words = [TAG_VOCAB[i] for i in row]
    if labels is None:
        return " ".join(words)
    return " ".join(w + ("/1" if l else "") for w, l in zip(words, labels))


# =============================================================================
# 自检 / 打印例子
# =============================================================================
if __name__ == "__main__":
    rng = np.random.default_rng(0)
    train, dev = make_dino_dataset(1500, 500, rng)
    print("恐龙名（前 10 个训练名字）：", ", ".join(train[:10]))
    lp = sum(dino_log_prob(s) for s in dev)
    n_chars = sum(len(s) + 1 for s in dev)
    print(f"  训练 {len(train)} 个，不同名字 {len(set(train))} 个；平均长度 {np.mean([len(s) for s in train]):.2f}；"
          f"dev 上真实分布的困惑度 {math.exp(-lp / n_chars):.3f}")
    print(f"  'stegosaurus' 的对数概率 {dino_log_prob('stegosaurus'):.3f}；'banana' 合法吗？{is_valid_dino('banana')}")
    X, y = make_agreement_batch(rng, 3, 6)
    for row, lab in zip(X, y):
        print("主谓一致：", agreement_text(row), "→", "were" if lab else "was")
    Xb, _ = make_agreement_batch(rng, 4000, 10)
    print(f"  间隔 10 时，含干扰名词的句子比例 {np.mean([any(t in _SING_IDS | _PLUR_IDS for t in r[4:-1]) for r in Xb]):.3f}")
    Xt, Yt, M = make_tagging_batch(rng, 3)
    for row, lab in zip(Xt, Yt):
        print("Teddy：", tagging_text(row, lab))