"""
=============================================================================
第 2 讲 · 代码 05：挑示例、排顺序 —— 检索相似的示例（KATE）与熵准则（Lu et al. 2022）
=============================================================================

对应讲义：lectures/lecture02_in_context_learning.md §7

【这个脚本要回答的问题】
    既然示例的选择和顺序影响这么大，能不能**有原则地**挑？两种代表性的做法：
    1. 按相似度检索示例（Liu et al. 2022，KATE）：每道测试题，挑和它最像的几个训练样本当示例。（演示 1）
    2. 不用标注数据挑顺序（Lu et al. 2022）：让模型自己生成一批"探测"输入，看每种顺序下模型在探测集上的预测分布——
       预测分布越均衡（熵越大），说明这种顺序越不偏，就选它。（演示 2：真实的迷你模型；演示 3：SimICL）

【怎么读输出】
    "配对检验"在同一组 300 道题上比较两种挑法（前置讲 P3 §3 的 McNemar）。
    GlobalE = 探测集上"预测标签的分布"的熵；LocalE = 探测集上"每条预测的概率分布"的熵的平均。

【运行】python3 05_selection_order.py     （约 40 秒；第一次运行还要训练 bursty 模型约 45 秒）
【依赖】PyTorch；import 同目录的 icl_lab.py，第 1 讲的 console.py、tiny_gpt.py，前置讲 P3 的 stats_kit.py。
"""

import itertools
import math
import os
import random
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "prereq3"))
from console import print_header, print_observations, print_table  # noqa: E402
from icl_lab import (TEMPLATES, VERBALIZERS, SimICL, accuracy, by_kind, correct_vector, demo_pool,  # noqa: E402
                     get_base_model, get_bursty_model, label_probs, render, render_english, sample_demos, test_set)
from stats_kit import mcnemar_exact  # noqa: E402


def bow_cosine(a, b):
    """
    两段文字的词袋余弦相似度：共同词数 / sqrt(词数a × 词数b)（每个词只算一次）。

    【为什么用这么简单的相似度】KATE 用的是 RoBERTa 句向量；这里用词袋，是为了一眼看出"像"是什么意思：
        共享方面词（soup、staff）和形容词。第 1 讲 §4 讲过词袋的盲点——它看不见 not 和 but 的作用。
    """
    A, B = set(a.split()), set(b.split())
    return len(A & B) / math.sqrt(len(A) * len(B)) if A and B else 0.0


def demo1_retrieval(models, tok, test, pool):
    """演示 1：随机示例 vs 最相似的示例 vs 最不相似的示例。"""
    print_header("演示 1：按相似度检索示例（每道题各挑各的；示例池 400 条任意长度的评论）")
    rows, res = [], {}
    for mname, m in models.items():
        for tname, vname in [("sentiment", "positive/negative"), ("sentiment", "great/awful")]:
            t, verb = TEMPLATES[tname], VERBALIZERS[vname]
            words = [verb["positive"], verb["negative"]]
            for k in [1, 2]:
                rng = random.Random(0)
                prompts = {"随机": [], "最相似": [], "最不相似": []}
                same_label = {"随机": 0, "最相似": 0, "最不相似": 0}
                for x, y, _ in test:
                    ranked = sorted(pool.all, key=lambda d: (bow_cosine(d[0], x), d[0]))
                    picks = {"随机": rng.sample(pool.all, k), "最相似": ranked[-k:], "最不相似": ranked[:k]}
                    for key, demos in picks.items():
                        prompts[key].append(render(t, verb, demos, x))        # 最相似的放在最后（离测试题最近）
                        same_label[key] += sum(1 for d in demos if d[1] == y) / k
                corr = {}
                for key, ps in prompts.items():
                    probs = label_probs(m, tok, ps, words)
                    corr[key] = correct_vector(probs, test)
                    res[(mname, tname, vname, k, key)] = (accuracy(probs, test), by_kind(probs, test))
                n01, n10, p = mcnemar_exact(corr["随机"], corr["最相似"])
                res[(mname, tname, vname, k, "p")] = (n01, n10, p)
                rows.append([mname, f"{tname} · {vname}", k] +
                            [f"{100 * res[(mname, tname, vname, k, key)][0]:.1f}%" for key in prompts] +
                            [f"{100 * same_label['最相似'] / len(test):.0f}%", f"{n10} : {n01}（p = {p:.3f}）"])
    print_table(["模型", "设计", "k", "随机", "最相似", "最不相似", "最相似示例和题目同标签的比例",
                 "随机 vs 最相似（只有随机对 : 只有最相似对，McNemar）"], rows)
    sig = [key for key, v in res.items() if key[-1] == "p" and v[2] < 0.05]
    better = [key for key in sig if res[key][0] > res[key][1]]          # n01 > n10：最相似对、随机错的题更多
    cells = [key[:-1] for key in res if key[-1] == "p"]
    far_best = sum(1 for c in cells if res[c + ("最不相似",)][0] >= max(res[c + ("随机",)][0], res[c + ("最相似",)][0]))
    print_observations([
        f"8 种（模型 × 设计 × k）里，最相似的示例和随机示例的差别只有 {len(sig)} 种在 0.05 水平上显著"
        f"（其中最相似更好的 {len(better)} 种：bursty 的 positive/negative、k = 2）。检索相似的示例**没有**稳定地帮到这两个迷你模型。",
        "最相似的示例和题目同标签的比例是 71%–72%，比随机示例的约 50% 高——如果模型会'照抄相似示例的标签'，这本该有帮助。"
        "它没有帮助，和 04 号脚本的结论一致：这两个模型不读示例里'输入 → 标签'的对应关系。",
        f"意外的是，'最不相似'的示例在 {far_best} / 8 种情形下最好（或并列最好）。词袋意义下最不相似的，往往是比较长、"
        "有两个分句的评论；对这两个小模型，示例的长度和句式也许比'像不像'更要紧——这只是一个猜测，这里没有验证。",
        "Liu et al.（2022）在 GPT-3 上报告 KATE 在情感分类、表格到文本生成、开放域问答上明显优于随机示例；一个合理的解释是"
        "大模型会从相似的示例里'借'答案的模式。**一种挑示例的方法好不好，取决于模型怎样使用示例**——要在你的模型上测。",
    ])
    return res


