"""
=============================================================================
第 12 讲 · 工具模块 2：checkers.py —— 模拟 NLI、claim 分解、groundedness 打分、SelfCheck、评估指标
=============================================================================

对应讲义：lectures/lecture12_grounding_hallucination.md §5（有参考的检测）、§6（无参考的检测）、§8（评估检测器）

-----------------------------------------------------------------------------
【这个模块提供什么 What's inside】
-----------------------------------------------------------------------------
    · nli()                 一个**模拟的** NLI 模型：给前提（段落里的事实）和假设（一个或几个 claim），
                            输出 (P(entail), P(contradict), P(neutral))。它有 4 个已知的弱点，每个都是一个旋钮：
                              1) 假设里有代词（没做 decontextualization）时，默认代词指前提里的那个人；
                              2) 名字只差一个中间名缩写时，有一定概率当成同一个人；
                              3) 数字差 1–2 时，有一定概率判成"支持"；
                              4) 前提太长（事实太多）时，有一定概率漏看；
                            另外，假设里有多个 claim 时，"部分支持"会被打一个**宽松的**中间分（partial_gamma）。
    · decompose()           把回答拆成 decontextualized 的原子 claim（可以模拟漏拆）。
    · claim_support()       一个 claim 对一组段落的支持分：先对每段算 entailment，再用 noisy-OR 或 max 汇总
                            （Fall 2023 Lecture 9 的 corroboration workflow）。
    · sentence_scores() / response_level_score()   句子级、整段回答级的同类打分（用来和 claim 级比较）。
    · selfcheck_scores()    SelfCheckGPT-NLI 风格的采样一致性分数（Manakul et al. 2023）；
                            selfcheck_matrix() 返回逐个采样的分数（用来扫描采样数 N）。
    · 指标：prf（precision / recall / F1）、roc_auc、wilson（比例的 95% 置信区间）、pearson、
            word_levenshtein（RARR 的 preservation）、mean。
    · 打印：print_header、table。

-----------------------------------------------------------------------------
【为什么 NLI 也要模拟】
-----------------------------------------------------------------------------
真实的 NLI 模型（TRUE 论文里最强的是在 ANLI 上训练的 T5-11B）很准，但并不完美，而且它们犯错的**方式**
已经被反复记录：代词和上下文（Choi et al. 2021 的 decontextualization）、数字、长前提、多 claim 的句子
（Fall 2023 Lecture 9 slides 的 "Failure Patterns"）。我们把这些方式写成旋钮，就能问：
**每一种弱点会让哪一类幻觉漏网？** 这是讲义 §5.4 的内容。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 checkers.py      # 自检：NLI 的几个典型输入、noisy-OR 的放大效应、指标函数（< 1 秒）
"""

import math
import random
import unicodedata
from dataclasses import dataclass

from fact_world import NUMERIC, Claim


# =============================================================================
# 打印与小工具
# =============================================================================
def print_header(title):
    """打印一个醒目的分隔标题（和第 1 讲的同名函数一样）。"""
    print("\n" + "=" * 100)
    print(title)
    print("=" * 100)


def dwidth(s):
    """
    字符串在终端里占几列：中文等"全角"字符占 2 列，英文占 1 列。

    【Python 小知识】unicodedata.east_asian_width(ch) 返回 "W"（wide）或 "F"（fullwidth）时是全角字符。
        str.ljust 按字符数补空格，中文会对不齐，所以表格要用这个函数算宽度。
    """
    return sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in s)


def pad(s, w, left=True):
    """按显示宽度补空格到 w 列（left=True 左对齐，否则右对齐）。"""
    gap = max(0, w - dwidth(s))
    return s + " " * gap if left else " " * gap + s


def mean(xs):
    """
    平均数。

    【Python 小知识】math.fsum 是"精确求和"：Python 3.12 起内置 sum() 对浮点数也做了补偿求和，
        和 3.9 的结果在最后一位可能不同；统一用 fsum，保证不同版本打印出的数字逐字相同。
    【边界】空列表返回 float('nan')。
    """
    xs = list(xs)
    return math.fsum(xs) / len(xs) if xs else float("nan")


