"""
=============================================================================
第 10 讲 · 代码 1：离策略评估（OPE）——重要性采样及其变体
=============================================================================

对应讲义：lectures/lecture10_offline_rl.md 第 1、2 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
我们只有行为策略 π_b 产生的历史数据，想知道另一个策略 π_e 表现如何。
如果直接对历史数据的回报取平均，得到的是 π_b 的价值，不是 π_e 的。

**重要性采样**（importance sampling, IS）的想法：给每条轨迹一个权重，
    "如果按 π_e 来走，这条轨迹出现的可能性是按 π_b 走的多少倍？"
    w(τ) = Π_t π_e(a_t|s_t) / π_b(a_t|s_t)
环境的转移概率在分子分母中相同、抵消掉了（和第 6 讲策略梯度中"转移概率消失"是同一个道理），
所以**不需要知道模型**。然后 J(π_e) ≈ (1/n) Σ_i w(τ_i) G(τ_i)。

本文件比较四种估计量：
    IS    ：普通重要性采样（无偏，方差可能极大）
    WIS   ：加权 / 自归一化重要性采样（有偏但一致，方差小得多）
    PDIS  ：逐决策重要性采样（利用"第 t 步的奖励只与前 t 步的动作有关"）
    WPDIS ：加权的逐决策重要性采样

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture10
    python 01_importance_sampling.py      （约 25 秒）
"""

import numpy as np

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

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


# =============================================================================
# 第 1 部分：四种重要性采样估计量
# =============================================================================
def importance_ratios(data, pi_e, pi_b):
    """
    每条轨迹每一步的重要性比 ρ_t = π_e(a_t|s_t) / π_b(a_t|s_t)。

    【返回】形状 (n, H) 的数组。

    【边界】如果某个 π_b(a|s) = 0 而 π_e(a|s) > 0，比值会是无穷大——
        这种 (s,a) 在数据中根本不会出现（π_b 从不选它），IS 对它"视而不见"，
        估计就会有偏（演示 3）。这就是 IS 需要的**覆盖假设**。
    """
    s = data[:, :, 0].astype(int)
    a = data[:, :, 1].astype(int)
    return pi_e[s, a] / pi_b[s, a]


def estimate_is(data, pi_e, pi_b):
    """
    普通（轨迹级）重要性采样：J ≈ (1/n) Σ_i w_i G_i，w_i = Π_t ρ_{i,t}。

    【为什么无偏？】E_{τ~π_b}[w(τ) G(τ)] = Σ_τ p_b(τ) · (p_e(τ)/p_b(τ)) · G(τ) = Σ_τ p_e(τ) G(τ) = J(π_e)。
    【问题】w 是 H 个比值的**乘积**。即使每个比值都不大，乘积也可能极大或极小——
        方差随视界 H 指数增长。
    """
    w = np.prod(importance_ratios(data, pi_e, pi_b), axis=1)
    G = data[:, :, 2].sum(axis=1)
    return float(np.mean(w * G))


def estimate_wis(data, pi_e, pi_b):
    """
    加权（自归一化）重要性采样：J ≈ Σ_i w_i G_i / Σ_i w_i。

    【直觉】把权重归一化成"加权平均"——估计值一定落在观察到的回报的范围内，
        不会因为一个巨大的权重而飞到天上去。
    【代价】有偏（分母也是随机的），但随着数据增多偏差趋于 0（一致）。实践中通常比 IS 好得多。
    """
    w = np.prod(importance_ratios(data, pi_e, pi_b), axis=1)
    G = data[:, :, 2].sum(axis=1)
    return float(np.sum(w * G) / np.sum(w))


def estimate_pdis(data, pi_e, pi_b):
    """
    逐决策重要性采样（per-decision IS）：J ≈ (1/n) Σ_i Σ_t (Π_{k≤t} ρ_{i,k}) r_{i,t}。

    【直觉】第 t 步的奖励 r_t 只取决于前 t 步的动作，与之后的动作无关——
        所以 r_t 只需要乘以"到 t 为止"的比值乘积，而不是整条轨迹的乘积。
        这和第 6 讲"利用因果性"（未来的动作影响不了过去的奖励）是同一个思想。
        仍然无偏，方差更小。
    """
    rho_cum = np.cumprod(importance_ratios(data, pi_e, pi_b), axis=1)   # 累积乘积：Π_{k≤t} ρ_k
    # 【Python 小知识】np.cumprod 沿某一维计算累积乘积：[a, b, c] → [a, a·b, a·b·c]
    return float(np.mean(np.sum(rho_cum * data[:, :, 2], axis=1)))


def estimate_wpdis(data, pi_e, pi_b):
    """
    加权的逐决策重要性采样：每个时刻 t 分别做自归一化。
        J ≈ Σ_t [ Σ_i ρ_{i,0:t} r_{i,t} / Σ_i ρ_{i,0:t} ]
    """
    rho_cum = np.cumprod(importance_ratios(data, pi_e, pi_b), axis=1)
    return float(np.sum(np.sum(rho_cum * data[:, :, 2], axis=0) / np.sum(rho_cum, axis=0)))


ESTIMATORS = {"IS": estimate_is, "WIS": estimate_wis, "PDIS": estimate_pdis, "WPDIS": estimate_wpdis}


