"""
=============================================================================
第 8 讲 · 共享模块：车道保持环境、专家、表格型模仿学习者
=============================================================================

【这个模块提供什么？】
    01（行为克隆）和 02（DAgger）两个文件共同使用的环境、专家和学习者。
    运行它们时当前目录必须是 code/lecture08/。

-----------------------------------------------------------------------------
【环境：有侧风的车道保持】
-----------------------------------------------------------------------------

    撞车 |  .  .  .  .  .  .  .  .  [中心]  .  .  .  .  .  .  .  .  | 撞车
        -8 -7 -6 -5 -4 -3 -2 -1    0     1  2  3  4  5  6  7    +8
                          └──── 车道（|x| ≤ 2）────┘

    - 状态 x：车相对车道中心的横向偏移，整数，-7 ~ +7
    - 动作 a：-1（向左打方向）、0（不动）、+1（向右打方向）
    - 每一步：x ← x + a + 风；风以 0.1 / 0.2 / 0.4 / 0.2 / 0.1 的概率为 -2 / -1 / 0 / +1 / +2
      （风有时比方向盘的修正能力更强，所以即使是专家，偶尔也会被吹离中心几格）
    - |x| ≥ 8：冲出路面，回合结束（撞车）
    - 代价：每一步如果车不在车道内（|x| > 2），代价 +1；撞车之后剩下的每一步都算代价 +1
      （所以一个回合的总代价 = "不在车道内的步数"，最大为回合长度 T）

    - 专家：总是往中心打方向：a = -sign(x)（x = 0 时不动）
      专家开车时，车几乎总在车道内，很少离开中心超过 2～3 格。

【为什么用这个例子？】
    它把"分布偏移"变得非常具体：专家的示范里几乎只有 |x| ≤ 3 的状态。
    学习者一旦因为风被吹到 |x| = 4、5……这些**从没见过**的状态，就不知道该怎么办了。
"""

import numpy as np

X_MAX = 7                  # 状态范围 -7 ~ +7；|x| >= 8 撞车
LANE = 2                   # |x| <= 2 算在车道内
ACTIONS = (-1, 0, 1)
N_STATES = 2 * X_MAX + 1   # 15 个状态
WIND = (-2, -1, 0, 1, 2)
WIND_PROBS = (0.1, 0.2, 0.4, 0.2, 0.1)


def idx(x):
    """状态 x（-7~7）→ 数组下标（0~14）。"""
    return x + X_MAX


def expert_action(x):
    """专家策略：往中心打方向。【Python 小知识】np.sign(x) 返回 -1、0 或 1。"""
    return int(-np.sign(x))


def step(x, a, rng):
    """
    环境的一步。返回 (新状态, 是否撞车)。
    【注意】这里的 a 是 -1/0/+1 本身（不是下标）。
    """
    wind = rng.choice(WIND, p=WIND_PROBS)
    x_new = x + a + wind
    crashed = abs(x_new) > X_MAX
    return x_new, crashed


def rollout(policy_fn, T, rng, x0=0, record=True):
    """
    用策略 policy_fn（输入 x，输出动作 -1/0/+1）开 T 步。

    【返回】(总代价, 途经的状态列表, 是否撞车)
        总代价 = 不在车道内的步数（撞车之后剩下的步数全部计入）。
    """
    x, cost, visited = x0, 0, []
    # 一次性采样整个回合的风（比每一步单独调用 rng.choice 快几十倍；结果的分布完全相同）
    winds = rng.choice(WIND, size=T, p=WIND_PROBS)
    for t in range(T):
        if record:
            visited.append(x)
        x = int(x + policy_fn(x) + winds[t])
        crashed = abs(x) > X_MAX
        if crashed:
            cost += T - t                   # 撞车：剩下的每一步都算代价
            return cost, visited, True
        cost += int(abs(x) > LANE)
    return cost, visited, False


class TabularImitator:
    """
    表格型模仿学习者：对每个见过的状态，记住专家在这个状态最常做的动作（多数投票）。

    【对没见过的状态怎么办？】
        default_action=0：什么都不做（不打方向）。这是一个"合理但无知"的默认行为——
        学习者从没见过这种情况，没有任何依据判断该往哪边打。

    【为什么用表格而不是神经网络？】
        这样"没见过就不会"的效应最直观。在高维问题中（比如图像输入），
        神经网络在远离训练数据的地方的输出同样不可靠——本质上是同一个问题。
    """

    def __init__(self, default_action=0):
        self.counts = np.zeros((N_STATES, 3))     # counts[状态, 动作下标]
        self.default_action = default_action

    def add(self, states, actions):
        """加入一批 (状态, 专家动作) 数据。"""
        for x, a in zip(states, actions):
            if abs(x) <= X_MAX:
                self.counts[idx(x), a + 1] += 1   # 动作 -1/0/+1 → 下标 0/1/2

    def __call__(self, x):
        """预测：见过的状态取多数票；没见过的状态用默认动作。"""
        c = self.counts[idx(x)]
        if c.sum() == 0:
            return self.default_action
        return int(np.argmax(c)) - 1

    def seen_states(self):
        """见过的状态集合（排好序的列表）。"""
        return [x for x in range(-X_MAX, X_MAX + 1) if self.counts[idx(x)].sum() > 0]


def evaluate(policy_fn, T, n_episodes, seed):
    """
    评估策略：跑 n_episodes 个长度为 T 的回合。
    【返回】(平均总代价, 撞车率, 状态访问分布（长度 N_STATES，和为 1）)
    """
    rng = np.random.default_rng(seed)
    costs, crashes = [], 0
    visits = np.zeros(N_STATES)
    for _ in range(n_episodes):
        cost, visited, crashed = rollout(policy_fn, T, rng)
        costs.append(cost)
        crashes += crashed
        for x in visited:
            visits[idx(x)] += 1
    return float(np.mean(costs)), crashes / n_episodes, visits / visits.sum()