"""
=============================================================================
第 11 讲 · 代码 1：贪心与 ε-贪心——为什么它们的遗憾是线性的
=============================================================================

对应讲义：lectures/lecture11_bandits_ucb.md 第 1–3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
多臂老虎机：有 K 台老虎机（"臂"），每台的中奖概率 μ_a 未知。一共可以拉 T 次，想让总奖励尽量多。
    - 一直拉"目前看起来最好"的臂（利用）？可能错过真正最好的臂。
    - 多试试别的臂（探索）？每次试差的臂都要付出代价。
这就是**探索与利用的两难**——第 1 讲 04 代码中我们第一次见到它，这一讲把它讲透。

怎么衡量一个策略好不好？**遗憾（regret）**：
    Regret(T) = T · μ* - E[总奖励] = Σ_a Δ_a · E[N_a(T)]
"如果一开始就知道哪个臂最好、一直拉它"与"实际做法"之间的期望差距。

本文件：
    1. 贪心：完全不探索 → 有相当比例的实验永远卡在次优臂上
    2. ε-贪心：一直以固定比例探索 → 遗憾随 T 线性增长；最佳的 ε 取决于 T
    3. 递减 ε-贪心、乐观初始化：能做得好，但参数必须"猜对"问题的难度

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture11
    python 01_greedy_epsilon.py      （约 35 秒）
"""

import numpy as np

from bandits import (DEFAULT_MEANS, decaying_epsilon_greedy, epsilon_greedy, greedy,
                     optimistic_greedy, pad, simulate)

np.set_printoptions(precision=3, suppress=True)

# 一个更难的老虎机：最好的两个臂只差 0.03
HARD_MEANS = np.array([0.2, 0.4, 0.5, 0.67, 0.7])


def fraction_wrong(result, means):
    """
    实验结束时，"被拉得最多的臂"不是最好的臂的实验所占的比例。

    【直觉】如果一个策略在 T 步之后最常拉的还不是最好的臂，说明它（至少到目前为止）"认错了"最好的臂。
        对于贪心这类策略，一旦认错，往往就永远纠正不过来——这样的实验会贡献线性增长的遗憾。
    """
    return float(np.mean(result["counts"].argmax(axis=1) != np.argmax(means)))