def evaluate_estimators(H, n_episodes, n_reps, seed, pi_e=None, pi_b=None):
    """
    重复 n_reps 次"收集 n_episodes 条历史数据 → 估计"，统计每个估计量的偏差、标准差、均方根误差。
    """
    pi_e = evaluation_policy() if pi_e is None else pi_e
    pi_b = behavior_policy() if pi_b is None else pi_b
    true_J = exact_value(pi_e, H)[0]
    rng = np.random.default_rng(seed)
    results = {name: [] for name in ESTIMATORS}
    for _ in range(n_reps):
        data, _ = sample_dataset(pi_b, n_episodes, H, rng)
        for name, fn in ESTIMATORS.items():
            results[name].append(fn(data, pi_e, pi_b))
    stats = {}
    for name, vals in results.items():
        vals = np.array(vals)
        stats[name] = (vals.mean() - true_J, vals.std(), np.sqrt(np.mean((vals - true_J) ** 2)))
    return true_J, stats


def main():
    pi_b, pi_e = behavior_policy(), evaluation_policy()

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：为什么不能直接用历史数据的平均回报？（H = 10）")
    print("=" * 78)
    H = 10
    J_b = exact_value(pi_b, H)[0]
    J_e = exact_value(pi_e, H)[0]
    data, _ = sample_dataset(pi_b, 1000, H, np.random.default_rng(0))
    print(f"  真实价值：行为策略 J(π_b) = {J_b:.3f}，评估策略 J(π_e) = {J_e:.3f}")
    print(f"  1000 条历史数据的平均回报 = {data[:, :, 2].sum(axis=1).mean():.3f}   ← 这是 π_b 的价值，不是 π_e 的")
    w = np.prod(importance_ratios(data, pi_e, pi_b), axis=1)
    print(f"  重要性权重 w = Π_t π_e/π_b 的统计：最小 {w.min():.2e}，中位数 {np.median(w):.3f}，"
          f"最大 {w.max():.1f}，平均 {w.mean():.3f}（理论上期望为 1）")
    print(f"  前 1% 的轨迹占了总权重的 {np.sort(w)[-10:].sum() / w.sum() * 100:.1f}%")
    print("  → 绝大多数轨迹的权重接近 0（它们\"不像\"π_e 会走出来的轨迹），极少数轨迹的权重巨大。")
    print("    IS 的估计几乎由这极少数轨迹决定——这就是它方差大的原因。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：四种估计量的比较（H = 10，每次 200 条历史数据，重复 500 次）")
    print("=" * 78)
    true_J, stats = evaluate_estimators(10, 200, 500, seed=1)
    print(f"  真实值 J(π_e) = {true_J:.3f}")
    print("   估计量     偏差      标准差     均方根误差")
    for name, (bias, std, rmse) in stats.items():
        print(f"   {name:<6s}   {bias:+7.3f}    {std:7.3f}     {rmse:7.3f}")
    print("  观察：")
    print("  - IS 和 PDIS 的偏差接近 0（无偏），但标准差很大。")
    print("  - PDIS 比 IS 好：每个奖励只乘以\"到那一步为止\"的比值。")
    print("  - WIS / WPDIS 有一点偏差（自归一化带来的），但方差小得多。")
    print("    综合误差（均方根误差）看，加权的版本通常更好——实践中常常首选 WPDIS / WIS。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：视界越长，重要性采样越难（每次 200 条数据，重复 300 次；表中为均方根误差）")
    print("=" * 78)
    print("     H     真实值      IS        WIS       PDIS      WPDIS")
    for H in [2, 5, 10, 20, 40]:
        true_J, stats = evaluate_estimators(H, 200, 300, seed=2)
        print(f"  {H:4d}   {true_J:7.2f}   " + "   ".join(f"{stats[n][2]:7.2f}" for n in ESTIMATORS))
    print("  → IS 的误差随视界急剧增大：权重是 H 个比值的乘积，方差随 H **指数**增长（\"视界诅咒\"）。")
    print("    WIS、WPDIS 缓解了这个问题，但长视界下所有重要性采样方法都会越来越困难。")
    print("    下一个文件的双重稳健估计利用一个模型来进一步降低方差。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 4：覆盖假设——如果行为策略从不做某个动作")
    print("=" * 78)
    pi_b_narrow = behavior_policy()
    pi_b_narrow[:2] = [1.0, 0.0]       # 病情危重（0、1）时，历史上从不积极治疗
    # π_e 在这些状态恰恰想积极治疗；用极小的数代替 0，避免除以 0（这些 (s,a) 在数据中本来就不会出现）
    pi_b_safe = np.maximum(pi_b_narrow, 1e-12)
    rng = np.random.default_rng(3)
    true_J = exact_value(pi_e, 10)[0]
    ests = {n: [] for n in ESTIMATORS}
    for _ in range(300):
        data, _ = sample_dataset(pi_b_narrow, 200, 10, rng)
        for name, fn in ESTIMATORS.items():
            ests[name].append(fn(data, pi_e, pi_b_safe))
    print(f"  假设历史上病情 ≤ 1 时 100% 保守治疗（从不积极治疗）。真实 J(π_e) = {true_J:.3f}")
    for name in ESTIMATORS:
        print(f"    {name:<6s} 估计的平均值 {np.mean(ests[name]):7.3f}（偏差 {np.mean(ests[name]) - true_J:+.3f}）")
    print("  → 所有估计量都有明显偏差，而且数据再多也消除不了。")
    print("    π_e 在病情危重时会积极治疗，但历史数据里**从来没有**这样的记录——")
    print("    没有任何数据能告诉我们\"积极治疗\"在那些状态会怎样。")
    print("    重要性采样要求：只要 π_e(a|s) > 0，就必须有 π_b(a|s) > 0（覆盖假设 / 支撑集假设）。")
    print("    这不是估计方法的缺陷，而是数据本身信息不足——任何方法都无能为力，除非做额外的假设。")


if __name__ == "__main__":
    main()