"""
=============================================================================
第 12 讲 · 代码 3：高斯过程上的采集函数 —— PI、EI、GP-UCB、Thompson 采样
=============================================================================

对应讲义：lectures/lecture12_thompson_acquisition.md 第 5 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
选项不再是有限的几个臂，而是一个连续的参数 x ∈ [0, 1]（例如生成时的温度、推荐的"新颖度"、一个剂量）。
效用 f(x) 未知，每次只能试一个 x、观测到带噪声的 f(x)。预算只有 25 次。怎样尽快找到 f 最大的 x？

贝叶斯优化：用高斯过程（第 6 讲）表示对 f 的信念，每一步用一个**采集函数**决定下一个 x。

真实的 f（"上帝视角"）有两个峰：x ≈ 0.25 处一个宽而矮的峰（0.6），x ≈ 0.8 处一个窄而高的峰（1.0）。

    演示 1：七种做法（随机、只看均值、PI、ξ 更大的 PI、EI、GP-UCB、Thompson 采样），40 次实验的平均结果。
    演示 2：一次实验里，PI 和 EI 把 25 次评估放在了哪里。
    演示 3：PI 为什么太贪心：在同一个后验下，比较"当前最好点旁边"和"没探索过的地方"的 PI、EI 值。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_gp_acquisition.py        # 约 10 秒

只需要 numpy。依赖本目录的 ts_tools.py。
"""

import numpy as np

from ts_tools import GP, acquisition, norm_cdf, norm_pdf, pad, print_header

GRID = np.linspace(0.0, 1.0, 201)
NOISE = 0.1


def f_true(x):
    """真实的效用：x = 0.25 处宽而矮的峰（0.6），x = 0.8 处窄而高的峰（1.0）。只在模拟里用。"""
    return 0.6 * np.exp(-((x - 0.25) / 0.12) ** 2) + 1.0 * np.exp(-((x - 0.8) / 0.05) ** 2)


F = f_true(GRID)
F_STAR = F.max()
SETTINGS = [("random", 0.01, "随机"), ("exploit", 0.01, "只看均值"), ("pi", 0.01, "PI（ξ=0.01）"),
            ("pi", 0.1, "PI（ξ=0.1）"), ("ei", 0.01, "EI"), ("ucb", 0.01, "GP-UCB"), ("ts", 0.01, "Thompson 采样")]


def run(acq, rep, xi=0.01, budget=25, init=3):
    """
    一次贝叶斯优化：先随机评估 init 个点，再按采集函数评估 budget 个点。

    【参数】xi：PI、EI 里的"至少要改进多少"ξ（越大越鼓励探索）
    【返回】(每一步后"已评估的点里最好的真实值"和全局最大值之差（数组）,
            最后按后验均值推荐的点的差距, 评估过的点的网格下标列表)
    【实现细节】"当前最好值"用已评估点上的**后验均值**的最大值（而不是带噪声的观测最大值）——
        有噪声时，观测最大值往往是被噪声抬高的，用它做基准会让 PI、EI 的行为失真。
    """
    rng = np.random.default_rng(rep)
    gp = GP(ell=0.08, sf=0.7, noise=NOISE)
    idx = list(rng.integers(0, len(GRID), init))
    y = list(F[idx] + NOISE * rng.normal(size=init))
    gaps = []
    for t in range(budget):
        mean, cov = gp.posterior(GRID, GRID[idx], np.array(y))
        k = acquisition(acq, mean, cov, mean[idx].max(), rng, t, xi=xi)
        idx.append(k)
        y.append(F[k] + NOISE * rng.normal())
        gaps.append(F_STAR - F[idx].max())
    mean, _ = gp.posterior(GRID, GRID[idx], np.array(y))
    return np.array(gaps), F_STAR - F[int(np.argmax(mean))], idx


def demo_compare(reps=40):
    print_header("演示 1：七种做法的比较")
    print(f"预算：先随机评估 3 个点，再评估 25 个点；观测噪声标准差 {NOISE}。{reps} 次独立实验的平均：\n")
    print(f"  {pad('采集函数', 16)}" + "".join(pad(f"{n} 次后差距", 13, "right") for n in [5, 10, 25])
          + pad("推荐点的差距", 13, "right") + pad("找到高峰", 9, "right") + pad("评估在 x>0.6", 13, "right"))
    for acq, xi, name in SETTINGS:
        res = [run(acq, rep, xi=xi) for rep in range(reps)]
        G = np.array([r[0] for r in res])
        found = np.mean([r[0][-1] < 0.1 for r in res])
        right = np.mean([np.mean(GRID[r[2][3:]] > 0.6) for r in res])
        print(f"  {pad(name, 16)}" + "".join(f"{G[:, n - 1].mean():>13.3f}" for n in [5, 10, 25])
              + f"{np.mean([r[1] for r in res]):>13.3f}{found:>9.2f}{right:>13.2f}")
    print("""
（“差距”：已评估的点里最好的真实值，离全局最大值 1.0 还差多少；“找到高峰”：25 次后差距 < 0.1 的实验比例。）

观察：
  · 只看均值（纯利用）：在第一个看起来不错的地方反复评估，三分之二的实验从没找到高峰。
  · PI（改进概率，ξ = 0.01）：只问“能不能比现在好一点点”，不问“能好多少”。当前最好点旁边“好一点点”的概率
    总是很高（接近一半），于是它在局部的峰上精雕细琢 —— 将近一半的实验困在 x ≈ 0.25 的矮峰上（演示 3）。
    把 ξ 调大（要求至少改进 0.1），PI 就愿意去别处看看了 —— 调好了 ξ 的 PI 在这里甚至是最好的。
    ξ 就是 PI 的“探索旋钮”，问题是事先不知道该调多大。
  · EI（期望改进）同时考虑“改进的概率”和“改进的幅度”，几乎每次都找到了高峰。
  · GP-UCB 和 Thompson 采样：用不确定性驱动探索，同样可靠。GP-UCB 前几步探索得多（5 次后差距最大），
    之后最好；Thompson 采样偶尔会在 25 次里还没完全“收尾”。
  · 随机：25 次里有约 30% 的实验没碰到那个窄峰。窄而高的最优点，正是需要“聪明地找”的情形。""")


