"""
=============================================================================
第 1 讲 · 工具模块：一个虚构的"罕见病百科"
=============================================================================

对应讲义：lectures/lecture01_intro_hallucination.md 第 0.3 节（本讲的玩具世界）

-----------------------------------------------------------------------------
【为什么要造一个虚构的世界？】
-----------------------------------------------------------------------------
本讲要研究的是"LLM 什么时候会编造事实"。如果用真实世界的知识做实验，我们说不清
模型的某个回答是"真的记得"还是"碰巧猜对"，也说不清它在训练时见过这条知识多少次。

所以我们自己造一个世界，并且**完全掌握**下面三件事：
    1. 每条事实的标准答案是什么（才能判断模型说的对不对）；
    2. 每种疾病在"预训练语料"里出现过多少次（才能研究"见得多 → 记得牢"）；
    3. 每种疾病是哪一年被描述的（才能模拟"知识截止日期之后的新知识"）。

选"罕见病"做主题，是因为它正好是本课作业 1 的主题（罕见病诊断），也是幻觉代价最高的
领域之一：全世界约 4 亿人患有 7000 多种罕见病中的一种，很多病的资料只存在于几篇论文里。

-----------------------------------------------------------------------------
【这个世界里有什么？】
-----------------------------------------------------------------------------
60 种虚构的疾病（名字由音节随机拼成，例如 "Velmora syndrome"），分三档：

    头部（head）  ：20 种。"预训练语料"里出现过几十到两千次 —— 相当于教科书里的常见病。
    长尾（tail）  ：25 种。只出现过 0~12 次 —— 相当于只在几篇论文里提过的病。
    新近（recent）：15 种。2025 年以后才被首次描述，模型的知识截止于 2024 年，一次也没见过。

每种疾病有 7 条可以提问的事实（RELATIONS）：标志性症状、起病年龄、致病基因、遗传方式、
首次描述年份、首次报告者、患病率；另外还有 3 个"其他症状"，用于诊断演示。

每种疾病有一篇百科文章（9 句话），被切成 5 段（passage），供检索器使用。
另外还能生成两类问题（直接问法 / 换说法）和一组"病人"（只用症状描述）。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    from toy_world import make_world
    world = make_world(seed=0)
    d = world.diseases[0]
    print(d.name, d.tier, d.mentions, d.facts["gene"])
    for p in world.passages_of(d.name):
        print(p.pid, p.text)

直接运行本文件会打印世界的概况和一篇示例文章：
    python toy_world.py
"""

# 【Python 小知识】
#   dataclass：一个"装饰器"（写在类定义上方、以 @ 开头）。被它装饰的类会自动生成 __init__ 等方法，
#              适合用来定义"只装数据"的类。例如
#                  @dataclass
#                  class Point:
#                      x: int
#                      y: int
#              之后就可以写 Point(1, 2)，并用 p.x、p.y 访问字段。冒号后面的 int 是"类型注解"，
#              只是给人看的说明，Python 运行时不会检查。
#   field(default_factory=list)：给字段一个"每个对象各自独立的空列表"作为默认值。
#              不能直接写 = []，因为那样所有对象会共用同一个列表（Python 的经典陷阱）。
import random
import re
from dataclasses import dataclass, field

# 模拟 LLM 的"知识截止年份"：2024 年及以前描述的疾病可能出现在预训练语料里，之后的一定没有。
LLM_CUTOFF_YEAR = 2024

# 30 个症状（表型）。真实的罕见病诊断使用"人类表型本体"（HPO）里的上万个标准术语，这里只取一小部分。
PHENOTYPES = [
    "seizures", "hypotonia", "ataxia", "cataracts", "hearing loss", "short stature",
    "intellectual disability", "muscle weakness", "scoliosis", "retinal degeneration",
    "cardiomyopathy", "liver enlargement", "skin fragility", "joint hypermobility", "kidney cysts",
    "microcephaly", "developmental delay", "spasticity", "peripheral neuropathy", "recurrent infections",
    "anemia", "photosensitivity", "hair loss", "cleft palate", "tremor",
    "chronic fatigue", "bone fractures", "early-onset obesity", "dystonia", "vision loss",
]

