"""
=============================================================================
CS486 · 第 2 讲 · 代码 01：亲手制造数据污染——分数被抬高了多少？怎样发现？
=============================================================================

【这个脚本要回答的问题】（讲义 §5.2）
    "数据污染"（data contamination）：基准测试的题目出现在了模型的训练数据里。
    对网页规模预训练的 LLM，这几乎无法完全避免（基准题会被贴到网上、被讨论、被复制）。
    它会让模型在这个基准上的分数虚高，从而让研究结论失真。

    我们在第 1 讲的实验台上亲手制造污染：一个 420 道题的"基准"，分成 7 组，每组 60 道：
        干净（0 次）/ 见过 1 次 / 4 次 / 16 次 / 64 次 / 16 次但都在训练早期 / 16 次但都在训练后期
    这些题被插进 2000 步的训练数据流里（只占全部训练题的约 5.5%）。

    演示 1：两个被污染的模型（直接回答、思维链）和一个**完全一样但没被污染**的对照模型（第 1 讲的
            direct_L2），在"原基准"和"孪生基准"（同分布的新题）上的分数。
    演示 2：剂量–反应：每组题的准确率——见过几次、什么时候见过，各有多大影响？
    演示 3：扰动检测：给题目改名（字母置换）、打乱陈述顺序，被"背下来"的题还能答对吗？
    演示 4：成员推断检测：只看模型的自信程度，能分辨哪些题被污染了吗？（AUC）

【运行】
    python3 01_contamination.py           # 第一次约 1.5 分钟（训练 2 个模型；若没有第 1 讲的缓存再加约 30 秒），之后约 10 秒
    python3 01_contamination.py --quick   # 冒烟测试，几秒钟，数字没有意义

【依赖】PyTorch + 本目录的 method_lab.py + 第 1 讲的 hop_world.py、tiny_reasoner.py、research_kit.py、console.py。
"""

import argparse
import os
import random
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

import method_lab as ml                         # noqa: E402
import tiny_reasoner as tr                      # noqa: E402
from console import pct, print_header, print_observations, print_table   # noqa: E402
from research_kit import wilson_interval        # noqa: E402

GROUP_NAMES = [g[0] for g in ml.CONTAM_GROUPS]


def acc_ci(results):
    """准确率和 Wilson 95% 区间，写成 "xx.x% [lo, hi]"。"""
    c = sum(r["correct"] for r in results)
    lo, hi = wilson_interval(c, len(results))
    return f"{pct(c / len(results))} [{pct(lo, 0)}, {pct(hi, 0)}]"


def group_acc(results, groups, name):
    """某一组题的准确率。"""
    ids = groups[name]
    return sum(results[i]["correct"] for i in ids) / len(ids)


def demo1_headline(models, bench, twin):
    print_header("演示 1：原基准 vs 孪生基准（同分布的新题）——'报告的分数'被抬高了多少？")
    rows, out = [], {}
    for label, (model, mode) in models.items():
        rb = tr.evaluate(model, bench, mode)
        rt = tr.evaluate(model, twin, mode)
        out[label] = (tr.accuracy(rb), tr.accuracy(rt), rb)
        rows.append([label, acc_ci(rb), acc_ci(rt), f"{100 * (tr.accuracy(rb) - tr.accuracy(rt)):+.1f}"])
    print_table(["模型", "原基准（420 道）", "孪生基准（420 道）", "差（点）"], rows, [30, 20, 20, 10])
    return out


def demo2_dose(bench, groups, headline):
    print_header("演示 2：剂量–反应——每组 60 道题的准确率")
    labels = list(headline)
    rows = []
    for name, reps, when in ml.CONTAM_GROUPS:
        rows.append([name] + [pct(group_acc(headline[lab][2], groups, name)) for lab in labels])
    print_table(["组"] + labels, rows, [22, 26, 24, 26])


def demo3_perturb(model, mode, bench, groups):
    print_header("演示 3：扰动检测（被污染的直接回答模型）——改名 / 打乱顺序之后还答得对吗？")
    rng = random.Random(9)
    renamed = [ml.rename_letters(p, rng) for p in bench]
    reordered = [ml.reorder_statements(p, rng) for p in bench]
    r0 = tr.evaluate(model, bench, mode)
    r1 = tr.evaluate(model, renamed, mode)
    r2 = tr.evaluate(model, reordered, mode)
    rows, out = [], {}
    for name, _, _ in ml.CONTAM_GROUPS:
        a = [group_acc(r, groups, name) for r in (r0, r1, r2)]
        out[name] = a
        rows.append([name] + [pct(x) for x in a])
    print_table(["组", "原题", "字母改名", "陈述换序"], rows, [22, 10, 10, 10])
    return out


def demo4_membership(model, mode, bench, groups):
    print_header("演示 4：成员推断——只看模型的自信程度，能分出被污染的题吗？（AUC，相对'干净'组）")
    max_p, gold_p = ml.answer_scores(model, bench, mode)
    clean = groups[GROUP_NAMES[0]]
    rows, out = [], {}
    for name in GROUP_NAMES[1:]:
        ids = groups[name]
        a1 = ml.auc([max_p[i] for i in ids], [max_p[i] for i in clean])
        a2 = ml.auc([gold_p[i] for i in ids], [gold_p[i] for i in clean])
        out[name] = (a1, a2)
        rows.append([name, f"{a1:.3f}", f"{a2:.3f}"])
    print_table(["组", "最大概率 AUC", "正确答案概率 AUC"], rows, [22, 14, 18])
    print("（0.5 = 和抛硬币一样分不开；1.0 = 完美分开。'最大概率'不需要知道答案，'正确答案概率'需要。）")
    return out


