"""
=============================================================================
第 6 讲 · 代码 4：高斯过程偏好学习 —— 不假设效用的形状，也知道自己哪里不知道
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 2 讲的线性 BT 假设"效用是特征的线性函数"。可是很多偏好不是线性的：回答太短不好、太长也不好，
最好的长度在中间；也许在某个更长的长度上还有第二个"甜点"（例如带一个完整的例子）。

设想一个一维的问题：x ∈ [0, 1] 是"回答长度"（归一化），真实效用 f*(x) 有一个主峰（x ≈ 0.3）
和一个小峰（x ≈ 0.78）。标注员只看过 x ≤ 0.6 的回答（数据收集时没生成过很长的回答）。

    演示 1：线性 BT、二次 BT、GP 三种模型的拟合；GP 的后验标准差在没有数据的地方变大。
    演示 2：长度尺度 ℓ 怎样选：边际似然（拉普拉斯近似的"证据"）vs 留出数据上的表现。
    演示 3：下一个问题该问什么？—— 后验不确定性告诉我们哪里最值得问（预告第 7、13 讲）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_gp_preference.py        # 约 2 秒

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

import numpy as np

from bayes_tools import fit_linear_bt, gp_pref_laplace, log_sigmoid, pad, print_header, sigmoid

GRID = np.linspace(0.0, 1.0, 61)          # 61 个网格点，间隔 1/60
SEEN = GRID <= 0.6 + 1e-9                 # 标注员看过的范围：x ≤ 0.6
SF = 1.5                                  # GP 的信号标准差（效用大概在 ±3 以内变化）


def true_utility(x):
    """
    真实效用：主峰在 x = 0.3（最喜欢的长度），x = 0.78 处有一个小峰，整体随长度略微下降。

    【边界】只在模拟里用（"上帝视角"），学习算法看不到它。
    """
    return 2.0 * np.exp(-((x - 0.3) / 0.12) ** 2) + 1.4 * np.exp(-((x - 0.78) / 0.08) ** 2) - 1.0 * x


def make_pairs(rng, n, allowed):
    """在 allowed（布尔数组）为 True 的网格点之间随机抽 n 对不同的点，按 BT 模型生成结果。"""
    idx = np.flatnonzero(allowed)
    a = rng.choice(idx, size=n)
    b = rng.choice(idx, size=n)
    keep = a != b
    a, b = a[keep], b[keep]
    f = true_utility(GRID)
    y = (rng.random(len(a)) < sigmoid(f[a] - f[b])).astype(int)
    return a, b, y


def poly_features(x, degree):
    """多项式特征 [x, x², …, x^degree]（不要常数项：BT 里常数会在差值中消掉）。"""
    return np.stack([x ** d for d in range(1, degree + 1)], axis=1)


def center_on_seen(f):
    """减去"看过的范围"内的平均值，让几条曲线可以放在一起比较（BT 只确定差值）。"""
    return f - f[SEEN].mean()


def centered_sd(cov):
    """
    "减去看过范围内的平均值"之后，每个网格点效用的后验标准差。

    【思路】centered = P f，P = I − 1·wᵀ（w 在看过的点上是 1/个数，其余为 0）。
        线性变换后的协方差是 P Cov Pᵀ。直接用 Cov 的对角线会把"整体平移"的不确定性也算进来 ——
        而整体平移只由先验决定，数据管不到，也不影响任何比较（代码 02 演示 4）。
    """
    G = cov.shape[0]
    w = SEEN / SEEN.sum()
    P = np.eye(G) - np.outer(np.ones(G), w)
    return np.sqrt(np.clip(np.diag(P @ cov @ P.T), 0.0, None))


def pair_metrics(f, a, b, y):
    """用效用曲线 f 预测比较 (a, b, y)：返回 (准确率, 对数损失)。"""
    d = f[a] - f[b]
    acc = float(np.mean((d > 0) == (y == 1)))
    ll = float(-np.mean(log_sigmoid(np.where(y == 1, d, -d))))
    return acc, ll


# =============================================================================
# 演示 1
# =============================================================================
def demo_fit(rng):
    print_header("演示 1：线性 BT、二次 BT、高斯过程 —— 谁能学出效用的形状？")
    a, b, y = make_pairs(rng, 300, SEEN)
    truth = center_on_seen(true_utility(GRID))
    w1 = fit_linear_bt(poly_features(GRID, 1), a, b, y)
    w2 = fit_linear_bt(poly_features(GRID, 2), a, b, y)
    lin = center_on_seen(poly_features(GRID, 1) @ w1)
    quad = center_on_seen(poly_features(GRID, 2) @ w2)
    gp = gp_pref_laplace(GRID, a, b, y, ell=0.08, sf=SF)
    gp_mean = center_on_seen(gp.mean)
    sd = centered_sd(gp.cov)
    print(f"{len(y)} 条比较，全部在 x ≤ 0.6 的范围内。GP 用 RBF 核，ℓ = 0.08，信号标准差 {SF}（演示 2 讲怎样选）。")
    print("所有曲线都减去了“x ≤ 0.6 范围内的平均值”（只有差值有意义），GP 的标准差也是对减去平均值之后的效用算的。\n")
    print(f"  {pad('x', 5)} {pad('真实', 7, 'right')} {pad('线性', 7, 'right')} {pad('二次', 7, 'right')} "
          f"{pad('GP 均值', 8, 'right')} {pad('GP ± 2 标准差', 17, 'right')}  {pad('有数据？', 8)}")
    for k in range(0, len(GRID), 5):
        band = f"[{gp_mean[k] - 2 * sd[k]:+.2f}, {gp_mean[k] + 2 * sd[k]:+.2f}]"
        seen = "是" if SEEN[k] else "否"
        print(f"  {GRID[k]:>5.2f} {truth[k]:>+7.2f} {lin[k]:>+7.2f} {quad[k]:>+7.2f} {gp_mean[k]:>+8.2f} {band:>17s}  {seen}")
    # 在新的比较上评估：看过的范围内 / 跨越到没看过的范围
    ta, tb, ty = make_pairs(rng, 4000, SEEN)
    ua, ub, uy = make_pairs(rng, 4000, np.ones(len(GRID), dtype=bool))
    print(f"\n  {pad('模型', 8)} {pad('x ≤ 0.6 内的准确率', 19, 'right')} {pad('全范围的准确率', 15, 'right')} "
          f"{pad('认为最好的长度', 15, 'right')}")
    f_true = true_utility(GRID)
    for name, f in [("真实效用", f_true), ("线性", lin), ("二次", quad), ("GP", gp_mean)]:
        acc_in, _ = pair_metrics(f, ta, tb, ty)
        acc_all, _ = pair_metrics(f, ua, ub, uy)
        print(f"  {pad(name, 8)} {acc_in:>19.3f} {acc_all:>15.3f} {GRID[np.argmax(f)]:>15.2f}")
    inside = float(sd[SEEN].mean())
    outside = float(sd[~SEEN].mean())
    print(f"""
  GP 的后验标准差：x ≤ 0.6 内平均 {inside:.2f}，x > 0.6 平均 {outside:.2f}

