"""
=============================================================================
第 10 讲 · 代码 4：机器人的轨迹偏好 —— 半空间、"体积移除"与信息增益
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
Sadigh 等（2017）的设定：一辆自动驾驶车（或机械臂）的一条轨迹用几个特征描述
（例如"离车道中心多远""速度多快""离前车多近"），人心中的奖励是特征的加权和 r = wᵀφ，
权重 w 只确定方向（长度和"人有多敏锐"混在一起），所以放在单位球面上。

每次给人看两条轨迹 a、b，问"哪条更好"。如果人总是答对，回答 "a 更好" 就告诉我们 wᵀ(φ_a − φ_b) > 0：
w 在一个**半空间**里。每个回答把可能的 w 的集合切掉一部分（"体积移除"）。

    · 体积移除（Sadigh 等 2017）：挑"预期能切掉最多体积"的问题 —— 两个答案各占一半时最好。
    · 信息增益（Bıyık 等 2019）：挑"答案和 w 之间的互信息最大"的问题（代码 01 的 BALD）。
    · 随机。

人的回答有噪声：P(a 更好) = σ(κ · wᵀ(φ_a − φ_b))。候选问题里混进了一些"两条几乎一样的轨迹"。

    演示 1：带噪声的人：三种挑法的误差（估计的 w 和真实 w 的夹角）；体积移除挑了多少"几乎一样"的问题？
    演示 2：几乎不出错的人（κ 很大）：结论变了吗？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_trajectory_preferences.py        # 约 25 秒

只需要 numpy。
"""

import numpy as np

D = 3                     # 特征维数
N_PARTICLES = 1000        # 用 1000 个"粒子"（单位球面上的点）表示 w 的后验
N_TRAJ = 60               # 轨迹库：60 条轨迹，任意两条可以组成一个问题
N_DUP = 60                # 再混进 60 个"两条几乎一样的轨迹"组成的问题
RULES = {"random": "随机", "volume": "体积移除", "infogain": "信息增益"}


def sigmoid(x):
    """σ(x)。"""
    return 0.5 * (1.0 + np.tanh(0.5 * x))


def entropy(p):
    """二元熵（nat）。"""
    p = np.clip(p, 1e-12, 1 - 1e-12)
    return -(p * np.log(p) + (1 - p) * np.log(1 - p))


def pad(text, width, align="left"):
    """按显示宽度补空格（汉字算 2 个宽度）。"""
    import unicodedata
    w = sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in text)
    spaces = " " * max(0, width - w)
    return text + spaces if align == "left" else spaces + text


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 78)
    print(f"  {title}")
    print("=" * 78)


def unit(v):
    """把向量（或每一行）缩放成长度 1。"""
    return v / np.linalg.norm(v, axis=-1, keepdims=True)


def run(rule, rep, kappa, steps=30):
    """
    一次实验：真实 w 随机，依次问 steps 个问题。

    【思路】
        · 后验用粒子表示：1000 个均匀分布在球面上的 w_s，每个带一个对数权重（先验均匀）。
          每得到一个回答，每个粒子的权重乘以"如果 w = w_s，人这样回答的概率"（第 6 讲的贝叶斯更新）。
        · 对每个候选问题 u = φ_a − φ_b：
            P(a 更好) = Σ_s 权重_s · σ(κ w_sᵀu)
            体积移除的得分 = 1 − P(a)² − P(b)² = 2 P(a) P(b)：预期切掉的后验质量（两个答案各一半时最大）
            信息增益的得分 = H(P(a)) − Σ_s 权重_s · H(σ(κ w_sᵀu))
    【返回】(每一步之后的夹角误差（度）数组, 挑了几个"几乎一样"的问题)
    """
    rng = np.random.default_rng(rep)
    w_true = unit(rng.normal(size=D))
    particles = unit(rng.normal(size=(N_PARTICLES, D)))
    logw = np.zeros(N_PARTICLES)
    phi = rng.normal(size=(N_TRAJ, D))
    i, j = np.triu_indices(N_TRAJ, k=1)
    U = np.vstack([phi[i] - phi[j], 0.02 * rng.normal(size=(N_DUP, D))])
    is_dup = np.r_[np.zeros(len(i)), np.ones(N_DUP)].astype(bool)
    P_ws = sigmoid(kappa * (particles @ U.T))            # (粒子, 候选)：每个粒子认为 a 更好的概率
    errors, n_dup = [], 0
    for _ in range(steps):
        wt = np.exp(logw - logw.max())
        wt /= wt.sum()
        pa = wt @ P_ws
        if rule == "random":
            k = int(rng.integers(len(U)))
        elif rule == "volume":
            k = int(np.argmax(pa * (1 - pa)))
        else:
            k = int(np.argmax(entropy(pa) - wt @ entropy(P_ws)))
        n_dup += int(is_dup[k])
        a_better = rng.random() < sigmoid(kappa * (w_true @ U[k]))
        logw += np.log(np.clip(P_ws[:, k] if a_better else 1 - P_ws[:, k], 1e-300, None))
        wt = np.exp(logw - logw.max())
        wt /= wt.sum()
        est = unit(wt @ particles)                       # 后验均值的方向
        errors.append(np.degrees(np.arccos(np.clip(est @ w_true, -1, 1))))
    return np.array(errors), n_dup