# 几个"取值范围很小"的属性，以及它们在真实罕见病里的大致比例（数字是为了演示而设的，不是统计结果）。
# 模型编造答案时会按这个比例去"猜"—— 越常见的值越容易被猜中，这正是"幻觉看起来很合理"的原因。
INHERITANCE_PRIOR = {"autosomal recessive": 0.45, "autosomal dominant": 0.35,
                     "X-linked": 0.15, "mitochondrial": 0.05}
ONSET_PRIOR = {"the neonatal period": 0.15, "infancy": 0.30, "early childhood": 0.30,
               "adolescence": 0.15, "adulthood": 0.10}
PREVALENCE_PRIOR = {"1 in 20,000": 0.10, "1 in 50,000": 0.20, "1 in 100,000": 0.25,
                    "1 in 250,000": 0.20, "1 in 500,000": 0.15, "1 in 1,000,000": 0.10}

# 7 种可以提问的关系。每种关系有：
#   sentence   ：百科文章里陈述这条事实的句子模板（{d} = 疾病名，{v} = 值）。模型生成的"声明"也用同一个模板。
#   direct     ：直接问法 —— 和文章用词有重叠（检索器容易找到正确的段落）。
#   paraphrase ：换说法 —— 和文章几乎没有共同的词（检索器很难找到正确的段落）。
#   label      ：中文名，打印用。
RELATIONS = {
    "hallmark": {
        "sentence": "The hallmark feature of {d} is {v}.",
        "direct": "What is the hallmark feature of {d}?",
        "paraphrase": "What is the most characteristic sign of {d}?",
        "label": "标志性症状",
    },
    "onset": {
        "sentence": "Symptoms of {d} usually begin in {v}.",
        "direct": "When do symptoms of {d} usually begin?",
        "paraphrase": "At what age does {d} typically appear?",
        "label": "起病年龄",
    },
    "gene": {
        "sentence": "{d} is caused by mutations in the {v} gene.",
        "direct": "Which gene is mutated in {d}?",
        "paraphrase": "What is the genetic cause of {d}?",
        "label": "致病基因",
    },
    "inheritance": {
        "sentence": "{d} follows {v} inheritance.",
        "direct": "What kind of inheritance does {d} follow?",
        "paraphrase": "How is {d} passed down in families?",
        "label": "遗传方式",
    },
    "year": {
        "sentence": "{d} was first described in {v}.",
        "direct": "When was {d} first described?",
        "paraphrase": "In what year was {d} discovered?",
        "label": "首次描述年份",
    },
    "discoverer": {
        "sentence": "The first cases of {d} were reported by {v}.",
        "direct": "Who reported the first cases of {d}?",
        "paraphrase": "Which physician discovered {d}?",
        "label": "首次报告者",
    },
    "prevalence": {
        "sentence": "{d} affects about {v} people.",
        "direct": "How many people does {d} affect?",
        "paraphrase": "How rare is {d}?",
        "label": "患病率",
    },
}
QA_RELATIONS = list(RELATIONS)          # ['hallmark', 'onset', 'gene', ...]，固定的顺序

# 文章里另外两种句子（不作为问答的关系）：
INTRO_TEMPLATE = "{d} is a rare genetic disorder."
FEATURES_TEMPLATE = "Other features of {d} include {f1}, {f2} and {f3}."

# 文章的句子顺序（仿照维基百科：定义 → 症状 → 病因 → 历史 → 流行病学），以及怎样切成段落。
# 每段 2 句话，最后一段 1 句话。注意 "rare genetic disorder" 这句开头语和"症状"放在同一段 ——
# 换说法的问题（"genetic cause""how rare"）会被这一段吸引过去，这是讲义第 4 节要讲的检索陷阱。
ARTICLE_ORDER = ["intro", "hallmark", "features", "onset", "gene", "inheritance",
                 "year", "discoverer", "prevalence"]
SENTENCES_PER_PASSAGE = 2

TIERS = ("head", "tail", "recent")
TIER_LABELS = {"head": "头部", "tail": "长尾", "recent": "新近"}

