"""
最小世界模型直觉（阶段 4 配套）
--------------------------------
真环境：1D 格子 0..N-1，动作 Left/Right，到达 N-1 得 +1。
世界模型：用有限真实数据学 P(s'|s,a) 计数表。
对比：
  1) 在「想象」里用模型贪心规划的回报估计
  2) 同一策略在真环境的实际回报
并展示想象地平线 H 变大时，模型误差如何让两者分叉。

运行：python toy_world_model.py
"""

from __future__ import annotations

import random
from collections import defaultdict

N = 8
ACTIONS = (-1, 1)
GAMMA = 0.95
SEED = 0


def env_step(s: int, a: int) -> tuple[int, float, bool]:
    ns = max(0, min(N - 1, s + a))
    if ns == N - 1:
        return ns, 1.0, True
    return ns, -0.05, False


def collect_real_data(rng: random.Random, episodes: int = 30):
    """偏右行为策略采集真实转移。"""
    data = []
    for _ in range(episodes):
        s = 0
        for _ in range(40):
            a = 1 if rng.random() < 0.7 else rng.choice(ACTIONS)
            ns, r, done = env_step(s, a)
            data.append((s, a, r, ns, done))
            s = ns
            if done:
                break
    return data


def fit_count_model(data):
    """
    世界模型：计数估计 P(s'|s,a) 与平均 r(s,a)。
    未见过的 (s,a) 用均匀兜底 → 会引入模型偏差。
    """
    trans = defaultdict(lambda: defaultdict(int))
    rew = defaultdict(list)
    for s, a, r, ns, _ in data:
        trans[(s, a)][ns] += 1
        rew[(s, a)].append(r)

    def sample_next(s, a, rng):
        key = (s, a)
        if key not in trans or not trans[key]:
            # OOD：瞎猜
            return rng.randint(0, N - 1), -0.05
        dist = trans[key]
        total = sum(dist.values())
        x = rng.randrange(total)
        acc = 0
        for ns, c in dist.items():
            acc += c
            if x < acc:
                mean_r = sum(rew[key]) / len(rew[key])
                return ns, mean_r
        ns = next(iter(dist))
        return ns, sum(rew[key]) / len(rew[key])

    coverage = {k: dict(v) for k, v in trans.items()}
    return sample_next, coverage


def greedy_right_policy(s: int) -> int:
    return 1


def imagine_return(sample_next, policy, H: int, rng: random.Random, s0: int = 0) -> float:
    s = s0
    g = 0.0
    disc = 1.0
    for _ in range(H):
        a = policy(s)
        ns, r = sample_next(s, a, rng)
        g += disc * r
        disc *= GAMMA
        s = ns
        if s == N - 1:
            break
    return g


def real_return(policy, max_steps: int = 40, s0: int = 0) -> float:
    s = s0
    g = 0.0
    disc = 1.0
    for _ in range(max_steps):
        a = policy(s)
        s, r, done = env_step(s, a)
        g += disc * r
        disc *= GAMMA
        if done:
            break
    return g


def main():
    rng = random.Random(SEED)
    data = collect_real_data(rng, episodes=25)
    sample_next, coverage = fit_count_model(data)

    print("=== 最小世界模型玩具 ===")
    print(f"真实转移条数: {len(data)}")
    print(f"覆盖的 (s,a) 数: {len(coverage)} / {N * len(ACTIONS)}")

    # 真实策略回报（多次平均）
    real_rets = [real_return(greedy_right_policy) for _ in range(30)]
    real_mean = sum(real_rets) / len(real_rets)
    print(f"\n真环境、始终向右策略 平均回报: {real_mean:.3f}")

    print("\n想象地平线 H | 想象回报均值 | |想象-真实|")
    print("-" * 48)
    for H in [1, 2, 4, 8, 16, 32]:
        imag = [
            imagine_return(sample_next, greedy_right_policy, H, random.Random(SEED + i))
            for i in range(50)
        ]
        m = sum(imag) / len(imag)
        print(f"  H={H:2d}        | {m:8.3f}      | {abs(m - real_mean):.3f}")

    print("\n解读:")
    print("  - H 很短：想象接近「短视」，误差小但用处有限。")
    print("  - H 变长：若模型有未见转移/计数噪声，误差会累积，想象回报偏离真实。")
    print("  - 模型基 RL 必须用真实交互校准；生物 oracle 当世界模型时同理。")
    print("  - 覆盖不足的 (s,a) 在本玩具里会均匀瞎猜 → 典型 OOD 模型风险。")


if __name__ == "__main__":
    main()