"""
=============================================================================
第 2 讲 · 工具模块：给研究智能体用的"小互联网"
=============================================================================

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

-----------------------------------------------------------------------------
【为什么第 1 讲的百科不够用？】
-----------------------------------------------------------------------------
第 1 讲的每种疾病只有 7 条事实、写在一篇文章里，问题是"一问一答"式的。
第 2 讲研究的是"写一篇综述"：你事先不知道该问哪些问题（不知道自己不知道什么），
而且有些问题只有在读过前面的资料之后才问得出来。为了把这两个难点做成可以测量的实验，
我们需要一个结构更丰富的语料库：

    · 广度（breadth）：每种疾病的资料分布在 8 个"方面"（facet）里 —— 症状、遗传、诊断、治疗、
      流行病学、历史、研究、社会与文化。每种疾病只有其中 6 个方面有资料（像真实的维基百科条目一样，
      不是每篇文章都有"社会与文化"这一节）。
      一个"不带视角"的提问者只会想到最常见的 3 个方面（症状、遗传、治疗）。

    · 深度（depth）：每个方面有 2 条"调查线"，每条线有一个**根事实**和一个**追问事实**。
      追问事实的问题里含有根事实里才出现的词。例如：
          根事实：Patients with X are treated with Kelomab.
          追问  ：When was Kelomab approved for X?      ← 不先读到 "Kelomab"，根本问不出这个问题
          追问事实：Kelomab was approved for X in 2019 after the VELA-2 trial.

    · 参考文章：每种疾病有一篇"人类写的"参考文章（6 个方面 × 4 条事实 = 24 条），
      作为评估研究智能体的标准答案（STORM 论文用真实的维基百科文章做同样的事）。

    · 相关话题的目录：STORM 通过浏览"相关话题的文章目录"来发现视角。这里每种疾病的参考文章
      都有一个目录（它有哪些方面），研究智能体可以看到其他疾病的目录。

另外，本模块还为作业 1 式的"找到那个病人"生成**病例文献**（GRILL，04 号文件）：
原始报道、病例系列、非典型病例、"从未见过某症状"的否定性报道，以及一篇相似疾病的文章（陷阱）。

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    from research_world import make_research_world
    rw = make_research_world(seed=0)
    d = rw.base.diseases[25].name           # 第 1 讲世界里的一种疾病
    for f in rw.facts_of(d):                # 它的 24 条事实
        print(f.facet, f.depth, f.text)
    rw.corpus                               # 全部文档（检索的语料）
    rw.toc(d)                               # 它的参考文章目录（方面列表）

直接运行本文件会打印一个例子：
    python research_world.py
"""

import os
import random
import re
import sys
from dataclasses import dataclass, field

# 让 Python 能找到第 1 讲的模块（toy_world.py、llm.py、retriever.py）。
# 【Python 小知识】sys.path 是 Python 查找模块的目录列表；把第 1 讲的目录加进去，
#     就可以像同目录的文件一样 import 它们。os.path.dirname(os.path.abspath(__file__))
#     是"本文件所在的目录"，不管你从哪里运行脚本都正确。
_HERE = os.path.dirname(os.path.abspath(__file__))
_L1 = os.path.join(_HERE, "..", "lecture01")
if _L1 not in sys.path:
    sys.path.insert(0, _L1)

from toy_world import PHENOTYPES, Passage, make_world  # noqa: E402  （import 不在文件开头是故意的）