def table(header, rows, widths=None):
    """
    打印一张对齐的表。

    【参数】header：列名列表；rows：每行一个列表（元素可以是字符串或数字）；widths：每列宽度（默认自动）。
    【思路】浮点数统一保留 3 位小数；第一列左对齐，其余右对齐。
    """
    def cell(x):
        if isinstance(x, float):
            return "nan" if math.isnan(x) else "%.3f" % x
        return str(x)
    srows = [[cell(x) for x in r] for r in rows]
    if widths is None:
        widths = [max(dwidth(str(h)), *(dwidth(r[i]) for r in srows)) + 2 for i, h in enumerate(header)]
    line = pad(str(header[0]), widths[0]) + "".join(pad(str(h), w, False) for h, w in zip(header[1:], widths[1:]))
    print(line)
    print("-" * dwidth(line))
    for r in srows:
        print(pad(r[0], widths[0]) + "".join(pad(x, w, False) for x, w in zip(r[1:], widths[1:])))


# =============================================================================
# 模拟 NLI
# =============================================================================
@dataclass
class NLIConfig:
    """
    模拟 NLI 的旋钮。

    【字段】
        p_miss            一般性的漏判：本来支持的 claim 被判成 neutral 的概率（措辞不同、需要一步推理……）
        p_name_confusion  名字只差中间名缩写（同名者）时，当成同一个人的概率
        p_num_blind       数字差 ≤ num_tol 时，判成"支持"的概率（NLI 对数字不敏感）
        num_tol           上面这条的容差（年）
        long_facts        前提里的事实数超过它，就算"长前提"
        p_long_miss       长前提里，本来支持的 claim 被漏看（判成 neutral）的概率
        partial_gamma     多 claim 假设的宽松程度：P(entail) 的基准 = 0.95 × (支持的比例)^gamma。
                          gamma = 1：按比例给分；gamma < 1：宽松（一半支持就给 0.67）；gamma 大：严格
        noise             输出概率上的噪声（半正态的标准差）
    """
    p_miss: float = 0.05
    p_name_confusion: float = 0.3
    p_num_blind: float = 0.5
    num_tol: int = 2
    long_facts: int = 4
    p_long_miss: float = 0.3
    partial_gamma: float = 0.5
    noise: float = 0.06


def _claim_status(world, prem_facts, h, decontext, rng, cfg):
    """
    判断一个 claim h 相对于前提事实的"真实 NLI 关系"（再叠加模拟 NLI 的弱点）。

    【返回】"support" / "contradict" / "neutral"
    【思路】遍历前提里属性相同的事实：
        · 主体是不是同一个人？做了 decontextualization 时按名字判断（同名者有 p_name_confusion 的概率被混淆）；
          没做时（代词 "This researcher"），NLI 默认代词指前提里的那个人 —— 于是**任何人**的同一属性都会被拿来比。
        · 值相同 → support；值不同 → contradict，但数字差得很少时有 p_num_blind 的概率误判成 support。
        · 只要有一条 support 就返回 support（前提里同时有支持和矛盾时，NLI 往往"看到想看的"）。
    """
    status = "neutral"
    for feid, frel, fval in prem_facts:
        if frel != h.relation:
            continue
        if not decontext:
            same = True
        elif feid == h.eid:
            same = True
        else:
            ns = world.entities[h.eid].namesake
            same = (feid == ns) and rng.random() < cfg.p_name_confusion
        if not same:
            continue
        if fval == h.value:
            status = "support"
            break
        if h.relation in NUMERIC and abs(fval - h.value) <= cfg.num_tol and rng.random() < cfg.p_num_blind:
            status = "support"
            break
        status = "contradict"
    if status == "support" and rng.random() < cfg.p_miss:
        status = "neutral"
    if status == "support" and len(prem_facts) > cfg.long_facts and rng.random() < cfg.p_long_miss:
        status = "neutral"
    return status


