"""
=============================================================================
第 12 讲 · 代码 1：Thompson 采样 —— 概率匹配、遗憾、批量反馈
=============================================================================

对应讲义：lectures/lecture12_thompson_acquisition.md 第 2、3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
Thompson 采样（Thompson 1933）：维护每个臂的后验，每一步从后验里抽一个样本，选样本里最好的臂。

    演示 1：概率匹配：Thompson 采样选每个臂的概率，正好等于"后验认为这个臂是最好的"的概率。
            和贪心、UCB1 在同一个后验下的选择比较。
    演示 2：遗憾：Thompson 采样、KL-UCB、UCB1、ε-贪心；遗憾 ÷ ln T 和 Lai–Robbins 下界比较。
    演示 3：批量 / 延迟反馈：奖励每 B 步才能看到一次（真实系统常常如此）。确定性的 UCB 类算法在一批里
            每一步都选同一个臂；Thompson 采样每一步重新抽样，天然地分散。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_thompson_bernoulli.py        # 约 25 秒

只需要 numpy。依赖本目录的 ts_tools.py（它会再导入第 11 讲的 bandit_tools.py）。
"""

import math

import numpy as np

from ts_tools import BERNOULLI_ALGOS, kl_ucb_index, lai_robbins_bound, pad, print_header, run_bernoulli


def demo_probability_matching(rng):
    print_header("演示 1：概率匹配 —— 选每个臂的概率 = 它是最好的后验概率")
    data = {"A": (6, 10), "B": (3, 4), "C": (20, 40)}       # (成功次数, 尝试次数)
    names = list(data)
    a = np.array([1 + s for s, n in data.values()], dtype=float)
    b = np.array([1 + n - s for s, n in data.values()], dtype=float)
    draws = rng.beta(a[None, :], b[None, :], size=(200000, 3))
    p_best = np.bincount(np.argmax(draws, axis=1), minlength=3) / len(draws)
    mu_hat = np.array([s / n for s, n in data.values()])
    n = np.array([n for _, n in data.values()], dtype=float)
    t = int(n.sum())
    ucb1 = mu_hat + np.sqrt(2 * math.log(t) / n)
    klucb = kl_ucb_index(mu_hat, n, t)
    print(f"三个臂到目前为止的数据（成功 / 尝试）：" + "，".join(f"{k} {s}/{m}" for k, (s, m) in data.items())
          + f"；先验 Beta(1, 1)。\n")
    print(f"  {pad('臂', 4)} {pad('后验均值', 9, 'right')} {pad('P(它是最好的)', 14, 'right')} {pad('UCB1 指数', 10, 'right')} "
          f"{pad('KL-UCB 指数', 12, 'right')}")
    for k in range(3):
        print(f"  {pad(names[k], 4)} {a[k] / (a[k] + b[k]):>9.3f} {p_best[k]:>14.3f} {ucb1[k]:>10.3f} {klucb[k]:>12.3f}")
    print(f"""
  下一步：贪心选 {names[int(np.argmax(mu_hat))]}（平均值最高）；UCB1 选 {names[int(np.argmax(ucb1))]}；KL-UCB 选 {names[int(np.argmax(klucb))]}；
  Thompson 采样以上表“P(它是最好的)”一列的概率随机选 —— 抽 20 万次，选中各臂的频率就是这一列。

观察：
  · Thompson 采样不需要设计“探索奖励”：一个臂只要有一定的可能是最好的，就有相应的机会被试。
    B 只试了 4 次、3 次成功：它可能很好，也可能只是运气；Thompson 采样按“它真的是最好的概率”分配机会。
    C 试了 40 次、平均 0.5，不太可能是最好的 —— 但仍有约 8% 的机会被试。
  · UCB 是确定性的：同样的数据，永远选同一个臂。Thompson 采样是随机的：这在“一批同时做决定”时很重要（演示 3）。
  · 这个原则叫“概率匹配”：按“某个行动是最优的”的后验概率去选它。""")


