"""
=============================================================================
第 7 讲 · 共享模块：一个可以精确计算一切的小链条 MDP
=============================================================================

【为什么需要它？】
    第 7 讲要讨论很多"估计量"的性质：Actor-Critic 的偏差、GAE 的偏差–方差权衡、
    性能差异引理、替代目标的近似误差……要看清这些性质，最好的办法是在一个
    **能精确算出真实答案**的小问题上做实验：精确的 J(θ)、V^π、Q^π、A^π、状态访问量 η、
    精确的策略梯度，统统可以用第 2 讲的线性代数算出来。

    本模块被 01、02、03 三个文件共同使用。
    【Python 小知识】同目录下的其他文件用 `from chain_mdp import ...` 导入，
    运行它们时当前目录必须是 code/lecture07/。

-----------------------------------------------------------------------------
【环境：回合制链条（第 6 讲 04 文件的加长版）】
-----------------------------------------------------------------------------

    [+1 终点] ← 0 — 1 — 2 — [3] — 4 — 5 — 6 → [+10 终点]
                              ▲
                             起点

    - N_STATES = 7 个非终止状态，从状态 3 出发
    - 动作：左 / 右；80% 按意图移动，20% 原地不动
    - 从最左边走出去 +1、从最右边走出去 +10，回合结束
    - 每一步有 -0.5 的时间成本
    - γ = 0.95
    - 策略：表格型 softmax，θ 的形状为 (7, 2)，共 14 个参数

    最优策略显然是"一直向右"，但从一个随机的初始策略出发，
    学习算法需要把"右边有 +10"这个信息传播回起点附近——这正是方差与偏差发挥作用的地方。
"""

import numpy as np

N_STATES = 7
START = 3
GAMMA = 0.95
STEP_COST = -0.5
LEFT, RIGHT = 0, 1
N_ACTIONS = 2


# =============================================================================
# 模型
# =============================================================================
def build_model():
    """
    返回 (P, Rsa)：
        P[s, a, s']：非终止状态之间的转移概率（走出边界的概率不在其中，所以行和可能 < 1）
        Rsa[s, a]  ：在 s 执行 a 的期望即时奖励
    """
    P = np.zeros((N_STATES, N_ACTIONS, N_STATES))
    Rsa = np.zeros((N_STATES, N_ACTIONS))
    for s in range(N_STATES):
        for a in (LEFT, RIGHT):
            target = s - 1 if a == LEFT else s + 1
            P[s, a, s] += 0.2
            Rsa[s, a] = STEP_COST
            if target < 0:
                Rsa[s, a] += 0.8 * 1.0
            elif target >= N_STATES:
                Rsa[s, a] += 0.8 * 10.0
            else:
                P[s, a, target] += 0.8
    return P, Rsa


P, RSA = build_model()


# =============================================================================
# 策略
# =============================================================================
def policy_probs(theta):
    """表格型 softmax：theta (N, 2) → π(a|s) (N, 2)。每行先减最大值，防止溢出。"""
    z = theta - theta.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)


def grad_log_pi(theta, s, a):
    """表格型 softmax 的得分函数：只有第 s 行非零，等于 e_a - π(·|s)（第 6 讲 2.2 节）。"""
    g = np.zeros_like(theta)
    g[s] = -policy_probs(theta)[s]
    g[s, a] += 1.0
    return g


# =============================================================================
# 精确计算（只用于评价，学习算法不能用）
# =============================================================================
def exact(pi):
    """
    对一个策略（概率矩阵 pi，形状 (N, 2)）精确计算：
        V   : V^π(s)
        Q   : Q^π(s, a)
        A   : 优势 A^π(s, a) = Q - V
        eta : 从起点出发的折扣访问量 η(s) = Σ_t γ^t Pr(s_t = s)
        J   : V^π(起点)
    返回字典。
    """
    P_pi = np.einsum("sa,sat->st", pi, P)
    R_pi = np.sum(pi * RSA, axis=1)
    I = np.eye(N_STATES)
    V = np.linalg.solve(I - GAMMA * P_pi, R_pi)
    Q = RSA + GAMMA * P @ V
    e0 = np.zeros(N_STATES)
    e0[START] = 1.0
    eta = np.linalg.solve(I - GAMMA * P_pi.T, e0)
    return {"V": V, "Q": Q, "A": Q - V[:, None], "eta": eta, "J": V[START]}


def J(theta):
    """目标函数 J(θ) = V^{π_θ}(起点)。"""
    return exact(policy_probs(theta))["J"]


def exact_gradient(theta):
    """
    策略梯度定理（第 6 讲 5.2 节），对表格型 softmax：
        ∂J/∂θ[s, b] = η(s) · π(b|s) · A^π(s, b)
    """
    pi = policy_probs(theta)
    ex = exact(pi)
    return ex["eta"][:, None] * pi * ex["A"]


# =============================================================================
# 采样
# =============================================================================
def sample_episode(theta, rng, max_steps=2000):
    """
    按策略采样一个回合。
    返回列表 [(s, a, r, s_next), ...]，s_next = None 表示回合在这一步结束。
    """
    pi = policy_probs(theta)
    s, episode = START, []
    for _ in range(max_steps):
        a = RIGHT if rng.random() < pi[s, RIGHT] else LEFT
        if rng.random() < 0.8:
            nxt = s - 1 if a == LEFT else s + 1
        else:
            nxt = s
        if nxt < 0:
            episode.append((s, a, STEP_COST + 1.0, None))
            return episode
        if nxt >= N_STATES:
            episode.append((s, a, STEP_COST + 10.0, None))
            return episode
        episode.append((s, a, STEP_COST, nxt))
        s = nxt
    return episode


if __name__ == "__main__":
    # 直接运行本模块：打印均匀随机策略和最优策略的精确信息
    np.set_printoptions(precision=3, suppress=True)
    for name, theta in [("均匀随机策略", np.zeros((N_STATES, 2))),
                        ("一直向右", np.tile([-5.0, 5.0], (N_STATES, 1)))]:
        ex = exact(policy_probs(theta))
        print(f"{name}：J = {ex['J']:.3f}")
        print(f"  V   = {ex['V']}")
        print(f"  η   = {ex['eta']}")