"""
=============================================================================
第 10 讲 · 代码 2：直接法与双重稳健估计——把模型和重要性采样结合起来
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
离策略评估有两大类方法，各有优缺点：

    直接法（Direct Method, DM）：先用数据学一个模型（或 Q 函数），再在模型上精确评估 π_e。
        第 3 讲的"确定性等价"就是这个思路；拟合 Q 评估（FQE）是它的函数逼近版本。
        方差小；但模型有误差时**有偏**，而且偏差不会随数据增多而自动消失（如果模型类选错了）。

    重要性采样（IS，01 文件）：不需要模型。
        无偏；但方差可能极大。

双重稳健（Doubly Robust, DR，Jiang & Li 2016）把两者结合：
    用模型给出一个"基准估计"，再用重要性采样**只修正模型的误差部分**。
    - 模型准确时，要修正的误差很小，方差很低；
    - 模型不准确时，重要性采样的修正项仍然保证**无偏**。
    "双重"的意思是：模型准确，或者行为策略已知（重要性比正确），只要满足一个就无偏。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture10
    python 02_doubly_robust.py      （约 30 秒）
"""

import numpy as np

from treatment_mdp import (N_STATES, N_ACTIONS, START, P, R, behavior_policy, evaluation_policy,
                           exact_value, sample_dataset)

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


# =============================================================================
# 第 1 部分：从数据估计模型（直接法）
# =============================================================================
def fit_model(data, next_states):
    """
    用数据"数"出一个表格模型（第 1 讲 7.7 节、第 3 讲确定性等价）：
        P̂(s'|s,a) = N(s,a,s') / N(s,a)，   R̂(s,a) = 在 (s,a) 观察到的奖励的平均

    【没见过的 (s,a) 怎么办？】这里假设"原地不动、奖励为 0"——一个任意的默认值。
        这正是直接法的软肋：数据没覆盖的地方，模型只能靠假设。
    """
    counts = np.zeros((N_STATES, N_ACTIONS, N_STATES))
    r_sum = np.zeros((N_STATES, N_ACTIONS))
    s = data[:, :, 0].astype(int).ravel()
    a = data[:, :, 1].astype(int).ravel()
    r = data[:, :, 2].ravel()
    np.add.at(counts, (s, a, next_states.ravel()), 1)
    np.add.at(r_sum, (s, a), r)
    n = counts.sum(axis=2)
    P_hat = np.zeros_like(counts)
    R_hat = np.zeros_like(r_sum)
    seen = n > 0
    P_hat[seen] = counts[seen] / n[seen][:, None]
    R_hat[seen] = r_sum[seen] / n[seen]
    for s_, a_ in zip(*np.where(~seen)):
        P_hat[s_, a_, s_] = 1.0                    # 没见过：假设原地不动
    return P_hat, R_hat


def wrong_model():
    """
    一个"故意错误"的模型：把两种治疗的效果对调了，而且把奖励整体高估了 1。
    用来检验：模型很差时，直接法和双重稳健分别会怎样。
    """
    return P[:, ::-1, :].copy(), R + 1.0


# =============================================================================
# 第 2 部分：估计量
# =============================================================================
def estimate_dm(model, pi_e, H):
    """直接法：在模型上精确计算 π_e 的价值。"""
    P_m, R_m = model
    return exact_value(pi_e, H, P_m, R_m)[0]


def estimate_dr(data, pi_e, pi_b, model, H):
    """
    双重稳健估计（Jiang & Li, 2016），对每条轨迹从后往前递推：

        DR_H = 0
        DR_t = V̂_t(s_t) + ρ_t · ( r_t + DR_{t+1} - Q̂_t(s_t, a_t) )

    估计值 = 所有轨迹 DR_0 的平均。其中 Q̂_t、V̂_t 是在模型上算出的 π_e 的价值（有限视界，依赖剩余步数）。

    【直觉】
        - V̂_t(s_t)：模型给出的"基准估计"；
        - r_t + DR_{t+1} - Q̂_t(s_t, a_t)：实际观察到的结果与模型预测之差——"模型的误差"；
        - 乘以 ρ_t：用重要性采样把这个误差从 π_b 的分布"搬到" π_e 的分布上。
        模型越准，括号里的误差越小，乘上一个大权重 ρ 也不会带来太大的方差。

    【为什么无偏？】（模型再差都成立，只要 ρ 是对的）
        对 a_t ~ π_b 求期望：E[ρ_t (r_t + DR_{t+1} - Q̂_t(s_t,a_t))] = E_{a~π_e}[r_t + DR_{t+1}] - E_{a~π_e}[Q̂_t(s_t,a)]，
        而 E_{a~π_e}[Q̂_t(s_t, a)] = V̂_t(s_t)，恰好与前面的 V̂_t 抵消，剩下 E_{π_e}[r_t + DR_{t+1}]——
        模型的影响在期望中被完全消掉了。
    """
    P_m, R_m = model
    _, Q_hat, V_hat = exact_value(pi_e, H, P_m, R_m)
    s = data[:, :, 0].astype(int)
    a = data[:, :, 1].astype(int)
    r = data[:, :, 2]
    rho = pi_e[s, a] / pi_b[s, a]
    dr = np.zeros(len(data))
    for t in range(H - 1, -1, -1):
        q = Q_hat[t][s[:, t], a[:, t]]
        v = V_hat[t][s[:, t]]
        dr = v + rho[:, t] * (r[:, t] + dr - q)
    return float(dr.mean())


