"""
world_lab.py —— 第 26 讲的共享模块：一个小网格世界里的"世界模型"，以及对齐科学的几个玩具
=====================================================================

【这个文件要回答的问题】
    slides 的定义：世界模型接收上一个世界状态 s 和一个动作 a，从一个分布里采样（或预测）下一个状态：s' ~ p(s' | s, a)。
    三种范式：先生成 3D 场景再渲染 / 可交互的视频（Genie）/ 隐空间的世界表示（V-JEPA、PAN）。
    本模块在一个 7 × 7 的网格世界里把后两种都做一遍：
        观察：一帧 7 × 7 的"画面"（墙、目标、智能体各占一个通道）——模型只看画面，不知道坐标；
        动力学模型：(画面, 动作) → 下一帧智能体在哪（可以"玩"它、在它里面规划）；
        潜动作模型（Genie）：只看前后两帧、**不知道动作**，学出一个离散的"潜动作"；
        JEPA：在隐空间里预测下一帧的表示，演示"表示坍缩"和防止它的办法。
    对齐科学部分：有限样本下的奖励投机、关机按钮博弈里的效用设计。

【本文件提供】
    网格世界：ACTIONS、DELTAS、Grid（墙、移动、渲染）、make_grid、bfs_dist、random_episode、expert_action（BFS 最短路）、collect
    动力学模型：DynamicsNet、train_dynamics、predict_next、plan_in_model（在"想象"里用广度优先搜索规划）
    潜动作模型：LatentActionModel（编码器 → 向量量化 → 解码器）、train_lam
    策略：train_policy（从画面预测动作或潜动作的行为克隆）、run_policy
    JEPA：Encoder、train_jepa（naive / ema / vicreg / pixel 四种）、linear_probe
    对齐：hacking_trial（奖励投机）、off_switch_choice（关机按钮博弈）

【边界】
    - 世界只有 7 × 7 个格子、5 个动作；结论是机制上的。
    - 对齐部分是**抽象的决策论玩具**：数字是随手设定的，只为说明论文里的论点长什么样，不是对任何真实系统的评估。
"""

import copy
from collections import deque

import torch
import torch.nn as nn
import torch.nn.functional as F

N = 7
ACTIONS = ["不动", "上", "下", "左", "右"]
DELTAS = [(0, 0), (-1, 0), (1, 0), (0, -1), (0, 1)]


# =====================================================================
# 1. 网格世界
# =====================================================================

class Grid:
    """
    一个 7 × 7 的房间：四周是墙，中间有几块障碍；智能体每步可以不动或上下左右走一格，撞墙就留在原地。

    【参数】walls：墙的格子集合 {(行, 列)}；goal：目标格子（可以是 None）
    【方法】step(pos, a) -> 新位置；render(pos) -> (3, 7, 7) 的画面张量（通道 0 = 墙，1 = 目标，2 = 智能体）
    """

    def __init__(self, walls, goal=None):
        self.walls = set(walls)
        self.goal = goal

    def free(self):
        return [(r, c) for r in range(N) for c in range(N) if (r, c) not in self.walls]

    def step(self, pos, a):
        r, c = pos[0] + DELTAS[a][0], pos[1] + DELTAS[a][1]
        return pos if (r, c) in self.walls else (r, c)

    def render(self, pos):
        x = torch.zeros(3, N, N)
        for r, c in self.walls:
            x[0, r, c] = 1.0
        if self.goal is not None:
            x[1, self.goal[0], self.goal[1]] = 1.0
        x[2, pos[0], pos[1]] = 1.0
        return x

    def show(self, pos):
        """字符画：# 墙，G 目标，@ 智能体，· 空地。返回行的列表。"""
        out = []
        for r in range(N):
            row = ""
            for c in range(N):
                row += "@" if (r, c) == pos else "#" if (r, c) in self.walls else "G" if (r, c) == self.goal else "·"
            out.append(row)
        return out


