"""
=============================================================================
第 10 讲 · 代码 3：在线 / 迭代地收集偏好 —— 在哪里采候选，比怎么挑更重要
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
代码 02 发现：从 SFT 的候选池里主动地挑，奖励模型整体更准了，可用它做 RLHF 的结果几乎没变好。
原因是 RLHF 会把策略推到 SFT 很少去的地方。这个文件把收集过程改成"一轮一轮"的：

    每一轮：用当前的奖励模型求 RLHF 策略（β = 0.3 的解析最优解）→ 采候选回答 → 挑 20 对请人比较 → 更新奖励模型

比较几种做法（都先用 20 对随机的 SFT 数据"热身"，再做 8 轮，共 180 对标注）：
    · 离线 + 随机    ：候选回答一直从 SFT 采，随机配对
    · 离线 + BALD    ：候选从 SFT 采，按信息增益批量挑
    · 在线 + 随机    ：候选从**当前策略**采，随机配对
    · 在线 + BALD    ：候选从当前策略采，按信息增益批量挑（每挑一对就"假装问了"更新协方差，第 7 讲演示 2）
    · 在线 + D 最优  ：同上，用 D 最优增量
    · 在线 + 双 Thompson：从奖励模型的后验里抽两个奖励函数，各自挑出它认为最好的回答，让人比较这两个
                         （Dwaracherla 等 2024 用于大模型的"探索"方法的简化版，第 12 讲详细讲 Thompson 采样）

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_online_collection.py        # 约 15 秒

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

import numpy as np

from active_tools import (ALL_GOLD, World, acquisition_scores, exact_stats, fit_laplace, optimal_logprobs, pad,
                          posterior_samples, print_header, sigmoid)

BETA = 0.3
STRATEGIES = {"offline_random": "离线 + 随机", "offline_bald": "离线 + BALD", "online_random": "在线 + 随机",
              "online_bald": "在线 + BALD", "online_dopt": "在线 + D 最优", "online_dts": "在线 + 双 Thompson"}


def pick_batch(rule, world, pool, w, cov, per, rng):
    """
    从候选回答 pool 里按信息增益或 D 最优**批量**挑 per 对。

    【思路】先随机组成 3000 个候选对；每挑一对，就用 Sherman–Morrison 假装已经问了它（信息不依赖答案，
        第 7 讲），更新协方差，再挑下一对 —— 避免一批里全是同一类问题。
    """
    ci = rng.integers(0, len(pool), 3000)
    cj = rng.integers(0, len(pool), 3000)
    keep = ci != cj
    Ap, Bp = pool[ci[keep]], pool[cj[keep]]
    Uc = world.X[Ap] - world.X[Bp]
    C = cov.copy()
    A, B = [], []
    for _ in range(per):
        if rule == "bald":
            sc = acquisition_scores("bald", Uc, w, C, rng, samples=posterior_samples(w, C, 64, rng))
        else:
            sc = acquisition_scores("dopt", Uc, w, C, rng)
        k = int(np.argmax(sc))
        A.append(Ap[k])
        B.append(Bp[k])
        u = Uc[k].copy()
        p = sigmoid(u @ w)
        Cu = C @ u
        C = C - np.outer(Cu, Cu) * p * (1 - p) / (1 + p * (1 - p) * (u @ Cu))
        Uc[k] = 0.0                                          # 挑过的对不再挑（特征差为 0 的对信息为 0）
    return np.array(A), np.array(B)


def pick_double_thompson(world, pool, w, cov, per, rng):
    """
    双 Thompson 采样：抽两个奖励函数 w₁、w₂ ~ 后验，各自在候选回答里挑最高分的，组成一对。

    【直觉】两个"可能的奖励函数"各自推荐最好的回答 —— 如果后验已经很确定，两者推荐同一个（换一个随机对手）；
        如果还不确定，它们推荐不同的回答，而这正是我们该搞清楚的"谁才是最好的"。
    """
    A, B = [], []
    for _ in range(per):
        S = posterior_samples(w, cov, 2, rng)
        i = int(np.argmax(world.X[pool] @ S[0]))
        j = int(np.argmax(world.X[pool] @ S[1]))
        if i == j:
            j = int(rng.integers(len(pool)))
        A.append(pool[i])
        B.append(pool[j])
    return np.array(A), np.array(B)


def run(world, strategy, rep, rounds=8, per=20, warmup=20, pool_size=200):
    """
    跑一次完整的收集过程，返回 dict：标注数 → (RLHF 后的真实质量, 候选回答的平均真实质量)。
    """
    rng = np.random.default_rng(rep)
    a = world.sample(world.ref_lp, warmup, rng)
    b = world.sample(world.ref_lp, warmup, rng)
    U = world.X[a] - world.X[b]
    Y = world.label(a, b, rng)
    w, cov = fit_laplace(U, Y)
    out = {warmup: (world.downstream(w, BETA)["gold"], np.nan)}
    for _ in range(rounds):
        current = optimal_logprobs(world.ref_lp, world.X @ w, BETA)      # 当前奖励模型对应的 RLHF 策略
        source = world.ref_lp if strategy.startswith("offline") else current
        pool = world.sample(source, pool_size, rng)
        if strategy.endswith("random"):
            A, B = pool[rng.integers(0, pool_size, per)], pool[rng.integers(0, pool_size, per)]
        elif strategy.endswith("bald"):
            A, B = pick_batch("bald", world, pool, w, cov, per, rng)
        elif strategy.endswith("dopt"):
            A, B = pick_batch("dopt", world, pool, w, cov, per, rng)
        else:
            A, B = pick_double_thompson(world, pool, w, cov, per, rng)
        U = np.vstack([U, world.X[A] - world.X[B]])
        Y = np.concatenate([Y, world.label(A, B, rng)])
        w, cov = fit_laplace(U, Y)
        out[len(Y)] = (world.downstream(w, BETA)["gold"], float(ALL_GOLD[pool].mean()))
    return out


def main():
    world = World()
    reps = 10
    best = exact_stats(optimal_logprobs(world.ref_lp, world.util, BETA), world.ref_lp)["gold"]
    print_header("在线 / 迭代地收集偏好：六种做法")
    print(f"每种做法重复 {reps} 次。表中是“用当前奖励模型做 RLHF（β = {BETA}）后的真实质量”；")
    print(f"完美奖励模型能达到 {best:+.3f}。\n")
    checkpoints = [20, 60, 100, 140, 180]
    results = {}
    for st in STRATEGIES:
        res = [run(world, st, rep) for rep in range(reps)]
        results[st] = res
    print(f"  {pad('标注数', 7)}" + "".join(pad(STRATEGIES[st], 20, "right") for st in STRATEGIES))
    for c in checkpoints:
        print(f"  {c:>7d}" + "".join(f"{np.mean([r[c][0] for r in results[st]]):>+20.3f}" for st in STRATEGIES))
    print("\n  候选回答本身的平均真实质量（最后一轮）：")
    print(f"  {pad('', 7)}" + "".join(f"{np.mean([r[180][1] for r in results[st]]):>+20.3f}" for st in STRATEGIES))
    print("""
观察：
  · 在线明显好于离线：候选从**当前策略**采，奖励模型就在“策略要去的地方”得到了数据 ——
    哪里被高估了，下一轮就会有比较来纠正它。几轮之后，在线的几种做法都接近完美奖励模型的水平。
  · 离线时，主动挑法（BALD）比随机稍好；在线时，随机配对反而和主动挑法一样好甚至更好。
    当前策略的回答已经集中在“高分区域”，在它们之间随机比较，问的正好是“好的回答里哪个更好”——
    这比“让参数整体更准”更贴近 RLHF 的需要。
  · 双 Thompson 这次不比随机配对强。它在真实大模型上的优势（Dwaracherla 等 2024）来自巨大的回答空间，
    这里每轮候选只有 200 个回答、特征只有 16 维，随机配对已经覆盖得不错。
  · 最后一行：在线做法的候选回答本身就好得多 —— 标注员看到的是“好回答 vs 好回答”，
    这也是 Llama 2 每周用最新的模型生成新回答、收集新偏好的原因（讲义第 5 节）。
  · 【边界】玩具世界里“当前策略”是解析解，现实中要跑 RL 或 DPO；每一轮的标注也有延迟和成本。""")


if __name__ == "__main__":
    main()