观察：
  · 线性 BT 只能说“越短越好”或“越长越好”，完全错过了主峰；二次 BT 能画出一个峰，但形状不对，
    而且在 x > 0.6 处信心十足地一路下降 —— 二次函数的形状替它做了外推。
  · GP 没有假设形状，只假设“相近的长度效用相近”，在有数据的范围内把主峰学得很准。
  · 在没有数据的 x > 0.6，GP 的均值慢慢回到先验（0 附近），标准差变大到接近先验的 {SF}：
    它**知道自己不知道**。真实的小峰（x ≈ 0.78）它当然学不到 —— 没有数据就没有信息 —— 但宽阔的不确定性带
    提醒我们：那里可能有东西，值得去问（演示 3）。
  · 这和第 2 讲“奖励模型的外推”是同一个问题：参数化的模型（线性、二次）在数据之外会“自信地胡说”。""")
    return a, b, y


# =============================================================================
# 演示 2
# =============================================================================
def demo_lengthscale(rng, data):
    print_header("演示 2：长度尺度 ℓ 怎样选 —— 边际似然 vs 留出数据")
    a, b, y = data
    ta, tb, ty = make_pairs(rng, 4000, SEEN)
    print(f"用演示 1 的 {len(y)} 条比较拟合，在 4000 条新的比较（x ≤ 0.6）上评估。\n")
    print(f"  {pad('ℓ', 6)} {pad('log 边际似然（拉普拉斯）', 25, 'right')} {pad('留出对数损失', 13, 'right')} "
          f"{pad('留出准确率', 11, 'right')} {pad('牛顿步数', 9, 'right')}")
    rows = []
    for ell in [0.01, 0.03, 0.05, 0.08, 0.15, 0.3, 1.0]:
        fit = gp_pref_laplace(GRID, a, b, y, ell=ell, sf=SF)
        acc, ll = pair_metrics(fit.mean, ta, tb, ty)
        rows.append((ell, fit.log_evidence, ll))
        print(f"  {ell:>6.2f} {fit.log_evidence:>25.2f} {ll:>13.4f} {acc:>11.3f} {fit.steps:>9d}")
    best_ev = max(rows, key=lambda r: r[1])[0]
    best_ll = min(rows, key=lambda r: r[2])[0]
    print(f"""
  边际似然最大的 ℓ = {best_ev}；留出对数损失最小的 ℓ = {best_ll}