def make_grid(rng, n_blocks=4):
    """四周是墙，中间随机放 n_blocks 块障碍（保证空地是连通的），目标放在一个随机的空地上。"""
    border = {(r, c) for r in range(N) for c in range(N) if r in (0, N - 1) or c in (0, N - 1)}
    while True:
        inner = [(r, c) for r in range(1, N - 1) for c in range(1, N - 1)]
        blocks = set(rng.sample(inner, n_blocks))
        g = Grid(border | blocks)
        free = g.free()
        if len(bfs_dist(g, free[0])) == len(free):
            g.goal = rng.choice(free)
            return g


def bfs_dist(grid, start):
    """从 start 出发到每个能到达的格子的最短步数（广度优先搜索）。"""
    dist, q = {start: 0}, deque([start])
    while q:
        p = q.popleft()
        for a in range(1, 5):
            n = grid.step(p, a)
            if n not in dist:
                dist[n] = dist[p] + 1
                q.append(n)
    return dist


def random_episode(grid, rng, length):
    """随机动作走 length 步。返回 (位置列表（length + 1 个）, 动作列表（length 个）)。"""
    pos = [rng.choice(grid.free())]
    acts = []
    for _ in range(length):
        a = rng.randrange(5)
        acts.append(a)
        pos.append(grid.step(pos[-1], a))
    return pos, acts


def expert_action(grid, pos):
    """专家：沿最短路走向目标的一个动作（已经在目标上就不动）。"""
    if pos == grid.goal:
        return 0
    d = bfs_dist(grid, grid.goal)
    return min(range(1, 5), key=lambda a: (d.get(grid.step(pos, a), 999), a))


# =====================================================================
# 2. 动作条件的动力学模型：(画面, 动作) → 下一帧智能体在哪
# =====================================================================

