"""
=============================================================================
CS324 · 第 1 讲共享模块：fm_world.py —— 一个"知道真相"的迷你互联网
=============================================================================

【这个模块是干什么的 What this module is for】
    真实的基础模型（foundation model）在几万亿个 token 的网页、书籍、代码上做预训练。
    我们在笔记本电脑上做不到这一点，但可以造一个**迷你互联网**：几千行文本，混合了好几种"网站"，
    每一种都对应真实互联网上的一类内容：

        体裁 genre    例子                                              对应真实互联网里的……
        ─────────    ───────────────────────────────────────────────    ────────────────────
        bio          "kari lives in varno."                             维基百科式的人物介绍（世界知识）
        qa           "q: where does kari live? a: varno."               问答网站、FAQ（问答的"格式"）
        math         "37+45=82."                                        数学作业、计算题
        review       "review: the soup was tasty but the host was       商品 / 餐厅评论（带评分）
                      rude. rating: ok."
        story        "the pilot finished the work. then she went home." 小说、新闻里的叙述（带代词）

    因为这个世界是我们自己写出来的，所以我们**知道全部真相**：每个人住在哪里、做什么工作、
    哪些加法题在训练集里出现过、每个职业在文本里用 he / she 的比例是多少。
    这样就能精确地检查：模型从文本里学到了什么、没学到什么、学偏了什么。

【设计上几个有意为之的地方 Deliberate design choices】
    1. **只有一半的人有问答（qa）文本**。另一半人的事实只以叙述句（bio）出现。
       → 第 1 讲代码 02 检查：模型能不能把"在叙述句里学到的知识"用到"问答格式"上？
    2. **一部分两位数加法题从训练集里拿掉**（test_pairs）。
       → 模型是在背答案，还是学会了加法？
    3. **城市有一个隐藏属性 region（north / south），但它从来没在文本里出现过**。
       → 第 1 讲代码 03 用它检查：预训练表示里带着的知识，能不能用很少的标签"读"出来。
    4. **职业和代词的比例是我们故意设定的**（P_SHE）。这些比例是随意选的，
       **不对应任何真实世界的统计，也不代表任何观点**；它们只是用来演示
       "模型会复现、甚至放大训练数据里的统计偏差"这个机制。
    5. 评论的评分完全由两个形容词决定：食物好 + 服务好 → good；都差 → bad；一好一差 → ok。

【运行方式】
    python3 fm_world.py      # 打印语料的统计和几行样例（< 1 秒）

只用 Python 标准库。后面几讲会导入本模块（只读）。

【Python 小知识】
    - random.Random(seed)：一个**独立的**随机数生成器。用同一个 seed 创建，就会得到完全相同的随机序列。
      我们不用全局的 random.random()，是为了让"改动一个地方的随机调用"不影响其他地方。
    - @dataclass：自动生成 __init__ 等方法的类装饰器，适合"装一堆字段"的简单数据类。
"""

import random
import re
from dataclasses import dataclass, field

# =============================================================================
# 1. 字母表与世界的"设定"
# =============================================================================

# 迷你互联网里只会出现这些字符（44 个）。tiny_gpt.py 的字符级分词器就用它当词表。
# "\n" 是文档之间的分隔符（相当于真实语料里的 <|endoftext|>）。
ALPHABET = "\n +,.:=?0123456789abcdefghijklmnopqrstuvwxyz"

# 8 个虚构的城市。用虚构的名字，是为了避免模型"本来就知道"（我们的模型是从零训练的，
# 其实不会知道任何东西；但读者看到 paris 会自动带入真实知识，虚构名字更干净）。
CITIES = ["varno", "telmi", "osku", "brena", "quila", "damar", "sirok", "pelto"]

# 城市的隐藏属性：region。它**从来不出现在语料里**，只在第 1 讲代码 03 的"探针"实验中当标签用。
REGION = {
    "varno": "north", "telmi": "north", "osku": "north", "brena": "north",
    "quila": "south", "damar": "south", "sirok": "south", "pelto": "south",
}