def main():
    ap = argparse.ArgumentParser(description="亲手制造数据污染")
    ap.add_argument("--quick", action="store_true", help="冒烟测试：每个模型只训练 60 步")
    args = ap.parse_args()
    t0 = time.perf_counter()

    print_header("准备：基准、污染日程与模型")
    bench, groups = ml.make_benchmark()
    twin = ml.twin_benchmark()
    rows = [[name, reps, {"all": "全程均匀", "early": "前 25% 的步数", "late": "最后 25% 的步数"}[when]]
            for name, reps, when in ml.CONTAM_GROUPS]
    print_table(["组（各 60 道）", "插入次数", "插入时间"], rows, [22, 10, 16], ["<", ">", "<"])
    direct, info = ml.contaminated_model("direct", quick=args.quick)
    cot, _ = ml.contaminated_model("cot", quick=args.quick)
    clean, _, _ = tr.get_model("direct_L2", quick=args.quick)
    print(f"  污染题共插入 {info['n_contam']} 次，占全部训练题的 "
          f"{pct(info['n_contam'] / (64 * (60 if args.quick else 2000)))}")
    models = {"直接回答（被污染）": (direct, "direct"),
              "直接回答（干净对照 direct_L2）": (clean, "direct"),
              "思维链（被污染）": (cot, "cot")}

    head = demo1_headline(models, bench, twin)
    demo2_dose(bench, groups, head)
    pert = demo3_perturb(direct, "direct", bench, groups)
    mem = demo4_membership(direct, "direct", bench, groups)

    if args.quick:
        print("\n（--quick 冒烟测试：模型只训练了 60 步，上面的数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    d_b, d_t, d_res = head["直接回答（被污染）"]
    c_b, c_t, _ = head["直接回答（干净对照 direct_L2）"]
    g = {name: group_acc(d_res, groups, name) for name in GROUP_NAMES}
    print_observations([
        f"被污染的直接回答模型在原基准上 {pct(d_b)}，在孪生基准上只有 {pct(d_t)}；没被污染的对照模型在两者上是 "
        f"{pct(c_b)} 和 {pct(c_t)}。只有约 5.5% 的训练题来自基准，就让'报告的分数'高了 "
        f"{100 * (d_b - d_t):.1f} 个点——而这个模型对这类题的真实能力（孪生基准）和对照模型一样，几乎是随机水平。",
        f"思维链模型被同样污染，在两个基准上都是 {pct(head['思维链（被污染）'][0])}：已经学会算法的模型，"
        "背不背题没有区别。污染最抬高的是**弱**模型的分数，于是它会**压缩**强弱模型之间的差距、扭曲排名。",
        f"剂量–反应（演示 2）：见过 1 / 4 / 16 / 64 次的题，准确率 {pct(g['见过 1 次'])} / {pct(g['见过 4 次'])} / "
        f"{pct(g['见过 16 次'])} / {pct(g['见过 64 次'])}（干净组 {pct(g[GROUP_NAMES[0]])}）。每组只有 60 道题，"
        "±10 个点以内的差别是噪声（比如'见过 1 次'比干净组还低）；但 16 次和 64 次的抬高远超噪声。"
        f"同样 16 次，集中在训练早期的只剩 {pct(g['16 次，都在训练早期'])}（接近干净组，被后面的训练'冲掉'了），"
        f"集中在后期的是 {pct(g['16 次，都在训练后期'])}。'重复次数'和'离训练结束有多近'都重要："
        "微调阶段混进的基准题比预训练早期的更危险。",
        f"扰动检测（演示 3）：见过 64 次的题，字母改名后从 {pct(pert['见过 64 次'][0])} 掉到 {pct(pert['见过 64 次'][1])}，"
        f"但只打乱陈述顺序时仍有 {pct(pert['见过 64 次'][2])}。模型记住的不是'这一串字符'，而是'这些字母、这些箭头 → 这个答案'"
        "的对应关系——注意力对陈述的排列顺序本来就不太敏感。所以扰动必须改掉模型赖以记忆的特征（这里是字母本身）"
        "才有检测力；只换顺序的扰动会漏掉污染。GSM-Symbolic 换人名和数字、GSM1k 重新出题，都是在改'记忆的钥匙'。",
        f"成员推断（演示 4）：不需要知道答案的'最大概率'只能分出见过 64 次的题（AUC {mem['见过 64 次'][0]:.3f}），"
        f"对见过 16 次、准确率已经被抬高的题只有 {mem['见过 16 次'][0]:.3f}；需要知道答案的'正确答案概率'灵敏得多"
        f"（16 次 {mem['见过 16 次'][1]:.3f}，后期 16 次 {mem['16 次，都在训练后期'][1]:.3f}），但对见过 1–4 次的题仍接近 0.5。"
        "基于置信度 / 损失的检测能抓到'背得熟'的题，抓不到轻度污染——重新出题（孪生基准）是更可靠的检查。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()