"""
=============================================================================
第 13 讲 · 代码 1：为什么在 MDP 中探索更难——RiverSwim 与 ε-贪心的失败
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
老虎机中，任何一个臂都可以**直接**拉。MDP 中，有些状态只有经过**一长串特定的动作**才能到达。
ε-贪心这种"每一步独立地随机一下"的探索，要碰巧连续做对很多步，概率随链的长度指数下降。

本文件：
    1. RiverSwim 的最优策略与几个简单策略的价值
    2. 随机游走需要多少步才能到达上游？——精确计算"期望到达时间"，看它如何随链长 N 指数增长
    3. 确定性等价（贪心地用估计的模型）、ε-贪心、ε-贪心 Q-learning 在 RiverSwim 上的表现

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture13
    python 01_why_exploration_is_hard.py      （约 25 秒）
"""

import numpy as np

from agents import summarize
from riverswim import H_DEFAULT, START, make_riverswim, pad, plan, policy_value, policy_value_probs

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


def hitting_time(P, probs, goal):
    """
    按一个固定的（平稳的）随机策略行动，从每个状态出发、第一次到达 goal 的**期望步数**。

    【参数】probs[s, a] = 在状态 s 选动作 a 的概率（不随时间变化）
    【方法】设 T(s) 为从 s 出发的期望到达时间，则
            T(goal) = 0，    T(s) = 1 + Σ_{s'} P_π(s' | s) · T(s')     （s ≠ goal）
        这是一个线性方程组，用 np.linalg.solve 直接求解。
        其中 P_π(s' | s) = Σ_a probs[s, a] · P[s, a, s'] 是"按这个策略走"时的状态转移概率。
    【边界】如果 goal 从某些状态根本到不了，方程组会奇异或得到负数/极大的值；RiverSwim 中不会出现这种情况。
    """
    N = P.shape[0]
    P_pi = np.einsum("sa,sat->st", probs, P)
    others = [s for s in range(N) if s != goal]
    M = np.eye(len(others)) - P_pi[np.ix_(others, others)]
    # 【Python 小知识】np.ix_(行下标, 列下标) 取出一个子矩阵：只保留"不是 goal"的行和列
    T = np.linalg.solve(M, np.ones(len(others)))
    out = np.zeros(N)
    out[others] = T
    return out


