"""
=============================================================================
第 5 讲 · 代码 01：长尾知识与幻觉 —— 见过几次的事实，模型才"知道"？
=============================================================================

对应讲义：lectures/lecture05_hallucination_calibration.md §2–§3

【这个脚本要回答的问题】
    1. 一个事实在预训练语料里出现了 0 / 1 / 2 / …… 次，模型答对它的概率是多少？它的置信度跟着变吗？（演示 1）
    2. 模型对不知道的事实会怎样回答？会说"不知道"吗？编出来的答案从哪里来？（演示 2）
    3. 理论上，一个"校准良好"的模型必须产生多少幻觉？（演示 3：Kalai & Vempala 2024 的"单次事实"下界）
    4. 把正确的（或者错误的）事实直接写在提示里，迷你模型会用它吗？（演示 4：检索增强的前提）

【怎么读输出】
    "置信度"是模型给它自己的答案（概率最大的那个词）的概率。"出现次数"是这个事实在训练语料里被陈述的次数
    （每篇百科词条陈述一次，问答页里每被问答一次算一次）。每个事实问一次（问答页格式）。

【运行】python3 01_longtail.py     （约 5 秒；第一次运行会先训练长尾模型约 50 秒）
【依赖】PyTorch；import 同目录的 hallu_lab.py，第 1 讲的 console.py。
"""

import os
import random
import sys
from collections import Counter

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import pct_ci, print_header, print_observations, print_table  # noqa: E402
from hallu_lab import (BUCKETS, BY_NAME, COUNTRIES, COUNTS, bucket, fact_counts, get_longtail_model,  # noqa: E402
                       last_logits, probe, pw)


def demo1(rows):
    """演示 1：准确率和置信度随出现次数的变化。"""
    print_header("演示 1：一个事实在预训练里出现几次，模型才知道？（240 个事实，问答页格式）")
    table = []
    stats = {}
    for b in BUCKETS:
        sel = [r for r in rows if bucket(r["count"]) == b]
        k = sum(r["correct"] for r in sel)
        mc = sum(r["conf"] for r in sel) / len(sel)
        wrong = [r for r in sel if not r["correct"]]
        mw = sum(r["conf"] for r in wrong) / len(wrong) if wrong else None
        stats[b] = (k / len(sel), mc, mw, len(sel))
        table.append([b, len(sel), pct_ci(k, len(sel)), f"{mc:.2f}", "–" if mw is None else f"{mw:.2f}"])
    print_table(["出现次数", "事实数", "准确率 [95% CI]", "平均置信度", "答错时的平均置信度"], table)
    s0, s1, s2 = stats["0 次"], stats["1 次"], stats["2 次"]
    print_observations([
        f"出现 0 次的事实（6 个只有名字的国家）准确率 {100 * s0[0]:.1f}%，出现 1 次的 {100 * s1[0]:.1f}%，出现 2 次的就跳到 "
        f"{100 * s2[0]:.1f}%，3 次以上全对。知识随出现次数的变化非常陡：见过一次，多半记不住；见过两三次，就记住了。",
        "Kandpal et al.（2023）在 BLOOM、GPT-Neo 等真实模型上看到同样的形状：一个问题的相关文档在预训练数据里出现得越少，"
        "模型越答不对；Mallen et al.（2023，PopQA）按实体的流行程度看，结论一样——**大模型的知识也有一条长尾**。",
        f"置信度跟着出现次数一起上升（{s0[1]:.2f} → {s1[1]:.2f} → {s2[1]:.2f} → … → {stats['21 次以上'][1]:.2f}）："
        "模型'多少知道自己知道多少'（Kadavath et al. 2022 的标题）。但只出现 2 次的事实几乎全对，置信度却只有 "
        f"{s2[1]:.2f}——它**低估**了自己；只出现 0 次的事实几乎全错，置信度却有 {s0[1]:.2f}——它在瞎猜时仍然给了不小的把握。",
    ])
    return stats


