"""
=============================================================================
第 6 讲 · 代码 1：价值迭代与策略迭代 —— 收缩映射、收敛速度，以及"奖励定义了最优"
=============================================================================

对应讲义：lectures/lecture06_rl_control.md 第 1–2 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
你大概已经会写 value iteration 和 policy iteration。这里不重复"怎么写"，而是用 Andrew Ng 课上的
4×3 网格世界（0.8 / 0.1 / 0.1 的打滑转移，+1 / −1 两个出口，每步 −0.02）回答几个"为什么"：

    演示 1：价值迭代的误差 ‖V_k − V*‖∞ 和收缩映射给的上界 γᵏ‖V_0 − V*‖∞ 比，谁快？
            为什么回合会结束时实际收敛远快于 γᵏ，而 agent "不想结束"时就恰好是 γᵏ？
            另外：贪心策略比价值函数早得多就变成最优的了。
    演示 2：策略迭代只需要几步。在一个随机 MDP 上把 γ 从 0.9 调到 0.999：
            VI 的迭代次数 ∝ 1/(1 − γ)，PI 几乎不变 —— PI 是 Bellman 方程上的 Newton 法。
    演示 3：只改每步的奖励 R(s)，最优策略怎样变化（"奖励定义了什么叫最优"）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_value_policy_iteration.py          # 约 1 秒

需要 numpy。依赖本目录的 rl_tools.py（它又导入 ../lecture01/linear_models.py）。
"""

import numpy as np

from rl_tools import (TabularMDP, greedy_policy, make_ng_grid, policy_evaluation, policy_iteration, print_header,
                      render_grid, value_iteration)


def side_by_side(blocks, titles, gap=6):
    """
    把几个"多行字符串块"横着并排打印（每块是一个字符串列表，比如 render_grid 的返回值）。

    【参数】blocks：list[list[str]]；titles：每块上方的标题；gap：块之间的空格数
    【边界】各块行数可以不同，短的块用空行补齐。标题只用 ASCII，避免中文宽度让列错位。
    """
    widths = [max(len(t), *(len(line) for line in b)) for b, t in zip(blocks, titles)]
    print((" " * gap).join(t.ljust(w) for t, w in zip(titles, widths)))
    for i in range(max(len(b) for b in blocks)):
        print((" " * gap).join((b[i] if i < len(b) else "").ljust(w) for b, w in zip(blocks, widths)))


def first_optimal_iteration(history, mdp, gamma, V_star):
    """
    在价值迭代的历史 V_0, V_1, … 中，找到"对 V_k 贪心得到的策略已经是最优策略"的第一个 k（并且之后一直最优）。

    【思路】不能直接比较动作（并列最优时 argmax 的选择可能不同），所以比较**策略的价值**：
        对 V_k 贪心得到 π_k，精确评估 V^{π_k}，和 V* 相差小于 1e-9 就算最优。
    【返回】int，或 None（直到最后都不是最优）
    """
    first = None
    for k, V in enumerate(history):
        pi = greedy_policy(V, mdp.P, mdp.R, gamma)
        optimal = np.max(np.abs(policy_evaluation(mdp.P, mdp.R, gamma, pi) - V_star)) < 1e-9
        if optimal and first is None:
            first = k
        elif not optimal:
            first = None                         # 之后又变回非最优：重新计数（要求"之后一直最优"）
    return first