def nli(world, prem_facts, hyp_claims, decontext, rng, cfg=None):
    """
    模拟 NLI：前提 = 一组事实，假设 = 一个或多个 claim。返回 (e, c, n) 三个概率（和为 1）。

    【参数】world：用来查同名者；prem_facts：[(eid, relation, value), ...]；hyp_claims：Claim 列表
        （空列表 = 不含任何事实的句子，比如客套话）；decontext：假设是否已经 decontextualized；
        rng：随机数生成器；cfg：NLIConfig。
    【思路】先对每个 claim 求关系，再汇总：
            fs = 支持的比例，fc = 矛盾的比例
            e_base = 0.95·fs^γ，c_base = 0.95·fc^γ（两者之和超过 0.95 时按比例缩小）
            n_base = 1 − e_base − c_base
        三个基准值各加一点半正态噪声后归一化。所以即使是 neutral，P(entail) 也有 0.05 左右的"杂音"——
        真实 NLI 模型也是这样，它会在 noisy-OR 汇总时累积起来（§5.3）。
    【返回】(e, c, n)
    【边界】hyp_claims 为空（客套话）时 fs = fc = 0 → 判 neutral，于是"Sure! I can help"会被当成
        "没有依据"（slides 的 precision issue）。
    【可信联系】真实系统里这一步是一个 NLI 模型（TRUE 论文推荐的那类）。
    """
    cfg = cfg or NLIConfig()
    if hyp_claims:
        st = [_claim_status(world, prem_facts, h, decontext, rng, cfg) for h in hyp_claims]
        fs = st.count("support") / len(st)
        fc = st.count("contradict") / len(st)
    else:
        fs = fc = 0.0
    g = cfg.partial_gamma
    e_base = 0.95 * fs ** g if fs > 0 else 0.0
    c_base = 0.95 * fc ** g if fc > 0 else 0.0
    tot = e_base + c_base
    if tot > 0.95:
        e_base, c_base = e_base * 0.95 / tot, c_base * 0.95 / tot
    n_base = max(0.0, 1.0 - e_base - c_base)
    e = e_base + 0.02 + abs(rng.gauss(0, cfg.noise))
    c = c_base + 0.02 + abs(rng.gauss(0, cfg.noise))
    n = n_base + 0.02 + abs(rng.gauss(0, cfg.noise))
    s = e + c + n
    return e / s, c / s, n / s


# =============================================================================
# claim 分解与 groundedness 打分
# =============================================================================
def decompose(response, rng, p_drop=0.0):
    """
    把回答拆成原子 claim（每个都是 decontextualized 的：主语是全名）。

    【思路】真实系统里这一步由 LLM 完成（"把下面的回答拆成独立的事实"），也可以退化成"按句子切"。
        这里 claim 结构是已知的，所以分解是精确的；p_drop 模拟抽取器**漏掉** claim 的概率
        —— 漏掉的 claim 永远不会被检查（§5.4 decomposition errors）。
    【返回】Claim 列表
    【边界】客套话没有 claim，自然被丢掉（相当于 slides 里的 "requires verification" 过滤）。
    """
    out = []
    for c in response.claims():
        if p_drop > 0 and rng.random() < p_drop:
            continue
        out.append(c)
    return out


def noisy_or(ps):
    """
    s = 1 − ∏(1 − p_j)：只要有一段支持，整体就支持（Fall 2023 Lecture 9 的 "OR operator"）。

    【思路】如果把每段的 entailment 概率看成"独立的证据"，noisy-OR 就是"至少一段支持"的概率。
    【边界】证据并不独立：k 段**无关**段落各有 0.05 的杂音，k = 10 时也会累积到 0.4 左右（§5.3）。
    """
    prod = 1.0
    for p in ps:
        prod *= (1.0 - p)
    return 1.0 - prod


def claim_support(world, claim, passages, rng, cfg=None, agg="noisy_or", decontext=True):
    """
    一个 claim 相对于一组段落的支持分（以及引用哪一段）。

    【参数】claim：Claim；passages：Passage 列表；agg："noisy_or" 或 "max"；decontext：claim 是否已经
        decontextualized（False 时模拟"直接拿带代词的句子去问 NLI"）。
    【返回】(score, best_pid, contra)：score ∈ [0,1] 是汇总后的支持分；best_pid 是 entailment 最高的段落
        编号（用来做引用 citation）；contra 是各段 P(contradict) 的最大值。
    【边界】passages 为空 → score = 0、best_pid = -1。
    """
    if not passages:
        return 0.0, -1, 0.0
    es, cs = [], []
    for p in passages:
        e, c, _ = nli(world, p.facts, [claim], decontext, rng, cfg)
        es.append(e)
        cs.append(c)
    best = max(range(len(es)), key=lambda j: (es[j], -j))
    score = noisy_or(es) if agg == "noisy_or" else max(es)
    return score, passages[best].pid, max(cs)