def report(kappa, reps):
    print(f"  {pad('挑法', 10)}" + "".join(pad(f"{n} 个问题后", 13, "right") for n in [5, 10, 20, 30])
          + pad("“几乎一样”的问题", 18, "right"))
    for rule, name in RULES.items():
        res = [run(rule, rep, kappa) for rep in range(reps)]
        E = np.array([r[0] for r in res]).mean(axis=0)
        print(f"  {pad(name, 10)}" + "".join(f"{E[n - 1]:>13.1f}" for n in [5, 10, 20, 30])
              + f"{np.mean([r[1] for r in res]):>18.1f}")


def main():
    reps = 30
    print_header("演示 1：会出错的人（κ = 5）")
    print(f"{D} 维特征，w 在单位球面上；{N_TRAJ} 条轨迹两两组成 {N_TRAJ * (N_TRAJ - 1) // 2} 个问题，"
          f"再混进 {N_DUP} 个“两条几乎一样的轨迹”的问题。")
    print(f"人回答 P(a 更好) = σ(κ wᵀ(φ_a − φ_b))。每种挑法重复 {reps} 次，报告估计方向和真实方向的平均夹角（度），")
    print("以及 30 个问题里平均挑了几个“几乎一样”的问题：\n")
    report(5.0, reps)
    print("""
观察：
  · 体积移除的大部分问题都是“两条几乎一样的轨迹”！对这种问题，无论 w 是什么，人都是在瞎猜（P = 0.5），
    于是“两个答案各占一半” —— 体积移除的得分 2P(a)P(b) 达到最大值 0.5。可是瞎猜的答案对 w 毫无信息。
    结果它比随机挑还差。
  · 信息增益没有这个问题：对“几乎一样”的问题，每个可能的 w 自己就是五五开，E[H(p)] = H(E[p])，信息增益为 0。
    它学得最快。Bıyık 等（2019）指出了这个问题，并且发现信息增益挑出的问题对人来说也更“容易回答”
    （两条轨迹差别明显）—— 难以区分的问题不仅没信息，还让人烦、让人更容易答错。""")

    print_header("演示 2：几乎不出错的人（κ = 50）")
    report(50.0, reps)
    print("""
观察：
  · 人几乎不出错时，“几乎一样”的轨迹也能给出有用的答案（只要有一点差别，人就能分辨），
    体积移除的表现接近信息增益 —— 体积移除最初就是在“人不出错”的假设下提出的。
  · 但它仍然挑了很多“几乎一样”的问题：在真实的人身上，这些问题恰恰是最难、最容易答错的。
  · 教训（和代码 02 的不确定性采样一样）：只看“答案有没有悬念”是不够的，要看“答案能不能告诉我们关于 w 的事”。
  · 同样的方法被用来个性化外骨骼的步态（Tucker 等 2020）：每次让使用者走两种步态、说哪个更舒服，
    几十次比较就能找到适合这个人的参数（第 13 讲的偏好贝叶斯优化）。""")


if __name__ == "__main__":
    main()