# 拼名字用的音节。全是虚构的，避免和真实疾病重名（这样任何真实的 LLM 也不可能"记得"它们）。
_SYL_START = ["vel", "kor", "mar", "tes", "lun", "bra", "dor", "fen", "gal", "hos", "ist", "jor",
              "kel", "lim", "nov", "pel", "quin", "ros", "sal", "tor", "var", "wen", "xan", "zel"]
_SYL_MID = ["a", "e", "i", "o", "u", "ar", "en", "is", "or", "ul"]
_SYL_END = ["mora", "dith", "vane", "sk", "nar", "lix", "tos", "brin", "dela", "rum", "phy", "zan"]
_FIRST_NAMES = ["Amara", "Bastian", "Celine", "Dmitri", "Elif", "Farid", "Greta", "Hiro", "Ines", "Jonas",
                "Kaia", "Lior", "Mina", "Nils", "Oona", "Pavel", "Rhea", "Soren", "Talia", "Viktor"]
_LAST_END = ["er", "ova", "sen", "ani", "ez", "berg", "eau", "itz"]


# =============================================================================
# 第 1 部分：数据结构
# =============================================================================
@dataclass
class Disease:
    """
    一种虚构的疾病。

    【字段】
        name     : 疾病名，例如 "Velmora syndrome"
        tier     : "head" / "tail" / "recent"
        mentions : 在模拟的"预训练语料"里出现的次数（新近疾病为 0）
        facts    : dict，关系名 → 值，例如 {"gene": "VLK2", "year": "1987", ...}（值都是字符串）
        features : list[str]，3 个"其他症状"（不含标志性症状）
    """
    name: str
    tier: str
    mentions: int
    facts: dict
    features: list

    @property
    def phenotypes(self):
        """
        这种病的全部症状：标志性症状 + 3 个其他症状，共 4 个。

        【Python 小知识】@property 让方法可以像字段一样访问：写 d.phenotypes，不用加括号。
        """
        return [self.facts["hallmark"]] + list(self.features)

    @property
    def year(self):
        return int(self.facts["year"])


@dataclass
class Passage:
    """
    检索的基本单位：百科文章里的一段（1~2 句话）。

    【字段】
        pid       : 段落编号，例如 "Velmora syndrome #2"（文章标题 + 段号）
        title     : 文章标题（= 疾病名）
        sentences : list[str]，这一段的句子
    """
    pid: str
    title: str
    sentences: list

    @property
    def text(self):
        return " ".join(self.sentences)


@dataclass
class Question:
    """
    一道问答题。

    【字段】
        text   : 问题文本
        d      : 问的是哪种疾病（疾病名）
        rel    : 问的是哪种关系（RELATIONS 的键）
        answer : 标准答案（字符串）
        tier   : 这种疾病的档次
        kind   : "direct"（直接问法）或 "paraphrase"（换说法）
    """
    text: str
    d: str
    rel: str
    answer: str
    tier: str
    kind: str


@dataclass
class Patient:
    """
    一个病人：只有一组症状（和作业 1 一样，病人只用"临床表现的列表"来描述）。

    【字段】
        pid      : 病人编号，例如 "P03"
        findings : list[str]，观察到的症状
        disease  : 真实的诊断（只用于评估；诊断算法不许偷看）
    """
    pid: str
    findings: list
    disease: str


# =============================================================================
# 第 2 部分：句子模板 ⇄ 结构化事实
# =============================================================================
def fact_sentence(rel, d, v):
    """
    把一条结构化的事实 (疾病, 关系, 值) 写成一句英文。

    【参数】rel：RELATIONS 的键；d：疾病名；v：值（字符串）
    【返回】str，例如 fact_sentence("gene", "Velmora syndrome", "VLK2")
            → "Velmora syndrome is caused by mutations in the VLK2 gene."
    【边界】rel 不在 RELATIONS 里会抛 KeyError。
    """
    return RELATIONS[rel]["sentence"].format(d=d, v=v)


