"""
=============================================================================
第 2 讲 · 代码 2：策略评估（Policy Evaluation）与动作价值函数 Q
=============================================================================

对应讲义：lectures/lecture02_mdp_planning.md 第 4、5 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
给定一个 MDP 和一个策略 π，**这个策略有多好？** 也就是计算 V^π(s)。

这是"评估"问题（讲义第 1 讲 6.3 节的"评估 vs 控制"）。它是后面一切的基础：
策略迭代要反复评估策略；价值迭代可以看成"评估只做一步"的策略迭代；
第 3 讲的 MC/TD 是"不知道模型时"的策略评估。

-----------------------------------------------------------------------------
【整体思路】
-----------------------------------------------------------------------------
关键观察（见 01_mdp_basics.py）：MDP + 固定策略 = MRP。
所以策略评估就是"计算一个 MRP 的价值"，第 1 讲的方法全都能用：

    方法 A  迭代法：反复做"贝尔曼期望备份"
            V_{k+1}(s) = Σ_a π(a|s) [ R(s,a) + γ Σ_{s'} P(s'|s,a) V_k(s') ]
    方法 B  解析解：解线性方程组 (I - γ P_π) V = R_π

然后引入**动作价值函数 Q^π(s, a)**：
    "在 s 先执行动作 a，之后一直按 π 行动"的期望回报。
它是下一个文件"策略改进"的关键工具。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 02_policy_evaluation.py

【数据约定】P: (N, A, N)，R: (N, A)，随机策略 pi: (N, A)，确定性策略 policy: (N,)
【下标约定】代码中的状态 i  <==>  讲义中的 s_{i+1}
"""

import numpy as np

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

TRY_LEFT, TRY_RIGHT = 0, 1
ARROWS = ["←", "→"]


# =============================================================================
# 第 0 部分：从 01_mdp_basics.py 复制来的基础函数（精简版，详细注释见 01 文件）
# =============================================================================
# 为了让每个文件都能单独运行，这里把需要的几个小函数复制了一份。
def build_mars_rover_mdp(success_prob=0.8):
    """火星车 MDP：7 个状态，2 个动作；成功概率 success_prob，否则原地不动；s1 奖励 +1，s7 奖励 +10。"""
    n, n_actions = 7, 2
    P = np.zeros((n, n_actions, n))
    for s in range(n):
        for a in range(n_actions):
            target = min(max(s + (-1 if a == TRY_LEFT else 1), 0), n - 1)
            P[s, a, target] += success_prob
            P[s, a, s] += 1.0 - success_prob
    R = np.tile(np.array([1.0, 0, 0, 0, 0, 0, 10.0])[:, None], (1, n_actions))
    return P, R


def deterministic_to_stochastic(policy, n_actions):
    """确定性策略 (N,) → 独热编码的随机策略矩阵 (N, A)。"""
    policy = np.asarray(policy, dtype=int)
    if ((policy < 0) | (policy >= n_actions)).any():
        raise ValueError(f"policy 中的动作必须在 0 到 {n_actions - 1} 之间，但收到了 {policy}")
    pi = np.zeros((len(policy), n_actions))
    pi[np.arange(len(policy)), policy] = 1.0
    return pi


def mdp_to_mrp(P, R, pi):
    """MDP + 策略 → MRP：P_π[s,t] = Σ_a π[s,a] P[s,a,t]，R_π[s] = Σ_a π[s,a] R[s,a]。"""
    return np.einsum("sa,sat->st", pi, P), np.einsum("sa,sa->s", pi, R)


def _check_inputs(P, R, pi, gamma, allow_gamma_one):
    """内部辅助函数：检查形状和 γ 的范围，出错时给出清楚的提示。"""
    n, n_actions, _ = P.shape
    if R.shape != (n, n_actions):
        raise ValueError(f"R 的形状应为 {(n, n_actions)}，但收到了 {R.shape}")
    if pi.shape != (n, n_actions):
        raise ValueError(f"pi 的形状应为 {(n, n_actions)}，但收到了 {pi.shape}。"
                         "（如果你传入的是确定性策略，请先用 deterministic_to_stochastic 转换）")
    if not np.allclose(pi.sum(axis=1), 1.0) or (pi < 0).any():
        raise ValueError("pi 的每一行必须是概率分布")
    upper_ok = gamma <= 1.0 if allow_gamma_one else gamma < 1.0
    if not (0.0 <= gamma and upper_ok):
        raise ValueError(f"gamma 必须在 {'[0, 1]' if allow_gamma_one else '[0, 1)'} 内，但收到了 {gamma}")


