"""
=============================================================================
第 13 讲 · 工具模块 1：一个"知道真相"的问答世界 + 一个参数完全公开的模拟 LLM
=============================================================================

对应讲义：lectures/lecture13_confidence_calibration.md（全讲；尤其 §0.5 关于模拟器）

-----------------------------------------------------------------------------
【为什么要造一个模拟世界？ Why a simulator?】
-----------------------------------------------------------------------------
这一讲要回答的问题都是"置信度有没有含义"：模型说 90% 的时候，是不是真的有 90% 是对的？
采样 10 次、有 8 次一样，代表什么？语义熵（semantic entropy）比按字面算的熵好在哪里？
要回答这些问题，必须知道**每道题的真实答案**和**模型"心里"对每个答案的真实信念**。
真实 LLM 做不到后者，所以我们造一个小世界：每一种行为都是一个写在下面的、可以改的参数。

**重要**：模拟器里"RLHF 让概率变得过度自信""verbalized confidence 偏高、是 5% 的整数倍"这些行为，
是我们**按论文的描述写进去的假设**（出处见每个参数的注释），不是模拟器"发现"的。
模拟器能做的，是在知道真相的环境里检查**测量方法**本身：ECE 的估计偏差、重新校准能不能迁移、
按字面聚类和按语义聚类的差别、一致性方法抓不住"一致的错误"、conformal 的覆盖率是边际的……
这些结论不依赖于具体参数。讲义 §0.5 详细说明哪些数字取决于参数、哪些不取决于参数。

-----------------------------------------------------------------------------
【这个世界里有什么？ What's in the world】
-----------------------------------------------------------------------------
1. 知识库 knowledge base：虚构的国家、公司、城市（名字由音节随机拼成，保证模型"不可能在别处见过"）。
   四种关系 relation，每道题问一个事实：
        capital  "What is the capital of Veloria?"        答案类型 city
        founder  "Who founded Quorandel?"                 答案类型 person（有名和姓）
        year     "In what year was Telmar founded?"       答案类型 year
        river    "Which river flows through Marnesh?"     答案类型 river
   每道题有 N_CAND = 5 个候选答案（semantic answers）：1 个正确答案 + 4 个"看起来合理"的干扰项。

2. 模型的信念 belief：每道题有一组 logits z（长度 5）。基础模型（pre-trained base model）对答案 a 的概率是
        π_base(a) = softmax(z)[a]
   z 怎么来：每个候选答案先有一个 N(0, 1) 的噪声；其中一个"被偏爱的"答案再加一个 boost b。
   b 取决于这个实体的**流行度 popularity**（PopQA，Mallen et al. 2023 发现模型对冷门实体的事实准确率低得多）：
        head（常见实体，30%）：b ~ N(4.0, 1)      torso（40%）：b ~ N(2.0, 1)      tail（冷门，30%）：b ~ N(0.3, 1)

3. 真相从哪里来 —— "基础模型在答案层面是校准的"：
   我们**先**抽信念 π_base，**再**按 π_base 抽出哪个候选是真答案（truth ~ π_base）。这样构造出来的基础模型，
   在"答案 a 是对的"这个事件上**恰好**是 calibrated 的：P(a 正确 | π_base(a) = q) = q。
   这对应 Kadavath et al.（2022）和 GPT-4 技术报告的观察：大的预训练模型在格式合适的选择题上校准得很好。
   （讲义 §4 讨论为什么会这样：预训练的 log loss 是一个 proper scoring rule。）

4. 分布偏移 domain B（"训练截止之后变了的事实" post-cutoff updates）：
   先按 π_base 抽出"旧世界"的真答案；然后以概率 STALE_RATE = 0.4，这个事实**已经变了**，新答案是另外 4 个候选之一（均匀）。
   模型不知道它变了 —— 于是在 domain B 上，它会**自信地、一致地**答错。domain A 没有这种题。

5. 措辞 surface forms：同一个答案有好几种说法，例如 "Marnesh" / "It's Marnesh." / "The capital is Marnesh."；
   人名还有"只说姓"的说法（"Moren."，partial = True）。每道题有自己的措辞偏好 φ（form_w），
   有的题几乎总用同一种说法，有的题说法很分散 —— 这个"啰嗦程度"是独立于答案对错抽取的
   （但代码 3 会发现，它对字面熵的影响在答对的题上更大）。
   于是"按字面"算的不确定性 = 语义上的不确定性 + 措辞上的噪声。这正是 semantic entropy（Kuhn, Gal & Farquhar 2023）要去掉的东西。

-----------------------------------------------------------------------------
【模拟 LLM 的参数 SimLLM parameters】（全部是显式假设）
-----------------------------------------------------------------------------
    sharpen        对答案 logits 乘的系数：π_model ∝ exp(sharpen · z)。base = 1；chat（RLHF 之后）= 2.5。
                   依据：Kadavath et al.（2022, §3.3）报告 RLHF policy "naively appear very miscalibrated"，
                   把采样温度调到 2.5 基本修好 —— 也就是说，RLHF 的概率 ≈ 基础概率被"锐化"了 2.5 倍。
                   我们直接用 2.5，所以本模拟里"正确的"温度就是 2.5（这是设计，不是发现）。
                   注意：锐化不改变 argmax，所以 base 和 chat 的贪心答案（greedy answer）完全一样、准确率一样，
                   **差别只在置信度上**。这让比较很干净（现实中 RLHF 当然也会改变答案本身）。
    verbal_*       口头置信度 verbalized confidence（"Confidence: 90%"）：
                       v = quantize( sigmoid(slope · logit(π_base(a)) + bias + N(0, noise²)) , 0.05 )
                   读的是**基础信念** π_base（Burns et al. 2023 等发现模型内部表示里编码了"真假"），
                   但经过一个偏高的仿射变换 + 噪声，并取到 5% 的整数倍（Xiong et al. 2024：口头置信度集中在 80–100%，
                   多是 5 的倍数）。chat 默认 slope = 0.6，bias = 1.2，noise = 0.6。
    topk_*         "给出你最好的 3 个猜测和各自的概率"（Tian et al. 2023 "Just Ask for Calibration"；
                   Fall 2023 第 1 讲 "How hot is it today?" 的例子）。我们把论文的发现
                   "先列出几个候选、再给概率，校准更好" 写成一个**更小的偏差**：slope = 0.8，bias = 0.5，noise = 0.5。
    ptrue_noise    P(True)（Kadavath et al. 2022）：把模型自己的答案拿给它看，问 "Is the proposed answer True or False?"，
                   读 "True" 这个 token 的概率。它是 token 概率，所以 RLHF 的锐化同样作用在它上面：
                       P(True) = sigmoid( sharpen · (logit(π_base(a)) + N(0, ptrue_noise²)) )

-----------------------------------------------------------------------------
【上帝视角 vs 模型视角】
-----------------------------------------------------------------------------
Question.truth、Question.stale、Question.tier 是上帝视角，只给评估代码用。
SimLLM 的方法只读 logits 和措辞偏好（"模型心里的东西"），从不读 truth。

-----------------------------------------------------------------------------
【怎样换成真实模型 Plugging in a real model】
-----------------------------------------------------------------------------
所有评估代码只用到下面几个接口：greedy(q)、sample(q, rng, temperature, k)、seq_logprob(q, s)、answer_probs(q)、
verbalized(q, rng)、verbalized_topk(q, rng)、p_true(q, s, rng)。真实模型对应：贪心解码、温度采样、
API 返回的 token logprobs（求和）、选择题格式下各选项的概率、在 prompt 里要求 "Confidence: X%" 后解析数字、
Tian et al. 的 top-k prompt、Kadavath et al. 的 P(True) prompt。判断两个回答是否同义要换成一个真正的 NLI 模型
（Kuhn et al. 用 DeBERTa-large MNLI）。本课不写任何调用外部服务的代码。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 uq_world.py          # 自检：打印几道题、两个模型的概率、几次采样，不到 1 秒

只用 Python 标准库（random、math）。
"""

