"""
=============================================================================
第 13 讲 · 共享模块：RiverSwim（逆流游泳）——一个著名的探索难题
=============================================================================

对应讲义：lectures/lecture13_mdp_exploration.md

-----------------------------------------------------------------------------
【环境描述】（Strehl & Littman, 2008 的 RiverSwim，这里做成有限视界的回合制版本）
-----------------------------------------------------------------------------
一条河被分成 N 段（状态 0, 1, ..., N-1）。智能体每个回合从最左边的状态 0（岸边）出发，走 H 步。

    状态：   [0] ─ [1] ─ [2] ─ [3] ─ [4] ─ [5]
             岸边                           上游
             小奖励 0.005                   大奖励 1

两个动作：
    动作 0（向左 / 顺流）：一定成功，移动到左边一格（在状态 0 则原地不动）。
                           在状态 0 选这个动作，得到奖励 0.005（岸边的一点小甜头）。
    动作 1（向右 / 逆流）：水流很急——
                           中间的状态：0.35 的概率前进一格，0.6 原地不动，0.05 被冲回一格；
                           状态 0：0.4 前进，0.6 原地不动；
                           状态 N-1：0.6 原地不动，0.4 被冲回一格；
                           在状态 N-1 选这个动作，得到奖励 1（上游的大奖）。

-----------------------------------------------------------------------------
【为什么它难？】
-----------------------------------------------------------------------------
    - 小奖励就在出发点，随便试试就能拿到；
    - 大奖励在最远端，需要**连续很多次**选择"向右"、并且运气不错，才能到达；
    - 随机地选动作时，"向左"一定成功而"向右"经常失败，所以随机游走强烈地偏向岸边——
      到达上游需要的时间随 N **指数**增长（01 文件）。
    ε-贪心拿到岸边的小奖励之后，就更不会往上游去了。

-----------------------------------------------------------------------------
【本模块提供】
-----------------------------------------------------------------------------
    make_riverswim(N)                  : 返回真实的转移 P[s, a, s'] 和奖励 R[s, a]
    plan(P, R, H)                      : 有限视界的动态规划（逆向归纳），返回 Q[h, s, a]、V[h, s]
    policy_value(P, R, H, policy)      : 精确计算一个（随时间变化的）确定性策略从起点出发的价值
    run_episode(...)                   : 按给定的策略在真实环境中走一个回合
"""

import unicodedata

import numpy as np

H_DEFAULT = 20          # 每个回合的步数
START = 0               # 起点：岸边


def make_riverswim(N=6):
    """
    构造 N 个状态的 RiverSwim。

    【返回】(P, R)
        P : 形状 (N, 2, N)，P[s, a, s'] = 在 s 做 a 之后到达 s' 的概率
        R : 形状 (N, 2)，R[s, a] = 期望奖励（这里奖励是伯努利随机的，R 是它的均值，见 run_episode）
    【边界】N ≥ 3（至少要有"岸边、中间、上游"）。
    """
    assert N >= 3
    P = np.zeros((N, 2, N))
    R = np.zeros((N, 2))
    for s in range(N):
        P[s, 0, max(s - 1, 0)] = 1.0                    # 向左：一定成功
        if s == 0:
            P[s, 1, 1] += 0.4
            P[s, 1, 0] += 0.6
        elif s == N - 1:
            P[s, 1, s] += 0.6
            P[s, 1, s - 1] += 0.4
        else:
            P[s, 1, s + 1] += 0.35
            P[s, 1, s] += 0.6
            P[s, 1, s - 1] += 0.05
    R[0, 0] = 0.005
    R[N - 1, 1] = 1.0
    return P, R