# =============================================================================
# 第 1 部分：贝尔曼期望备份（单个状态）
# =============================================================================
def bellman_expectation_backup_state(P, R, pi, V, s, gamma):
    """
    对**一个状态** s 做一次贝尔曼期望备份，返回更新后的 V(s)。

    【公式】（讲义 4.2 节）
        新 V(s) = Σ_a π(a|s) · [ R(s,a) + γ · Σ_{s'} P(s'|s,a) · V(s') ]
                  ─────────   ──────────────────────────────────────
                  按策略对动作   "执行 a" 的价值 = 即时奖励 + 折扣后的下一状态平均价值
                  取平均

    【为什么叫"备份"（backup）？】
        因为信息是从"后继状态" s' **往回**传给 s 的：先知道 s' 有多好，才能算出 s 有多好。
        画成树就是：s 在上面，下面分叉出各个动作 a，每个动作再分叉出各个 s'；
        我们把叶子上的 V(s') "备份"到根节点 s。讲义里有这张图。

    【参数】
        P, R  : MDP。
        pi    : (N, A) 随机策略。
        V     : 长度 N 的数组，当前的价值估计。
        s     : int，要更新的状态。
        gamma : float，折扣因子。

    【返回】
        float，状态 s 的新价值（**不会修改**传入的 V）。

    【用途】
        这个函数写得非常"慢"（逐个状态），是为了和讲义里的手算过程一一对应。
        真正的算法里我们用下面向量化的版本一次更新所有状态。
    """
    n_actions = P.shape[1]
    new_value = 0.0
    for a in range(n_actions):
        # P[s, a] @ V = Σ_{s'} P(s'|s,a) V(s')：执行 a 之后，下一状态的平均价值
        value_of_a = R[s, a] + gamma * (P[s, a] @ V)
        new_value += pi[s, a] * value_of_a
    return float(new_value)


# =============================================================================
# 第 2 部分：方法 A —— 迭代法策略评估
# =============================================================================
def policy_evaluation_iterative(P, R, pi, gamma, tol=1e-10, max_iters=100_000,
                                in_place=False, order=None):
    """
    迭代法策略评估：从 V_0 = 0 出发，反复做贝尔曼期望备份，直到收敛。

    【两种更新方式】
        同步更新（in_place=False，默认，教科书标准写法）：
            用**上一轮**的 V_k 计算所有状态的 V_{k+1}，全部算完再一起替换。
            需要两份数组（旧的和新的）。
                V_new = R_π + γ P_π V_old
        原地更新（in_place=True，也叫 Gauss-Seidel 式）：
            按某个顺序逐个更新状态，更新 s 时**立刻使用**本轮已经更新过的其他状态的新值。
            只需要一份数组，而且新信息传播得更快。
            **更新顺序会影响传播速度**：如果奖励在右边（s7），从右往左扫，
            一轮之内信息就能从 s7 一路传到 s1；从左往右扫则每轮只能传一格。
            两种方式收敛到**同一个** V^π。
            但要注意：信息传得快，**不一定**意味着总轮数更少。如果某个状态有"自环"
            （比如 s7 执行 TryRight 永远停在 s7），它自己的值每轮只能逼近真实值 γ 倍，
            无论按什么顺序更新都一样，总轮数就被它卡住了（演示 3 会看到这个现象）。

    【参数】
        P, R      : MDP。
        pi        : (N, A) 随机策略。
        gamma     : float，[0, 1]。γ=1 允许传入，但对没有终止状态的 MDP 通常不收敛。
        tol       : float，停止阈值：一整轮中所有状态的最大变化量 < tol 时停止。
        max_iters : int，最多迭代轮数（"一轮" = 所有状态都更新一次）。
        in_place  : bool，是否原地更新。
        order     : 状态更新顺序（只对 in_place=True 有意义），例如 [6,5,4,3,2,1,0]。
                    默认 None 表示 0, 1, ..., N-1。

    【返回】
        (V, n_iters, converged)：价值、实际轮数、是否收敛。

    【边界】
        - converged=False 时 V 不可信，调用者必须检查。
        - order 必须是 0..N-1 的一个排列（每个状态恰好出现一次），否则 ValueError。
    """
    _check_inputs(P, R, pi, gamma, allow_gamma_one=True)
    n = P.shape[0]
    P_pi, R_pi = mdp_to_mrp(P, R, pi)
    V = np.zeros(n)

    if order is None:
        order = list(range(n))
    # sorted(order) == list(range(n)) 检查 order 是不是 0..N-1 的一个排列
    if sorted(order) != list(range(n)):
        raise ValueError(f"order 必须是 0 到 {n - 1} 的一个排列，但收到了 {order}")

    for k in range(1, max_iters + 1):
        if not in_place:
            # 同步更新：右边全部使用旧的 V，算完后整体替换
            V_new = R_pi + gamma * (P_pi @ V)
            delta = np.abs(V_new - V).max()
            V = V_new
        else:
            # 原地更新：逐个状态更新，每次都用"最新"的 V
            delta = 0.0
            for s in order:
                old = V[s]
                V[s] = R_pi[s] + gamma * (P_pi[s] @ V)   # 这里的 V 已经包含本轮更新过的状态
                delta = max(delta, abs(V[s] - old))
        if delta < tol:
            return V, k, True
    return V, max_iters, False