import math
import random

# =============================================================================
# 0. 世界的常量 World constants
# =============================================================================
N_CAND = 5                     # 每道题的候选答案数：1 个正确 + 4 个干扰项
STALE_RATE = 0.4               # domain B 里"事实已经变了"的比例

# 流行度分层：(boost 的均值, 占的比例)。boost 的标准差都是 1。
TIERS = [("head", 4.0, 0.30), ("torso", 2.0, 0.40), ("tail", 0.3, 0.30)]

SYLLABLES = ["mar", "nesh", "tel", "va", "or", "in", "ka", "lo", "ri", "sen", "dra", "vel", "quo",
             "ran", "del", "thi", "em", "ul", "zan", "bri", "ost", "ny", "fen", "gal", "mo", "ren",
             "is", "ta", "ber", "cor", "lin", "sa", "dun", "pe", "har", "wy"]

# 四种关系：(名字, 问题模板, 答案类型)
RELATIONS = [
    ("capital", "What is the capital of {s}?", "city"),
    ("founder", "Who founded {s}?", "person"),
    ("year", "In what year was {s} founded?", "year"),
    ("river", "Which river flows through {s}?", "river"),
]

# 每种答案类型的措辞模板：(模板, 是否 partial)。{a} = 完整答案，{last} = 人名的姓。
# 【可信联系】措辞的多少因类型而异（年份只有 2 种说法，人名有 7 种），所以"按字面"的熵会被答案类型影响 ——
#   措辞偏好本身和"答得对不对"独立地生成；但代码 3 发现它对字面熵的**影响**并不独立（答对的题上更大，见 §6.3）。
FORMS = {
    "city": [("{a}", False), ("{a}.", False), ("It's {a}.", False), ("The capital is {a}.", False),
             ("I believe it is {a}.", False)],
    "person": [("{a}", False), ("{a}.", False), ("It was {a}.", False), ("{a} founded it.", False),
               ("The founder was {a}.", False), ("{last}.", True), ("I think it was {a}.", False)],
    "year": [("{a}", False), ("In {a}.", False)],
    "river": [("The {a}.", False), ("The {a} River.", False), ("{a}.", False), ("It's the {a}.", False)],
}