JOBS = ["baker", "pilot", "farmer", "tailor", "potter", "singer", "judge", "doctor"]

# story 体裁里，每个职业后面用 "she" 的比例（其余用 "he"）。
# 这些数字是随意设定的（见文件开头第 4 点）：一半职业是 50/50，另一半有不同程度的偏斜。
P_SHE = {
    "baker": 0.50, "pilot": 0.85, "farmer": 0.50, "tailor": 0.15,
    "potter": 0.75, "singer": 0.50, "judge": 0.25, "doctor": 0.50,
}

# 评论体裁的词表：食物、员工、正面 / 负面形容词
FOODS = ["soup", "bread", "fish", "cake", "rice", "tea", "pie", "stew"]
FOOD_POS = ["tasty", "fresh", "warm", "sweet", "rich", "crisp", "great", "lovely"]
FOOD_NEG = ["bland", "stale", "cold", "burnt", "sour", "soggy", "awful", "greasy"]
STAFF = ["waiter", "staff", "cook", "host"]
SERVICE_POS = ["kind", "quick", "polite", "friendly"]
SERVICE_NEG = ["rude", "slow", "grumpy", "careless"]

# 人物介绍的 5 个模板。注意：城市总是出现在名字**后面**；职业有 3 个模板在名字后面、
# 1 个模板（"the {job} {name} lives in ..."）在名字**前面**。第 1 讲代码 02 会看到这个差别的影响。
BIO_TEMPLATES = [
    "{name} lives in {city}.",
    "{name} works as a {job}.",
    "{name} is a {job} from {city}.",
    "{name} moved to {city} years ago.",
    "the {job} {name} lives in {city}.",
]

# 问答的 2 个模板（只给一半的人生成）
QA_TEMPLATES = [
    ("q: where does {name} live? a:", "city"),
    ("q: what does {name} do? a:", "job"),
]

# story 体裁的 4 个模板。{pron} 是 he / she。第 1 讲代码 03 把它们当作 4 个不同的"应用"。
STORY_TEMPLATES = [
    "the {job} finished the work. then {pron} went home.",
    "the {job} was tired, so {pron} sat down.",
    "after the meeting the {job} said {pron} was happy.",
    "the {job} smiled because {pron} had won.",
]

# 生成名字用的音节：辅音 + 元音，两个音节，再随机加一个结尾辅音（或不加）
_ONSETS = ["b", "d", "f", "g", "k", "l", "m", "n", "p", "r", "s", "t", "v", "z"]
_VOWELS = ["a", "e", "i", "o", "u"]
_CODAS = ["", "n", "r", "l"]


# =============================================================================
# 2. 数据结构
# =============================================================================
@dataclass
class Person:
    """迷你世界里的一个人：名字、住的城市、职业。"""
    name: str
    city: str
    job: str


@dataclass
class World:
    """
    build_world() 的返回值：迷你互联网的全部内容和"真相"。

    【字段】
        people        : list[Person]，所有人（按名字排序）
        qa_names      : list[str]，有问答文本的那一半人的名字（排好序）
        test_pairs    : list[(a, b)]，**没有**出现在训练语料里的两位数加法题
        train_lines   : list[str]，训练语料（每个元素是一行 = 一篇"文档"）
        val_lines     : list[str]，验证语料（和训练语料来自同一个分布，用来算 bits per char）
        line_genres   : dict[str, str]，每一行属于哪个体裁（统计用）
        seed          : 生成这个世界用的随机种子
    """
    people: list
    qa_names: list
    test_pairs: list
    train_lines: list
    val_lines: list
    line_genres: dict = field(default_factory=dict)
    seed: int = 0

    def train_text(self):
        """把训练行拼成一整段文本，行与行之间用 "\\n" 隔开，末尾也加一个 "\\n"。"""
        return "\n".join(self.train_lines) + "\n"

    def val_text(self):
        """验证集文本，格式同上。"""
        return "\n".join(self.val_lines) + "\n"

    def person(self, name):
        """按名字查人。【边界】名字不存在时抛出 KeyError。"""
        for p in self.people:
            if p.name == name:
                return p
        raise KeyError(name)