# =============================================================================
# 第 3 部分：方法 B —— 解析解
# =============================================================================
def policy_evaluation_exact(P, R, pi, gamma):
    """
    解析法策略评估：V^π = (I - γ P_π)^{-1} R_π（用 np.linalg.solve 求解，不显式求逆）。

    【参数】
        gamma : float，必须在 [0, 1) 内（γ=1 时 I - P_π 奇异，见第 1 讲 8.7 节）。

    【返回】
        长度 N 的数组，精确的 V^π。

    【边界】
        - 计算量 O(N³)，状态数很大时不可行，那时用迭代法。
    """
    _check_inputs(P, R, pi, gamma, allow_gamma_one=False)
    n = P.shape[0]
    P_pi, R_pi = mdp_to_mrp(P, R, pi)
    return np.linalg.solve(np.eye(n) - gamma * P_pi, R_pi)


# =============================================================================
# 第 4 部分：动作价值函数 Q^π
# =============================================================================
def q_from_v(P, R, V, gamma):
    """
    由状态价值 V 计算动作价值 Q：
        Q(s, a) = R(s, a) + γ · Σ_{s'} P(s'|s, a) · V(s')

    【直觉】
        Q^π(s, a) 回答的问题是："如果我现在**不按策略**，而是先做 a，之后再回到策略 π，
        能拿多少？"——它让我们可以**比较同一个状态下不同动作的好坏**，
        这正是改进策略所需要的信息（见 03_policy_iteration.py）。

    【注意】
        这个计算**需要知道模型 P 和 R**。这就是为什么第 4 讲（无模型控制）要直接学习 Q，
        而不是先学 V 再换算——没有模型就没法从 V 换算出 Q。

    【参数】
        V : 长度 N 的数组。可以是任何策略的 V^π，也可以是任意的价值估计。

    【返回】
        形状 (N, A) 的数组。

    【Python 小知识】
        np.einsum("sat,t->sa", P, V) 计算 结果[s,a] = Σ_t P[s,a,t] V[t]，
        等价于 P @ V（numpy 会把 V 乘到 P 的最后一维上），两种写法都可以。
    """
    return R + gamma * np.einsum("sat,t->sa", P, V)


def v_from_q(Q, pi):
    """
    由动作价值 Q 和策略 π 计算状态价值：V(s) = Σ_a π(a|s) Q(s, a)。

    【直觉】状态的价值 = 按策略选动作时，各动作价值的加权平均。
    对确定性策略，就是 V(s) = Q(s, π(s))。
    """
    return np.sum(pi * Q, axis=1)