def plan(P, R, H, q_max=None):
    """
    有限视界的动态规划（逆向归纳）：从最后一步往前算。

        V[H, s] = 0
        Q[h, s, a] = R[s, a] + Σ_{s'} P[s, a, s'] · V[h+1, s']
        V[h, s] = max_a Q[h, s, a]

    【参数】
        P, R  : 模型（可以是真实的，也可以是估计的、乐观的、从后验中抽出来的）
        H     : 视界
        q_max : 如果给出，把 Q[h] 裁剪到不超过 q_max[h]。
                乐观算法加了探索奖励之后，Q 可能超过"理论上可能的最大值"（每步最多得 1，剩 H-h 步），
                裁剪能避免这种无意义的过度乐观。
    【返回】Q（形状 (H, N, 2)）、V（形状 (H+1, N)）
    【RL 联系】这就是第 2 讲的价值迭代，只是有限视界下不需要迭代到收敛：从后往前算 H 步就是精确解。
        最优策略是随时间变化的：policy[h, s] = argmax_a Q[h, s, a]。
    """
    N, A = R.shape
    Q = np.zeros((H, N, A))
    V = np.zeros((H + 1, N))
    for h in range(H - 1, -1, -1):
        Q[h] = R + P @ V[h + 1]
        # 【Python 小知识】P 形状 (N, A, N)，V[h+1] 形状 (N,)；P @ V[h+1] 对最后一维求和，得到 (N, A)
        if q_max is not None:
            Q[h] = np.minimum(Q[h], q_max[h])
        V[h] = Q[h].max(axis=1)
    return Q, V


def policy_value(P, R, H, policy, start=START):
    """
    精确计算一个随时间变化的确定性策略 policy[h, s] 在真实模型 (P, R) 下、从 start 出发的期望总奖励。

    【实现】与 plan() 相同的逆向归纳，只是每一步不取 max，而是按策略选定的动作。
    """
    N = R.shape[0]
    V = np.zeros(N)
    idx = np.arange(N)
    for h in range(H - 1, -1, -1):
        a = policy[h]
        V = R[idx, a] + P[idx, a] @ V
    return float(V[start])


def policy_value_probs(P, R, H, probs, start=START):
    """
    精确计算一个**随机**策略的价值。probs[h, s, a] = 第 h 步在状态 s 选动作 a 的概率。

    【用途】ε-贪心这样的策略不是确定性的：以 1-ε 的概率选贪心动作，以 ε 的概率随机选。
        逆向归纳时，每一步的价值是各动作价值按概率的加权平均：
            V_h(s) = Σ_a probs[h, s, a] · (R[s, a] + Σ_{s'} P[s, a, s'] V_{h+1}(s'))
    """
    N = R.shape[0]
    V = np.zeros(N)
    for h in range(H - 1, -1, -1):
        Q = R + P @ V
        V = np.sum(probs[h] * Q, axis=1)
    return float(V[start])


def run_episode(P, R, H, choose, rng, start=START):
    """
    在真实环境中走一个回合。

    【参数】
        choose : 函数 choose(h, s) → 动作，由算法提供（可以带随机性，比如 ε-贪心）
    【返回】(transitions, total)
        transitions : 列表 [(h, s, a, r, s'), ...]
        total       : 本回合的**期望**奖励之和 Σ R[s, a]（用期望奖励而不是抽到的 0/1 来记分，
                      去掉奖励本身的随机性，和第 11 讲的"伪遗憾"是同一个思路）
    【奖励】实际观察到的奖励 r 是伯努利随机的：以概率 R[s, a] 为 1，否则为 0。
        所以算法必须通过多次尝试才能知道每个 (s, a) 的奖励有多大——和老虎机一样。
    """
    s = start
    transitions, total = [], 0.0
    for h in range(H):
        a = choose(h, s)
        r = float(rng.random() < R[s, a])
        s2 = int(rng.choice(len(R), p=P[s, a]))
        transitions.append((h, s, a, r, s2))
        total += R[s, a]
        s = s2
    return transitions, total


def pad(text, width):
    """在右边补空格，使字符串在终端里的显示宽度为 width（中文字符算 2 格），用来对齐表格。"""
    w = sum(2 if unicodedata.east_asian_width(ch) in "WF" else 1 for ch in text)
    return text + " " * max(0, width - w)


if __name__ == "__main__":
    P, R = make_riverswim(6)
    Q, V = plan(P, R, H_DEFAULT)
    print("最优价值 V*(起点) =", round(V[0, START], 3))
    print("第 0 步的最优动作（0 = 左，1 = 右）：", Q[0].argmax(axis=1))
    always_left = np.zeros((H_DEFAULT, 6), dtype=int)
    print("\"一直向左\"的价值 =", round(policy_value(P, R, H_DEFAULT, always_left), 3))