def main():
    N, H = 6, H_DEFAULT
    P, R = make_riverswim(N)

    # ---------------------------------------------------------------------
    print("=" * 78)
    print(f"演示 1：RiverSwim（N = {N} 个状态，每回合 H = {H} 步，从岸边的状态 0 出发）")
    print("=" * 78)
    Q, V = plan(P, R, H)
    opt = Q.argmax(axis=2)
    print(f"  最优价值 V*(起点) = {V[0, START]:.3f}")
    print("  最优策略（行 = 第几步 h，列 = 状态；→ = 向右逆流，← = 向左顺流）：")
    for h in [0, 5, 10, 14, 15, 16, 19]:
        print(f"    h = {h:2d}：" + " ".join("→" if a == 1 else "←" for a in opt[h]))
    left = np.zeros((H, N), dtype=int)
    uniform = np.full((H, N, 2), 0.5)
    print(f"  \"一直向左\"（守着岸边的小奖励）的价值 = {policy_value(P, R, H, left):.3f}")
    print(f"  完全随机的价值 = {policy_value_probs(P, R, H, uniform):.3f}")
    print("  观察：")
    print("  - 最优策略是一路向右逆流而上：虽然前十几步没有任何奖励，但到达上游后每一步都能拿到 1。")
    print("  - 快结束时（从 h = 14 开始），离上游太远的状态已经来不及游过去了，最优策略改为向左拿岸边的小奖励——")
    print("    有限视界下，最优策略会随剩余时间变化。")
    print(f"  - \"一直向左\"的价值只有最优的 {policy_value(P, R, H, left) / V[0, START] * 100:.0f}%，但它的奖励来得最快、最容易被发现。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：从岸边出发，期望要走多少步才能第一次到达上游？（精确计算）")
    print("=" * 78)
    print("    N     一直向右     完全随机     ε-贪心（ε=0.3，贪心=向左）   ε-贪心（ε=0.1，贪心=向左）")
    prev = None
    for n in [4, 6, 8, 10]:
        Pn, _ = make_riverswim(n)
        right = np.tile([0.0, 1.0], (n, 1))
        rand = np.full((n, 2), 0.5)
        eps3 = np.tile([1 - 0.15, 0.15], (n, 1))        # ε = 0.3：以 0.3 的概率随机，随机时一半概率向右
        eps1 = np.tile([1 - 0.05, 0.05], (n, 1))
        t = [hitting_time(Pn, p, n - 1)[START] for p in (right, rand, eps3, eps1)]
        ratio = "" if prev is None else f"   （随机游走：多 2 个状态，时间多 {t[1] / prev:.0f} 倍）"
        print(f"   {n:3d}   {t[0]:9.1f}   {t[1]:12.3g}      {t[2]:16.3g}               {t[3]:16.3g}{ratio}")
        prev = t[1]
    print("  观察：")
    print("  - \"一直向右\"的到达时间随 N **线性**增长（每一格平均要 1/0.35 ≈ 3 步）。")
    print("  - 随机游走的到达时间随 N **指数**增长：在中间的状态，随机选动作时每一步")
    print("    向右前进的概率 = 0.5 × 0.35 = 0.175，向左后退的概率 = 0.5 + 0.5 × 0.05 = 0.525——")
    print("    后退是前进的 3 倍。要到达上游，相当于一个强烈偏向左边的随机游走要\"逆流\"走 N-1 格。")
    print("    每多 2 个状态，时间约乘以 3² = 9。")
    print("  - ε-贪心更糟：一旦贪心的动作是\"向左\"（因为岸边的小奖励最先被发现），")
    print("    向右的概率只有 ε/2，到达时间是天文数字：N = 10、ε = 0.1 时约 3.7 × 10¹⁵ 步。")
    print("  - （N 再大，这些数字会超过双精度浮点数能可靠表示的范围，线性方程组的解不再准确，所以只算到 N = 10。）")
    print("  - 这就是 MDP 中探索的核心困难：**需要连续做对很多步才能发现的东西，随机探索几乎不可能发现**。")
    print("    老虎机里不存在这个问题——每个臂都可以直接拉。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：几种\"朴素\"算法在 RiverSwim 上的表现（2000 个回合，20 个随机种子）")
    print("=" * 78)
    print(f"  最优价值 V* = {V[0, START]:.3f}；每个回合的遗憾最多约 {V[0, START]:.1f}（完全学不会时）")
    print("   " + pad("算法", 36) + "累积遗憾（第 100 / 500 / 2000 回合）  最后 100 回合  从没到过上游  第一次到达上游")
    print("   " + pad("", 36) + "                                     每回合遗憾    的种子比例    的回合（中位数）")
    rows = [("确定性等价（贪心地用估计的模型）", "ce", None),
            ("确定性等价 + ε-贪心（ε = 0.1）", "ce_eps", 0.1),
            ("确定性等价 + ε-贪心（ε = 0.3）", "ce_eps", 0.3),
            ("Q-learning + ε-贪心（ε = 0.1）", "q_eps", 0.1),
            ("Q-learning + ε-贪心（ε = 0.3）", "q_eps", 0.3)]
    for label, name, param in rows:
        s = summarize(name, P, R, H, 2000, range(20), param)
        print("   " + pad(label, 36) + "".join(f"{x:9.1f}" for x in s["cum"])
              + f"            {s['last']:6.3f}        {s['never'] * 100:5.0f}%          {s['first_median']:6.1f}")
    print("  观察：")
    print("  - 所有这些算法的遗憾都在线性增长，最后每个回合的遗憾接近 V*——它们学到的基本上是\"守着岸边\"。")
    print("  - 确定性等价：没见过的 (s, a) 被当作\"原地不动、没有奖励\"，所以它从不主动去尝试——和老虎机中的贪心一样。")
    print("  - 加上 ε-贪心也几乎没用：ε-贪心的随机性是\"每一步独立\"的，没法把智能体送到很远的地方（演示 2）。")
    print("  - 注意\"第一次到达上游\"：很多种子**到达过**上游，却仍然没学会！偶尔到达一两次，")
    print("    留下的数据太少（比如只知道\"在上游向右拿到了 1\"，却不知道怎样可靠地游过去），")
    print("    估计出的模型/Q 值仍然认为\"游过去不值得\"，于是再也不去了。**探索需要持续，而不是碰运气。**")


if __name__ == "__main__":
    main()