def probe_set(model, tok, n=60, seed=0):
    """
    让模型自己"写"一批评论当探测集（Lu et al. 2022 的做法：不用任何标注数据）。

    【思路】从 "<doc> reviews : review :" 开始，温度 1 采样，写到 "rating" 或 "<doc>" 停下，去掉末尾的标点。
    【返回】n 条评论正文（可能有语法不通的，这没关系：探测集只用来看"预测分布是否均衡"）
    """
    from tiny_gpt import generate
    out = []
    for s in range(n):
        g = generate(model, tok, "<doc> reviews : review :", max_new_tokens=14, temperature=1.0,
                     stop=("rating", "<doc>"), seed=seed + s)
        out.append(g.replace("rating", "").replace("<doc>", "").strip().rstrip(" ."))
    return out


def entropy(ps):
    """一个概率分布的熵（自然对数）。"""
    return -sum(p * math.log(p) for p in ps if p > 0)


def score_orders(get_probs, demos_list, test, probe):
    """
    对每一种顺序，算测试集准确率、探测集上的 GlobalE 和 LocalE。

    【参数】get_probs(demos, 输入列表) → 概率列表；demos_list：所有候选顺序
    【返回】[{"order", "acc", "globalE", "localE"}, ...]
    """
    out = []
    for demos in demos_list:
        acc = accuracy(get_probs(demos, [x for x, _, _ in test]), test)
        pp = get_probs(demos, probe)
        frac = sum(1 for p in pp if p[0] > p[1]) / len(pp)
        out.append({"order": "".join("P" if d[1] == "positive" else "N" for d in demos), "acc": acc,
                    "globalE": entropy([frac, 1 - frac]), "localE": sum(entropy(p) for p in pp) / len(pp)})
    return out


def summarize_orders(scored):
    """按 GlobalE / LocalE 取前 4 名的平均准确率，和全体平均、最好、最差比较。"""
    accs = [r["acc"] for r in scored]
    top_g = sorted(scored, key=lambda r: -r["globalE"])[:4]
    top_l = sorted(scored, key=lambda r: -r["localE"])[:4]
    return {"mean": sum(accs) / len(accs), "best": max(accs), "worst": min(accs),
            "globalE": sum(r["acc"] for r in top_g) / 4, "localE": sum(r["acc"] for r in top_l) / 4,
            "g_last": "".join(r["order"][-1] for r in top_g)}