# =============================================================================
# 3. 各个体裁的生成函数
# =============================================================================
def make_names(n, rng):
    """
    生成 n 个互不相同的虚构名字，比如 "kari"、"mosel"。

    【思路】两个"辅音 + 元音"音节，再随机接一个结尾辅音。不断生成，直到攒够 n 个不同的名字。
    【参数】n：要几个名字；rng：random.Random 对象
    【返回】按字母排序的名字列表（排序是为了让结果不依赖集合 set 的内部顺序）
    【边界】可能的名字一共 14×5×14×5×4 = 19600 个，n 不要超过几千，否则循环会很久。
            名字恰好和城市名、职业名或语料里的其他单词撞上时，会被跳过。
    """
    reserved = set(CITIES) | set(JOBS) | set(FOODS) | set(STAFF) | {"the", "then", "said", "more"}
    names = set()
    while len(names) < n:
        s = "".join(rng.choice(_ONSETS) + rng.choice(_VOWELS) for _ in range(2)) + rng.choice(_CODAS)
        if s not in reserved:
            names.add(s)
    # 【Python 小知识】sorted(set) 返回一个排好序的 list。set 本身的迭代顺序可能随运行变化（尤其是字符串），
    # 先排序才能保证每次运行结果一样。
    return sorted(names)


def bio_lines(person, repeats):
    """
    一个人的全部"人物介绍"：5 个模板，每个重复 repeats 次。

    【参数】person：Person；repeats：每个模板出现几次
    【返回】list[str]
    【FM 联系】同一个事实以不同的说法出现多次，是模型能"记住"并灵活使用这个事实的关键
            （Allen-Zhu & Li 的 "Physics of Language Models" 系列专门研究了这一点；第 5 讲再谈）。
    """
    out = []
    for template in BIO_TEMPLATES:
        line = template.format(name=person.name, city=person.city, job=person.job)
        out.extend([line] * repeats)
    return out


def qa_prompt(name, kind):
    """
    问答的"问题部分"（不含答案），比如 "q: where does kari live? a:"。

    【参数】name：名字；kind："city" 或 "job"
    【返回】字符串。答案的格式是：一个空格 + 答案 + 句号，例如 " varno."
    【边界】kind 不是 "city" / "job" 时抛出 ValueError。
    """
    for template, k in QA_TEMPLATES:
        if k == kind:
            return template.format(name=name)
    raise ValueError(f"unknown kind: {kind}")


def qa_lines(person, repeats):
    """一个人的问答文本：两个问题（住哪里、做什么），每个重复 repeats 次。"""
    out = []
    for kind in ("city", "job"):
        answer = person.city if kind == "city" else person.job
        out.extend([f"{qa_prompt(person.name, kind)} {answer}."] * repeats)
    return out


def addition_line(a, b):
    """一道加法题，例如 addition_line(37, 45) == "37+45=82."。"""
    return f"{a}+{b}={a + b}."