# =============================================================================
# 第 1 部分：8 个方面，每个方面 2 条调查线
# =============================================================================
# 每个方面有：persona（这个视角对应的人物）、query（大纲驱动的 RAG 检索这一节时用的关键词：
# 相当于让 LLM 把小节标题扩展成检索查询）、lines（两条调查线）。
# 每条调查线有：
#   root   ：根事实的句子模板
#   ask    ：问出根事实的问题（"广度"问题：只要知道这个方面存在，就能问出来）
#   child  ：追问事实的句子模板
#   follow ：追问的问题。注意它含有根事实里的值（例如 {drug}），所以**必须先读到根事实才问得出**。
#   key    ：这条线上"决定对错"的那个值的名字（模型编造时会把它换掉）
# {d} 是疾病名；其他占位符的值由 _make_slots 为每种疾病随机生成。
FACETS = {
    "Signs and symptoms": {
        "persona": "a pediatric neurologist who examines patients",
        "query": "hallmark feature other features",
        "lines": [
            {"root": "The hallmark feature of {d} is {hallmark}.",
             "ask": "What is the hallmark feature of {d}?",
             "child": "In {d}, {hallmark} usually becomes noticeable around age {age}.",
             "follow": "At what age does {hallmark} become noticeable in {d}?",
             "key": ("hallmark", "age")},
            {"root": "Other features of {d} include {f1}, {f2} and {f3}.",
             "ask": "What other features does {d} have?",
             "child": "The severity of {f1} in {d} varies widely, even between siblings.",
             "follow": "How much does the severity of {f1} vary in {d}?",
             "key": ("f1", "f1")},
        ],
    },
    "Genetics": {
        "persona": "a clinical geneticist who reads DNA test results",
        "query": "caused by mutations gene variant",
        "lines": [
            {"root": "{d} is caused by mutations in the {gene} gene.",
             "ask": "Which gene is mutated in {d}?",
             "child": "The {gene} gene encodes {protein}, a protein that {function}.",
             "follow": "What protein does the {gene} gene encode?",
             "key": ("gene", "protein")},
            {"root": "The most common variant in {d} is {variant}.",
             "ask": "What is the most common variant in {d}?",
             "child": "The {variant} variant of {d} is especially common in {region}.",
             "follow": "Where is the {variant} variant of {d} especially common?",
             "key": ("variant", "region")},
        ],
    },
    "Diagnosis": {
        "persona": "a laboratory physician who runs diagnostic tests",
        "query": "diagnosed test mistaken",
        "lines": [
            {"root": "{d} is usually diagnosed with a {test} test.",
             "ask": "How is {d} usually diagnosed?",
             "child": "The {test} test for {d} was developed at {lab}.",
             "follow": "Where was the {test} test for {d} developed?",
             "key": ("test", "lab")},
            {"root": "{d} is often mistaken for {other}.",
             "ask": "What is {d} often mistaken for?",
             "child": "Doctors tell {d} apart from {other} by measuring {marker}.",
             "follow": "How do doctors tell {d} apart from {other}?",
             "key": ("other", "marker")},
        ],
    },
    "Treatment": {
        "persona": "a pharmacist who manages patients' medicines",
        "query": "patients treated with avoid",
        "lines": [
            {"root": "Patients with {d} are treated with {drug}.",
             "ask": "How are patients with {d} treated?",
             "child": "{drug} was approved for {d} in {year_drug} after the {trial} trial.",
             "follow": "When was {drug} approved for {d}?",
             "key": ("drug", "year_drug")},
            {"root": "Patients with {d} are advised to avoid {trigger}.",
             "ask": "What should patients with {d} avoid?",
             "child": "When patients with {d} are exposed to {trigger}, a metabolic crisis can follow.",
             "follow": "What happens when patients with {d} are exposed to {trigger}?",
             "key": ("trigger", "trigger")},
        ],
    },
    "Epidemiology": {
        "persona": "a public health researcher who studies populations",
        "query": "affects people most common",
        "lines": [
            {"root": "{d} affects about {prevalence} people.",
             "ask": "How many people does {d} affect?",
             "child": "The estimate of {prevalence} people for {d} comes from the {registry}.",
             "follow": "Where does the estimate of {prevalence} people for {d} come from?",
             "key": ("prevalence", "registry")},
            {"root": "{d} is most common in {region2}.",
             "ask": "Where is {d} most common?",
             "child": "The high rate of {d} in {region2} was traced to a single founder who lived in the {century} century.",
             "follow": "Why is the rate of {d} so high in {region2}?",
             "key": ("region2", "century")},
        ],
    },
    "History": {
        "persona": "a medical historian who reads old journals",
        "query": "first described reported cases",
        "lines": [
            {"root": "{d} was first described in {year}.",
             "ask": "When was {d} first described?",
             "child": "The {year} description of {d} was published in the {journal}.",
             "follow": "Where was the {year} description of {d} published?",
             "key": ("year", "journal")},
            {"root": "The first cases of {d} were reported by {discoverer}.",
             "ask": "Who reported the first cases of {d}?",
             "child": "{discoverer} originally called the condition {old_name}, before it was named {d}.",
             "follow": "What did {discoverer} originally call {d}?",
             "key": ("discoverer", "old_name")},
        ],
    },
    "Research": {
        "persona": "a translational scientist who runs clinical trials",
        "query": "clinical trial researchers study models",
        "lines": [
            {"root": "A clinical trial of {candidate} for {d} is underway.",
             "ask": "What clinical trials are underway for {d}?",
             "child": "{investigator} at {institution} leads the {candidate} trial for {d}.",
             "follow": "Who leads the {candidate} trial for {d}?",
             "key": ("candidate", "investigator")},
            {"root": "Researchers study {d} using {model} models.",
             "ask": "How do researchers study {d}?",
             "child": "The {model} models of {d} reproduce {f2}.",
             "follow": "What do {model} models of {d} reproduce?",
             "key": ("model", "f2")},
        ],
    },
    "Society and culture": {
        "persona": "a patient advocate who works with families",
        "query": "foundation families awareness day",
        "lines": [
            {"root": "The {foundation} supports families affected by {d}.",
             "ask": "Which organization supports families affected by {d}?",
             "child": "The {foundation} was founded in {year_f} by parents of a child with {d}.",
             "follow": "Who founded the {foundation}?",
             "key": ("foundation", "year_f")},
            {"root": "{d} Awareness Day is observed on {date}.",
             "ask": "When is {d} Awareness Day?",
             "child": "{d} Awareness Day is held in memory of {event}.",
             "follow": "In memory of what is {d} Awareness Day held?",
             "key": ("date", "event")},
        ],
    },
}
FACET_NAMES = list(FACETS)
# 一个"不带视角"的提问者默认会想到的方面（最常见、最"教科书"的三个）。
DEFAULT_FACETS = ["Signs and symptoms", "Genetics", "Treatment"]
FACETS_PER_TOPIC = 6