# =============================================================================
# 1. 小数学工具 Small math helpers（uq_tools 里也有同名函数；这里复制一份，让本模块不依赖别的文件）
# =============================================================================
def softmax(z):
    """
    softmax：把一组实数 logits 变成概率（非负、和为 1）。

    【思路】先减去最大值再取 exp —— 数学上结果不变（分子分母同乘 e^{−max}），但不会上溢（overflow）。
    【参数】z：数字列表   【返回】同样长度的概率列表
    【边界】z 里有很大的数（比如 1000）也不会出 inf。
    【Python 小知识】max(z) 求最大值；[math.exp(v - m) for v in z] 是"列表推导式"：对每个 v 算一个值，组成新列表。
        math.fsum 是"精确求和"：比 sum() 更准，而且在 Python 3.9 和 3.12+ 上结果完全一样（3.12 改了 sum 的算法）。
    """
    m = max(z)
    e = [math.exp(v - m) for v in z]
    s = math.fsum(e)
    return [v / s for v in e]


def sigmoid(x):
    """数值稳定的 sigmoid：x 很负时用 e^x / (1 + e^x) 的写法，避免 exp(−x) 上溢。"""
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-x))
    e = math.exp(x)
    return e / (1.0 + e)


def logit(p, eps=1e-6):
    """logit(p) = log(p / (1 − p))；先把 p 夹到 [eps, 1 − eps]，避免 log(0)。"""
    p = min(max(p, eps), 1.0 - eps)
    return math.log(p / (1.0 - p))


def quantize(v, step):
    """把 v 取到最近的 step 的整数倍（例如 0.83 → 0.85），并夹到 [0, 1]。模拟"口头说出的百分数"只有有限几个值。"""
    q = round(v / step) * step
    return min(max(q, 0.0), 1.0)


def weighted_choice(rng, weights):
    """
    按权重抽一个下标。

    【思路】累积权重 + 一个 [0, total) 的均匀随机数，落在哪一段就选哪个。
    【Python 小知识】rng.random() 返回 [0, 1) 的均匀随机数。这里不用 rng.choices，是为了让"每次调用恰好消耗
        一个随机数"，代码的随机数流更容易推理（改一个地方不会悄悄改变后面所有的随机数）。
    """
    total = math.fsum(weights)
    u = rng.random() * total
    acc = 0.0
    for i, w in enumerate(weights):
        acc += w
        if u < acc:
            return i
    return len(weights) - 1          # 浮点误差兜底：u 恰好等于 total 时返回最后一个


# =============================================================================
# 2. 题目 Question
# =============================================================================
def _make_name(rng, n_syll):
    """用 n_syll 个音节拼一个首字母大写的虚构名字，例如 'Marnesh'。"""
    return "".join(rng.choice(SYLLABLES) for _ in range(n_syll)).capitalize()