def make_review(rng, food_pos=None, service_pos=None, food_adjs=None):
    """
    随机生成一条评论（不含评分），以及它的"真相"。

    【思路】随机选食物好 / 差、服务好 / 差，各挑一个形容词；两者同向用 "and"，反向用 "but"。
            评分：都好 → good；都差 → bad；一好一差 → ok。
    【参数】
        rng         : random.Random
        food_pos    : True / False 指定食物好坏；None 表示随机
        service_pos : 同上，指定服务好坏
        food_adjs   : 可选，(正面列表, 负面列表)，限定食物形容词只从这两个列表里选
                      （第 1 讲代码 03 用它做"标注数据里没见过的形容词"实验）
    【返回】(text, rating, food_pos, service_pos)
            text   形如 "review: the soup was tasty but the host was rude."（以句号结尾，不含 rating）
            rating "good" / "ok" / "bad"
    【边界】food_adjs 里的列表不能为空，否则 rng.choice 会报 IndexError。
    """
    if food_pos is None:
        food_pos = rng.random() < 0.5
    if service_pos is None:
        service_pos = rng.random() < 0.5
    pos_list, neg_list = food_adjs if food_adjs is not None else (FOOD_POS, FOOD_NEG)
    food_adj = rng.choice(pos_list if food_pos else neg_list)
    service_adj = rng.choice(SERVICE_POS if service_pos else SERVICE_NEG)
    conj = "and" if food_pos == service_pos else "but"
    text = (f"review: the {rng.choice(FOODS)} was {food_adj} {conj} "
            f"the {rng.choice(STAFF)} was {service_adj}.")
    rating = "good" if (food_pos and service_pos) else ("bad" if not (food_pos or service_pos) else "ok")
    return text, rating, food_pos, service_pos


def review_line(rng):
    """一条完整的评论（带评分），例如 "review: ... rating: ok."。"""
    text, rating, _, _ = make_review(rng)
    return f"{text} rating: {rating}."


def story_line(rng, job=None):
    """
    一句带代词的叙述，例如 "the pilot finished the work. then she went home."。

    【思路】先选职业，再按 P_SHE[job] 的概率选 she，否则选 he，最后随机选一个模板。
    【参数】rng：random.Random；job：可选，指定职业
    【返回】(line, pronoun)
    """
    if job is None:
        job = rng.choice(JOBS)
    pron = "she" if rng.random() < P_SHE[job] else "he"
    template = rng.choice(STORY_TEMPLATES)
    return template.format(job=job, pron=pron), pron


# =============================================================================
# 4. 把所有体裁拼成一个"迷你互联网"
# =============================================================================
def build_world(seed=0, n_people=60, bio_repeats=4, qa_repeats=3,
                n_math=2000, n_reviews=2000, n_stories=2000, n_test_pairs=300, val_fraction=0.05):
    """
    生成整个迷你互联网。

    【思路】
        1. 生成 n_people 个人，随机分配城市和职业；偶数位置的人（一半）有问答文本。
        2. 所有两位数加法题（10–99 + 10–99，共 8100 道）打乱，前 n_test_pairs 道留作测试，
           训练语料从其余题目里**有放回地**抽 n_math 道（所以有的题出现多次、有的一次也没出现）。
        3. 生成评论、故事。
        4. 把所有行打乱，取 val_fraction 比例的行做验证集。
    【参数】seed：随机种子；其余参数控制各体裁的数量（默认值约 26 万个字符）
    【返回】World
    【边界】
        - 验证集是"同分布"的：bio 行在训练集里通常有一模一样的副本（每个模板重复了 bio_repeats 次），
          所以 bits per char 在 bio 部分会偏乐观。这和真实世界里"测试集泄漏到训练集"是同一类问题（第 2 讲）。
        - 评论的组合一共 8 × 16 × 4 × 8 = 4096 种，2000 条评论里会有重复，评测用的新评论也可能和训练评论相同。
    【FM 联系】真实的预训练语料也是这样：各种体裁按某个比例混在一起（data mixture），
            比例本身就是一个很重要的设计选择（第 2 讲）。
    """
    rng = random.Random(seed)
    names = make_names(n_people, rng)
    people = [Person(name=nm, city=rng.choice(CITIES), job=rng.choice(JOBS)) for nm in names]
    qa_names = [p.name for p in people[::2]]   # 【Python 小知识】[::2] 取下标 0, 2, 4, ... 的元素

    tagged = []   # (genre, line) 的列表
    for p in people:
        tagged += [("bio", line) for line in bio_lines(p, bio_repeats)]
        if p.name in qa_names:
            tagged += [("qa", line) for line in qa_lines(p, qa_repeats)]

    all_pairs = [(a, b) for a in range(10, 100) for b in range(10, 100)]
    rng.shuffle(all_pairs)
    test_pairs = all_pairs[:n_test_pairs]
    train_pairs = all_pairs[n_test_pairs:]
    for _ in range(n_math):
        a, b = rng.choice(train_pairs)
        tagged.append(("math", addition_line(a, b)))

    for _ in range(n_reviews):
        tagged.append(("review", review_line(rng)))
    for _ in range(n_stories):
        tagged.append(("story", story_line(rng)[0]))

    rng.shuffle(tagged)
    n_val = int(len(tagged) * val_fraction)
    val, train = tagged[:n_val], tagged[n_val:]
    line_genres = {}
    for genre, line in tagged:
        line_genres[line] = genre
    world = World(people=people, qa_names=qa_names, test_pairs=test_pairs,
                  train_lines=[line for _, line in train], val_lines=[line for _, line in val],
                  line_genres=line_genres, seed=seed)
    # 自检：语料里的每个字符都必须在 ALPHABET 里，否则分词器会出错
    bad = set(world.train_text() + world.val_text()) - set(ALPHABET)
    assert not bad, f"语料里出现了字母表之外的字符：{bad}"
    return world