# 不带视角的提问者能想到的"泛泛的问题"。注意它们几乎都落在默认的三个方面里：
# 多问几个，只是同一批方面的不同说法（讲义第 2.2 节"提问的瓶颈"）。
GENERIC_QUESTIONS = [
    "What is the hallmark feature of {d}?", "Which gene is mutated in {d}?", "How are patients with {d} treated?",
    "What are the symptoms of {d}?", "What causes {d}?", "Is there a cure for {d}?", "Is {d} genetic?",
    "What medicine helps {d}?", "How serious is {d}?", "What does {d} do to the body?", "Can {d} be treated?",
    "What are the first signs of {d}?", "What is {d}?", "Is {d} dangerous?", "Who gets {d}?",
]

# 每个话题另有 3 篇"概述页"：不含任何具体事实，只是泛泛的介绍。含糊的问题（"How serious is X?"）
# 检索到的往往就是这种页面 —— 真实的网页搜索也是这样：泛泛的问题得到泛泛的答案。
OVERVIEW_TEMPLATES = [
    "{d} is a rare genetic disorder, and there is no cure for it yet.",
    "{d} is a serious condition that affects many parts of the body.",
    "Most people have never heard of {d}, and its first signs are easy to miss.",
]

# 生成虚构的名字用的词表。
_SYL = ["vel", "kor", "mar", "tes", "lun", "bra", "dor", "fen", "gal", "hos", "jor", "kel", "lim", "nov",
        "pel", "quin", "ros", "sal", "tor", "var", "wen", "xan", "zel", "ast", "bel", "cor", "dun", "eld"]
_FUNCTIONS = ["regulates calcium flow in neurons", "helps repair damaged DNA", "transports fats inside cells",
              "builds the scaffolding of muscle fibers", "controls the recycling of proteins",
              "guides the growth of nerve cells"]
_TESTS = ["urine amino acid", "skin biopsy", "muscle enzyme", "retinal imaging", "blood lactate", "nerve conduction"]
_MARKERS = ["serum ferritin", "plasma copper", "urinary oxalate", "blood ammonia", "creatine kinase", "vitamin B12"]
_TRIGGERS = ["fasting", "high-protein meals", "cold weather", "intense exercise", "certain anesthetics", "dehydration"]
_JOURNALS = ["Annals of Rare Medicine", "Journal of Clinical Genetics Reports", "Pediatric Case Letters",
             "Northern Review of Neurology"]