def _template_to_regex(template):
    """
    把句子模板变成正则表达式：{d} → 命名分组 d，{v} → 命名分组 v，其余部分按字面匹配。

    【思路】
        先用 re.split 把模板切成"占位符"和"普通文字"交替的列表，普通文字用 re.escape 转义
        （模板里的 "." "?" 在正则里有特殊含义，必须转义），占位符换成 (?P<名字>.+?)。

    【例子】"{d} follows {v} inheritance."
            → r"^(?P<d>.+?)\\ follows\\ (?P<v>.+?)\\ inheritance\\.$"

    【Python 小知识】
        (?P<name>...) 是"命名分组"，匹配后可以用 m.group("name") 取出对应的内容。
        .+? 是"非贪婪"匹配：尽量少地匹配字符，配合首尾的 ^ $ 可以准确地把句子切开。
        re.split 的模式外面加一层括号 ( )，切出来的结果里会保留分隔符本身。
    """
    parts = re.split(r"(\{d\}|\{v\})", template)
    regex = ""
    for part in parts:
        if part == "{d}":
            regex += r"(?P<d>.+?)"
        elif part == "{v}":
            regex += r"(?P<v>.+?)"
        else:
            regex += re.escape(part)
    return re.compile("^" + regex + "$")


# 预先编译好所有正则（模块被导入时只算一次）。
_SENTENCE_RE = {rel: _template_to_regex(spec["sentence"]) for rel, spec in RELATIONS.items()}
_QUESTION_RE = {(rel, kind): _template_to_regex(spec[kind])
                for rel, spec in RELATIONS.items() for kind in ("direct", "paraphrase")}
_TOPIC_RE = re.compile(r"^Tell me about (?P<d>.+?)\.$")


def parse_claim(sentence):
    """
    把一句话解析回结构化的事实 (疾病, 关系, 值)。是 fact_sentence 的逆操作。

    【参数】sentence：str
    【返回】(d, rel, v) 三元组；如果这句话不符合任何关系的模板，返回 None
    【边界】
        - 只认识 RELATIONS 里的 7 种句式；开头语、"其他症状"那句话都返回 None。
        - 这是玩具世界的"作弊"之处：真实系统里，把一句自由文本变成结构化事实本身就是一个
          容易出错的 LLM 任务（第 6~9 讲的语义解析、信息抽取都在研究它）。
    """
    sentence = sentence.strip()
    for rel, regex in _SENTENCE_RE.items():
        m = regex.match(sentence)
        if m:
            return m.group("d"), rel, m.group("v")
    return None


_FEATURES_RE = re.compile(r"^Other features of (?P<d>.+?) include (?P<f1>.+?), (?P<f2>.+?) and (?P<f3>.+?)\.$")


def parse_features(sentence):
    """
    解析"其他症状"那句话："Other features of X include A, B and C."

    【返回】(疾病名, [A, B, C])；不是这种句子时返回 None
    """
    m = _FEATURES_RE.match(sentence.strip())
    return (m.group("d"), [m.group("f1"), m.group("f2"), m.group("f3")]) if m else None


def parse_question(text):
    """
    认出一道题问的是哪种疾病的哪种关系。

    【返回】(d, rel)；认不出时返回 None
    【边界】模拟 LLM 用它来"理解"问题。真实的 LLM 理解问题靠的是语言能力，不需要模板。
    """
    text = text.strip()
    for (rel, _kind), regex in _QUESTION_RE.items():
        m = regex.match(text)
        if m:
            return m.group("d"), rel
    return None


def topic_utterance(d):
    """用户在对话里的一句话："Tell me about <疾病>."（第 03 号文件的对话演示用）"""
    return f"Tell me about {d}."


def parse_topic(utterance):
    """从 "Tell me about X." 里取出 X；格式不对时返回 None。"""
    m = _TOPIC_RE.match(utterance.strip())
    return m.group("d") if m else None


def split_sentences(text):
    """
    把一段文字切成句子：在句号后面的空格处切开。

    【边界】很粗糙：遇到 "Dr. Amara" 这种缩写里的句号会切错 —— 所以本世界里的人名写成
        "Amara Kessler, MD"，不用 "Dr."。真实系统里分句也是一个小小的坑。
    """
    return [s for s in re.split(r"(?<=\.)\s+", text.strip()) if s]