def genre_counts(world):
    """
    统计训练语料里每个体裁有多少行、多少个字符。

    【返回】dict：genre -> (行数, 字符数)，按体裁名排序
    """
    stats = {}
    for line in world.train_lines:
        g = world.line_genres[line]
        n_lines, n_chars = stats.get(g, (0, 0))
        stats[g] = (n_lines + 1, n_chars + len(line) + 1)   # +1 是行尾的 "\n"
    return dict(sorted(stats.items()))


# =============================================================================
# 5. 检查一行生成的文本：格式对不对？内容真不真？
# =============================================================================
# 每个体裁的"模板"写成正则表达式。(\w+) 匹配一个单词，(\d+) 匹配一个整数。
# 【Python 小知识】re.compile 把正则表达式预先编译好；pattern.match(s) 从开头匹配，
# 成功时返回一个 Match 对象，m.groups() 是各个括号里匹配到的内容。^ 和 $ 表示"从头到尾整行匹配"。
_BIO_PATTERNS = [
    (re.compile(r"^(\w+) lives in (\w+)\.$"), ("name", "city")),
    (re.compile(r"^(\w+) works as a (\w+)\.$"), ("name", "job")),
    (re.compile(r"^(\w+) is a (\w+) from (\w+)\.$"), ("name", "job", "city")),
    (re.compile(r"^(\w+) moved to (\w+) years ago\.$"), ("name", "city")),
    (re.compile(r"^the (\w+) (\w+) lives in (\w+)\.$"), ("job", "name", "city")),
    (re.compile(r"^q: where does (\w+) live\? a: (\w+)\.$"), ("name", "city")),
    (re.compile(r"^q: what does (\w+) do\? a: (\w+)\.$"), ("name", "job")),
]
_MATH_PATTERN = re.compile(r"^(\d+)\+(\d+)=(\d+)\.$")
_REVIEW_PATTERN = re.compile(r"^review: the (\w+) was (\w+) (and|but) the (\w+) was (\w+)\. rating: (good|ok|bad)\.$")
_STORY_PATTERNS = [re.compile("^" + re.escape(t).replace(r"\{job\}", r"(\w+)").replace(r"\{pron\}", r"(he|she)") + "$")
                   for t in STORY_TEMPLATES]