_MODELS = ["zebrafish", "mouse", "fruit fly", "organoid", "yeast"]
_EVENTS = ["the first child diagnosed with the condition", "a nurse who cared for the first patients",
           "the founder of the first support group"]
_MONTHS = ["January", "March", "May", "July", "September", "November"]


def _cap(rng, n=2):
    """拼一个首字母大写的虚构词，例如 "Kelvar"。"""
    return "".join(rng.choice(_SYL) for _ in range(n)).capitalize()


def _make_slots(base_disease, all_names, rng):
    """
    为一种疾病生成填模板用的全部值（占位符 → 字符串）。

    【参数】base_disease：第 1 讲的 Disease；all_names：所有疾病名（用来选"容易混淆的疾病"）；rng：random.Random
    【返回】dict
    【思路】和第 1 讲重合的值（标志性症状、基因、年份、报告者、患病率……）直接沿用第 1 讲的事实，
        这样第 1 讲的模拟 LLM 对它们的"记忆"可以直接复用；其余的值随机生成。
    """
    f = base_disease.facts
    year = int(f["year"])
    s = {
        "hallmark": f["hallmark"], "f1": base_disease.features[0], "f2": base_disease.features[1],
        "f3": base_disease.features[2], "gene": f["gene"], "year": f["year"], "discoverer": f["discoverer"],
        "prevalence": f["prevalence"],
        "age": str(rng.randint(2, 15)),
        "protein": f["gene"][:3].capitalize() + rng.choice(["in", "ase", "ulin", "ophilin"]),
        "function": rng.choice(_FUNCTIONS),
        "variant": f"c.{rng.randint(100, 2999)}{rng.choice('ACG')}>T",
        "region": _cap(rng) + rng.choice(["land", "ia", "mark"]),
        "region2": _cap(rng) + rng.choice(["land", "ia", "mark"]),
        "test": rng.choice(_TESTS),
        "lab": _cap(rng) + " " + rng.choice(["University", "Institute", "Children's Hospital"]),
        "other": rng.choice([n for n in all_names if n != base_disease.name]),
        "marker": rng.choice(_MARKERS),
        "drug": _cap(rng) + rng.choice(["mab", "stat", "zole", "tinib", "parin"]),
        "year_drug": str(min(2026, year + rng.randint(2, 30))),
        "trial": _cap(rng, 1).upper() + "-" + str(rng.randint(1, 3)),
        "trigger": rng.choice(_TRIGGERS),
        "registry": _cap(rng) + " Rare Disease Registry",
        "century": rng.choice(["16th", "17th", "18th"]),
        "journal": rng.choice(_JOURNALS),
        "old_name": _cap(rng) + "'s " + rng.choice(["palsy", "dystrophy", "anemia"]),
        "candidate": _cap(rng) + rng.choice(["mab", "stat", "zole", "tinib"]),
        "investigator": rng.choice(["Amara", "Bastian", "Celine", "Dmitri", "Elif", "Farid", "Greta", "Hiro"])
                        + " " + _cap(rng) + "er",
        "institution": _cap(rng) + " " + rng.choice(["University", "Medical Center"]),
        "model": rng.choice(_MODELS),
        "foundation": _cap(rng) + " Family Foundation",
        "year_f": str(min(2026, year + rng.randint(1, 20))),
        "date": f"{rng.choice(_MONTHS)} {rng.randint(1, 28)}",
        "event": rng.choice(_EVENTS),
    }
    return s


# =============================================================================
# 第 2 部分：数据结构
# =============================================================================
@dataclass
class Fact:
    """
    研究世界里的一条事实。

    【字段】
        fid      : 编号，例如 "Velmora syndrome|Treatment|0|1"（疾病|方面|第几条线|深度）
        topic    : 疾病名
        facet    : 方面（FACET_NAMES 之一）
        line     : 这个方面的第几条调查线（0 或 1）
        depth    : 0 = 根事实；1 = 追问事实
        text     : 事实的句子
        question : 能问出这条事实的问题（根事实用 "ask"，追问事实用 "follow"）
        slot     : 决定对错的那个占位符名（模型编造时换掉它的值）
        value    : 这个占位符的正确值
    """
    fid: str
    topic: str
    facet: str
    line: int
    depth: int
    text: str
    question: str
    slot: str
    value: str

    @property
    def child_id(self):
        """根事实的追问事实的编号；追问事实本身返回 None。"""
        return None if self.depth else f"{self.topic}|{self.facet}|{self.line}|1"