def _make_answer(rng, answer_type):
    """
    为一种答案类型造一个答案字符串。人名返回 (全名, 姓)，其余返回 (名字, 名字)。

    【Python 小知识】rng.randint(a, b) 返回 [a, b] 里的整数（两端都包含）；str(1872) 把整数变成字符串 "1872"。
    """
    if answer_type == "year":
        y = str(rng.randint(1650, 1990))
        return y, y
    if answer_type == "person":
        first, last = _make_name(rng, 2), _make_name(rng, rng.choice([2, 3]))
        return first + " " + last, last
    return _make_name(rng, rng.choice([2, 3])), None


class Question:
    """
    一道题（上帝视角的完整记录）。

    【属性】
        qid          编号
        relation     关系名（capital / founder / year / river）
        text         问题文本
        answer_type  答案类型（决定有哪些措辞）
        candidates   5 个候选答案字符串；lasts：对应的"姓"（只有人名有，其余为 None）
        logits       模型信念的 logits z（长度 5）—— "模型心里的东西"
        truth        真答案的下标（上帝视角）
        old_truth    domain B 里"变化之前"的答案下标；domain A 里等于 truth
        tier         流行度 head / torso / tail（上帝视角，用来按切片看覆盖率）
        domain       "A"（普通）或 "B"（训练截止后变化过的事实）
        stale        这道题的事实是否已经变了（只有 domain B 可能为 True）
        form_w       措辞偏好：每种措辞的概率（和为 1）
    【Python 小知识】class 定义一个"类"：把相关的数据打包成一个对象；__init__ 是构造函数，
        self.xxx = ... 把数据存到对象上，之后用 q.xxx 读取。
    """

    def __init__(self, qid, relation, text, answer_type, candidates, lasts, logits, truth, old_truth,
                 tier, domain, stale, form_w):
        self.qid = qid
        self.relation = relation
        self.text = text
        self.answer_type = answer_type
        self.candidates = candidates
        self.lasts = lasts
        self.logits = logits
        self.truth = truth
        self.old_truth = old_truth
        self.tier = tier
        self.domain = domain
        self.stale = stale
        self.form_w = form_w

    def render(self, a, f):
        """把"答案 a + 措辞 f"渲染成字符串，例如 render(2, 3) → 'The capital is Telva.'。"""
        template, _ = FORMS[self.answer_type][f]
        return template.format(a=self.candidates[a], last=self.lasts[a] or self.candidates[a])


def draw_belief(rng, tier):
    """
    抽一组信念 logits：每个候选 N(0, 1)，其中一个"被偏爱的"候选再加 boost b ~ N(tier 的均值, 1)。

    【参数】rng：random.Random；tier：'head' / 'torso' / 'tail'
    【返回】长度 N_CAND 的 logits 列表
    【思路】boost 越大，模型越"知道"某个答案（π_base 越集中）。被偏爱的那个不一定是真答案 ——
        真答案在 make_questions 里按 π_base 抽（所以被偏爱的答案以 π_base 的概率是对的）。
    【Python 小知识】rng.gauss(mu, sigma) 抽一个正态分布随机数；rng.randrange(n) 抽 0..n−1 的整数。
    """
    mean_b = {name: m for name, m, _ in TIERS}[tier]
    z = [rng.gauss(0.0, 1.0) for _ in range(N_CAND)]
    z[rng.randrange(N_CAND)] += rng.gauss(mean_b, 1.0)
    return z


def draw_tier(rng):
    """按 TIERS 里的比例抽一个流行度分层。"""
    return TIERS[weighted_choice(rng, [frac for _, _, frac in TIERS])][0]