# =============================================================================
# 演示 1：价值迭代 vs 收缩映射的上界
# =============================================================================
def demo_value_iteration():
    print_header("演示 1：Value iteration —— 实际误差 vs 收缩映射的上界 γ^k‖V_0 − V*‖")
    gamma = 0.99
    rows = {}
    for step_reward in [-0.02, 0.02]:
        mdp = make_ng_grid(step_reward=step_reward)
        V_star, _ = value_iteration(mdp.P, mdp.R, gamma, tol=1e-14)
        pi_star = greedy_policy(V_star, mdp.P, mdp.R, gamma)
        hist = []
        _, n_iters = value_iteration(mdp.P, mdp.R, gamma, tol=1e-10, history=hist)
        errs = np.array([np.max(np.abs(V - V_star)) for V in hist])
        # 最优策略下"γ × 转移矩阵"的谱半径：VI 后期每一步误差大约乘以这个数
        P_pi = mdp.P[np.arange(mdp.S), pi_star]
        rho = float(np.max(np.abs(np.linalg.eigvals(gamma * P_pi))))
        rows[step_reward] = (mdp, V_star, pi_star, errs, n_iters, rho,
                             first_optimal_iteration(hist, mdp, gamma, V_star))

    mdp, V_star, pi_star, errs, n_iters, rho, k_opt = rows[-0.02]
    print(f"Ng 的 4x3 网格世界：R(s) = −0.02（非终止状态），γ = {gamma}，转移 0.8 / 0.1 / 0.1。V* 和 π*：\n")
    side_by_side([render_grid(mdp, values=V_star), render_grid(mdp, pi=pi_star)], ["V*", "pi*"])

    ks = [1, 2, 5, 10, 15, 20, 30, 50, 100, 200]
    print(f"\n‖V_k − V*‖∞（从 V_0 = 0 出发）。bound = γ^k ‖V_0 − V*‖∞：\n")
    print(f"{'k':>5}" + "".join(f"{'R=' + format(r, '+.2f') + ' err':>16}{'bound':>11}" for r in rows))
    for k in ks:
        line = f"{k:>5}"
        for r, (_, _, _, e, _, _, _) in rows.items():
            err_s = f"{e[k]:.2e}" if k < len(e) else "(converged)"
            line += f"{err_s:>16}{gamma ** k * e[0]:>11.2e}"
        print(line)
    print(f"\n{'step reward':>12}{'VI iters (tol 1e-10)':>22}{'rho(γ·P_pi*)':>15}{'greedy optimal from k =':>25}")
    for r, (_, _, _, e, n, rh, ko) in rows.items():
        print(f"{r:>+12.2f}{n:>22}{rh:>15.3f}{str(ko):>25}")
    k_pos = rows[0.02][6]
    n_neg, n_pos = rows[-0.02][4], rows[0.02][4]
    print(f"""
观察 Observations：
  · Bellman 最优算子是 γ-收缩，所以误差至少按 γ^k 衰减（bound 列）。γ = 0.99 时这个上界要 ~2300 步
    才能降到 1e-10；可是 R = −0.02 时 VI 只用了 {n_neg} 步。
  · 原因：回合会结束。最优策略下 agent 很快走到出口，"还没结束"的概率每步都在缩小，
    后期每步误差大约乘以 ρ(γ·P_π*) = {rho:.3f}，而不是 γ = 0.99。γ^k 是最坏情况的保证。
  · R = +0.02 时活着本身有奖励，最优策略是永远不进出口（看演示 3）：回合永远不结束，
    前几步的误差恰好等于上界 γ^k，后期的速率 ρ = {rows[0.02][5]:.3f} 也接近 γ，VI 需要 {n_pos} 步。
  · 贪心策略比价值函数早得多就对了：R = −0.02 时从第 {k_opt} 步起，对 V_k 贪心的策略已经是最优策略，
    而 V_k 本身还差得远（第 {k_opt} 步误差 {errs[k_opt]:.1e}）。我们要的是策略，价值只是中间产物 ——
    这也是 policy iteration 的出发点。（R = +0.02 时很多动作并列最优，从第 {k_pos} 步起策略就最优了。）""")


# =============================================================================
# 演示 2：策略迭代只需要几步
# =============================================================================
def make_random_mdp(S, A, n_next, rng):
    """
    随机生成一个"没有终止状态"的 MDP：每个 (s, a) 随机连到 n_next 个下一状态，概率来自 Dirichlet 分布。

    【参数】S、A：状态数、动作数；n_next：每个 (s, a) 的可能后继个数（越小越"稀疏"）；rng
    【返回】TabularMDP。奖励 R(s) ~ N(0, 1)。
    【numpy 小知识】rng.dirichlet(np.ones(k)) 生成 k 个和为 1 的非负数（"随机概率向量"）；
        rng.choice(S, size=k, replace=False) 不放回地抽 k 个不同的状态。
    """
    P = np.zeros((S, A, S))
    for s in range(S):
        for a in range(A):
            nxt = rng.choice(S, size=n_next, replace=False)
            P[s, a, nxt] = rng.dirichlet(np.ones(n_next))
    return TabularMDP(P, rng.normal(size=S), np.zeros(S, dtype=bool), start=0)