观察：
  · ℓ 太小（0.01）：相邻网格点几乎不相关，每个点只能靠自己的那几条比较，相当于第 5 讲“每个选项一个分数”
    的 BT 加上弱先验 —— 数据被摊得太薄，留出表现差。
  · ℓ 太大（1.0）：函数被迫非常平滑，接近一条直线或一个宽包，画不出主峰的形状。
  · 边际似然（“证据”）只用训练数据，就挑出了和留出数据差不多的 ℓ。它自动权衡“拟合得好”和“模型太灵活”
    （贝叶斯的奥卡姆剃刀），不需要像第 5 讲那样做交叉验证。
  · 【边界】拉普拉斯近似的证据只是近似；数据很少时，它和交叉验证可能选出不同的 ℓ。""")


# =============================================================================
# 演示 3
# =============================================================================
def ucb_vs_best(fit):
    """
    对每个网格点 x，计算"和当前最好的点比，乐观地看能好多少"：E[f(x) − f(best)] + 2·sd[f(x) − f(best)]。

    【思路】差值的后验方差 = Var(x) + Var(best) − 2 Cov(x, best)。差值不受"整体平移"影响，不用再减平均值。
    【返回】(当前最好的下标, 均值差数组, 差值标准差数组, 乐观值数组)
    """
    best = int(np.argmax(fit.mean))
    var_diff = np.clip(np.diag(fit.cov) + fit.cov[best, best] - 2 * fit.cov[:, best], 0.0, None)
    mean_diff = fit.mean - fit.mean[best]
    sd_diff = np.sqrt(var_diff)
    return best, mean_diff, sd_diff, mean_diff + 2 * sd_diff


def demo_next_query(rng, data):
    print_header("演示 3：下一个问题该问什么？（预告第 7、13 讲）")
    a, b, y = data
    fit = gp_pref_laplace(GRID, a, b, y, ell=0.08, sf=SF)
    best, mean_diff, sd_diff, ucb = ucb_vs_best(fit)
    f = true_utility(GRID)
    print(f"当前后验均值最高的长度：x = {GRID[best]:.2f}。拿其他长度和它比，“乐观地看能好多少”：\n")
    print(f"  {pad('x', 6)} {pad('均值差', 8, 'right')} {pad('差值的标准差', 13, 'right')} {pad('均值差 + 2 标准差', 18, 'right')} "
          f"{pad('真实效用差', 11, 'right')}")
    for x in [0.2, 0.45, 0.6, 0.7, 0.78, 0.9, 1.0]:
        k = int(round(x * 60))
        print(f"  {GRID[k]:>6.2f} {mean_diff[k]:>+8.2f} {sd_diff[k]:>13.2f} {ucb[k]:>+18.2f} {f[k] - f[best]:>+11.2f}")

    # 一个小实验：再问 40 个问题，每次都拿"乐观值最高的长度"和"当前最好的长度"比较
    budget = 40
    results = {}
    for name in ["按乐观值提问", "随机提问（全范围）"]:
        qa, qb, qy = list(a), list(b), list(y)
        asked = []
        local = np.random.default_rng(7)
        fit_t = fit
        for _ in range(budget):
            best_t, _, _, ucb_t = ucb_vs_best(fit_t)
            if name == "按乐观值提问":
                cand = int(np.argmax(np.where(np.arange(len(GRID)) == best_t, -np.inf, ucb_t)))
                pa, pb = cand, best_t
            else:
                pa, pb = local.choice(len(GRID), size=2, replace=False)
            asked.append(GRID[pa])
            qa.append(pa)
            qb.append(pb)
            qy.append(int(local.random() < sigmoid(f[pa] - f[pb])))
            fit_t = gp_pref_laplace(GRID, np.array(qa), np.array(qb), np.array(qy), ell=0.08, sf=SF)
        results[name] = (fit_t, np.array(asked))
    print(f"\n  再问 {budget} 个问题之后（每个问题是一对长度的比较），x > 0.6 区域里学到的效用（减去 x ≤ 0.6 的平均值）：\n")
    print(f"  {pad('x', 6)} {pad('真实', 7, 'right')}" + "".join(pad(n, 20, "right") for n in results))
    truth = center_on_seen(f)
    for x in [0.65, 0.7, 0.75, 0.78, 0.82, 0.9, 1.0]:
        k = int(round(x * 60))
        row = f"  {GRID[k]:>6.2f} {truth[k]:>+7.2f}"
        for name, (ft, _) in results.items():
            m = center_on_seen(ft.mean)
            sdc = centered_sd(ft.cov)
            row += pad(f"{m[k]:+.2f} ± {2 * sdc[k]:.2f}", 20, "right")
        print(row)
    print()
    for name, (ft, asked) in results.items():
        mae = float(np.mean(np.abs(center_on_seen(ft.mean)[~SEEN] - truth[~SEEN])))
        print(f"  {pad(name, 20)}：x > 0.6 区域的平均绝对误差 {mae:.2f}；问到 x > 0.6 的次数 {int(np.sum(asked > 0.6))} / {budget}，"
              f"其中 0.7~0.86 之间 {int(np.sum((asked >= 0.7) & (asked <= 0.86)))} 次")
    print("""