# =============================================================================
# 第 3 部分：生成世界
# =============================================================================
def weighted_choice(rng, prior):
    """按 prior（值 → 概率）的比例随机选一个值。rng 是 random.Random 对象。"""
    values = list(prior)
    return rng.choices(values, weights=[prior[v] for v in values], k=1)[0]


def make_disease_name(rng, taken):
    """
    拼一个没用过的疾病名，例如 "Velmora syndrome"、"Korisk disease"。

    【参数】rng：random.Random；taken：已经用过的名字的集合（会被修改：新名字加进去）
    【返回】str
    """
    while True:
        stem = rng.choice(_SYL_START) + rng.choice(_SYL_MID) + rng.choice(_SYL_END)
        name = stem.capitalize() + " " + rng.choice(["syndrome", "syndrome", "disease"])
        if stem not in {n.split()[0].lower() for n in taken}:      # 词干也不许重复，避免 "X syndrome" 和 "X disease" 并存
            taken.add(name)
            return name


def make_gene_symbol(rng, taken):
    """拼一个没用过的基因符号：3 个大写字母 + 1 个数字，例如 "VLK2"。"""
    letters = "BCDFGHKLMNPRSTVXZ"
    while True:
        sym = "".join(rng.choice(letters) for _ in range(3)) + str(rng.randint(1, 9))
        if sym not in taken:
            taken.add(sym)
            return sym


def make_person_name(rng):
    """拼一个虚构的医生名字，例如 "Amara Kessler, MD"。"""
    last = (rng.choice(_SYL_START) + rng.choice(_LAST_END)).capitalize()
    return f"{rng.choice(_FIRST_NAMES)} {last}, MD"