def demo_policy_iteration():
    print_header("演示 2：Policy iteration 只要几步 —— 它是 Bellman 方程上的 Newton 法")
    mdp = make_ng_grid()
    gamma = 0.99
    V_star, _ = value_iteration(mdp.P, mdp.R, gamma, tol=1e-14)
    hist = []
    policy_iteration(mdp.P, mdp.R, gamma, pi0=np.full(mdp.S, 2), history=hist)   # 初始策略：全部向下（S）
    print("(a) 4x3 网格，初始策略 = 全部向下（v）。每一轮：精确评估 V^π，再贪心改进。\n")
    print(f"{'iter':>5}{'#actions changed':>18}{'V^pi(start)':>13}{'‖V^pi − V*‖∞':>15}")
    prev = None
    for k, (pi, V) in enumerate(hist):
        live = ~mdp.terminal
        changed = "-" if prev is None else str(int(np.sum(pi[live] != prev[live])))
        print(f"{k:>5}{changed:>18}{V[mdp.start]:>13.4f}{np.max(np.abs(V - V_star)):>15.2e}")
        prev = pi
    blocks = [render_grid(mdp, pi=pi) for pi, _ in hist]
    print()
    side_by_side(blocks, [f"pi_{k}" for k in range(len(hist))], gap=4)

    print("\n(b) 随机 MDP：S = 200 个状态，A = 4，每个 (s, a) 随机连到 5 个后继，没有终止状态。")
    print("    VI 停止条件 ‖V_{k+1} − V_k‖∞ < 1e-8；PI 从 \"全选动作 0\" 出发。\n")
    rng = np.random.default_rng(12)
    rmdp = make_random_mdp(200, 4, 5, rng)
    print(f"{'gamma':>7}{'VI iters':>10}{'PI iters':>10}{'‖V_VI − V*‖∞':>15}   PI errors ‖V^pi_k − V*‖∞, k = 0, 1, ...")
    vi_iters, pi_iters = [], []
    for gamma in [0.9, 0.99, 0.999]:
        V_vi, k_vi = value_iteration(rmdp.P, rmdp.R, gamma, tol=1e-8)
        h = []
        _, V_pi, k_pi = policy_iteration(rmdp.P, rmdp.R, gamma, history=h)
        errs = "  ".join(f"{np.max(np.abs(V - V_pi)):.0e}" for _, V in h[:-1])
        print(f"{gamma:>7}{k_vi:>10}{k_pi:>10}{np.max(np.abs(V_vi - V_pi)):>15.1e}   {errs}  0")
        vi_iters.append(k_vi)
        pi_iters.append(k_pi)
    print(f"""
观察 Observations：
  · (a) PI 在网格上 {len(hist) - 1} 轮改进就找到最优策略（第 {len(hist)} 轮评估只是确认"没有动作需要再改"）。
    确定性策略一共有 4^9 = 262144 个，PI 只看了其中 {len(hist)} 个。每一轮的价值都不降低（policy improvement theorem）。
  · (b) VI 的迭代次数随 1/(1 − γ) 增长：{vi_iters[0]} → {vi_iters[1]} → {vi_iters[2]}（γ 每靠近 1 十倍，步数约多十倍），
    而 PI 始终是 {min(pi_iters)}–{max(pi_iters)} 轮。PI 的误差序列每一轮都"缩小很多倍"，越到后面缩得越快 ——
    这是 Newton 法的特征（第 1 讲 §3.7）：PI 恰好是对 Bellman 方程 V = TV 做 Newton 法，而 VI 是不动点迭代。
  · 代价：PI 每轮要解一个 S×S 的线性方程组（O(S³)），VI 每步只要 O(S²A)。S 很大时 PI 每轮太贵，
    实践中常用折中的 modified policy iteration（评估时只做几次迭代，不精确求解）。
  · VI 的停止条件不等于误差：‖V_(k+1) − V_k‖ < ε 只保证 ‖V − V*‖ ≤ γε/(1 − γ)。表中第 4 列随 γ 变大，
    正是这个 1/(1 − γ) 放大。""")


# =============================================================================
# 演示 3：奖励定义了"最优"
# =============================================================================
def demo_step_reward():
    print_header("演示 3：只改每步的奖励 R(s) —— 最优策略怎样变化（γ = 0.99）")
    gamma = 0.99
    rewards = [-2.0, -0.4, -0.04, -0.02, 0.0, 0.02]
    blocks, titles, starts = [], [], []
    for r in rewards:
        mdp = make_ng_grid(step_reward=r)
        V, _ = value_iteration(mdp.P, mdp.R, gamma, tol=1e-12)
        pi = greedy_policy(V, mdp.P, mdp.R, gamma)
        blocks.append(render_grid(mdp, pi=pi))
        titles.append(f"R={r:+.2f}")
        starts.append(V[mdp.start])
    side_by_side(blocks, titles, gap=4)
    print("\n" + "    ".join(f"V*(1,1)={v:+.3f}".ljust(len(b[0])) for v, b in zip(starts, blocks)))
    print("""
（终止状态画成它的奖励 +1 / -1，墙画成 ##；R = 0 和 +0.02 时有些格子的几个动作并列最优，箭头只是其中一个。）

观察 Observations：
  · R = −2：活着太痛苦，从 (3,2)、(4,1) 直接冲向 −1 出口也比多走几步划算。
  · R = −0.4 → −0.04 → −0.02：每步越便宜，越愿意绕远路避开 −1。(2,1) 在 −0.04 时从向右变成向左；
    (3,1) 在 CS229 课上的 −0.02 时从向上变成向左 —— 绕一大圈走上面那条路。(3,2) 一直选向上
    （只有 10% 的概率滑进 −1）。
  · R = 0：时间不再有成本（只有 γ 的折扣），agent 在 (3,2) 选择"往墙上撞"（向左）、在 (4,1) 选择向下撞边界：
    这样永远不会滑进 −1，只靠 10% 的"滑动"慢慢离开。极度谨慎。
  · R = +0.02：活着就有奖励，最好永远不结束 —— 连 +1 出口也躲着走。
  · 教训：奖励函数不是"告诉 agent 怎么做"，而是"定义什么叫好"。一个看起来无害的小数字能完全改变行为。
    讲义 §6 的 RL debugging 第三问"是不是奖励函数错了？"就是在问这件事。
    CS234 第 2 讲的 05_gridworld.py 在同一个网格上把 γ 和噪声也扫了一遍。""")


def main():
    demo_value_iteration()
    demo_policy_iteration()
    demo_step_reward()


if __name__ == "__main__":
    main()