def demo_regret(rng):
    print_header("演示 2：遗憾 —— Thompson 采样、KL-UCB、UCB1、ε-贪心")
    checks = [1000, 10000, 20000]
    for title, means, runs in [("容易：5 个臂，点击率 0.3、0.4、0.5、0.6、0.7", np.array([0.3, 0.4, 0.5, 0.6, 0.7]), 100),
                               ("难：10 个臂，最好的 0.5，其余 0.45", np.array([0.5] + [0.45] * 9), 100)]:
        coef = lai_robbins_bound(list(means), math.e)              # ln e = 1：下界的系数
        print(f"\n【{title}】每种算法跑 {runs} 次、每次 2 万步。Lai–Robbins 系数 {coef:.2f}\n")
        print(f"  {pad('算法', 18)}" + "".join(pad(f"T={c}", 9, "right") for c in checks)
              + pad("遗憾/ln T（T=2万）", 20, "right") + pad("最后 10% 选中最好的", 20, "right"))
        for algo, name in BERNOULLI_ALGOS.items():
            out = run_bernoulli(algo, means, checks[-1], runs, rng)
            r = out["regret"]
            print(f"  {pad(name, 18)}" + "".join(f"{r[c - 1]:>9.1f}" for c in checks)
                  + f"{r[-1] / math.log(checks[-1]):>20.2f}{out['best_frac'].mean():>20.3f}")
    print("""
观察：
  · 容易的问题上，Thompson 采样和 KL-UCB 的遗憾最小，“遗憾 / ln T”接近（甚至低于）Lai–Robbins 系数 ——
    它们在渐近意义上是最优的（Agrawal & Goyal 2012；Kaufmann 等 2012；Garivier & Cappé 2011）。
    UCB1 的置信半径来自 Hoeffding 不等式，偏宽，探索过多。
  · 难的问题上，所有算法在 2 万步里都还没分清（差距只有 0.05）。Thompson 采样的遗憾和 ε-贪心差不多
    （ε = 0.1 在这里恰好合适），但最后选中最好臂的比例最高；UCB1 探索得太多，遗憾最大。
  · 遗憾 / ln T 低于 Lai–Robbins 系数并不矛盾：下界是 T → ∞ 的结论（第 11 讲）。
  · Thompson 采样的实现只有一行（从 Beta 分布抽样），没有需要调的常数；先验可以放进领域知识。""")


def demo_batch(rng):
    print_header("演示 3：批量 / 延迟反馈")
    means = np.array([0.3, 0.4, 0.5, 0.6, 0.7])
    batches = [1, 20, 100, 500]
    print("容易的问题，1 万步。奖励每 B 步才能看到一次（B 步之内算法拿不到新信息）。报告平均累计遗憾（100 次实验）：\n")
    print(f"  {pad('算法', 18)}" + "".join(pad(f"B={b}", 10, "right") for b in batches))
    for algo in ["ucb1", "kl_ucb", "thompson"]:
        row = ""
        for b in batches:
            out = run_bernoulli(algo, means, 10000, 100, rng, batch=b)
            row += f"{out['regret'][-1]:>10.1f}"
        print(f"  {pad(BERNOULLI_ALGOS[algo], 18)}{row}")
    print("""
观察：
  · 真实系统里，反馈常常是成批到来的：推荐系统每小时才更新一次模型；标注团队一次交回几百条偏好（第 10 讲）。
  · UCB 类算法是确定性的：一批之内信息不变，它就把整批都押在同一个臂上 —— 批越大，越像“每批只做一次决定”。
  · Thompson 采样每一步重新抽样，一批之内自然地分散到几个“可能最好”的臂上，对批量和延迟更稳健：
    B = 100 时 KL-UCB 的遗憾几乎翻倍，Thompson 采样只增加了约 14%（Chapelle & Li 2011 在雅虎的实验里强调了这一点）。
  · B = 500 时所有算法都明显变差 —— 1 万步只有 20 次“带着新信息”的决定 —— 但 Thompson 采样仍然最低。
    UCB1 的相对增幅不大，是因为它本来就探索过多。
  · 第 10 讲的“批量主动学习”是同一个问题：一批之内要多样，而不是重复问同一个问题。""")


def main():
    rng = np.random.default_rng(12)
    demo_probability_matching(rng)
    demo_regret(rng)
    demo_batch(rng)


if __name__ == "__main__":
    main()