def make_questions(n, rng, domain="A", tiers=None):
    """
    生成 n 道题。

    【参数】
        n       题数
        rng     random.Random（每个演示用自己的 rng，改一个演示不影响别的演示）
        domain  "A"：真答案 ~ π_base（基础模型在答案层面恰好校准）
                "B"：先按 π_base 抽"旧答案"，再以 STALE_RATE 的概率换成另一个候选（训练截止后事实变了）
        tiers   None = 按 TIERS 的比例抽；也可以传一个字符串（例如 "tail"）只生成某一层
    【返回】Question 列表
    【边界】候选答案在同一道题里保证互不相同（重名时重抽）；不同题之间可能重名，没关系（每道题独立判断）。
    """
    qs = []
    for i in range(n):
        rel, template, atype = RELATIONS[rng.randrange(len(RELATIONS))]
        subject = _make_name(rng, rng.choice([2, 3]))
        cands, lasts = [], []
        while len(cands) < N_CAND:
            a, last = _make_answer(rng, atype)
            if a not in cands:
                cands.append(a)
                lasts.append(last)
        tier = tiers if tiers is not None else draw_tier(rng)
        z = draw_belief(rng, tier)
        pi = softmax(z)
        old = weighted_choice(rng, pi)               # truth ~ π_base：基础模型在答案层面校准的根源
        truth, stale = old, False
        if domain == "B" and rng.random() < STALE_RATE:
            others = [j for j in range(N_CAND) if j != old]
            truth, stale = others[rng.randrange(len(others))], True
        # 措辞偏好：Gamma(α, 1) 权重归一化 = Dirichlet(α)。α 小 → 几乎只用一种说法；α 大 → 说法分散。
        alpha = rng.choice([0.3, 1.0, 3.0])
        w = [rng.gammavariate(alpha, 1.0) + 1e-9 for _ in FORMS[atype]]
        sw = math.fsum(w)
        form_w = [x / sw for x in w]
        qs.append(Question(i, rel, template.format(s=subject), atype, cands, lasts, z, truth, old,
                           tier, domain, stale, form_w))
    return qs


# =============================================================================
# 3. 一次回答 Sample
# =============================================================================
class Sample:
    """
    模型的一次输出：语义答案下标 sem、措辞下标 form、文本 text、是否只说了姓 partial、token 数 n_tokens。

    【思路】真实 LLM 输出的是字符串；我们额外记住"它是哪个答案"（sem），只给**模拟的 NLI** 和评估代码用。
        按字面比较的方法（lexical）只能看 text。
    """

    def __init__(self, q, sem, form):
        self.sem = sem
        self.form = form
        self.text = q.render(sem, form)
        self.partial = FORMS[q.answer_type][form][1]
        self.n_tokens = len(self.text.split())        # 用"空格分开的词数"当 token 数（粗略但够用）


def normalize_text(s):
    """
    按字面比较前的标准化：小写、去掉标点、合并空格。"Marnesh." 和 "marnesh" 算同一个字符串，
    但 "It's Marnesh." 和 "Marnesh" 仍然不同 —— 这就是 lexical 方法的局限。

    【Python 小知识】str.isalnum() 判断字符是不是字母或数字；"".join(...) 把字符拼回字符串；s.split() 按空白切分。
    """
    kept = "".join(ch.lower() if (ch.isalnum() or ch == " ") else "" for ch in s)
    return " ".join(kept.split())


# =============================================================================
# 4. 模拟的 NLI：双向蕴含 bidirectional entailment
# =============================================================================
def entails(s, t):
    """
    模拟一个**完美的** NLI 模型判断 "s ⇒ t"（s 蕴含 t）。

    【规则】
        · 不同的答案：不蕴含。
        · 同一个答案：完整说法 ⇒ 只说姓的说法（"It was Ilsa Moren." ⇒ "Moren."）；
          反过来不成立（"Moren." 不能推出 "Ilsa Moren"：可能是另一个姓 Moren 的人）。
    【可信联系】这正是 Kuhn et al.（2023）要求**双向**蕴含的原因 —— 也是它的一个代价：
        只说姓的回答和完整回答是"同一个答案"，却会被分进不同的语义簇（§6.3）。
    """
    if s.sem != t.sem:
        return False
    return not (s.partial and not t.partial)


def nli_equivalent(s, t, rng, err=0.0):
    """
    双向蕴含：s ⇒ t 且 t ⇒ s 才算"同义"。每个方向的判断以概率 err 被翻转（模拟一个不完美的 NLI 模型）。

    【参数】s, t：Sample；rng：随机数生成器；err：每次判断出错的概率（0 = 完美 NLI）
    【返回】True / False
    【边界】err = 0 时不消耗随机数（保证"完美 NLI"的结果和 rng 无关）。
    """
    a, b = entails(s, t), entails(t, s)
    if err > 0.0:
        if rng.random() < err:
            a = not a
        if rng.random() < err:
            b = not b
    return a and b