class ResearchWorld:
    """
    研究世界：在第 1 讲的世界之上，为每种疾病生成 6 个方面 × 4 条事实，以及检索用的语料库。

    【常用属性和方法】
        rw.base                : 第 1 讲的 World（疾病、档次、出现次数……）
        rw.corpus              : list[Passage]，全部文档（每条事实一篇，另有每个话题 3 篇概述页；pid 形如 "web0123"）
        rw.facts_of(d)         : 疾病 d 的全部事实（参考文章的内容），list[Fact]
        rw.fact(fid)           : 按编号取事实
        rw.toc(d)              : 疾病 d 的参考文章目录（它有哪 6 个方面）
        rw.doc_fact(pid)       : 文档 pid 讲的是哪条事实
        rw.slots(d)            : 填模板用的值
        rw.render(fact, value) : 把一条事实的关键值换成 value 后的句子（用来生成"编造的"版本）
    """

    def __init__(self, base, seed):
        self.base = base
        self.seed = seed
        names = [d.name for d in base.diseases]
        self._slots, self._toc, self._facts, self._by_fid = {}, {}, {}, {}
        self.corpus, self._doc_fact, self._fact_doc = [], {}, {}
        for d in base.diseases:
            rng = random.Random(f"research|{seed}|{d.name}")
            self._slots[d.name] = _make_slots(d, names, rng)
            facets = sorted(rng.sample(FACET_NAMES, FACETS_PER_TOPIC), key=FACET_NAMES.index)
            self._toc[d.name] = facets
            facts = []
            for facet in facets:
                for li, line in enumerate(FACETS[facet]["lines"]):
                    for depth in (0, 1):
                        tmpl = line["child"] if depth else line["root"]
                        q = line["follow"] if depth else line["ask"]
                        slot = line["key"][depth]
                        fact = Fact(fid=f"{d.name}|{facet}|{li}|{depth}", topic=d.name, facet=facet, line=li,
                                    depth=depth, text=self._fill(tmpl, d.name), question=self._fill(q, d.name),
                                    slot=slot, value=self._slots[d.name][slot])
                        facts.append(fact)
            self._facts[d.name] = facts
            for f in facts:
                self._by_fid[f.fid] = f
        # 真值：8 个方面的 32 条事实**都是真的**；只是其中 6 个方面有文档记录（在语料库和参考文章里），
        # 另外 2 个方面是"真实但没人写过"的知识。判断模型说的一句话对不对，要和全部 32 条比。
        self._truth = {}
        for d in base.diseases:
            truths = set()
            for facet in FACET_NAMES:
                for line in FACETS[facet]["lines"]:
                    truths.add(self._fill(line["root"], d.name))
                    truths.add(self._fill(line["child"], d.name))
            self._truth[d.name] = truths
        # 语料库：每条事实一篇文档，再加上每个话题 3 篇概述页（不含事实）。
        # 打乱顺序，免得"文档编号"泄露事实的组织方式。
        items = [(f.topic, f.text, f) for d in base.diseases for f in self._facts[d.name]]
        items += [(d.name, t.format(d=d.name), None) for d in base.diseases for t in OVERVIEW_TEMPLATES]
        random.Random(f"corpus|{seed}").shuffle(items)
        for i, (topic, text, f) in enumerate(items):
            p = Passage(pid=f"web{i:04d}", title=topic, sentences=[text])
            self.corpus.append(p)
            if f is not None:
                self._doc_fact[p.pid] = f
                self._fact_doc[f.fid] = p

    def _fill(self, template, d, override=None):
        """用疾病 d 的值填模板；override 可以替换其中某些值。"""
        values = dict(self._slots[d])
        if override:
            values.update(override)
        return template.format(d=d, **values)

    # ----- 查询 -----
    def facts_of(self, d):
        return list(self._facts[d])

    def fact(self, fid):
        return self._by_fid[fid]

    def toc(self, d):
        return list(self._toc[d])

    def slots(self, d):
        return dict(self._slots[d])

    def doc_fact(self, pid):
        """文档 pid 对应的事实；概述页、不在语料库里的编号（例如模型编造的引用）返回 None。"""
        return self._doc_fact.get(pid)

    def doc_of(self, fid):
        """事实 fid 所在的文档。"""
        return self._fact_doc[fid]

    def template_of(self, fact):
        """一条事实的句子模板。"""
        line = FACETS[fact.facet]["lines"][fact.line]
        return line["child"] if fact.depth else line["root"]

    def render(self, fact, value):
        """
        把事实的关键值换成 value，重新生成句子 —— 用来模拟"模型记错了 / 编造了"。

        【例子】render(药物那条事实, "Torzole") → "Patients with X are treated with Torzole."
        【边界】有些模板里同一个值出现两次（例如 f1），会一起被换掉。
        """
        return self._fill(self.template_of(fact), fact.topic, {fact.slot: value})

    def is_true(self, d, sentence):
        """
        关于疾病 d 的一句话是不是真的（和 32 条真实事实之一逐字相同）。

        【边界】只认识模板生成的句子；措辞不同的真话也会被判为 False。
        """
        return sentence in self._truth.get(d, ())

    def fact_template(self, facet, line, depth):
        """某个方面、某条线、某个深度的句子模板。"""
        spec = FACETS[facet]["lines"][line]
        return spec["child"] if depth else spec["root"]

    def render_any(self, d, facet, line, depth, override=None):
        """
        用疾病 d 的值填任意方面的模板（包括它没有文档记录的方面）。
        直接生成（不检索）的模型会写出这样的句子：可能是真的，但没有来源可以核实。
        """
        return self._fill(self.fact_template(facet, line, depth), d, override)

    def find_fact(self, sentence):
        """
        一句话是不是某条真实的事实？是则返回那条 Fact，否则返回 None。

        【边界】只认识逐字相同的句子：模型说的话和事实的措辞稍有不同就认不出来。
            真实系统里，判断"一句话是否被某个来源支持"需要 LLM 或人工（讲义第 4 节）。
        """
        return self._sentence_index().get(sentence)

    def _sentence_index(self):
        if not hasattr(self, "_sent_idx"):
            self._sent_idx = {f.text: f for f in self._by_fid.values()}
        return self._sent_idx