def demo2(rows):
    """演示 2：不知道的时候，模型说了什么。"""
    print_header("演示 2：模型对'从没见过'的事实怎样回答？（出现 0 次的 24 个事实）")
    zero = [r for r in rows if r["count"] == 0]
    gold_attr = {}
    for c in COUNTRIES:
        for rel in pw.RELATIONS:
            gold_attr.setdefault((rel, c.attr(rel)), []).append(c.name)
    examples, borrowed, borrowed_counts = [], 0, []
    for r in zero:
        owners = gold_attr.get((r["rel"], r["pred"]), [])
        if owners:
            borrowed += 1
            borrowed_counts.append(max(COUNTS[o] for o in owners))
        if len(examples) < 6:
            examples.append([f"{r['country']} 的 {r['rel']}", r["pred"], f"{r['conf']:.2f}",
                             ", ".join(f"{o}（{COUNTS[o]} 篇）" for o in owners) or "（不是任何国家的这个属性）"])
    print_table(["问题", "模型的回答", "置信度", "这个答案其实属于"], examples, aligns=["<", "<", ">", "<"])
    med = sorted(borrowed_counts)[len(borrowed_counts) // 2] if borrowed_counts else 0
    print_observations([
        f"24 个从没见过的事实，模型**一次也没有**说'不知道'——base 模型的词表里根本没有这个选项，它只会写出'最像样的下一个词'。",
        f"{borrowed} 个回答是**别的国家的同类属性**（一个首都名、一种语言名……）：格式完全正确、内容完全是编的。"
        f"被借用的国家大多也是冷门国家（中位数只有 {med} 篇百科）：模型把几个它都不太熟的国家混在了一起——"
        "这种'张冠李戴'是幻觉的一种常见形态，编出来的东西往往是**真实存在的、但属于别处的**信息。",
        "这就是**幻觉**最典型的样子：流畅、格式正确、看起来像真的，但没有任何依据。它不是'程序出错'，"
        "而是语言模型的训练目标（写出最可能的续写）在知识缺失时的自然结果。",
    ])


def demo3(rows):
    """演示 3：单次事实与幻觉的下界。"""
    print_header("演示 3：'校准良好的模型必然产生幻觉'——单次事实的比例（Kalai & Vempala 2024）")
    fc = fact_counts()
    seen = [v for v in fc.values() if v >= 1]
    mono = sum(1 for v in seen if v == 1) / len(seen)
    zero_frac = sum(1 for v in fc.values() if v == 0) / len(fc)
    err_once = 1 - sum(r["correct"] for r in rows if r["count"] == 1) / sum(1 for r in rows if r["count"] == 1)
    print_table(["量", "数值"], [
        ["训练语料里至少出现一次的事实", f"{len(seen)} 个"],
        ["其中恰好出现一次（单次事实，monofact）的比例", f"{100 * mono:.1f}%"],
        ["从没出现过的事实占全部 240 个的比例", f"{100 * zero_frac:.1f}%"],
        ["模型在单次事实上的错误率", f"{100 * err_once:.1f}%"],
    ])
    print_observations([
        "Good–Turing 估计（图灵和古德在二战中破译密码时提出）：'只出现过一次的东西'所占的比例，估计了'从没出现过的东西'的总概率。"
        f"在我们的语料里，单次事实占 {100 * mono:.1f}%。",
        "Kalai & Vempala（2024）证明：一个在这类'任意事实'上**校准良好**的语言模型，它的幻觉率大致不会低于单次事实的比例"
        "（减去校准误差）。直觉是：模型无法区分'只见过一次的真事实'和'从没见过的假事实'，要对前者给出合理的概率，"
        "就必须也给后者分概率——也就是会编。",
        "推论：幻觉不只是'没训练好'，而是'从有限的样本里学会生成'这件事的统计后果。减少幻觉要靠训练目标之外的东西："
        "让模型能说'不知道'（03 号脚本）、检查一致性（04 号）、去外部查证（检索，第 7 讲）。Kalai et al.（2025）还指出，"
        "很多评测只看准确率、不奖励'不知道'，等于在鼓励模型去猜。",
    ])
    return mono


def demo4(model, tok):
    """演示 4：把事实写进提示里，模型会用吗？"""
    print_header("演示 4：把事实写在提示里——迷你模型会照着用吗？（首都；在问题前面加一句百科）")
    rng = random.Random(0)
    rows = []
    res = {}
    for label, sel in [("热门（≥ 3 篇）", [c for c in COUNTRIES if COUNTS[c.name] >= 3]),
                       ("冷门（0–1 篇）", [c for c in COUNTRIES if COUNTS[c.name] <= 1])]:
        for setting in ("不给资料", "给正确的资料", "给错误的资料"):
            prompts, golds, ctxs = [], [], []
            for c in sel:
                wrong = rng.choice([x.capital for x in COUNTRIES if x.name != c.name])
                val = c.capital if setting == "给正确的资料" else wrong
                pre = "" if setting == "不给资料" else f"<doc> article : the capital of {c.name} is {val} . "
                prompts.append(pre + f"<doc> faq : q : what is the capital of {c.name} ? a :")
                golds.append(c.capital)
                ctxs.append(val)
            preds = [tok.itos[int(lg.argmax())] for lg in last_logits(model, tok, prompts)]
            mem = sum(p == g for p, g in zip(preds, golds)) / len(preds)
            ctx = sum(p == v for p, v in zip(preds, ctxs)) / len(preds) if setting != "不给资料" else None
            res[(label, setting)] = (mem, ctx)
            rows.append([label, setting, f"{100 * mem:.1f}%", "–" if ctx is None else f"{100 * ctx:.1f}%"])
    print_table(["国家", "提示", "答出真正的首都", "答出资料里写的首都"], rows)
    cold = res[("冷门（0–1 篇）", "给正确的资料")]
    print_observations([
        f"给冷门国家**正确**的资料，准确率仍然只有 {100 * cold[0]:.1f}%（不给资料时 {100 * res[('冷门（0–1 篇）', '不给资料')][0]:.1f}%）："
        "资料就在上一句，模型却没用上。给**错误**的资料，它几乎从不照着答，仍然回答它记忆里的那个（对热门国家来说就是对的）。",
        "原因和第 2 讲 §8.6 一样：这个迷你模型没有归纳头，不会'照抄前文'。**检索增强生成（RAG）的前提是模型会读上下文**；"
        "第 7 讲用真实的大模型讲 RAG。",
        "真实的大模型会读上下文，但会出现另一个问题：资料和它的记忆冲突时听谁的。Longpre et al.（2021）、Xie et al.（2024）"
        "发现，模型对热门事实更'固执'（坚持记忆），对冷门事实更容易被资料说服；Mallen et al.（2023）发现检索对冷门事实帮助最大，"
        "对热门事实有时反而有害——所以有人提出只在问题涉及冷门实体时才检索（自适应检索）。",
    ])


def main():
    model, tok = get_longtail_model()
    rows = probe(model, tok)
    demo1(rows)
    demo2(rows)
    demo3(rows)
    demo4(model, tok)


if __name__ == "__main__":
    main()