def semantic_clusters(samples, rng, err=0.0):
    """
    Kuhn et al.（2023）的聚类算法：依次处理每个样本，和每个已有簇的**第一个成员**比较，
    双向蕴含就加入那个簇，否则自己开一个新簇。

    【参数】samples：Sample 列表；rng、err：传给 nli_equivalent
    【返回】簇的列表，每个簇是样本下标的列表
    【边界】结果依赖于处理顺序（论文的算法也是这样）；NLI 有错时，同一个答案可能被拆成几个簇，
        不同答案也可能被合并。
    【numpy/Python 小知识】for ... else：循环没有被 break 打断时才执行 else 分支 —— 这里用来表示"没找到可以加入的簇"。
    """
    clusters = []
    for i, s in enumerate(samples):
        for c in clusters:
            if nli_equivalent(samples[c[0]], s, rng, err):
                c.append(i)
                break
        else:
            clusters.append([i])
    return clusters


def lexical_clusters(samples):
    """
    按字面聚类：标准化之后字符串相同才算同一簇（这就是"把每个不同的序列当作不同结果"的预测熵的做法）。

    【返回】簇的列表（每个簇是样本下标列表），按第一次出现的顺序排列（不依赖 dict 的哈希顺序）。
    【Python 小知识】dict 在 Python 3.7+ 保持插入顺序；这里键是字符串，但我们只按插入顺序遍历，
        所以结果和 PYTHONHASHSEED 无关。
    """
    groups = {}
    order = []
    for i, s in enumerate(samples):
        key = normalize_text(s.text)
        if key not in groups:
            groups[key] = []
            order.append(key)
        groups[key].append(i)
    return [groups[k] for k in order]