# =============================================================================
# 第 5 部分：用蒙特卡洛做"交叉验证"
# =============================================================================
def monte_carlo_policy_value(P, R, pi, gamma, start_state, n_episodes, horizon, rng):
    """
    用蒙特卡洛模拟估计 V^π(start_state)，用来独立检验上面算法的结果。

    【为什么要做这个？】
        写 RL 代码时，用两种**完全不同**的方法算同一个量、看结果是否一致，
        是发现 bug 最有效的手段之一。解析解和迭代法都依赖同一个 P_π，
        如果 mdp_to_mrp 写错了，它俩会"一起错"；而蒙特卡洛是直接模拟交互，不经过 P_π。

    【返回】
        (估计值, 标准误差)
    """
    n, n_actions, _ = P.shape
    returns = np.zeros(n_episodes)
    for i in range(n_episodes):
        s, G, discount = start_state, 0.0, 1.0
        for _ in range(horizon):
            a = rng.choice(n_actions, p=pi[s])
            G += discount * R[s, a]
            discount *= gamma
            s = rng.choice(n, p=P[s, a])
        returns[i] = G
    return float(returns.mean()), float(returns.std(ddof=1) / np.sqrt(n_episodes))


# =============================================================================
# 第 6 部分：演示
# =============================================================================
def print_row(label, values, width=16):
    """打印一行价值表。label 只用 ASCII 或放在固定宽度内，避免中文对齐问题。"""
    print(f"  {label:<{width}s}" + "".join(f"{v:8.2f}" for v in values))