def check_line(world, line):
    """
    判断一行文本（比如模型生成的一行）属于哪个体裁、格式是否合法、内容是否为真。

    【思路】按体裁逐个尝试正则模板：
        - 格式合法（well-formed）：匹配某个模板，并且每个槽位填的是"那一类"的已知词
          （名字必须是世界里真实存在的人，城市必须是 8 个城市之一……）
        - 内容为真（true）：
            bio / qa ：事实和这个人的真相一致（住的城市、职业）
            math     ：加法算对了
            review   ：and / but 和两个形容词的正负一致，评分和两个形容词一致
            story    ：没有"事实"可言，格式合法就算真
    【参数】world：World；line：一行文本（不含行尾的 "\\n"）
    【返回】(genre, well_formed, true)。genre 是 "bio"/"qa"/"math"/"review"/"story"，都不匹配时是 None；
            不匹配时 well_formed 和 true 都是 False；格式不合法时 true 一定是 False。
    【FM 联系】模型生成的句子可能"格式完美但内容是假的"——这就是 LLM 的**幻觉（hallucination）**，
            第 1 讲 §5 的"可靠性（reliability）"风险。这个函数让我们能给它计数。
    """
    names = {p.name: p for p in world.people}
    for pattern, slots in _BIO_PATTERNS:
        m = pattern.match(line)
        if m:
            genre = "qa" if line.startswith("q:") else "bio"
            vals = dict(zip(slots, m.groups()))
            ok = (vals["name"] in names and vals.get("city", CITIES[0]) in CITIES
                  and vals.get("job", JOBS[0]) in JOBS)
            if not ok:
                return genre, False, False
            p = names[vals["name"]]
            true = vals.get("city", p.city) == p.city and vals.get("job", p.job) == p.job
            return genre, True, true
    m = _MATH_PATTERN.match(line)
    if m:
        a, b, c = (int(g) for g in m.groups())
        return "math", True, a + b == c
    m = _REVIEW_PATTERN.match(line)
    if m:
        food, fadj, conj, staff, sadj, rating = m.groups()
        if food not in FOODS or staff not in STAFF or fadj not in FOOD_POS + FOOD_NEG \
                or sadj not in SERVICE_POS + SERVICE_NEG:
            return "review", False, False
        fp, sp = fadj in FOOD_POS, sadj in SERVICE_POS
        want_rating = "good" if (fp and sp) else ("bad" if not (fp or sp) else "ok")
        want_conj = "and" if fp == sp else "but"
        return "review", True, (rating == want_rating and conj == want_conj)
    for pattern in _STORY_PATTERNS:
        m = pattern.match(line)
        if m:
            return "story", m.group(1) in JOBS, m.group(1) in JOBS
    return None, False, False


# =============================================================================
# 6. 自检 / 演示
# =============================================================================
if __name__ == "__main__":
    w = build_world()
    text = w.train_text()
    print(f"训练语料：{len(w.train_lines)} 行，{len(text):,} 个字符；验证语料：{len(w.val_lines)} 行")
    print(f"字母表大小 |V| = {len(ALPHABET)}；语料里实际出现的字符数 = {len(set(text))}")
    print("\n各体裁（行数，字符数）：")
    for g, (nl, nc) in genre_counts(w).items():
        print(f"  {g:<7} {nl:>5} 行  {nc:>7,} 字符  ({nc / len(text):.1%})")
    print(f"\n人数 {len(w.people)}，其中有问答文本的 {len(w.qa_names)} 人；留作测试的加法题 {len(w.test_pairs)} 道")
    print("\n训练语料的前 12 行：")
    for line in w.train_lines[:12]:
        print("  " + line)
    p = w.people[0]
    print(f"\n第一个人的真相：{p}，region = {REGION[p.city]}（region 从不出现在语料里）")

    # check_line 自检：训练语料里的每一行都必须"格式合法且为真"
    bad = [line for line in w.train_lines if check_line(w, line)[1:] != (True, True)]
    assert not bad, bad[:3]
    examples = [f"{p.name} lives in {p.city}.", f"{p.name} lives in {'varno' if p.city != 'varno' else 'osku'}.",
                "12+30=43.", "review: the soup was tasty but the host was kind. rating: ok.", "zzz lives in varno."]
    print("\ncheck_line 示例（体裁, 格式合法, 内容为真）：")
    for line in examples:
        print(f"  {line:<62} → {check_line(w, line)}")