观察：
  · 在有数据的范围内，其他长度和当前最好的点的差距已经比较确定，“乐观地看”也不会更好。
  · 没有数据的一侧，均值虽然低，但不确定性大，“乐观地看”可能更好 —— 最乐观的是离数据最远的 x = 1.0。
  · 按乐观值提问时（打印每一步问的 x 可以看到）：第一个问题落在最远的 x = 1.0，随后集中在小峰附近
    （0.77~0.88，0.7~0.86 之间共问了 10 次，随机提问 7 次）。大约 16 个问题之后，那一侧“乐观地看”
    也不可能比 0.3 好了，于是转而反复比较 0.27 和 0.30，细化“最好的长度究竟在哪” —— 先探索、再利用。
    问完以后，它学到的 x > 0.6 一侧的形状更接近真实（平均绝对误差约为随机提问的一半）。
  · 但它的不确定性带仍然很宽，甚至比随机提问还宽。原因：每个问题都是“拿某个长度和当前最好的 x = 0.3 比”，
    而 0.3 几乎总是赢 —— 答案几乎是预先知道的，每个回答带来的信息很少（第 7 讲：一个比较的信息量
    正比于 p(1 − p)，p 接近 0 或 1 时接近 0）。随机提问常常比较两个都在未知区域的长度，信息量反而更大。
    好的提问策略要同时考虑“哪里可能更好”和“哪个问题信息量大”（第 7、13 讲）。
  · 这个“均值 + 若干倍标准差”的规则就是 UCB（上置信界），第 11~13 讲的老虎机和偏好贝叶斯优化会系统地讲；
    第 7 讲从 Fisher 信息的角度讲“哪个问题信息量最大”。只有知道“哪里不确定”，才能聪明地提问。""")


def main():
    rng = np.random.default_rng(64)
    data = demo_fit(rng)
    demo_lengthscale(rng, data)
    demo_next_query(rng, data)


if __name__ == "__main__":
    main()