def main():
    P, R = build_mars_rover_mdp(success_prob=0.8)
    n, n_actions, _ = P.shape
    gamma = 0.9
    header = "  " + " " * 16 + "".join(f"{f's{j + 1}':>8s}" for j in range(n))

    always_left = deterministic_to_stochastic(np.zeros(n, dtype=int), n_actions)
    always_right = deterministic_to_stochastic(np.ones(n, dtype=int), n_actions)
    random_pi = np.full((n, n_actions), 0.5)

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：手算贝尔曼备份（复现讲义 4.3 节，策略 = 一直向右，γ = 0.9）")
    print("=" * 78)
    V0 = np.zeros(n)
    V1 = np.array([bellman_expectation_backup_state(P, R, always_right, V0, s, gamma) for s in range(n)])
    V2 = np.array([bellman_expectation_backup_state(P, R, always_right, V1, s, gamma) for s in range(n)])
    V3 = np.array([bellman_expectation_backup_state(P, R, always_right, V2, s, gamma) for s in range(n)])
    print(f"  V_2(s6) = 0 + 0.9 × (0.2 × V_1(s6) + 0.8 × V_1(s7)) = 0.9 × (0.2×0 + 0.8×10) = {V2[5]:.2f}")
    print(f"  V_2(s7) = 10 + 0.9 × (1.0 × V_1(s7)) = 10 + 9 = {V2[6]:.2f}")
    print(f"  V_3(s5) = 0 + 0.9 × (0.2 × V_2(s5) + 0.8 × V_2(s6)) = 0.9 × 0.8 × 7.2 = {V3[4]:.3f}\n")

    V_exact_right = policy_evaluation_exact(P, R, always_right, gamma)
    print(header)
    print_row("V_0", V0)
    print_row("V_1", V1)
    print_row("V_2", V2)
    print_row("V_3", V3)
    print_row("... V_inf", V_exact_right)
    print("\n  观察：奖励的信息每迭代一次向左\"传播\"一格——V_2 时 s6 才知道右边有 +10，")
    print("        V_3 时 s5 才知道。这就是\"备份\"：价值从后继状态一步步往回传。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print(f"演示 2：评估三个策略（γ = {gamma}，解析解）")
    print("=" * 78)
    print(header)
    for name, pi in [("always-left", always_left), ("always-right", always_right), ("random", random_pi)]:
        print_row(name, policy_evaluation_exact(P, R, pi, gamma))
    print("\n  观察：一直向右在**每一个**状态都最好，连 s1 都是——")
    print("        在 s1 原地拿 +1 的总价值只有 1/(1-0.9) = 10，不如花几步走到 s7 去拿 +10。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：迭代法 vs 解析解；同步更新 vs 原地更新（策略 = 一直向右）")
    print("=" * 78)
    left_to_right = list(range(n))              # [0, 1, ..., 6]
    right_to_left = list(range(n - 1, -1, -1))  # [6, 5, ..., 0]
    schemes = [
        ("sync", dict(in_place=False)),
        ("in-place s1->s7", dict(in_place=True, order=left_to_right)),
        ("in-place s7->s1", dict(in_place=True, order=right_to_left)),
    ]
    # 【Python 小知识】dict(in_place=True, ...) 创建一个字典；
    #   调用函数时写 **kwargs 会把字典"展开"成关键字参数，
    #   即 f(**{"in_place": True}) 等价于 f(in_place=True)。
    print("  (a) 只迭代 1 轮后的 V——看信息传播得有多远：")
    print(header)
    for name, kwargs in schemes:
        V_one, _, _ = policy_evaluation_iterative(P, R, always_right, gamma, max_iters=1, **kwargs)
        print_row(name, V_one)
    print("  → 同步更新和 s1→s7 的原地更新，1 轮后只有 s7 知道 +10；")
    print("    s7→s1 的原地更新先算 s7，再用 s7 的新值算 s6，再用 s6 的新值算 s5……")
    print("    1 轮之内 +10 的信息就传遍了所有状态。")

    print("\n  (b) 收敛到 tol=1e-10 所需的总轮数：")
    for name, kwargs in schemes:
        V_it, k_it, _ = policy_evaluation_iterative(P, R, always_right, gamma, **kwargs)
        print(f"    {name:<18s}：{k_it:4d} 轮，与解析解最大差 {np.abs(V_it - V_exact_right).max():.1e}")
    print("  → 三者都收敛到同一个 V^π，但总轮数一样！原因是 s7 的\"自环\"：")
    print("    V(s7) ← 10 + 0.9 × V(s7)，它自己每轮只能把误差缩小到 0.9 倍，和更新顺序无关。")
    print("    所以\"信息传得更快\"不一定等于\"总轮数更少\"——瓶颈在最慢收敛的那个部分。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 4：动作价值 Q^π——看出\"哪里可以改进\"（策略 = 一直向左）")
    print("=" * 78)
    V_left = policy_evaluation_exact(P, R, always_left, gamma)
    Q_left = q_from_v(P, R, V_left, gamma)
    print(header)
    print_row("V^pi(s)", V_left)
    print_row("Q^pi(s, Left)", Q_left[:, TRY_LEFT])
    print_row("Q^pi(s, Right)", Q_left[:, TRY_RIGHT])
    better = np.flatnonzero(Q_left[:, TRY_RIGHT] > V_left + 1e-9)
    print(f"\n  检查 V = Σ_a π(a|s) Q(s,a)：最大差 {np.abs(v_from_q(Q_left, always_left) - V_left).max():.1e} ✓")
    print(f"  在这些状态中 Q(s, Right) > V(s)：{['s' + str(s + 1) for s in better]}")
    print("  → 在这些状态\"偏离一次策略、改成向右\"会更好。把这些状态的动作改成向右，")
    print("    就得到一个更好的策略——这正是下一个文件\"策略改进\"的思路。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 5：用蒙特卡洛交叉验证（随机策略，V^π(s4)）")
    print("=" * 78)
    rng = np.random.default_rng(0)
    V_rand = policy_evaluation_exact(P, R, random_pi, gamma)
    est, se = monte_carlo_policy_value(P, R, random_pi, gamma, start_state=3,
                                       n_episodes=3000, horizon=150, rng=rng)
    print(f"  解析解：V^π(s4) = {V_rand[3]:.3f}")
    print(f"  蒙特卡洛（3000 条轨迹）：{est:.3f} ± {2 * se:.3f}（约 95% 置信区间）")
    print("  → 两种完全独立的方法结果一致，说明 mdp_to_mrp 和策略评估的实现是对的。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 6：边界情况")
    print("=" * 78)
    try:
        policy_evaluation_exact(P, R, always_right, gamma=1.0)
    except ValueError as e:
        print(f"  [解析解 γ=1]     ValueError: {e}")
    V_g1, k_g1, ok = policy_evaluation_iterative(P, R, always_right, gamma=1.0, max_iters=500)
    print(f"  [迭代法 γ=1]     {k_g1} 轮后 converged={ok}，V(s7) = {V_g1[6]:.0f}（每轮 +10，永远不收敛）")
    try:
        policy_evaluation_exact(P, R, np.ones(n, dtype=int), gamma)   # 忘了转换确定性策略
    except (ValueError, IndexError) as e:
        print(f"  [忘记转换策略]   {type(e).__name__}: {e}")
    V_g0 = policy_evaluation_exact(P, R, random_pi, gamma=0.0)
    print(f"  [γ=0]            V^π = {V_g0}（等于即时奖励，任何策略都一样）")


if __name__ == "__main__":
    main()