"""
=============================================================================
第 15 讲 · 代码 2：模型过时了怎么办？——阻塞迷宫、捷径迷宫与 Dyna-Q+
=============================================================================

对应讲义：lectures/lecture15_model_based_rl.md 第 3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
模型是从过去的经验中学来的。如果世界变了，模型就过时了，基于它的规划可能把智能体带向错误的方向。

两种变化（Sutton & Barto 例 8.2、8.3）：
    - **阻塞**：原来的路被堵住了，另一边开了新路。模型说"这里能走"，实际走过去撞墙——
      **变坏的变化很容易发现**：按照旧计划行动，马上就会碰壁，模型随之被纠正。
    - **捷径**：原来的路还在，另一边多开了一条更短的路。模型说"那边是墙"，而智能体按照旧计划走得好好的，
      根本不会去那边看——**变好的变化很难发现**。

Dyna-Q+ 的对策：规划时给"很久没在真实环境中试过"的 (s, a) 加探索奖励 κ·sqrt(τ)，
鼓励智能体定期回去验证模型——这又是第 11–13 讲的"乐观探索"。

环境是确定性的，模型记住最近一次的结果就是准确的，所以这里用学习率 α = 1。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture15
    python 02_changing_env.py      （约 40 秒）
"""

import numpy as np

from maze import blocking_mazes, pad, run_dyna, shortcut_mazes


def greedy_path_length(maze, Q, limit=100):
    """
    按最终 Q 表的贪心策略，从起点走到终点需要多少步（不加任何随机性）。
    走 limit 步还没到（比如在墙边来回打转），返回 None。
    【用途】判断智能体最终学到的是哪条路：阻塞迷宫的新路是 16 步；捷径迷宫的捷径是 10 步、原路是 16 步。
    """
    s = maze.idx(maze.start)
    for k in range(1, limit + 1):
        q = Q[s]
        s, _, done = maze.step(s, int(np.argmax(q)))
        if done:
            return k
    return None


def run_experiment(before, after, change_at, total, agents, n_seeds, checkpoints):
    """对每个智能体跑 n_seeds 次，返回各检查点的平均累计奖励，以及最终贪心路径长度的分布。"""
    results = {}
    for name, n, kappa in agents:
        cums, paths = [], []
        for sd in range(n_seeds):
            out = run_dyna(lambda t: before if t < change_at else after, n, np.random.default_rng(sd),
                           max_steps=total, alpha=1.0, kappa=kappa)
            cums.append([out["cum_reward"][c - 1] for c in checkpoints])
            paths.append(greedy_path_length(after, out["Q"]))
        results[name] = (np.mean(cums, axis=0), paths)
    return results


def describe_paths(paths):
    """把最终贪心路径长度统计成"几步：几个种子"的字符串。"""
    from collections import Counter
    c = Counter("到不了" if p is None else f"{p}步" for p in paths)
    return "，".join(f"{k}×{v}" for k, v in sorted(c.items(), key=lambda kv: (kv[0] == "到不了", kv[0])))