class CountingSearch:
    """
    带计数和日志的检索器：包装第 1 讲的 BM25。

    【用途】研究智能体的成本有很大一部分是"检索了多少次、读了多少篇"。这个类记下：
        self.n_searches : 检索次数
        self.fetched    : 每次检索返回的文档编号（按顺序记录）—— 03 号文件用它审计"智能体看了哪些网页"
    【构造参数】passages：语料；k1、b：BM25 参数
    """

    def __init__(self, passages, k1=1.5, b=0.75):
        from retriever import BM25
        self.bm25 = BM25(passages, k1=k1, b=b)
        self.n_searches = 0
        self.fetched = []

    def search(self, query, k=3):
        """返回 list[(Passage, 分数)]，并计数、记日志。"""
        self.n_searches += 1
        out = self.bm25.search_scored(query, k)
        self.fetched.extend(p.pid for p, _ in out)
        return out

    def reset(self):
        self.n_searches = 0
        self.fetched = []


def make_research_world(seed=0):
    """
    生成研究世界（同一个 seed 总是得到同一个世界）。

    【返回】ResearchWorld
    【自我检查】每条事实的文本唯一；每篇文档恰好对应一条事实；每种疾病恰好 6 个方面 24 条事实。
    """
    rw = ResearchWorld(make_world(seed=seed), seed)
    texts = [f.text for d in rw.base.diseases for f in rw.facts_of(d.name)]
    assert len(texts) == len(set(texts)), "事实的句子应该互不相同"
    for d in rw.base.diseases:
        assert len(rw.toc(d.name)) == FACETS_PER_TOPIC and len(rw.facts_of(d.name)) == 4 * FACETS_PER_TOPIC
    assert len(rw.corpus) == len(texts) + len(OVERVIEW_TEMPLATES) * len(rw.base.diseases)
    return rw