def demo_where(rep=6):
    print_header("演示 2：一次实验里，评估都放在了哪里")
    print(f"把 [0, 1] 分成 10 段，统计第 {rep} 次实验（一次 PI 被困住的实验）里 25 次评估落在每一段的次数"
          "（真实的高峰在 0.75~0.85）：\n")
    print(f"  {pad('采集函数', 16)}" + "".join(f"{a / 10:>5.1f}~" for a in range(10)) + pad("最终差距", 10, "right"))
    for acq, xi, name in SETTINGS[2:]:
        gaps, _, idx = run(acq, rep, xi=xi)
        counts = np.histogram(GRID[idx[3:]], bins=10, range=(0, 1))[0]
        print(f"  {pad(name, 16)}" + "".join(f"{c:>6d}" for c in counts) + f"{gaps[-1]:>10.3f}")
    print("""
观察：
  · PI（ξ = 0.01）把大部分评估都花在了 0.2~0.3（矮峰）上，高峰所在的 0.7~0.9 一次都没去。
  · EI、GP-UCB、ξ 更大的 PI 先在各处看看，然后集中到 0.7~0.9 的高峰上。
  · 这一次 Thompson 采样也被困住了：它是随机的，偶尔会连续抽到“最大值在矮峰附近”的函数样本。
    演示 1 里它有约 7% 的实验没找到高峰 —— 随机性让它在平均意义上探索得很好，但单次运行有波动。""")


def demo_why_pi():
    print_header("演示 3：PI 为什么太贪心")
    gp = GP(ell=0.08, sf=0.7, noise=NOISE)
    X = np.array([0.15, 0.22, 0.25, 0.3, 0.4])            # 已经在矮峰附近评估了几次
    y = f_true(X)
    mean, cov = gp.posterior(GRID, X, y)
    mean = np.where(np.abs(mean) < 5e-4, 0.0, mean)          # 远处的均值是 1e-6 量级，避免显示成 −0.000
    sd = np.sqrt(np.diag(cov))
    best = mean[np.searchsorted(GRID, X)].max()
    print(f"已经在 x = {X.tolist()} 评估过（都在矮峰附近），当前最好的后验均值 {best:.3f}。比较几个候选点：\n")
    print(f"  {pad('x', 6)} {pad('后验均值', 9, 'right')} {pad('后验标准差', 11, 'right')} {pad('PI（ξ=0.01）', 13, 'right')} "
          f"{pad('EI（ξ=0.01）', 13, 'right')} {pad('真实值', 7, 'right')}")
    for x in [0.24, 0.27, 0.6, 0.8, 0.95]:
        k = int(round(x * 200))
        z = (mean[k] - best - 0.01) / sd[k]
        pi = norm_cdf(z)
        ei = (mean[k] - best - 0.01) * norm_cdf(z) + sd[k] * norm_pdf(z)
        print(f"  {x:>6.2f} {mean[k]:>9.3f} {sd[k]:>11.3f} {float(pi):>13.3f} {float(ei):>13.4f} {F[k]:>7.3f}")
    print("""
观察：
  · 在当前最好点旁边（x = 0.24、0.27），后验均值和“当前最好”几乎一样、标准差很小：PI ≈ 0.4~0.5，
    “比现在好一点点”的可能性很大（PI 约 0.4）。在没探索过的地方（x = 0.8），均值低、标准差大：PI 只有约 0.2。
    PI 选“旁边”，于是一直在矮峰上打转。
  · EI 看的是“期望能改进多少”：旁边的点即使能改进，也只是一点点（EI 很小）；远处虽然概率低，可一旦改进就是
    一大截，标准差大让 EI 反而更高。所以 EI 会去远处看看。
  · 这是采集函数设计的核心：**探索的价值来自“可能有大的改进”，而不只是“可能有改进”**。
    GP-UCB 用“均值 + 若干倍标准差”、Thompson 采样用“随机的一条可能的函数”，以不同的方式表达了同一件事。""")


def main():
    demo_compare()
    demo_where()
    demo_why_pi()


if __name__ == "__main__":
    main()