"""
=============================================================================
第 10 讲 · 代码 2：主动 vs 随机地挑回答对，训练奖励模型
=============================================================================

对应讲义：lectures/lecture10_active_llm.md 第 3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
设定（"离线候选池"）：从 SFT 策略采 300 个回答，随机组成约 6000 个候选对。标注预算有限（最多 160 对），
每次从候选池里挑一对交给专家标注员，拿到答案后更新奖励模型，再挑下一对。

六种挑法：随机、不确定性（p 最接近 0.5）、方差最大、D 最优、信息增益（BALD）、集成分歧。
两个评估：
    · 奖励模型和标注员方向的一致率（在全部候选对上）—— "奖励模型整体准不准"
    · 用这个奖励模型做 RLHF（β = 0.3 的解析最优解）后的真实质量 —— "拿它优化出的策略好不好"

    演示 1：六种挑法的学习曲线。
    演示 2：不确定性采样的陷阱：它挑了多少"两个回答特征完全相同"的对？去掉这些对之后呢？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_active_pairs.py        # 约 20 秒

需要 numpy 和 PyTorch（只用来训练第 8 讲的 SFT 策略）。依赖本目录的 active_tools.py。
"""

import numpy as np

from active_tools import (RULES, World, acquisition_scores, bootstrap_ensemble, exact_stats, fit_laplace,
                          optimal_logprobs, pad, posterior_samples, print_header)

BETA = 0.3
CHECKPOINTS = [20, 40, 80, 160]


def make_pool(world, rep, drop_identical=False):
    """
    候选池：从 SFT 采 300 个回答，随机组成约 6000 个候选对。

    【参数】drop_identical：True 时去掉"两个回答特征完全相同"的对（例如只是词序不同）
    【返回】(A, B, Uc)：候选对的两个回答下标、特征差
    """
    rng = np.random.default_rng(1000 + rep)
    pool = world.sample(world.ref_lp, 300, rng)
    ci = rng.integers(0, 300, 6000)
    cj = rng.integers(0, 300, 6000)
    keep = ci != cj
    A, B = pool[ci[keep]], pool[cj[keep]]
    Uc = world.X[A] - world.X[B]
    if drop_identical:
        nz = np.abs(Uc).sum(axis=1) > 0
        A, B, Uc = A[nz], B[nz], Uc[nz]
    return A, B, Uc


def run_active(world, rule, rep, budget=CHECKPOINTS[-1], warmup=10, drop_identical=False):
    """
    依次挑 budget 对去问。前 warmup 对随机挑（冷启动：一开始所有候选的分数都一样）。

    【返回】dict：检查点 → (一致率, RLHF 后的真实质量, 已选的对里"特征完全相同"的比例)
    """
    A, B, Uc = make_pool(world, rep, drop_identical)
    rng = np.random.default_rng(rep)
    available = np.ones(len(A), dtype=bool)
    chosen, labels = [], []
    d = Uc.shape[1]
    w, cov = np.zeros(d), 4.0 * np.eye(d)                  # 先验 N(0, 4I)
    out = {}
    for t in range(budget):
        idx = np.flatnonzero(available)
        r = "random" if t < warmup else rule
        samples = posterior_samples(w, cov, 64, rng) if r == "bald" else None
        ens = bootstrap_ensemble(Uc[chosen], np.array(labels), 5, rng) if r == "qbc" else None
        k = idx[int(np.argmax(acquisition_scores(r, Uc[idx], w, cov, rng, samples=samples, ensemble=ens)))]
        available[k] = False
        chosen.append(k)
        labels.append(int(world.label(A[[k]], B[[k]], rng)[0]))
        w, cov = fit_laplace(Uc[chosen], np.array(labels))
        if t + 1 in CHECKPOINTS:
            identical = float(np.mean(np.abs(Uc[chosen]).sum(axis=1) == 0))
            out[t + 1] = (world.agreement(w, A, B), world.downstream(w, BETA)["gold"], identical)
    return out