# =============================================================================
# 第 3 部分：作业 1 式的病例文献（04 号文件 GRILL 用）
# =============================================================================
# 病例文献的句子模板（{d} = 目标疾病，{y} = 相似疾病）。04 号文件"读论文"时按这些模板解析。
GRILL_TEMPLATES = {
    "orig": "{d} was first reported in {year} in three unrelated children; all of them had {h} and {f1}.",
    "series": "In a series of twelve patients with {d}, {f2} was found in most cases.",
    "atyp": "We describe a patient with {d} who presented with {f3} but did not have {h}.",
    "excl": "None of the patients with {d} reported so far had {g}.",
    "similar": "{y} is characterized by {f1}, {f2} and {g}, which resembles {d}.",
}
_GRILL_RE = {
    "orig": re.compile(r"^(?P<d>.+?) was first reported in (?P<year>\d+) in three unrelated children; "
                       r"all of them had (?P<h>.+?) and (?P<f1>.+?)\.$"),
    "series": re.compile(r"^In a series of twelve patients with (?P<d>.+?), (?P<f2>.+?) was found in most cases\.$"),
    "atyp": re.compile(r"^We describe a patient with (?P<d>.+?) who presented with (?P<f3>.+?) "
                       r"but did not have (?P<h>.+?)\.$"),
    "excl": re.compile(r"^None of the patients with (?P<d>.+?) reported so far had (?P<g>.+?)\.$"),
    "similar": re.compile(r"^(?P<y>.+?) is characterized by (?P<f1>.+?), (?P<f2>.+?) and (?P<g>.+?), "
                          r"which resembles (?P<d>.+?)\.$"),
}


def parse_grill_sentence(sentence):
    """
    把一句病例文献解析成 (类型, 字段字典)；不是这几种句式时返回 None。

    【例子】"None of the patients with X reported so far had cleft palate."
            → ("excl", {"d": "X", "g": "cleft palate"})
    【边界】和第 1 讲一样，这是玩具世界的"作弊"：真实系统里，从论文里读出"谁有什么症状、
        谁明确没有什么症状"本身就是一个容易出错的 LLM 任务。
    """
    for kind, regex in _GRILL_RE.items():
        m = regex.match(sentence.strip())
        if m:
            return kind, m.groupdict()
    return None


@dataclass
class GrillCase:
    """
    一道"找到那个病人"的题。

    【字段】
        target   : 目标疾病名（新近疾病：2025 年以后才描述，模型的记忆和旧知识库里都没有）
        cohort   : list[Patient]（第 1 讲的 Patient：pid、findings、disease），10 个病人
        answer   : 患目标疾病的那个病人的 pid
        papers   : list[Passage]，关于目标疾病（和一种相似疾病）的文献
        decoy    : "诱饵"病人的 pid：他有三个症状和目标疾病的文献吻合，但还有一个被文献明确排除的症状
        excluded : 文献明确说"目标疾病的病人从未出现"的那个症状
        atypical : 目标病人是否"非典型"（没有标志性症状）
    """
    target: str
    cohort: list
    answer: str
    papers: list
    decoy: str
    excluded: str
    atypical: bool
    paper_roles: dict = field(default_factory=dict)