def sentence_scores(world, response, passages, rng, cfg=None, agg="noisy_or"):
    """
    句子级打分：每句话（可能含 0、1、2 个 claim）对每段算 NLI，再汇总。

    【思路】带代词的句子**没有** decontextualized（这正是"按句子切"会遇到的问题）。
    【返回】[(Sentence, score), ...]
    """
    out = []
    for s in response.sentences:
        es = [nli(world, p.facts, s.claims, not s.uses_pronoun, rng, cfg)[0] for p in passages]
        score = (noisy_or(es) if agg == "noisy_or" else max(es)) if es else 0.0
        out.append((s, score))
    return out


def response_level_score(world, response, passages, rng, cfg=None):
    """
    整段回答级打分：把所有检索段落拼成一个长前提，把整段回答当成一个假设，问一次 NLI。

    【思路】这是最省事的做法（"直接问上下文是否蕴含整个回答"，Fall 2024 slides 也提到可以跳过分解）。
        它的问题：前提很长（漏看）、假设里有多个 claim（部分支持被宽松地打分），而且只给一个数，
        没法告诉你**哪一句**有问题。
    """
    facts = [f for p in passages for f in p.facts]
    return nli(world, facts, response.claims(), True, rng, cfg)[0]


# =============================================================================
# SelfCheckGPT 风格的采样一致性
# =============================================================================
def selfcheck_scores(world, claims, samples, rng, cfg=None):
    """
    SelfCheckGPT-NLI：对主回答里的每个 claim，看它和 N 个随机采样的回答是否矛盾。

    【思路】Manakul et al.（2023）的 NLI 变体：对每个采样 s，score_s = P(contradict) / (P(entail) + P(contradict))，
        再对 N 个采样取平均。分数越高越可能是幻觉。直觉：模型**知道**的事实，每次采样都会说一样的话；
        **编**出来的事实，每次编得不一样。
    【参数】claims：主回答的 claim 列表；samples：Response 列表（同一个问题、较高温度的采样）。
    【返回】与 claims 等长的分数列表。
    【边界】某个采样没有提到这个属性 → NLI 输出 neutral，e 和 c 都很小，比值接近 0.5（一个噪声项）；
        这是 SelfCheck 的真实性质：采样越"啰嗦"、覆盖的属性越全，分数越干净。
    【可信联系】它检测的是**不一致**，不是**错误**：每次都错得一样的 claim（记错的事实）得分很低（§6.1）。
    """
    return [mean(row) for row in selfcheck_matrix(world, claims, samples, rng, cfg)]


def selfcheck_matrix(world, claims, samples, rng, cfg=None):
    """
    selfcheck_scores 的"展开版"：返回 claims × samples 的矩阵，第 i 行第 j 列是 claim i 对采样 j 的
    P(contradict) / (P(entail) + P(contradict))。

    【用途】想看"采样数 N 的影响"时，只要对每行的前 N 个取平均（不必重新采样）。
    """
    cfg = cfg or NLIConfig()
    sample_facts = [[(c.eid, c.relation, c.value) for c in s.claims()] for s in samples]
    out = []
    for h in claims:
        row = []
        for facts in sample_facts:
            e, c, _ = nli(world, facts, [h], True, rng, cfg)
            row.append(c / (e + c))
        out.append(row)
    return out


# =============================================================================
# 评估指标
# =============================================================================
def prf(flags, labels):
    """
    precision / recall / F1。

    【参数】flags：检测器是否报警（bool 列表）；labels：真相（True = 真的是幻觉）。
    【返回】(precision, recall, f1, n_flagged)；没有报警时 precision = nan。
    """
    tp = sum(1 for f, y in zip(flags, labels) if f and y)
    fp = sum(1 for f, y in zip(flags, labels) if f and not y)
    fn = sum(1 for f, y in zip(flags, labels) if (not f) and y)
    p = tp / (tp + fp) if tp + fp else float("nan")
    r = tp / (tp + fn) if tp + fn else float("nan")
    f1 = 2 * p * r / (p + r) if (tp and not math.isnan(p) and not math.isnan(r)) else 0.0
    return p, r, f1, tp + fp


