"""
=============================================================================
第 6 讲 · 代码 2：最简单的策略梯度——老虎机上的 REINFORCE 与基线
=============================================================================

对应讲义：lectures/lecture06_policy_gradient.md 第 4、7 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
在只有一步的问题（多臂老虎机）上理解策略梯度的核心思想。
没有状态、没有时间，只剩下最本质的东西：

    目标：J(θ) = E_{a~π_θ}[r(a)]           （期望奖励）
    梯度：∇J(θ) = E_{a~π_θ}[r · ∇log π_θ(a)]   （似然比技巧，讲义 4.2 节）
    更新：θ ← θ + α · r · ∇log π_θ(a)       （用一个样本估计梯度）

直觉："试一个动作；如果奖励高，就增加它的概率，奖励越高增加得越多。"

然后我们会看到一个问题：如果所有奖励都是正的（比如都在 10 左右），
那**每个**被选中的动作概率都会增加——只是增加得多少不同。梯度估计的方差很大。
解决办法是减去一个**基线** b：θ ← θ + α (r - b) ∇log π_θ(a)。
"比平均好就增加概率，比平均差就减少概率。"

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 02_bandit_policy_gradient.py      （约 5 秒）
"""

import numpy as np

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


def softmax(h):
    z = h - h.max()
    e = np.exp(z)
    return e / e.sum()


# =============================================================================
# 第 1 部分：老虎机上的 REINFORCE
# =============================================================================
def reinforce_bandit(true_means, n_steps, alpha, use_baseline, rng, noise_std=1.0):
    """
    在多臂老虎机上运行 REINFORCE（Sutton & Barto 2.8 节称之为"梯度老虎机算法"）。

    【参数化】每个臂一个偏好 θ_a，π(a) = softmax(θ)_a。
        此时得分函数 ∇_θ log π(a) = e_a - π（e_a 是第 a 个位置为 1 的向量），
        即：被选中的臂的偏好 +(1 - π(a))，其他臂的偏好 -π(b)。

    【更新】
        θ ← θ + α (r - b) (e_a - π)
        - 不用基线：b = 0
        - 用基线：  b = 到目前为止所有奖励的平均值（用增量平均计算）

    【参数】
        true_means   : 每个臂的真实平均奖励。
        n_steps      : 步数（拉多少次）。
        alpha        : 学习率。
        use_baseline : 是否使用基线。
        noise_std    : 奖励噪声的标准差。

    【返回】
        长度 n_steps 的布尔数组：每一步是否选中了最优臂。

    【边界】
        - 更新时，基线必须用**这一步之前**的平均值（不包含当前奖励 r）。
          如果把当前的 r 也算进基线，基线就和当前动作有关了，梯度估计会有偏。
    """
    n_arms = len(true_means)
    best = int(np.argmax(true_means))
    theta = np.zeros(n_arms)
    baseline, n = 0.0, 0
    chose_best = np.zeros(n_steps, dtype=bool)
    for t in range(n_steps):
        pi = softmax(theta)
        a = int(rng.choice(n_arms, p=pi))
        r = rng.normal(true_means[a], noise_std)
        b = baseline if use_baseline else 0.0
        grad_log_pi = -pi.copy()
        grad_log_pi[a] += 1.0                      # e_a - π
        theta += alpha * (r - b) * grad_log_pi
        n += 1
        baseline += (r - baseline) / n             # 更新基线（在用完之后）
        chose_best[t] = a == best
    return chose_best


# =============================================================================
# 第 2 部分：直接测量梯度估计的方差
# =============================================================================
def gradient_estimates(theta, true_means, n_samples, baseline, rng, noise_std=1.0):
    """
    在固定的 θ 处，抽取 n_samples 个单样本梯度估计 (r - b)(e_a - π)。

    【返回】形状 (n_samples, n_arms) 的数组，每行是一个梯度估计。
    """
    pi = softmax(theta)
    n_arms = len(true_means)
    arms = rng.choice(n_arms, size=n_samples, p=pi)
    rewards = rng.normal(np.asarray(true_means)[arms], noise_std)
    grads = -np.tile(pi, (n_samples, 1))
    grads[np.arange(n_samples), arms] += 1.0
    return (rewards - baseline)[:, None] * grads