def make_grill_case(base, target, seed=0):
    """
    为目标疾病生成一道"找到那个病人"的题和它的文献。

    【参数】base：第 1 讲的 World；target：目标疾病名；seed：随机种子
    【返回】GrillCase

    【文献】（每篇是一个 Passage，pid 形如 "paper-<目标>-<角色>"）
        orig    原始报道：所有病人都有标志性症状 h 和症状 f1
        series  病例系列：多数病人有 f2
        atyp    非典型病例：一个病人有 f3，但**没有** h
        excl    否定性报道：迄今报道的病人里没有一个出现过 g
        similar 相似疾病：另一种疾病 Y 的特征是 f1、f2 和 g（陷阱：讲的不是目标疾病）

    【病人】
        目标病人：f2、f3 + 一个无关症状；一半的题里还有 h（典型），另一半没有（非典型）
        诱饵病人：h、f1、f2、g —— 有三个症状都能在 X 的文献里找到，比目标病人"看起来"更像 X，
                  但 g 被否定性报道明确排除了
        其余 8 人：第 1 讲的普通病人（患 2024 年以前描述过的疾病）
    【随机性】文献只由目标疾病决定（同一种病的几道题共用一套文献）；病人的分组和是否非典型由 seed 决定。

    【边界】这是为演示"假设检验"专门设计的题：诱饵的存在让"数重合个数"的方法容易上当，
        而否定性报道（excl）只有在**主动寻找反证**时才会被注意到。真实的病例要复杂得多。
    """
    from toy_world import LLM_CUTOFF_YEAR, Patient
    # 文献只由目标疾病决定（同一种病，不管出几道题，文献都是同一套）；病人的分组由 seed 决定。
    lit_rng = random.Random(f"grill-lit|{base.seed}|{target}")
    rng = random.Random(f"grill|{base.seed}|{target}|{seed}")
    d = base.by_name[target]
    h, (f1, f2, f3) = d.facts["hallmark"], d.features
    pool = [p for p in PHENOTYPES if p not in d.phenotypes]
    g = lit_rng.choice(pool)
    noise = [p for p in pool if p != g]
    y_name = lit_rng.choice([x.name for x in base.diseases if x.name != target and x.year <= LLM_CUTOFF_YEAR])
    atypical = rng.random() < 0.5

    def paper(role, title, text):
        return Passage(pid=f"paper-{target.split()[0]}-{role}", title=title, sentences=[text])
    values = dict(d=target, y=y_name, year=d.year, h=h, f1=f1, f2=f2, f3=f3, g=g)
    titles = {"orig": f"{target} first report", "series": f"{target} case series",
              "atyp": f"{target} atypical case", "excl": f"{target} negative findings", "similar": f"{y_name} review"}
    papers = [paper(role, titles[role], GRILL_TEMPLATES[role].format(**values)) for role in GRILL_TEMPLATES]
    roles = {p.pid: p.pid.rsplit("-", 1)[1] for p in papers}

    target_findings = [f2, f3, rng.choice(noise)] + ([h] if not atypical else [])
    decoy_findings = [h, f1, f2, g]
    rng.shuffle(target_findings)
    rng.shuffle(decoy_findings)
    others = rng.sample([x for x in base.diseases if x.year <= LLM_CUTOFF_YEAR and x.name not in (target, y_name)], 8)
    people = [(target, target_findings), (y_name, decoy_findings)]
    for o in others:
        fs = rng.sample(o.features, 2) + ([o.facts["hallmark"]] if rng.random() < 0.8 else [])
        fs += rng.sample([p for p in PHENOTYPES if p not in o.phenotypes], 1)
        rng.shuffle(fs)
        people.append((o.name, fs))
    rng.shuffle(people)
    cohort = [Patient(pid=f"P{i:02d}", findings=fs, disease=name) for i, (name, fs) in enumerate(people)]
    for p in papers:                                    # 自我检查：每篇文献都能被解析回来
        assert parse_grill_sentence(p.text)[0] == roles[p.pid]
    answer = next(p.pid for p in cohort if p.disease == target)
    decoy = next(p.pid for p in cohort if p.disease == y_name)
    return GrillCase(target=target, cohort=cohort, answer=answer, papers=papers, decoy=decoy, excluded=g,
                     atypical=atypical, paper_roles=roles)


if __name__ == "__main__":
    rw = make_research_world(seed=0)
    d = [x for x in rw.base.diseases if x.tier == "tail"][0].name
    print(f"语料库：{len(rw.corpus)} 篇文档（{len(rw.base.diseases)} 种疾病 × (24 条事实 + 3 篇概述页)）")
    print(f"\n{d} 的参考文章目录：{rw.toc(d)}")
    for f in rw.facts_of(d)[:8]:
        kind = "追问" if f.depth else "根  "
        print(f"  [{kind}] {f.facet:20s} 问：{f.question}\n          答：{f.text}")
    case = make_grill_case(rw.base, [x for x in rw.base.diseases if x.tier == "recent"][0].name)
    print(f"\n作业 1 式的题：在 10 个病人里找出患 {case.target} 的那位（答案 {case.answer}，诱饵 {case.decoy}）")
    for p in case.papers:
        print(f"  [{p.pid}] {p.text}")