def estimate_wdr(data, pi_e, pi_b, model, H):
    """
    加权双重稳健估计（WDR，Thomas & Brunskill, 2016）：把 DR 中的重要性权重换成**自归一化**的权重，
    就像 WIS 之于 IS（01 文件）。

        w_{i,t} = ρ_{i,0:t} / Σ_j ρ_{j,0:t}          （第 t 步的归一化累积权重；w_{i,-1} = 1/n）
        WDR = Σ_t [ Σ_i w_{i,t} r_{i,t}  -  Σ_i ( w_{i,t} Q̂_t(s_{i,t}, a_{i,t}) - w_{i,t-1} V̂_t(s_{i,t}) ) ]

    【效果】有一点偏差（来自归一化），但方差通常比 DR 小得多。
    （Thomas 与 Brunskill——CS234 的主讲人——提出的这个估计量，是实践中最常用的 OPE 方法之一。）
    """
    P_m, R_m = model
    _, Q_hat, V_hat = exact_value(pi_e, H, P_m, R_m)
    s = data[:, :, 0].astype(int)
    a = data[:, :, 1].astype(int)
    r = data[:, :, 2]
    n = len(data)
    rho_cum = np.cumprod(pi_e[s, a] / pi_b[s, a], axis=1)
    w = rho_cum / rho_cum.sum(axis=0, keepdims=True)
    w_prev = np.hstack([np.full((n, 1), 1.0 / n), w[:, :-1]])
    est = 0.0
    for t in range(H):
        q = Q_hat[t][s[:, t], a[:, t]]
        v = V_hat[t][s[:, t]]
        est += np.sum(w[:, t] * r[:, t]) - np.sum(w[:, t] * q - w_prev[:, t] * v)
    return float(est)


def estimate_wpdis(data, pi_e, pi_b):
    """加权逐决策重要性采样（与 01 文件相同），作为对照。"""
    s = data[:, :, 0].astype(int)
    a = data[:, :, 1].astype(int)
    rho_cum = np.cumprod(pi_e[s, a] / pi_b[s, a], axis=1)
    return float(np.sum(np.sum(rho_cum * data[:, :, 2], axis=0) / np.sum(rho_cum, axis=0)))


def run(H, n_episodes, n_reps, seed, pi_b=None):
    """重复实验，返回每个估计量的 (偏差, 标准差, 均方根误差)。"""
    pi_e = evaluation_policy()
    pi_b = behavior_policy() if pi_b is None else pi_b
    pi_b_safe = np.maximum(pi_b, 1e-12)
    true_J = exact_value(pi_e, H)[0]
    rng = np.random.default_rng(seed)
    names = ["WPDIS", "DM（学到的模型）", "DM（错误模型）", "DR（学到的模型）", "DR（错误模型）",
             "WDR（学到的模型）", "WDR（错误模型）", "DR（真实模型）"]
    ests = {n: [] for n in names}
    for _ in range(n_reps):
        data, nxt = sample_dataset(pi_b, n_episodes, H, rng)
        learned = fit_model(data, nxt)
        ests["WPDIS"].append(estimate_wpdis(data, pi_e, pi_b_safe))
        ests["DM（学到的模型）"].append(estimate_dm(learned, pi_e, H))
        ests["DM（错误模型）"].append(estimate_dm(wrong_model(), pi_e, H))
        ests["DR（学到的模型）"].append(estimate_dr(data, pi_e, pi_b_safe, learned, H))
        ests["DR（错误模型）"].append(estimate_dr(data, pi_e, pi_b_safe, wrong_model(), H))
        ests["WDR（学到的模型）"].append(estimate_wdr(data, pi_e, pi_b_safe, learned, H))
        ests["WDR（错误模型）"].append(estimate_wdr(data, pi_e, pi_b_safe, wrong_model(), H))
        ests["DR（真实模型）"].append(estimate_dr(data, pi_e, pi_b_safe, (P, R), H))
    out = {}
    for n, v in ests.items():
        v = np.array(v)
        out[n] = (v.mean() - true_J, v.std(), np.sqrt(np.mean((v - true_J) ** 2)))
    return true_J, out