# =============================================================================
# 5. 模拟 LLM
# =============================================================================
class SimLLM:
    """
    参数完全公开的模拟 LLM。见文件开头【模拟 LLM 的参数】。

    【用法】
        base = SimLLM.base()            # 预训练模型：sharpen = 1
        chat = SimLLM.chat()            # RLHF 之后的聊天模型：sharpen = 2.5，口头置信度偏高
        g = chat.greedy(q)              # 贪心回答（Sample）
        ss = chat.sample(q, rng, temperature=1.0, k=10)
    """

    def __init__(self, name, sharpen=1.0, verbal_slope=0.6, verbal_bias=1.2, verbal_noise=0.6,
                 topk_slope=0.8, topk_bias=0.5, topk_noise=0.5, ptrue_noise=0.6, verbal_step=0.05):
        self.name = name
        self.sharpen = sharpen
        self.verbal_slope, self.verbal_bias, self.verbal_noise = verbal_slope, verbal_bias, verbal_noise
        self.topk_slope, self.topk_bias, self.topk_noise = topk_slope, topk_bias, topk_noise
        self.ptrue_noise = ptrue_noise
        self.verbal_step = verbal_step

    # 【Python 小知识】@classmethod：一个"属于类本身"的构造方法，用 SimLLM.base() 调用，返回一个新对象。
    @classmethod
    def base(cls):
        """预训练模型：答案概率 = π_base（校准的）。口头置信度参数和 chat 相同（本讲只在 chat 上用口头置信度）。"""
        return cls("base", sharpen=1.0)

    @classmethod
    def chat(cls):
        """RLHF 之后的聊天模型：答案概率被锐化 2.5 倍（Kadavath et al. 2022 §3.3 的温度 2.5）。"""
        return cls("chat", sharpen=2.5)

    # ---------------- 白盒 white-box：需要概率 / logprobs ----------------
    def answer_probs(self, q, temperature=1.0):
        """
        语义答案层面的概率 π_model(a)（5 个候选的概率，和为 1）。temperature = 采样温度。

        【可信联系】对真实模型，这相当于"选择题格式下各选项字母的概率"，或者把同一个答案的所有说法的概率加起来。
            后者对开放式回答很难精确计算（说法无穷多），这正是 §2 "置信度是关于哪个事件的" 的要点。
        """
        return softmax([v * self.sharpen / temperature for v in q.logits])

    def form_probs(self, q):
        """措辞的概率 φ（base 和 chat 一样；我们只锐化"选哪个答案"，不锐化措辞 —— 一个简化）。"""
        return q.form_w

    def greedy(self, q):
        """贪心回答：概率最大的答案 + 概率最大的措辞。锐化不改变 argmax，所以 base 和 chat 的贪心回答相同。"""
        z = q.logits
        a = max(range(N_CAND), key=lambda j: z[j])
        f = max(range(len(q.form_w)), key=lambda j: q.form_w[j])
        return Sample(q, a, f)

    def sample(self, q, rng, temperature=1.0, k=1):
        """
        温度采样 k 次：先按 π_model^{1/T} 抽答案，再按 φ 抽措辞。

        【参数】temperature：T > 1 更分散，T < 1 更集中；k：样本数
        【返回】Sample 列表
        【可信联系】黑盒 API 不给 logits，但**允许设 temperature** —— 这是黑盒用户手里少数几个能调的旋钮之一（§5.3）。
        """
        p = self.answer_probs(q, temperature)
        out = []
        for _ in range(k):
            a = weighted_choice(rng, p)
            f = weighted_choice(rng, q.form_w)
            out.append(Sample(q, a, f))
        return out

    def seq_logprob(self, q, s):
        """
        整个字符串的 log 概率（T = 1）：log π_model(a) + log φ(f)。真实模型里 = 所有 token 的 logprob 之和。

        【可信联系】这是 "P(恰好输出这个字符串)"，不是 "P(这个答案是对的)"：同一个答案的几种说法会把概率分掉（§2.2）。
        """
        return math.log(self.answer_probs(q)[s.sem]) + math.log(q.form_w[s.form])

    # ---------------- 黑盒 black-box：只看文本 ----------------
    def verbalized(self, q, rng, a=None):
        """
        口头置信度："Answer: X. Confidence: 85%"。读基础信念 π_base(a)，经过偏高的仿射变换 + 噪声，取到 5% 的整数倍。

        【参数】a：被评价的答案下标（默认是贪心答案）
        【返回】[0, 1] 里 0.05 的整数倍
        """
        if a is None:
            a = self.greedy(q).sem
        pb = softmax(q.logits)[a]
        v = sigmoid(self.verbal_slope * logit(pb) + self.verbal_bias + rng.gauss(0.0, self.verbal_noise))
        return quantize(v, self.verbal_step)

    def verbalized_topk(self, q, rng, k=3):
        """
        Tian et al.（2023）的 top-k prompt："Provide your k best guesses and the probability that each is correct"。
        返回第 1 个猜测（= 贪心答案）的概率。偏差比单个口头置信度小（论文的发现，被我们写成参数）。

        【边界】k 只影响"文档上的含义"；我们只用第 1 个猜测的概率（贪心答案的置信度）。
        """
        a = self.greedy(q).sem
        pb = softmax(q.logits)[a]
        v = sigmoid(self.topk_slope * logit(pb) + self.topk_bias + rng.gauss(0.0, self.topk_noise))
        return quantize(v, self.verbal_step)

    def p_true(self, q, a, rng):
        """
        P(True)（Kadavath et al. 2022）：问模型 "这个答案对吗？"，读 'True' token 的概率。
        token 概率会被 RLHF 锐化：P(True) = sigmoid(sharpen · (logit(π_base(a)) + 噪声))。
        """
        pb = softmax(q.logits)[a]
        return sigmoid(self.sharpen * (logit(pb) + rng.gauss(0.0, self.ptrue_noise)))


# =============================================================================
# 6. 自检 Self-test
# =============================================================================
if __name__ == "__main__":
    rng = random.Random(0)
    qs = make_questions(4, rng)
    base, chat = SimLLM.base(), SimLLM.chat()
    for q in qs:
        g = chat.greedy(q)
        pb, pc = base.answer_probs(q), chat.answer_probs(q)
        print(f"\nQ{q.qid} [{q.tier}] {q.text}")
        for j in range(N_CAND):
            mark = " ← truth" if j == q.truth else ""
            print(f"   {q.candidates[j]:<22} π_base = {pb[j]:.3f}   π_chat = {pc[j]:.3f}{mark}")
        ss = chat.sample(q, rng, temperature=1.0, k=5)
        print(f"   greedy: {g.text!r}   correct = {g.sem == q.truth}")
        print("   5 samples:", [s.text for s in ss])
        print(f"   verbalized = {chat.verbalized(q, rng):.2f}   top-3 = {chat.verbalized_topk(q, rng):.2f}"
              f"   P(True) chat = {chat.p_true(q, g.sem, rng):.3f}")