def roc_auc(scores, labels):
    """
    ROC AUC = P(随机一个正例的分数 > 随机一个负例的分数)，平局算 1/2（Mann–Whitney 统计量）。

    【思路】按分数排序后给每个值分配平均秩（相同分数取平均秩），正例秩和减去最小可能值，再除以正负对数。
    【边界】没有正例或没有负例时返回 nan。
    【Python 小知识】sorted(range(n), key=lambda i: scores[i]) 得到"按分数从小到大"的下标顺序。
    """
    n = len(scores)
    order = sorted(range(n), key=lambda i: scores[i])
    ranks = [0.0] * n
    i = 0
    while i < n:
        j = i
        while j + 1 < n and scores[order[j + 1]] == scores[order[i]]:
            j += 1
        avg = (i + j) / 2.0 + 1.0
        for k in range(i, j + 1):
            ranks[order[k]] = avg
        i = j + 1
    n_pos = sum(1 for y in labels if y)
    n_neg = n - n_pos
    if n_pos == 0 or n_neg == 0:
        return float("nan")
    rank_sum = math.fsum(r for r, y in zip(ranks, labels) if y)
    return (rank_sum - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg)


def wilson(k, n, z=1.96):
    """
    比例 k/n 的 Wilson 95% 置信区间。

    【思路】比"p ± 1.96·√(p(1−p)/n)"更好：p 接近 0 或 1、n 很小时也不会越出 [0, 1]。
    【返回】(lo, hi)；n = 0 时返回 (nan, nan)。
    """
    if n == 0:
        return float("nan"), float("nan")
    p = k / n
    d = 1 + z * z / n
    centre = (p + z * z / (2 * n)) / d
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return max(0.0, centre - half), min(1.0, centre + half)


def pearson(xs, ys):
    """皮尔逊相关系数；任一序列方差为 0 时返回 nan。"""
    mx, my = mean(xs), mean(ys)
    sxy = math.fsum((x - mx) * (y - my) for x, y in zip(xs, ys))
    sxx = math.fsum((x - mx) ** 2 for x in xs)
    syy = math.fsum((y - my) ** 2 for y in ys)
    if sxx == 0 or syy == 0:
        return float("nan")
    return sxy / math.sqrt(sxx * syy)


def word_levenshtein(a, b):
    """
    两段文字之间的**词级**编辑距离（插入、删除、替换各算 1）。

    【思路】经典动态规划，只保留上一行，内存 O(len(b))。RARR 用的是字符级的 Levenshtein；
        这里用词级只是为了快（传记只有几十个词），结论不变。
    【Python 小知识】a.split() 按空白切词；prev[:] 之类的切片会复制列表。
    """
    x, y = a.split(), b.split()
    prev = list(range(len(y) + 1))
    for i in range(1, len(x) + 1):
        cur = [i] + [0] * len(y)
        for j in range(1, len(y) + 1):
            cost = 0 if x[i - 1] == y[j - 1] else 1
            cur[j] = min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + cost)
        prev = cur
    return prev[len(y)]


# =============================================================================
# 自检
# =============================================================================
if __name__ == "__main__":
    from fact_world import World
    w = World()
    rng = random.Random(1)
    e = next(x for x in w.entities if x.namesake >= 0)
    ns = w.entities[e.namesake]
    prem = [(e.eid, "born_year", e.facts["born_year"]), (e.eid, "born_city", e.facts["born_city"])]
    print("前提：", e.name, prem)
    tests = [
        ("对的年份", Claim(e.eid, "born_year", e.facts["born_year"], "x", e.name), True),
        ("差 1 年", Claim(e.eid, "born_year", e.facts["born_year"] + 1, "x", e.name), True),
        ("差 10 年", Claim(e.eid, "born_year", e.facts["born_year"] + 10, "x", e.name), True),
        ("同名者的值、带代词", Claim(ns.eid, "born_year", e.facts["born_year"], "x", ns.name), False),
        ("同名者的值、全名", Claim(ns.eid, "born_year", e.facts["born_year"], "x", ns.name), True),
        ("前提没提到的属性", Claim(e.eid, "award", e.facts["award"], "x", e.name), True),
    ]
    for label, h, dec in tests:
        avg = [mean(nli(w, prem, [h], dec, rng)[i] for _ in range(400)) for i in range(3)]
        print("  " + pad(label, 20) + "  E=%.3f C=%.3f N=%.3f" % tuple(avg))
    print("noisy-OR of k copies of 0.06:", ["k=%d: %.2f" % (k, noisy_or([0.06] * k)) for k in (1, 4, 10, 20)])
    print("roc_auc([1,2,3,4],[0,0,1,1]) =", roc_auc([1, 2, 3, 4], [0, 0, 1, 1]),
          " wilson(5,10) = (%.3f, %.3f)" % wilson(5, 10),
          " levenshtein =", word_levenshtein("a b c d", "a x c d e"))