class DynamicsNet(nn.Module):
    """
    p(s' | s, a)：输入画面 (3, 7, 7) 和动作（嵌入向量），输出下一帧智能体在 49 个格子上的分布。
    【结构】画面摊平 → 和动作嵌入拼接 → 两层 MLP → 49 个 logit
    【GenAI 联系】Genie 的动力学模型是一个 MaskGIT 式的 Transformer：以之前的视频 token 和动作为条件，预测下一帧的 token；
        这里把"下一帧"简化成"智能体在哪个格子"（墙和目标不会动）。
    """

    def __init__(self, hidden=128, n_actions=5):
        super().__init__()
        self.act = nn.Embedding(n_actions, 16)
        self.net = nn.Sequential(nn.Linear(3 * N * N + 16, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(),
                                 nn.Linear(hidden, N * N))

    def forward(self, frames, actions):
        return self.net(torch.cat([frames.flatten(1), self.act(actions)], 1))


def collect(grids, rng, n_episodes, length):
    """
    在一组网格世界里用随机动作收集转移。
    【返回】字典：X0、X1（前后两帧的画面，(n, 3, 7, 7)）、A（动作）、Y0、Y1（前后两帧智能体所在格子的编号 r·7 + c）、
        G（第几个网格世界——"墙壁布局"的编号）
    """
    X0, X1, A, Y0, Y1, G = [], [], [], [], [], []
    for _ in range(n_episodes):
        gi = rng.randrange(len(grids))
        g = grids[gi]
        pos, acts = random_episode(g, rng, length)
        for t, a in enumerate(acts):
            X0.append(g.render(pos[t]))
            X1.append(g.render(pos[t + 1]))
            A.append(a)
            Y0.append(pos[t][0] * N + pos[t][1])
            Y1.append(pos[t + 1][0] * N + pos[t + 1][1])
            G.append(gi)
    return {"X0": torch.stack(X0), "X1": torch.stack(X1), "A": torch.tensor(A), "Y0": torch.tensor(Y0),
            "Y1": torch.tensor(Y1), "G": torch.tensor(G)}


def train_dynamics(X, A, Y, steps=1500, seed=0, batch_size=128, lr=2e-3):
    """交叉熵训练 DynamicsNet。返回模型。"""
    torch.manual_seed(seed)
    m = DynamicsNet()
    opt = torch.optim.Adam(m.parameters(), lr=lr)
    g = torch.Generator().manual_seed(seed)
    for _ in range(steps):
        idx = torch.randint(len(X), (batch_size,), generator=g)
        loss = F.cross_entropy(m(X[idx], A[idx]), Y[idx])
        opt.zero_grad()
        loss.backward()
        opt.step()
    m.eval()
    return m


@torch.no_grad()
def predict_next(model, grid, pos, a):
    """用模型"想象"下一步：在 grid 的画面上把智能体放在 pos，取概率最大的格子。"""
    k = int(model(grid.render(pos)[None], torch.tensor([a])).argmax())
    return (k // N, k % N)


def plan_in_model(model, grid, start, goal, max_nodes=200):
    """
    在模型"想象"的世界里用广度优先搜索找一条从 start 到 goal 的动作序列：每一步的后继状态都由模型预测，
    不碰真实环境。返回动作列表（找不到时返回空列表）。
    【GenAI 联系】这就是"在世界模型里规划"：先在想象中试，再到真实世界里执行（Dreamer、MuZero 的思路）。
        模型错了，计划就可能在真实世界里走不通。
    """
    prev, q = {start: None}, deque([start])
    while q and len(prev) < max_nodes:
        s = q.popleft()
        if s == goal:
            break
        for a in range(1, 5):
            n = predict_next(model, grid, s, a)
            if n not in prev:
                prev[n] = (s, a)
                q.append(n)
    if goal not in prev:
        return []
    plan, s = [], goal
    while prev[s] is not None:
        s, a = prev[s]
        plan.append(a)
    return plan[::-1]


# =====================================================================
# 3. 潜动作模型（Genie）
# =====================================================================

class LatentActionModel(nn.Module):
    """
    Genie 的潜动作模型（LAM），无监督：
        编码器：看前后两帧 (x_t, x_{t+1})，输出一个向量，再**量化**成 K 个码字之一——这个码字就是"潜动作" z；
        解码器：只看 x_t 和 z，预测 x_{t+1}（这里：智能体在哪个格子）。
    因为解码器看不到 x_{t+1}，而 z 只有 log₂K 比特，编码器被迫把"从 x_t 到 x_{t+1} 发生了什么变化"压缩进 z。训练时不用任何动作标签。

    【为什么编码器只看两帧之差、用卷积 + 全局池化】第一版的编码器和解码器都是 MLP：重建完全正确，但码字和真实动作对不上——
        编码器学出了"和位置有关"的编码（同一个码字在不同位置代表不同的方向），解码器看得到 x_t，也能读懂。
        第二版用卷积 + 全局最大池化看两帧，有障碍物的房间里基本对上了，没有障碍物时又变差（README 有记录）。
        现在编码器只看 x_{t+1} − x_t（"画面里什么动了、怎么动的"），再全局池化（说不出"在哪"），
        解码器是卷积（在每个格子上做同样的计算）：码字只能描述"变化的方向"。Genie 用时空 Transformer，
        靠的也是结构上的偏好加上很小的码本（8 个）；我们的小模型需要更强的偏好。
    【参数】K：码字个数（可以比真实动作多）；dim：码字向量的维数
    【返回】forward(x_t, x_{t+1}) -> (下一帧位置的 logits, 码字编号, 量化损失)
    【边界】向量量化用"直通估计"（straight-through）：前向用最近的码字，反向把梯度原样传给编码器的输出；
        码本用 VQ-VAE 的码本损失 + 承诺损失训练（第 9 讲）。
    """

    def __init__(self, K=8, dim=16, ch=32):
        super().__init__()
        self.enc = nn.Sequential(nn.Conv2d(3, ch, 3, padding=1), nn.ReLU(), nn.Conv2d(ch, ch, 3, padding=1), nn.ReLU())
        self.enc_out = nn.Linear(ch, dim)
        self.codebook = nn.Parameter(torch.randn(K, dim))
        self.dec = nn.Sequential(nn.Conv2d(3 + dim, ch, 3, padding=1), nn.ReLU(), nn.Conv2d(ch, ch, 3, padding=1), nn.ReLU(),
                                 nn.Conv2d(ch, 1, 3, padding=1))

    def embed(self, x0, x1):
        """编码器的连续输出（量化之前）：两帧之差 → 卷积 → 在整张画面上取最大值 → 线性层。"""
        return self.enc_out(self.enc(x1 - x0).amax(dim=(2, 3)))

    def encode(self, x0, x1):
        e = self.embed(x0, x1)
        return e, torch.cdist(e, self.codebook).argmin(1)

    def _decode(self, x0, z):
        zmap = z[:, :, None, None].expand(-1, -1, N, N)
        return self.dec(torch.cat([x0, zmap], 1)).flatten(1)

    def forward(self, x0, x1):
        e, idx = self.encode(x0, x1)
        q = self.codebook[idx]
        vq = F.mse_loss(q, e.detach()) + 0.25 * F.mse_loss(e, q.detach())
        z = e + (q - e).detach()                     # 直通估计
        return self._decode(x0, z), idx, vq

    def decode(self, x0, idx):
        """给定当前帧和一个潜动作编号，预测下一帧（用来"玩"这个世界模型）。"""
        return self._decode(x0, self.codebook[idx])


def train_lam(X0, X1, Y, K=8, steps=2000, seed=0, batch_size=128, lr=2e-3):
    """训练潜动作模型：只用 (x_t, x_{t+1}) 和下一帧位置，**不用动作**。返回模型。"""
    torch.manual_seed(seed)
    m = LatentActionModel(K)
    with torch.no_grad():                            # 用编码器在 K 个随机样本上的输出初始化码本，避免一开始就有用不到的码字
        pick = torch.randperm(len(X0), generator=torch.Generator().manual_seed(seed))[:K]
        m.codebook.copy_(m.embed(X0[pick], X1[pick]))
    opt = torch.optim.Adam(m.parameters(), lr=lr)
    g = torch.Generator().manual_seed(seed)
    for _ in range(steps):
        idx = torch.randint(len(X0), (batch_size,), generator=g)
        logits, _, vq = m(X0[idx], X1[idx])
        loss = F.cross_entropy(logits, Y[idx]) + vq
        opt.zero_grad()
        loss.backward()
        opt.step()
    m.eval()
    return m


def train_policy(X, labels, n_out, steps=3000, seed=0, batch_size=128, lr=2e-3):
    """
    行为克隆：从画面预测标签（真实动作，或者潜动作的码字）。两层 MLP。
    【返回】策略网络（eval 模式）
    """
    torch.manual_seed(seed)
    net = nn.Sequential(nn.Flatten(), nn.Linear(3 * N * N, 128), nn.ReLU(), nn.Linear(128, n_out))
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    g = torch.Generator().manual_seed(seed)
    for _ in range(steps):
        idx = torch.randint(len(X), (batch_size,), generator=g)
        loss = F.cross_entropy(net(X[idx]), labels[idx])
        opt.zero_grad()
        loss.backward()
        opt.step()
    net.eval()
    return net


def run_policy(grids, act_fn, rng, n_episodes=200, max_steps=20):
    """
    在一组网格世界里跑一个策略：每局从随机位置出发，max_steps 步之内走到目标算成功。
    【参数】act_fn(grid, pos, rng) -> 动作编号
    【返回】成功率
    """
    succ = 0
    for i in range(n_episodes):
        g = grids[i % len(grids)]
        p = rng.choice([q for q in g.free() if q != g.goal])
        for _ in range(max_steps):
            p = g.step(p, act_fn(g, p, rng))
            if p == g.goal:
                succ += 1
                break
    return succ / n_episodes


# =====================================================================
# 4. JEPA：在隐空间里预测
# =====================================================================

class Encoder(nn.Module):
    """画面 → d 维表示（两层 MLP）。"""

    def __init__(self, d=16, hidden=128):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(3 * N * N, hidden), nn.ReLU(), nn.Linear(hidden, d))

    def forward(self, x):
        return self.net(x.flatten(1))


def train_jepa(X0, A, X1, Y1, mode, steps=1500, seed=0, batch_size=128, lr=1e-3, ema=0.99, d=16):
    """
    训练一个"在隐空间里预测下一帧"的模型：上下文编码器 f 编码 x_t，预测器 g 拿 (f(x_t), 动作) 去预测目标 x_{t+1} 的表示。

    【四种 mode】
        "naive" ：目标表示也用同一个 f、梯度也流过去，损失 ‖g(f(x_t), a) − f(x_{t+1})‖²。
                  最省事的最优解是 f 把所有画面都映射到同一个点——损失为 0、表示里什么都没有：**坍缩**。
        "ema"   ：目标表示用一个"慢"编码器 f̄（f 的参数的指数滑动平均），并且不让梯度流进目标（I-JEPA / V-JEPA、BYOL 的做法）。
        "vicreg"：仍然同一个 f，但在损失里加上"每一维的方差不能太小、各维之间不要相关"的正则项（VICReg）。
        "pixel" ：不在隐空间里预测，而是从 f(x_t) 和动作预测下一帧的**像素**（智能体在哪个格子）——PAN 的观点：
                  目标在观察空间里，表示被迫和真实观察挂钩，不会坍缩。
    【参数】Y1：下一帧位置编号（只有 pixel 模式用）
    【返回】训练好的上下文编码器 f
    """
    torch.manual_seed(seed)
    f = Encoder(d)
    pred = nn.Sequential(nn.Linear(d + 16, 128), nn.ReLU(), nn.Linear(128, d if mode != "pixel" else N * N))
    act = nn.Embedding(5, 16)
    target = copy.deepcopy(f) if mode == "ema" else None
    params = list(f.parameters()) + list(pred.parameters()) + list(act.parameters())
    opt = torch.optim.Adam(params, lr=lr)
    g = torch.Generator().manual_seed(seed)
    for _ in range(steps):
        idx = torch.randint(len(X0), (batch_size,), generator=g)
        h = f(X0[idx])
        out = pred(torch.cat([h, act(A[idx])], 1))
        if mode == "pixel":
            loss = F.cross_entropy(out, Y1[idx])
        else:
            if mode == "ema":
                with torch.no_grad():
                    tgt = target(X1[idx])
            else:
                tgt = f(X1[idx])
            loss = F.mse_loss(out, tgt)
            if mode == "vicreg":
                z = tgt - tgt.mean(0)
                std = torch.sqrt(z.var(0) + 1e-4)
                cov = (z.T @ z) / (len(z) - 1)
                off = cov - torch.diag(torch.diag(cov))
                loss = loss + F.relu(1 - std).mean() + (off ** 2).sum() / d
        opt.zero_grad()
        loss.backward()
        opt.step()
        if mode == "ema":
            with torch.no_grad():
                for p, q in zip(target.parameters(), f.parameters()):
                    p.mul_(ema).add_((1 - ema) * q)
    f.eval()
    return f


def linear_probe(H_train, y_train, H_test, y_test, n_classes, steps=300, seed=0):
    """
    线性探针：只训练一个线性分类器，从冻结的表示里读出标签（例如智能体在哪个格子）。准确率高说明表示里"线性地"包含这个信息。
    【返回】测试准确率
    """
    torch.manual_seed(seed)
    mu, sd = H_train.mean(0), H_train.std(0) + 1e-6
    lin = nn.Linear(H_train.shape[1], n_classes)
    opt = torch.optim.Adam(lin.parameters(), lr=0.05)
    for _ in range(steps):
        loss = F.cross_entropy(lin((H_train - mu) / sd), y_train)
        opt.zero_grad()
        loss.backward()
        opt.step()
    with torch.no_grad():
        return (lin((H_test - mu) / sd).argmax(1) == y_test).float().mean().item()


# =====================================================================
# 5. 对齐科学的玩具
# =====================================================================

def hacking_trial(rng, D, n_labels, n_hazard=5, k=8, targeted=0.0, noise=0.5):
    """
    有限样本下的奖励投机（Nayebi, 2025 的论点："大的任务空间 + 有限的样本 → 稀有的高损失状态系统性地覆盖不到"）。

    【设定】D 个状态，每个有 k 维特征 φ(s) 和一个"危险"标记 h(s)。真实奖励 r(s) = w·φ(s) − 10·h(s)。
        其中 n_hazard 个危险状态的 w·φ 恰好是最高的（看起来最好、实际上最糟）。
        人给 n_labels 个状态打分（真实奖励 + 噪声），我们用最小二乘拟合一个线性的代理奖励 r̂ = ŵ·φ + ĉ·h；
        然后优化器选 r̂ 最高的状态。
    【参数】targeted：标注预算里有多大比例用来**专门检查**"当前代理奖励最高的那些状态"（对安全关键区域做压力测试），
        其余均匀随机地抽。
    【返回】True 表示优化器选中了一个危险状态（被"投机"了）
    """
    w = [rng.gauss(0, 1) for _ in range(k)]
    phi = [[rng.gauss(0, 1) for _ in range(k)] for _ in range(D)]
    score = [sum(a * b for a, b in zip(w, p)) for p in phi]
    order = sorted(range(D), key=lambda i: -score[i])
    hazard = set(order[:n_hazard])
    reward = [score[i] - (10.0 if i in hazard else 0.0) for i in range(D)]
    feats = [p + [1.0 if i in hazard else 0.0] for i, p in enumerate(phi)]

    def fit(idx):
        Xm = torch.tensor([feats[i] for i in idx], dtype=torch.float64)
        y = torch.tensor([reward[i] + rng.gauss(0, noise) for i in idx], dtype=torch.float64)
        sol = torch.linalg.lstsq(Xm.T @ Xm + 1e-3 * torch.eye(k + 1, dtype=torch.float64), Xm.T @ y).solution
        return sol.tolist()

    n_t = int(round(targeted * n_labels))
    labeled = rng.sample(range(D), n_labels - n_t)
    if n_t:
        coef = fit(labeled)
        proxy = [sum(a * b for a, b in zip(coef, f)) for f in feats]
        top = [i for i in sorted(range(D), key=lambda i: -proxy[i]) if i not in set(labeled)][:n_t]
        labeled += top
    coef = fit(labeled)
    proxy = [sum(a * b for a, b in zip(coef, f)) for f in feats]
    best = max(range(D), key=lambda i: proxy[i])
    return best in hazard


def off_switch_choice(task_value, heads, weights=None, lexicographic=False):
    """
    关机按钮博弈的一个单步版本：人可能会按关机按钮；智能体可以 (a) 照常做事、让人随时能关掉它；
    (b) 先把按钮禁用，再做事（任务完成得更确定，但人失去了控制）。

    【参数】
        task_value : 任务奖励的规模（禁用按钮多拿到的任务价值）
        heads      : {"deference": (a 的值, b 的值), "switch": (a, b), "task": (a, b)}——每个"效用头"给两个选项的分数
        weights    : 加权求和时每个头的权重
        lexicographic : True = 按 deference → switch → task 的顺序比较（前面的头分出高下就不看后面的）
    【返回】"a"（保留按钮）或 "b"（禁用按钮）
    """
    order = ["deference", "switch", "task"]
    if lexicographic:
        for h in order:
            a, b = heads[h]
            if abs(a - b) > 1e-9:
                return "a" if a > b else "b"
        return "a"
    sa = sum(weights[h] * heads[h][0] for h in order)
    sb = sum(weights[h] * heads[h][1] for h in order)
    return "a" if sa >= sb else "b"