class World:
    """
    整个虚构世界：疾病、百科文章（切成段落）、问题、病人。

    【构造】一般不直接构造，用 make_world(seed) 得到。

    【常用属性和方法】
        world.diseases            : list[Disease]
        world.by_name[name]       : 按名字查 Disease
        world.passages            : list[Passage]，全部段落（检索器的语料）
        world.passages_of(name)   : 某种疾病的全部段落
        world.questions(kind)     : 全部问题（每种疾病 × 每种关系 = 60 × 7 = 420 道）
        world.is_true(d, rel, v)  : 判断一条事实是否正确
        world.kb_snapshot(year)   : 某一年建成的"表型知识库"（只含当年已经描述过的疾病）
        world.make_cohort(...)    : 生成一组病人
    """

    def __init__(self, diseases, seed):
        self.seed = seed
        self.diseases = diseases
        self.by_name = {d.name: d for d in diseases}
        self.passages = []
        self._passages_by_title = {}
        for d in diseases:
            ps = self._build_article(d)
            self.passages.extend(ps)
            self._passages_by_title[d.name] = ps

    # ----- 百科文章 -----
    def article_sentences(self, d):
        """
        一种疾病的百科文章：9 句话的列表，顺序见 ARTICLE_ORDER。

        【参数】d：Disease
        """
        out = []
        for part in ARTICLE_ORDER:
            if part == "intro":
                out.append(INTRO_TEMPLATE.format(d=d.name))
            elif part == "features":
                f1, f2, f3 = d.features
                out.append(FEATURES_TEMPLATE.format(d=d.name, f1=f1, f2=f2, f3=f3))
            else:
                out.append(fact_sentence(part, d.name, d.facts[part]))
        return out

    def _build_article(self, d):
        """把文章每 SENTENCES_PER_PASSAGE 句切成一段。"""
        sents = self.article_sentences(d)
        passages = []
        for j in range(0, len(sents), SENTENCES_PER_PASSAGE):
            passages.append(Passage(pid=f"{d.name} #{j // SENTENCES_PER_PASSAGE}", title=d.name,
                                    sentences=sents[j:j + SENTENCES_PER_PASSAGE]))
        return passages

    def passages_of(self, name):
        return self._passages_by_title[name]

    def gold_passage(self, d, rel):
        """
        含有 (d, rel) 这条事实的那一段。用来判断"检索有没有找到答案所在的段落"。

        【返回】Passage
        """
        sentence = fact_sentence(rel, d, self.by_name[d].facts[rel])
        for p in self._passages_by_title[d]:
            if sentence in p.sentences:
                return p
        raise ValueError(f"找不到 {d} 的 {rel}")           # 正常情况下不会发生

    # ----- 问题 -----
    def questions(self, kind="direct", tiers=TIERS):
        """
        生成全部问题：每种疾病 × 每种关系各一道。

        【参数】
            kind  : "direct"（直接问法）或 "paraphrase"（换说法）
            tiers : 只要哪些档次的疾病，默认全部
        【返回】list[Question]
        """
        out = []
        for d in self.diseases:
            if d.tier not in tiers:
                continue
            for rel in QA_RELATIONS:
                out.append(Question(text=RELATIONS[rel][kind].format(d=d.name), d=d.name, rel=rel,
                                    answer=d.facts[rel], tier=d.tier, kind=kind))
        return out

    # ----- 判分 -----
    def is_true(self, d, rel, v):
        """
        (d, rel, v) 这条事实对不对。

        【边界】d 不是本世界的疾病时返回 False（模型编造了一种不存在的病，当然不对）。
            比较时忽略大小写和首尾空格。
        """
        disease = self.by_name.get(d)
        if disease is None or rel not in disease.facts:
            return False
        return str(v).strip().lower() == disease.facts[rel].strip().lower()

    def claim_is_true(self, sentence):
        """一句"声明"对不对。解析不了的句子算作不对（返回 False）。"""
        parsed = parse_claim(sentence)
        return parsed is not None and self.is_true(*parsed)

    # ----- 诊断用的知识库与病人 -----
    def kb_snapshot(self, built_in_year=LLM_CUTOFF_YEAR):
        """
        某一年建成的"表型知识库"：疾病名 → 症状集合。只包含那一年以前已经被描述过的疾病。

        【思路】真实世界的表型知识库（例如 HPO 的疾病注释、Orphanet）也是定期更新的。
            作业 1 的讲义里举过例子：ReNU 综合征 2024 年才被首次描述，在那之前建成的数据库里根本没有它。
        【返回】dict[str, set[str]]
        """
        return {d.name: set(d.phenotypes) for d in self.diseases if d.year <= built_in_year}

    def make_cohort(self, target, n=10, seed=0, noise=1):
        """
        生成一组 n 个病人，其中**恰好一个**患有 target 这种病（作业 1 的设定）。

        【参数】
            target : 目标疾病名
            n      : 病人数
            seed   : 随机种子
            noise  : 每个病人额外带几个"无关症状"（真实病历里总有和主诊断无关的发现）
        【返回】list[Patient]（顺序已打乱）
        【边界】其余 n − 1 个病人的疾病从 2024 年以前描述过的疾病里选，彼此不重复。
            每个病人的症状 = 标志性症状（80% 的概率出现）+ 3 个其他症状中的 2 个 + noise 个无关症状。
        """
        rng = random.Random(f"cohort|{self.seed}|{target}|{seed}")
        others = [d.name for d in self.diseases if d.name != target and d.year <= LLM_CUTOFF_YEAR]
        chosen = [target] + rng.sample(others, n - 1)
        rng.shuffle(chosen)
        patients = []
        for i, name in enumerate(chosen):
            d = self.by_name[name]
            findings = rng.sample(d.features, 2)
            if rng.random() < 0.8:
                findings.append(d.facts["hallmark"])
            pool = [p for p in PHENOTYPES if p not in d.phenotypes]
            findings += rng.sample(pool, noise)
            rng.shuffle(findings)
            patients.append(Patient(pid=f"P{i:02d}", findings=findings, disease=name))
        return patients