def print_table(true_J, out, only=None):
    print(f"  真实值 J(π_e) = {true_J:.3f}")
    print("   估计量                 偏差      标准差    均方根误差")
    for n, (b, s, e) in out.items():
        if only is not None and n not in only:
            continue
        # 名字里有中文，用固定格式对齐：先英文缩写，再补空格
        pad = " " * (22 - len(n) - sum(1 for ch in n if ord(ch) > 127))
        print(f"   {n}{pad}{b:+7.3f}   {s:7.3f}    {e:7.3f}")


def main():
    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：各种估计量的比较（H = 10，每次 200 条历史数据，重复 300 次）")
    print("=" * 78)
    true_J, out = run(10, 200, 300, seed=0)
    print_table(true_J, out)
    print("  观察：")
    print("  - DM（学到的模型）误差最小：这个问题只有 5 个状态、2 个动作，200 条轨迹足以把模型估计得很准。")
    print("    直接法把**所有**轨迹的信息汇集到一个模型里（第 3 讲：确定性等价的数据效率最高）；")
    print("    而重要性采样类方法只能通过每条轨迹的权重来利用数据，还会被权重放大奖励的噪声。")
    print("  - DM（错误模型）：严重有偏（+3.6），而且它根本没用数据——数据再多也纠正不了。")
    print("  - DR（错误模型）：**几乎无偏**（+0.18）！重要性采样的修正项把模型的错误纠正了回来。")
    print("    WDR（错误模型）有一点偏差（-0.9，来自自归一化），但远小于 DM 的 +3.6，而且随数据增多而消失（演示 2）。")
    print("    这就是双重稳健的价值：一份\"保险\"——模型错了，估计也不至于有偏。")
    print("  - 即使用**真实模型**，DR 的标准差仍有 3.6：模型只能预测\"期望\"奖励，")
    print("    每一步奖励的随机噪声仍会被重要性权重放大。WDR 的自归一化权重大大缓解了这一点。")
    print("  - 在大规模问题中（状态连续、用神经网络做函数逼近），模型**一定**有误差，")
    print("    DM 的偏差可能很大，这时 DR / WDR 的\"保险\"就非常值得。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：数据量的影响（H = 10；表中为均方根误差，每种设置重复 200 次）")
    print("=" * 78)
    print("   数据条数     WPDIS    DM(学到)   DM(错误)   DR(错误)   WDR(学到)  WDR(错误)")
    for n in [20, 50, 200, 1000]:
        _, out = run(10, n, 200, seed=1)
        vals = [out[k][2] for k in ["WPDIS", "DM（学到的模型）", "DM（错误模型）", "DR（错误模型）",
                                    "WDR（学到的模型）", "WDR（错误模型）"]]
        print(f"   {n:6d}     " + "   ".join(f"{v:7.3f}" for v in vals))
    print("  → 数据越多，除了\"DM（错误模型）\"以外，所有估计量都越来越准；")
    print("    错误模型的 DM 停在一个固定的偏差上，而 DR / WDR 即使用错误的模型，误差也在持续下降。")
    print("    这就是\"双重稳健\"：模型对，或者重要性比对，只要满足一个就行。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：覆盖不足时，所有方法都会失败（病情 ≤ 1 时历史上从不积极治疗）")
    print("=" * 78)
    pi_b_narrow = behavior_policy()
    pi_b_narrow[:2] = [1.0, 0.0]
    true_J, out = run(10, 200, 200, seed=2, pi_b=pi_b_narrow)
    print_table(true_J, out, only=["WPDIS", "DM（学到的模型）", "DR（学到的模型）", "WDR（学到的模型）", "DR（真实模型）"])
    print("  → 所有从数据中学习的方法都有明显偏差：")
    print("    DM 在没见过的 (状态, 动作) 上只能用任意的默认假设（这里是\"原地不动、奖励为 0\"）；")
    print("    重要性采样的修正也无法触及这些 (状态, 动作)——它们在数据中根本不出现。")
    print("    只有\"真实模型\"的 DR 是无偏的：模型本身补上了数据缺失的部分（双重稳健中\"模型正确\"的那一半）。")
    print("    但真实模型恰恰是我们没有的。")
    print("    **数据中没有的信息，任何估计方法都变不出来**；这是离线 RL 最根本的限制。")


if __name__ == "__main__":
    main()