def main():
    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：阻塞迷宫（第 1000 步时，原来的路被堵住，另一边开了新路；共 3000 步，20 个种子）")
    print("=" * 78)
    before, after = blocking_mazes()
    print("  变化之前：                  变化之后：")
    for a, b in zip(before.render().split("\n"), after.render().split("\n")):
        print(f"    {a}          {b}")
    print(f"  最短路径：变化前 {before.shortest_path()} 步，变化后 {after.shortest_path()} 步")
    cps = [1000, 1500, 2000, 3000]
    agents = [("Q-learning（n = 0）", 0, 0.0), ("Dyna-Q（n = 10）", 10, 0.0), ("Dyna-Q+（n = 10，κ = 1e-4）", 10, 1e-4)]
    res = run_experiment(before, after, 1000, 3000, agents, 20, cps)
    print("\n   " + pad("智能体", 30) + "累计到达终点的次数（第 1000 / 1500 / 2000 / 3000 步）  变化后的 2000 步内")
    for name, (cum, paths) in res.items():
        print("   " + pad(name, 30) + "".join(f"{x:9.1f}" for x in cum) + f"              {cum[3] - cum[0]:6.1f}")
    print("\n   最终贪心策略从起点到终点的步数（变化后的迷宫，最短 16 步）：")
    for name, (cum, paths) in res.items():
        print("     " + pad(name, 30) + describe_paths(paths))
    print("  观察：")
    print("  - 墙变化之后（第 1000 步），所有智能体都有一段时间几乎拿不到奖励：按照旧的模型和 Q 值，")
    print("    它们一次次地走向已经被堵住的右边缺口。")
    print("  - \"变坏\"的变化会被发现：撞墙的经验更新了模型，规划不再把智能体引向右边的缺口。")
    print("    但发现\"旧路不通\"之后，还得找到**新路**——左边缺口那一带以前很少去，Q 值还是 0。")
    print("    Dyna-Q 只能靠 ε-贪心的随机探索慢慢摸过去：2000 步之后，20 个种子中只有 8 个学会了新路。")
    print("  - Dyna-Q+ 恢复得快得多（变化后拿到的奖励是 Dyna-Q 的 3 倍），19 个种子学会了最短的新路。")
    print("    它的探索奖励让它更系统地探索整个迷宫——在变化之前就拿到了更多奖励，变化之后也更快找到新路。")
    print("  - 普通 Q-learning 在 3000 步内连原来的迷宫都没学会（只到达终点几次）——没有模型，学得太慢了。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：捷径迷宫（第 3000 步时，右边多开了一条更短的路，原路仍然可以走；共 6000 步，20 个种子）")
    print("=" * 78)
    before, after = shortcut_mazes()
    print("  变化之前：                  变化之后：")
    for a, b in zip(before.render().split("\n"), after.render().split("\n")):
        print(f"    {a}          {b}")
    print(f"  最短路径：变化前 {before.shortest_path()} 步，变化后 {after.shortest_path()} 步")
    cps = [3000, 4000, 5000, 6000]
    agents = [("Q-learning（n = 0）", 0, 0.0), ("Dyna-Q（n = 50）", 50, 0.0),
              ("Dyna-Q+（n = 50，κ = 1e-3）", 50, 1e-3), ("Dyna-Q+（n = 50，κ = 1e-2）", 50, 1e-2)]
    res = run_experiment(before, after, 3000, 6000, agents, 20, cps)
    print("\n   " + pad("智能体", 30) + "累计到达终点的次数（第 3000 / 4000 / 5000 / 6000 步）  变化后的 3000 步内")
    for name, (cum, paths) in res.items():
        print("   " + pad(name, 30) + "".join(f"{x:9.1f}" for x in cum) + f"              {cum[3] - cum[0]:6.1f}")
    print("\n   最终贪心策略从起点到终点的步数（捷径 10 步，原路 16 步）：")
    for name, (cum, paths) in res.items():
        print("     " + pad(name, 30) + describe_paths(paths))
    print("  观察：")
    print("  - Dyna-Q 从不发现捷径（20 个种子全都还在走原路）：它的模型说\"右边是墙\"，")
    print("    而沿着原路走得好好的，没有任何理由去右边看看。")
    print("    ε-贪心的随机探索要连续好几步走到右边那个缺口，概率很小（第 13 讲：随机探索很难走到远处）。")
    print("    **模型\"变好\"的错误不会自己暴露出来**——你永远不会发现一个从没去验证过的机会。")
    print("  - Dyna-Q+（κ = 1e-3）：很久没去过的地方，探索奖励 κ·sqrt(τ) 越来越大，")
    print("    规划最终会把智能体\"派\"过去看看——于是发现了捷径，变化后拿到的奖励明显多于 Dyna-Q。")
    print("    有 5 个种子的贪心路径\"到不了\"终点：Dyna-Q+ 的 Q 值里混进了探索奖励，")
    print("    在最后这一刻，贪心策略正被\"派去\"验证某个很久没去的角落，而不是直奔终点——")
    print("    它的 Q 值衡量的是\"利用 + 探索\"的综合价值，而不只是到达终点的价值。")
    print("  - 一个意外：普通 Q-learning 有 13 个种子学会了捷径！原因是它学得太慢——到第 3000 步时")
    print("    连原来的路都还没学会（只到过终点 5 次左右），还在四处乱走，于是等于在**新迷宫**上从头学起。")
    print("    它没有\"过时的知识\"需要克服，但总奖励远少于 Dyna-Q+。")
    print("  - κ 太大（1e-2）：探索奖励压倒了真实奖励（真实奖励只有 1，而 sqrt(τ) 可以到几十），")
    print("    智能体一直在\"验证\"各个角落，几乎不去终点。和第 11 讲 UCB 的 c 太大、第 13 讲 β 太大一样：探索过度。")


if __name__ == "__main__":
    main()