"""
=============================================================================
第 5 讲 · 代码 04：用"多问几遍"检测幻觉 —— SelfCheckGPT 与语义熵的思路
=============================================================================

对应讲义：lectures/lecture05_hallucination_calibration.md §6

【这个脚本要回答的问题】
    很多 API 不给 token 的概率。拿不到置信度时，怎样判断模型是不是在编？SelfCheckGPT（Manakul et al. 2023）的想法是：
    **对同一个问题采样好几次**——模型真知道的事，每次都说一样的；编出来的，每次编得不一样。
    1. 采样 2 / 5 / 10 / 20 次，"和自己一致的程度"能多好地区分对错？和直接用概率比呢？（演示 1）
    2. 一致性检测会漏掉什么？"每次都编同一个答案"的幻觉有多少？（演示 2）
    3. 在"教过它不知道的事实"的模型上，一致性检测还管用吗？（演示 3）

【怎么读输出】
    一致性 = k 次采样里，和贪心答案相同的比例；样本熵 = k 次采样的答案分布的熵（语义熵的最简版本：答案都是一个词，
    "意思相同"就是"字面相同"）。AUROC 衡量这个分数区分"答对 / 答错"的能力（0.5 = 瞎猜）。

【运行】python3 04_selfcheck.py     （约 10 秒；需要 02 号脚本训练的缓存）
【依赖】PyTorch；import 同目录的 hallu_lab.py，第 1 讲的 console.py。
"""

import math
import os
import sys
from collections import Counter

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import print_header, print_observations, print_table  # noqa: E402
from hallu_lab import (BY_NAME, COUNTRIES, HEAD_TRAIN, QA_TASKS, TAIL_TRAIN, auroc, chat_probe, chat_question,  # noqa: E402
                       faq_prompt, get_longtail_model, get_sft_model, probe, sample_first)

KS = [2, 5, 10, 20]


def consistency_scores(model, tok, rows, prompt_of, k_max=20, temperature=1.0):
    """
    每个事实采样 k_max 次（第一个词），返回 {k: (一致性列表, 负的样本熵列表)}，k 取 KS。

    【思路】一致性 = 前 k 个样本里和贪心答案相同的比例；样本熵 = −Σ p log p（p 是前 k 个样本里每个答案的频率），
        取负号是为了"分数越高越可信"，和置信度方向一致。
    """
    samples = [sample_first(model, tok, prompt_of(r), k=k_max, temperature=temperature, seed=1) for r in rows]
    out = {}
    for k in KS:
        agree, negent = [], []
        for r, s in zip(rows, samples):
            ss = s[:k]
            agree.append(sum(x == r["pred"] for x in ss) / k)
            cnt = Counter(ss)
            negent.append(sum((v / k) * math.log(v / k) for v in cnt.values()))
        out[k] = (agree, negent)
    return out, samples


def demo1(model, tok, rows):
    """演示 1：一致性 vs 概率。"""
    print_header("演示 1：采样 k 次的一致性能区分对错吗？（长尾模型，240 个事实，温度 1）")
    ys = [r["correct"] for r in rows]
    scores, samples = consistency_scores(model, tok, rows, lambda r: faq_prompt(BY_NAME[r["country"]], r["rel"]))
    table = [["概率（需要 logprobs）", "1 次", f"{auroc([r['conf'] for r in rows], ys):.3f}", "–"]]
    for k in KS:
        a, e = scores[k]
        table.append([f"采样一致性（k = {k}）", f"{k + 1} 次（贪心 + k 次采样）", f"{auroc(a, ys):.3f}", f"{auroc(e, ys):.3f}"])
    print_table(["方法", "每个问题的生成次数", "AUROC（一致性 / 概率）", "AUROC（负样本熵）"], table)
    print_observations([
        f"只采样 2 次，一致性的 AUROC 就有 {auroc(scores[2][0], ys):.3f}；采样 10 次，{auroc(scores[10][0], ys):.3f}——"
        f"接近直接用概率（{auroc([r['conf'] for r in rows], ys):.3f}）。拿不到概率的黑盒 API，可以用'多问几遍'代替。",
        "原因很直接：模型知道的事实，下一个词的分布集中在正确答案上，采样几次都是它；不知道的事实，分布是散的，每次采到的不一样。"
        "一致性是在用采样**估计**那个分布有多集中——和概率是同一个信息，只是更贵、更粗。",
        "代价：k 次采样就是 k 倍的生成费用。SelfCheckGPT（Manakul et al. 2023）对长篇回答逐句检查和其他样本是否一致；"
        "语义熵（Kuhn et al. 2023；Farquhar et al. 2024）先把'意思相同'的回答归为一类再算熵——"
        "长回答里同一个意思有很多种说法，字面比较会把它们当成不一致。",
    ])
    return samples