def demo2_orders(models, tok, test, pool):
    """演示 2：迷你 GPT 上用熵准则挑顺序。"""
    print_header("演示 2：用熵准则挑顺序（sentiment · great/awful；每组 4 个示例的 24 种顺序，5 组；探测集 = 模型自己写的 60 条评论）")
    t, verb = TEMPLATES["sentiment"], VERBALIZERS["great/awful"]
    words = [verb["positive"], verb["negative"]]
    rows, res = [], {}
    for mname, m in models.items():
        probe = probe_set(m, tok)

        def get_probs(demos, xs, m=m):
            return label_probs(m, tok, [render(t, verb, demos, x) for x in xs], words)
        sums = []
        for s in range(5):
            demos = sample_demos(pool, random.Random(100 + s), pattern="PPNN")
            sm = summarize_orders(score_orders(get_probs, [list(p) for p in itertools.permutations(demos)], test, probe))
            sums.append(sm)
            rows.append([mname, s, f"{100 * sm['worst']:.1f}%", f"{100 * sm['mean']:.1f}%", f"{100 * sm['best']:.1f}%",
                         f"{100 * sm['globalE']:.1f}%（最后一个：{sm['g_last']}）", f"{100 * sm['localE']:.1f}%"])
        res[mname] = {key: sum(x[key] for x in sums) / len(sums) for key in ("worst", "mean", "best", "globalE", "localE")}
        res[mname]["probe"] = probe[:4]
    print_table(["模型", "示例组", "最差顺序", "24 种的平均", "最好顺序", "GlobalE 前 4 名的平均", "LocalE 前 4 名的平均"], rows)
    for mname in models:
        print(f"  {mname} 自己写的探测评论（前 4 条）：{res[mname]['probe']}")
    b, u = res["base"], res["bursty"]
    print_observations([
        f"base 模型：GlobalE 前 4 名平均 {100 * b['globalE']:.1f}%，比 24 种顺序的平均（{100 * b['mean']:.1f}%）还差，"
        f"接近最差的顺序（{100 * b['worst']:.1f}%）。",
        f"bursty 模型：GlobalE 前 4 名平均 {100 * u['globalE']:.1f}%，远高于 24 种的平均（{100 * u['mean']:.1f}%），"
        f"几乎追上最好的顺序（{100 * u['best']:.1f}%）。",
        "两个模型上，熵准则挑中的都是'最后一个示例是好评'的顺序（表里的 PPPP）。02 号脚本演示 2 发现：这类顺序对 bursty 更好、"
        "对 base 更差。同一个准则、同样的选择，一个模型受益，另一个受害。",
        "熵准则的假设是'预测越均衡越好'，而'均衡'是对探测集说的：模型自己写的评论好坏各占多少、像不像真实的测试题，都会影响结论。"
        f"LocalE（每条预测本身的不确定性）在两个模型上都不如 GlobalE（base {100 * b['localE']:.1f}%、bursty {100 * u['localE']:.1f}%）。",
        "Lu et al.（2022）在 GPT 系列模型、11 个文本分类任务上报告，用熵准则挑顺序带来约 13% 的相对提升。"
        "那些模型的顺序效应主要表现为'偏向某个标签'，正是 GlobalE 能检测到的那种（演示 3）。",
    ])
    return res


def demo3_sim(test, pool):
    """演示 3：SimICL-large 上的熵准则。"""
    print_header("演示 3：SimICL-large 上用熵准则挑顺序（4 个示例 × 24 种顺序 × 5 组；探测集 = 示例池里 60 条不带标签的评论）")
    sim = SimICL("large")
    rng0 = random.Random(0)
    test_sub = rng0.sample(test, 200)
    probe = [x for x, _, _ in random.Random(1).sample(pool.all, 60)]

    def get_probs(demos, xs):
        return [sim.label_probs(render_english("review", ("positive", "negative"), demos, x)) for x in xs]
    rows, sums = [], []
    for s in range(5):
        demos = sample_demos(pool, random.Random(100 + s), pattern="PPNN")
        sm = summarize_orders(score_orders(get_probs, [list(p) for p in itertools.permutations(demos)], test_sub, probe))
        sums.append(sm)
        rows.append([s, f"{100 * sm['worst']:.1f}%", f"{100 * sm['mean']:.1f}%", f"{100 * sm['best']:.1f}%",
                     f"{100 * sm['globalE']:.1f}%", f"{100 * sm['localE']:.1f}%"])
    print_table(["示例组", "最差顺序", "24 种的平均", "最好顺序", "GlobalE 前 4 名的平均", "LocalE 前 4 名的平均"], rows)
    m = {key: sum(x[key] for x in sums) / len(sums) for key in ("worst", "mean", "best", "globalE", "localE")}
    print_observations([
        f"SimICL-large 的顺序效应来自近因和顺序噪声两项'加在分数上的偏差'：24 种顺序平均 {100 * m['mean']:.1f}%，"
        f"最差 {100 * m['worst']:.1f}%、最好 {100 * m['best']:.1f}%。",
        f"GlobalE 前 4 名平均 {100 * m['globalE']:.1f}%：比平均好一点，避开了最差的那几种（偏差把预测推向一边时，探测集上的"
        f"预测分布不均衡，熵就低），但离最好的顺序还有距离。LocalE 前 4 名 {100 * m['localE']:.1f}%，反而低于平均。",
        "两个演示合起来的结论：一种无监督的挑选准则，只有在它的假设（这里是'好的顺序 = 不偏的顺序'）符合你的模型时才有效。"
        "**用一小份带标注的开发集验证一下**，比相信任何一个准则都可靠。",
    ])


def main():
    test, pool = test_set(), demo_pool()
    base, tok = get_base_model()
    bursty, _ = get_bursty_model()
    models = {"base": base, "bursty": bursty}
    demo1_retrieval(models, tok, test, pool)
    demo2_orders(models, tok, test, pool)
    demo3_sim(test, pool)


if __name__ == "__main__":
    main()