"""
=============================================================================
第 10 讲 · 共享模块：一个"治疗决策"MDP（用于离策略评估）
=============================================================================

【场景】
    一位病人的健康状况用 0~4 表示（0 = 危重，4 = 健康），从 2 出发。
    每个时间步，医生选择"保守治疗"或"积极治疗"：

        保守治疗（动作 0）：好转 +1 的概率 0.2，不变 0.5，恶化 -1 的概率 0.3
        积极治疗（动作 1）：好转 +1 的概率 0.6，不变 0.2，恶化 -1 的概率 0.2
                           但有 0.5 的副作用代价
    每一步的奖励 = 当前健康状况 - 治疗代价。共 H 步（有限视界，不打折）。

    - 行为策略 π_b（医院过去的做法，产生历史数据）：不管病情如何，70% 保守、30% 积极
    - 评估策略 π_e（想评估的新方案）：病情不好（≤ 2）时 90% 积极治疗，否则 90% 保守治疗

    离策略评估的问题：**只用 π_b 产生的历史记录，估计 π_e 的期望总奖励**——
    不能拿病人去试验新方案。

【本模块提供】
    - 精确的 J(π)（有限视界动态规划），作为"标准答案"
    - 用行为策略采样历史数据
    - 所有文件共用的常量与小工具

    01、02 文件共用本模块。运行它们时当前目录必须是 code/lecture10/。
"""

import numpy as np

N_STATES = 5
N_ACTIONS = 2
START = 2
TREAT_COST = 0.5

# 转移概率 P[s, a, s']
P = np.zeros((N_STATES, N_ACTIONS, N_STATES))
for s in range(N_STATES):
    for a, (up, stay, down) in enumerate([(0.2, 0.5, 0.3), (0.6, 0.2, 0.2)]):
        P[s, a, min(s + 1, N_STATES - 1)] += up
        P[s, a, s] += stay
        P[s, a, max(s - 1, 0)] += down

# 期望即时奖励 R[s, a] = 健康状况 - 治疗代价
R = np.array([[s - (TREAT_COST if a == 1 else 0.0) for a in range(N_ACTIONS)] for s in range(N_STATES)])


def behavior_policy():
    """π_b：不管病情，70% 保守、30% 积极。形状 (N_STATES, N_ACTIONS)。"""
    return np.tile([0.7, 0.3], (N_STATES, 1))


def evaluation_policy():
    """π_e：病情 ≤ 2 时 90% 积极治疗，否则 90% 保守治疗。"""
    pi = np.zeros((N_STATES, N_ACTIONS))
    for s in range(N_STATES):
        pi[s] = [0.1, 0.9] if s <= 2 else [0.9, 0.1]
    return pi


def exact_value(pi, H, P_model=P, R_model=R):
    """
    有限视界动态规划（第 2 讲 8.7 节）精确计算 J(π) = E[Σ_{t=0}^{H-1} r_t]（从 START 出发）。

    同时返回每个时刻的 Q_t 和 V_t（双重稳健估计要用）：
        V_H = 0；Q_t(s,a) = R(s,a) + Σ_s' P(s'|s,a) V_{t+1}(s')；V_t(s) = Σ_a π(a|s) Q_t(s,a)

    【参数】P_model、R_model：默认用真实模型；也可以传入估计的模型（直接法 DM）。
    【返回】(J, Q, V)：Q 形状 (H, N, A)，V 形状 (H+1, N)。
    """
    V = np.zeros((H + 1, N_STATES))
    Q = np.zeros((H, N_STATES, N_ACTIONS))
    for t in range(H - 1, -1, -1):
        Q[t] = R_model + P_model @ V[t + 1]
        V[t] = np.sum(pi * Q[t], axis=1)
    return V[0, START], Q, V


def sample_dataset(pi, n_episodes, H, rng):
    """
    按策略 pi 采样 n_episodes 条长度为 H 的轨迹。

    【返回】形状 (n_episodes, H, 3) 的数组，最后一维是 (s, a, r)；另外返回下一状态数组 (n_episodes, H)。
        奖励在这里带有噪声：真实奖励 + N(0, 1)——病人的反馈本来就有随机性。
    """
    data = np.zeros((n_episodes, H, 3))
    next_states = np.zeros((n_episodes, H), dtype=int)
    for i in range(n_episodes):
        s = START
        for t in range(H):
            a = int(rng.random() < pi[s, 1])
            r = R[s, a] + rng.normal(0, 1.0)
            s_next = int(rng.choice(N_STATES, p=P[s, a]))
            data[i, t] = (s, a, r)
            next_states[i, t] = s_next
            s = s_next
    return data, next_states