def demo_rules(world, reps=8):
    print_header("演示 1：六种挑法的学习曲线")
    best = exact_stats(optimal_logprobs(world.ref_lp, world.util, BETA), world.ref_lp)["gold"]
    sft = exact_stats(world.ref_lp, world.ref_lp)["gold"]
    print(f"每种挑法重复 {reps} 次（每次新的候选池）。参照：SFT 策略的真实质量 {sft:+.3f}；用**完美**奖励模型做 RLHF"
          f"（β = {BETA}）的真实质量 {best:+.3f}。\n")
    results = {}
    for rule in RULES:
        res = [run_active(world, rule, rep) for rep in range(reps)]
        results[rule] = {c: np.mean([r[c] for r in res], axis=0) for c in CHECKPOINTS}
    for title, col, fmt in [("奖励模型和标注员方向的一致率", 0, "{:>18.3f}"),
                            (f"用它做 RLHF（β = {BETA}）后的真实质量", 1, "{:>+18.3f}"),
                            ("已选的对里“两个回答特征完全相同”的比例", 2, "{:>18.2f}")]:
        print(f"  {title}")
        print(f"  {pad('标注数', 7)}" + "".join(pad(RULES[r], 18, "right") for r in RULES))
        for c in CHECKPOINTS:
            print(f"  {c:>7d}" + "".join(fmt.format(results[r][c][col]) for r in RULES))
        print()
    print("""观察：
  · 一致率：D 最优、BALD、集成分歧这几种“看参数不确定性”的挑法，学得比随机快 —— 80 对左右就达到
    随机挑 160 对的水平。方差最大（只看方向、不管答案有没有把握）不如它们。
  · 不确定性采样（只看 p 是否接近 0.5）在前 80 对几乎原地踏步：它挑的大多是“两个回答特征完全相同”的对
    （第三张表）—— 永远五五开，问了也白问（代码 01 演示 3：BALD = 0）。
  · 可是“用它做 RLHF 后的真实质量”几乎**没有**区别：主动挑法让奖励模型在 SFT 的回答上整体更准，
    但 RLHF 关心的是奖励模型在“策略要去的地方”准不准 —— 那些回答在 SFT 的候选池里很少（代码 01 演示 2）。
    同样的预算，“在哪里采候选”可能比“怎么挑”更重要（代码 03）。
  · 【边界】玩具世界的特征只有 16 维，随机挑 160 对已经学得不错，主动学习的空间有限；
    真实奖励模型的维数高得多，Dwaracherla 等（2024）报告主动探索在大模型上能省下大量标注。""")
    return results


def demo_identical(world, reps=8):
    print_header("演示 2：去掉“特征完全相同”的对之后，不确定性采样怎样？")
    print(f"  {pad('标注数', 7)} {pad('不确定性（原候选池）', 21, 'right')} {pad('不确定性（去掉相同的对）', 25, 'right')} "
          f"{pad('随机', 7, 'right')}")
    orig = [run_active(world, "uncertainty", rep) for rep in range(reps)]
    dropped = [run_active(world, "uncertainty", rep, drop_identical=True) for rep in range(reps)]
    rnd = [run_active(world, "random", rep) for rep in range(reps)]
    for c in CHECKPOINTS:
        print(f"  {c:>7d} {np.mean([r[c][0] for r in orig]):>21.3f} {np.mean([r[c][0] for r in dropped]):>25.3f} "
              f"{np.mean([r[c][0] for r in rnd]):>7.3f}")
    print("""
（表中是和标注员方向的一致率。）

观察：
  · 去掉特征完全相同的对，不确定性采样不再原地踏步，但仍然不比随机好多少：剩下的“p ≈ 0.5”的对
    常常是“两个回答的特征差和 w 几乎垂直”（第 7 讲练习 5），对学 w 的帮助有限。
  · 真实世界里也有“特征完全相同”的对：同一个模型对同一个提示词生成两次几乎一样的回答。
    只看“模型有没有把握”的挑法会把预算浪费在它们身上；信息增益（BALD）会自动避开。""")


def main():
    world = World()
    demo_rules(world)
    demo_identical(world)


if __name__ == "__main__":
    main()