def main():
    means = DEFAULT_MEANS
    gaps = means.max() - means

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：问题设定，以及贪心为什么会失败（T = 1000，重复 1000 次）")
    print("=" * 78)
    print(f"  5 个臂的真实成功率 μ（智能体不知道）：{means}")
    print(f"  每个臂的差距 Δ_a = μ* - μ_a：            {gaps}")
    print("  遗憾 Regret(T) = Σ_a Δ_a · N_a(T)：每拉一次差距为 Δ 的臂，就\"亏\"Δ。")
    print(f"  完全随机地拉：每步平均亏 {gaps.mean():.2f}，T = 1000 时遗憾约 {gaps.mean() * 1000:.0f}。")
    res = simulate(greedy(), means, T=1000, n_runs=1000, seed=0)
    most = res["counts"].argmax(axis=1)
    print(f"\n  贪心（每个臂先拉一次，之后总拉经验均值最高的臂）的平均遗憾：{res['regret'].mean():.1f}")
    print("  实验结束时\"被拉得最多的臂\"的分布：")
    for a in range(len(means)):
        print(f"    臂 {a}（μ = {means[a]:.1f}）：{np.mean(most == a) * 100:5.1f}% 的实验")
    stuck = np.where(most != 4)[0][0]
    print(f"\n  举一个卡住的例子（第 {stuck} 次实验）：前 20 步拉的臂 = {res['first_pulls'][stuck]}")
    print(f"    结束时各臂被拉的次数 = {res['counts'][stuck].astype(int)}")
    print("  → 贪心有相当比例的实验永远卡在了次优的臂上。原因：最好的臂（臂 4）在最初一两次拉动中")
    print("    运气不好（得到 0），经验均值被低估；而贪心从此不再拉它，也就永远没有机会纠正这个低估。")
    print("    被低估的臂得不到数据，数据不来，低估就一直持续——这就是\"不探索\"的致命之处。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：ε-贪心——遗憾随 T 线性增长（重复 200 次）")
    print("=" * 78)
    cps = [100, 1000, 10000, 100000]
    agents = [("贪心", greedy())] + [(f"ε = {e}", epsilon_greedy(e)) for e in [0.01, 0.05, 0.1, 0.2]]
    print("   " + pad("策略", 10) + "   T=100   T=1000  T=10000 T=100000   R(10⁵)/R(10⁴)  后 9 万步每步遗憾  ε·平均差距")
    table = {}
    for name, pol in agents:
        res = simulate(pol, means, T=100000, n_runs=200, seed=1, checkpoints=cps)
        r = res["regret"].mean(axis=0)
        table[name] = r
        slope = (r[3] - r[2]) / 90000
        eps = 0.0 if name == "贪心" else float(name.split("=")[1])
        print("   " + pad(name, 10) + "".join(f"{x:9.1f}" for x in r)
              + f"   {r[3] / r[2]:10.1f}        {slope:10.4f}      {eps * gaps.mean():10.4f}")
    for k, T in enumerate(cps):
        best = min(table, key=lambda n: table[n][k])
        print(f"   T = {T:>6d} 时遗憾最小的是：{best}")
    print("  观察：")
    print("  - 贪心和 ε ≥ 0.05 的 ε-贪心，从 T = 10⁴ 到 10⁵ 遗憾增长了 6.6～10 倍——遗憾与 T 成正比，即**线性遗憾**。")
    print("  - ε-贪心后期的每步遗憾 ≈ ε × 平均差距：就算已经知道哪个臂最好，每一步仍以 ε 的概率随机拉，")
    print("    随机拉到的臂平均亏\"平均差距\"那么多。探索的比例永远不减少，代价就永远在累积。")
    print("  - ε = 0.01 只增长了 2.1 倍，看起来不像线性。原因：它探索得很少，前 1 万步里还有不少实验没找到最好的臂，")
    print("    这一阶段的遗憾占了大头。但它后 9 万步的每步遗憾（0.0029）已经接近 ε × 平均差距（0.0022），")
    print("    只要 T 再长下去，遗憾最终也按 T 线性增长——只是进入\"线性阶段\"更晚。")
    print("  - 最佳的 ε 取决于 T：T 小时要多探索（尽快找到好臂），T 大时要少探索（长期的探索代价更重要）。")
    print("    可 T 往往事先不知道；而且无论选哪个 ε，只要 T 足够长，遗憾都是线性的。")
    print("  - 贪心的线性遗憾来自\"卡住\"的那部分实验，ε-贪心的线性遗憾来自\"永不停止的探索\"。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：递减 ε-贪心与乐观初始化——参数要\"猜对\"问题的难度（T = 10⁵，重复 200 次）")
    print("=" * 78)
    agents = [
        ("ε = 0.01（对照）", epsilon_greedy(0.01)),
        ("递减 ε：c=1, d=0.1", decaying_epsilon_greedy(1, 0.1)),
        ("递减 ε：c=1, d=0.3", decaying_epsilon_greedy(1, 0.3)),
        ("递减 ε：c=5, d=0.1", decaying_epsilon_greedy(5, 0.1)),
        ("乐观初始化：伪计数 1", optimistic_greedy(1.0, 1)),
        ("乐观初始化：伪计数 10", optimistic_greedy(1.0, 10)),
    ]
    for label, m in [("容易的老虎机（最小差距 0.10）", means), ("困难的老虎机（最小差距 0.03）", HARD_MEANS)]:
        print(f"\n  {label}：μ = {m}")
        print("   " + pad("策略", 22) + "  T=1000  T=10000 T=100000   认错最好臂的实验比例")
        for name, pol in agents:
            res = simulate(pol, m, T=100000, n_runs=200, seed=2, checkpoints=[1000, 10000, 100000])
            r = res["regret"].mean(axis=0)
            print("   " + pad(name, 22) + "".join(f"{x:9.1f}" for x in r) + f"        {fraction_wrong(res, m) * 100:5.1f}%")
    print("\n  观察：")
    print("  - 递减 ε-贪心的探索概率 ε_t = c·K/(d²·t) 按 1/t 减小，总探索次数只按 ln T 增长。")
    print("    理论保证要求 d ≤ 最小差距；但\"安全\"的参数（d = 0.1，尤其 c = 5）探索得太多，遗憾反而很大。")
    print("  - 在容易的老虎机上，\"大胆\"的参数（d = 0.3）和乐观初始化（伪计数 10）表现极好；")
    print("    换到困难的老虎机上，它们有相当比例的实验认错了最好的臂——那些实验的遗憾会线性增长。")
    print("  - 问题出在：这些方法的参数隐含了\"问题有多难\"（差距有多大）的假设，而差距恰恰是未知的。")
    print("    我们想要一个**不需要知道差距**、能自动适应问题难度的方法——这就是下一个文件的 UCB。")


if __name__ == "__main__":
    main()