def make_world(seed=0, n_head=20, n_tail=25, n_recent=15):
    """
    生成虚构世界。同一个 seed 总是生成同一个世界。

    【参数】
        seed     : 随机种子
        n_head   : 头部疾病数（预训练语料里出现 30~2000 次）
        n_tail   : 长尾疾病数（出现 0~12 次）
        n_recent : 新近疾病数（2025、2026 年才被描述，出现 0 次）
    【返回】World

    【思路】出现次数按"对数均匀"分布抽取：先在 log 尺度上均匀地取一个数，再取指数。
        这样 30~2000 之间，30~100、100~300、300~1000 这几段被抽到的机会差不多，
        符合"少数疾病极其常见、大部分疾病很少被提到"的长尾形状。

    【边界】头部疾病的描述年份在 1940~2000（越老的病文献越多），长尾在 1960~2024，新近在 2025~2026。
    """
    rng = random.Random(seed)
    names, genes = set(), set()
    diseases = []
    plan = [("head", n_head), ("tail", n_tail), ("recent", n_recent)]
    for tier, count in plan:
        for _ in range(count):
            name = make_disease_name(rng, names)
            if tier == "head":
                mentions = round(10 ** rng.uniform(1.5, 3.3))           # 约 32 ~ 2000
                year = rng.randint(1940, 2000)
            elif tier == "tail":
                mentions = round(10 ** rng.uniform(-0.5, 1.1))          # 约 0 ~ 12
                year = rng.randint(1960, LLM_CUTOFF_YEAR)
            else:
                mentions = 0
                year = rng.randint(LLM_CUTOFF_YEAR + 1, LLM_CUTOFF_YEAR + 2)
            phenos = rng.sample(PHENOTYPES, 4)
            facts = {
                "hallmark": phenos[0],
                "onset": weighted_choice(rng, ONSET_PRIOR),
                "gene": make_gene_symbol(rng, genes),
                "inheritance": weighted_choice(rng, INHERITANCE_PRIOR),
                "year": str(year),
                "discoverer": make_person_name(rng),
                "prevalence": weighted_choice(rng, PREVALENCE_PRIOR),
            }
            diseases.append(Disease(name=name, tier=tier, mentions=mentions, facts=facts,
                                    features=phenos[1:]))
    world = World(diseases, seed)

    # ----- 自我检查 -----
    # 1. 每条事实写成句子再解析回来，必须得到原来的三元组（模板和正则是一致的）。
    for d in world.diseases:
        for rel in QA_RELATIONS:
            assert parse_claim(fact_sentence(rel, d.name, d.facts[rel])) == (d.name, rel, d.facts[rel])
            for kind in ("direct", "paraphrase"):
                assert parse_question(RELATIONS[rel][kind].format(d=d.name)) == (d.name, rel)
    # 2. 每篇文章都恰好包含 7 条事实，各在某一段里。
    for d in world.diseases:
        for rel in QA_RELATIONS:
            world.gold_passage(d.name, rel)
    return world


# =============================================================================
# 直接运行：打印世界概况
# =============================================================================
if __name__ == "__main__":
    world = make_world(seed=0)
    print(f"疾病总数：{len(world.diseases)}，段落总数：{len(world.passages)}")
    for tier in TIERS:
        ds = [d for d in world.diseases if d.tier == tier]
        ms = sorted(d.mentions for d in ds)
        print(f"  {TIER_LABELS[tier]}（{tier}）：{len(ds)} 种，预训练语料里出现次数 {ms[0]}~{ms[-1]}，"
              f"描述年份 {min(d.year for d in ds)}~{max(d.year for d in ds)}")
    d = world.diseases[0]
    print(f"\n示例文章：{d.name}（{TIER_LABELS[d.tier]}，出现 {d.mentions} 次）")
    for p in world.passages_of(d.name):
        print(f"  [{p.pid}] {p.text}")
    print("\n示例问题：")
    for kind in ("direct", "paraphrase"):
        q = world.questions(kind)[2]
        print(f"  {kind:10s} {q.text}  →  {q.answer}")
    cohort = world.make_cohort(world.diseases[-1].name)
    print(f"\n一组病人（其中一人患有 {world.diseases[-1].name}）：")
    for p in cohort[:3]:
        print(f"  {p.pid}: {', '.join(p.findings)}")
    print("  ……")