def main():
    rng = np.random.default_rng(0)
    true_means = np.array([0.2, 0.5, 1.0, 0.0, 0.3])    # 臂 2 最好

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：一步 REINFORCE 更新做了什么？")
    print("=" * 78)
    theta = np.zeros(5)
    pi = softmax(theta)
    print(f"  初始 θ = {theta}，π = {pi}")
    for a, r in [(2, 1.5), (3, -0.8)]:
        g = -pi.copy()
        g[a] += 1
        new_theta = theta + 0.5 * r * g
        print(f"  若选中臂 {a}、奖励 r = {r:+.1f}：θ 变为 {new_theta}，π 变为 {softmax(new_theta)}")
    print("  → 奖励为正：被选中的臂概率上升，其他臂下降；奖励为负：被选中的臂概率下降。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：学习曲线——有无基线、奖励是否整体平移（选中最优臂的比例，200 次运行平均）")
    print("=" * 78)
    n_runs, n_steps = 200, 1000
    checkpoints = [10, 50, 100, 250, 500, 1000]
    print("  设置                               " + "".join(f"{f'第{c}步':>8s}" for c in checkpoints))
    for shift in [0.0, 10.0]:
        for use_baseline in [False, True]:
            results = np.mean([reinforce_bandit(true_means + shift, n_steps, 0.1, use_baseline, rng)
                               for _ in range(n_runs)], axis=0)
            label = f"奖励 +{shift:<4.0f}{'有基线' if use_baseline else '无基线'}"
            # 用窗口平均让曲线更平滑
            smooth = [results[max(0, c - 10):c].mean() for c in checkpoints]
            print(f"  {label:<24s}          " + "".join(f"{x * 100:7.1f}%" for x in smooth))
    print("\n  观察：")
    print("  - 奖励在 0 附近时，有无基线差别不大。")
    print("  - 把所有奖励都加 10（问题本质没变，最优臂还是同一个）：")
    print("    **无基线**的版本几乎学不会！因为每次的奖励都约为 10，被选中的臂的概率总是被大幅提高，")
    print("    \"选中了就加分\"压倒了\"好的多加、差的少加\"的信号；而且大步长让策略很快锁定在")
    print("    某个随机早期选中的臂上。")
    print("    **有基线**的版本完全不受影响：r - b 只看\"比平均好多少\"，与奖励的整体平移无关。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：直接测量——基线不改变梯度的期望，但大幅减小方差")
    print("=" * 78)
    theta = np.array([0.1, 0.3, 0.2, -0.2, 0.0])
    pi = softmax(theta)
    true_grad = pi * (true_means - pi @ true_means)      # 精确梯度：∂J/∂θ_a = π_a (r_a - J)
    print(f"  π = {pi}")
    print(f"  精确梯度 ∇J(θ)       = {true_grad}")
    for shift in [0.0, 10.0]:
        J = pi @ (true_means + shift)
        for b, name in [(0.0, "b = 0"), (J, "b = J(θ)")]:
            g = gradient_estimates(theta, true_means + shift, 200_000, b, rng)
            print(f"  奖励 +{shift:<4.0f} {name:<9s}：样本均值 {g.mean(axis=0)}，"
                  f"总方差 {g.var(axis=0).sum():8.3f}")
    print("  → 四种情况的样本均值都约等于精确梯度（无偏）；")
    print("    但奖励 +10 且无基线时，方差大了约 100 倍。用 b ≈ J(θ)（平均奖励）作基线，")
    print("    方差回到与\"奖励在 0 附近\"时相同的水平。")


if __name__ == "__main__":
    main()