def demo2(rows, samples):
    """演示 2：一致的幻觉。"""
    print_header("演示 2：一致性检测会漏掉什么？——'每次都编同一个答案'（k = 20）")
    table = []
    for lo, hi, label in [(0.8, 1.01, "高（≥ 0.8）"), (0.4, 0.8, "中（0.4–0.8）"), (0.0, 0.4, "低（< 0.4）")]:
        sel = [(r, s) for r, s in zip(rows, samples) if lo <= sum(x == r["pred"] for x in s) / 20 < hi]
        wrong = [r for r, _ in sel if not r["correct"]]
        table.append([label, len(sel), len(wrong), f"{100 * len(wrong) / max(1, len(sel)):.1f}%"])
    print_table(["一致性", "事实数", "其中答错（幻觉）", "错误率"], table)
    high_wrong = table[0][2]
    total_wrong = sum(1 for r in rows if not r["correct"])
    print_observations([
        f"一致性低的那一档错误率 {table[2][3]}，高的那一档只有 {table[0][3]}：在这个 base 模型上，一致性是一个很好的信号。",
        f"但高一致性的那一档里仍有 {high_wrong} 个错误答案（全部 {total_wrong} 个错误里的 {100 * high_wrong / total_wrong:.0f}%）："
        "模型**每次都编同一个**答案。这种'自信的、稳定的'幻觉在一致性上和真知识看起来一模一样，检测不出来。"
        "base 模型上它很少；演示 3 会看到，某些微调会让它大量增加。",
        "这对应第 4 讲的'系统性错误'：投票修不好它们，一致性也查不出它们。要发现它们，需要模型之外的信息：检索资料来核对"
        "（第 7 讲）、用另一个独立的模型交叉检查、或者让模型逐条验证自己的说法（第 10 讲的验证链）。",
    ])


def demo3():
    """演示 3：教过不知道的事实之后。"""
    print_header("演示 3：在指令微调过的模型上做一致性检测（对话格式；没进过微调数据的国家）")
    trained = {c.name for c in HEAD_TRAIN + TAIL_TRAIN}
    facts = [(c, rel) for c in COUNTRIES if c.name not in trained for rel in QA_TASKS]
    table = []
    res = {}
    for variant, label in [("known", "只用知道的事实"), ("teach", "也教它不知道的事实")]:
        m, t = get_sft_model(variant)
        from instruct_lab import render_chat
        rows = chat_probe(m, t, facts)
        ys = [r["correct"] for r in rows]
        scores, samples = consistency_scores(m, t, rows, lambda r: render_chat(chat_question(BY_NAME[r["country"]], r["rel"], 0)))
        a10 = scores[10][0]
        stable_wrong = sum(1 for r, s in zip(rows, samples) if not r["correct"] and sum(x == r["pred"] for x in s) / 20 >= 0.8)
        res[variant] = (auroc([r["conf"] for r in rows], ys), auroc(a10, ys), stable_wrong, sum(not y for y in ys))
        table.append([label, f"{100 * sum(ys) / len(ys):.1f}%", f"{res[variant][0]:.3f}", f"{res[variant][1]:.3f}",
                      f"{stable_wrong} / {res[variant][3]}"])
    print_table(["微调配方", "准确率", "AUROC（概率）", "AUROC（一致性，k = 10）", "稳定的幻觉 / 全部错误"], table)
    k, t = res["known"], res["teach"]
    print_observations([
        f"只用知道的事实微调：一致性的 AUROC {k[1]:.3f}；也教了它不知道的事实之后，{t[1]:.3f}。稳定的幻觉（一致性 ≥ 0.8 的错误答案）"
        f"从 {k[2]} 个变成 {t[2]} 个。",
        "用没有依据的事实做微调，让模型在不知道的问题上也给出**稳定**的答案（02 号脚本看到它同时变得更自信）。"
        "幻觉越稳定，靠模型自己的信号（概率、一致性）就越难发现——这是 Gekhman et al.（2024）的警告在检测这一侧的样子。",
    ])


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


if __name__ == "__main__":
    main()