# ============================================================
#  ReflectDrive-2 · AutoEdit 训练阶段伪代码
#  对应论文 Section 3.2 / 4.3 / 4.4
# ============================================================
#
#  符号说明（与论文保持一致）
#  z0      : 干净的连续坐标路径点序列  [(x1,y1),...,(x8,y8)]
#  x0      : z0 离散化后的 token 序列  [x1,y1,...,x8,y8]  长度 L=16
#  x̃0      : 扰动轨迹的 token 序列
#  c       : 多模态上下文 (视觉tokens, 导航指令, 自车状态)
#  pθ      : 共享的条件 token 模型（Transformer 主干）
#   assessments  : 掩码符号，用于草稿阶段
#  t       : 掩码比例 ∈ [0, 1]，决定盖住多少 token
# ============================================================


import random
import math


# ----------------------------------------------------------
# 工具函数
# ----------------------------------------------------------

def tokenize(waypoints):
    """
    把连续坐标路径点量化为离散 token 序列。
    每个路径点 (x, y) → 两个整数 token (x_bin, y_bin)。
    8 个路径点 → 长度 16 的 token 序列。
    """
    tokens = []
    for (x, y) in waypoints:
        tokens.append(quantize(x, axis='longitudinal'))  # 纵向坐标 bin
        tokens.append(quantize(y, axis='lateral'))        # 横向坐标 bin
    return tokens  # 长度 L = 16


def detokenize(tokens):
    """tokenize 的逆操作：token 序列 → 连续坐标路径点。"""
    waypoints = []
    for i in range(0, len(tokens), 2):
        x = dequantize(tokens[i],   axis='longitudinal')
        y = dequantize(tokens[i+1], axis='lateral')
        waypoints.append((x, y))
    return waypoints


# ----------------------------------------------------------
# Step 1 · 结构感知扰动（Structure-Aware Perturbation, SAP）
#          论文公式 (5)(6)
# ----------------------------------------------------------

def structure_aware_perturbation(z0, beta_range=(0.7, 1.3), alpha_max=0.1):
    """
    对干净轨迹 z0 施加两类结构化扰动，模拟驾驶中最常见的错误模式。

    Args:
        z0         : 干净路径点列表，[(x1,y1),...,(x8,y8)]
        beta_range : 纵向进度缩放系数范围 (βmin, βmax)
        alpha_max  : 横向旋转角度上限 αmax（弧度）

    Returns:
        z̃0 : 扰动后的路径点列表
    """

    # ── 扰动类型随机选一种（或两种都加）──────────────────────

    perturbation_type = random.choice(['longitudinal', 'lateral', 'both'])

    z_tilde = [wp for wp in z0]   # 先复制一份

    # ── 纵向进度扰动：沿弧长缩放 ────────────────────────────
    #    论文公式 (5)：z̃i = Interp(z0, β·dᵢ)
    #    β < 1 → 进度不足（刹车过早）
    #    β > 1 → 进度超前（刹车过晚 / 超速）
    if perturbation_type in ('longitudinal', 'both'):

        beta = random.uniform(*beta_range)

        # 计算每个路径点的弧长累积值
        arc_lengths = compute_arc_lengths(z0)          # [d1, d2, ..., d8]

        # 用缩放后的弧长重新插值轨迹
        scaled_arc = [beta * d for d in arc_lengths]   # β·dᵢ
        z_tilde = interpolate_by_arc(z0, scaled_arc)   # 沿原轨迹重新采样


    # ── 横向方向扰动：在自车坐标系中旋转 ───────────────────
    #    论文公式 (6)：z̃i = R(α)·zᵢ
    #    α 为随机旋转角，产生整体侧向偏移，但保持轨迹平滑性
    if perturbation_type in ('lateral', 'both'):

        alpha = random.uniform(-alpha_max, alpha_max)

        cos_a, sin_a = math.cos(alpha), math.sin(alpha)

        z_tilde = [
            (cos_a * x - sin_a * y,
             sin_a * x + cos_a * y)
            for (x, y) in z_tilde
        ]

    return z_tilde   # z̃0：扰动后的连续坐标路径点


# ----------------------------------------------------------
# Step 2 · AutoEdit 损失（LSAP）
#          论文公式 (7)(8)
# ----------------------------------------------------------

def compute_autoedit_loss(model, x0, context):
    """
    训练 AutoEdit：给定扰动 token 序列，预测干净 token 序列。

    注意：AutoEdit 不使用 assessments！
    输入是完整的（扰动后的）具体 token，目标是干净的具体 token。
    这与草稿阶段的掩码填空训练完全不同。

    Args:
        model   : 共享的条件 token 模型 pθ
        x0      : 干净轨迹的 token 序列，长度 L=16
        context : 多模态上下文 c = (视觉tokens, 导航, 自车状态)

    Returns:
        L_SAP : AutoEdit 纠错损失（标量）
    """

    L = len(x0)   # = 16

    # 2-a. 生成扰动后的连续坐标路径点
    z0     = detokenize(x0)                        # token → 连续坐标
    z_tilde = structure_aware_perturbation(z0)      # 施加纵向/横向扰动

    # 2-b. 把扰动路径点重新离散化为 token 序列
    x_tilde = tokenize(z_tilde)                    # 扰动 token 序列 x̃0

    # 2-c. 用模型对扰动 token 序列做预测
    #      输入：x̃0（扰动后的完整具体 token）+ 多模态上下文 c
    #      输出：每个位置上的替换 token 概率分布
    #      论文公式 (7)：qθ(· | x̃0, c) = softmax( hθ(x̃0, c) )
    logits = model(x_tilde, context)               # shape: [L, vocab_size]
    probs  = softmax(logits, dim=-1)               # qθ(· | x̃0, c)

    # 2-d. 计算交叉熵损失（对所有 L 个位置平均）
    #      论文公式 (8)：L_SAP = -E[ 1/L · Σ log qθ(x⁰ᵢ | x̃0, c) ]
    #      目标是干净 token x0，而非扰动 token x̃0
    L_SAP = 0.0
    for i in range(L):
        clean_token_id = x0[i]                     # 第 i 位的干净 token
        L_SAP += -math.log(probs[i][clean_token_id] + 1e-9)

    L_SAP /= L    # 对序列长度取平均

    return L_SAP


# ----------------------------------------------------------
# Step 3 · 草稿阶段掩码扩散损失（LDLM）
#          论文公式 (1)
# ----------------------------------------------------------

def compute_dlm_loss(model, x0, context):
    """
    训练草稿生成：随机盖住若干 token，让模型预测所有位置的原始 token。

    与 AutoEdit 损失的关键区别：
      - AutoEdit：输入是"有规律错误"的具体 token
      - DLM    ：输入是"随机盖 assessments"的不完整序列

    Args:
        model   : 共享的条件 token 模型 pθ
        x0      : 干净轨迹的 token 序列
        context : 多模态上下文 c

    Returns:
        L_DLM : 掩码扩散损失（标量）
    """

    L = len(x0)   # = 16

    # 3-a. 随机采样掩码比例 t ∈ [0, 1]
    t = random.uniform(0, 1)

    # 3-b. 按比例 t 独立随机将每个 token 替换为 assessments
    x_masked = []
    for token in x0:
        if random.random() < t:
            x_masked.append(MASK_TOKEN_ID)   # 盖住
        else:
            x_masked.append(token)           # 保留

    # 3-c. 模型对所有位置（包括未盖住的）做预测
    #      注意：论文选择"全位置监督"，而非仅对 assessments 位置监督
    #      实验证明全位置监督让训练更稳定、草稿更连贯
    logits = model(x_masked, context)        # shape: [L, vocab_size]
    probs  = softmax(logits, dim=-1)

    # 3-d. 对所有 L 个位置计算交叉熵（目标均为干净 token x0）
    #      论文公式 (1)：L_DLM = -E[ 1/L · Σ log pθ(x⁰ᵢ | xt, c) ]
    L_DLM = 0.0
    for i in range(L):
        L_DLM += -math.log(probs[i][x0[i]] + 1e-9)

    L_DLM /= L

    return L_DLM


# ----------------------------------------------------------
# Step 4 · 可行驶区域场损失（Lfield）
#          论文公式 (11)(12)(13)(14)
# ----------------------------------------------------------

def compute_field_loss(model, x0, context, dac_cost_field):
    """
    空间惩罚：让模型在概率层面就"远离"不可行驶区域。

    Args:
        model          : 共享的条件 token 模型 pθ
        x0             : 干净轨迹的 token 序列
        context        : 多模态上下文 c
        dac_cost_field : 可行驶区域代价场 C ∈ R^[H×W]
                         非可行驶区域 = 高代价，可行驶区域 = 0

    Returns:
        L_field : 可行驶区域场损失（标量）
    """

    # 4-a. 用带随机掩码的输入让模型输出 logits（与 DLM 阶段共用）
    logits = model(apply_random_mask(x0), context)   # shape: [L, vocab_size]

    L_field = 0.0

    # 4-b. 对每个路径点 t（共 8 个）计算空间惩罚
    for t in range(8):

        x_logit = logits[2*t]       # 第 t 个路径点的纵向 logit
        y_logit = logits[2*t + 1]   # 第 t 个路径点的横向 logit

        p_x = softmax(x_logit)      # 纵向坐标的概率分布 p^(t)_x  shape:[W]
        p_y = softmax(y_logit)      # 横向坐标的概率分布 p^(t)_y  shape:[H]

        # 4-c. 构造联合空间分布（外积）
        #      论文公式 (11)：p^(t)_xy[i,j] = p^(t)_x[i] · p^(t)_y[j]
        #      假设纵横坐标独立——近似但够用
        p_xy = outer_product(p_x, p_y)   # shape: [H, W]

        # 4-d. 用代价场加权的对数障碍函数计算惩罚
        #      论文公式 (12)：L_field = Σ_t Σ_{i,j} -log(1 - p^(t)_xy[i,j]) · C[i,j]
        #
        #      直觉：
        #        · C[i,j] 大（靠近不可行驶区域）→ 这一项权重大
        #        · p_xy[i,j] 大（模型对这个位置置信高）→ 惩罚也大
        #        · 两个大值相乘 → 强烈惩罚"高置信度踩线"
        for i in range(H):
            for j in range(W):
                cost = dac_cost_field[i][j]
                if cost > 0:
                    penalty = -math.log(1 - p_xy[i][j] + 1e-9) * cost
                    L_field += penalty

    return L_field


# ----------------------------------------------------------
# Step 5 · 总监督损失 & 训练主循环
#          论文公式 (15)
# ----------------------------------------------------------

def supervised_training(model, dataloader, lambda_SAP=1.0, lambda_field=0.1):
    """
    监督微调（SFT）主循环：同时优化三个目标。

    总损失：L_sup = L_DLM + λ_SAP · L_SAP + λ_field · L_field
    论文公式 (15)

    三个损失的分工：
      L_DLM   → 教模型从 assessments 还原干净轨迹（起草能力）
      L_SAP   → 教模型将扰动 token 映射回干净 token（纠错能力）
      L_field → 在概率层面惩罚踩入不可行驶区域（空间约束）
    """

    optimizer = AdamW(model.parameters())

    for batch in dataloader:

        x0         = batch['trajectory_tokens']   # 干净轨迹 token 序列
        context    = batch['context']             # 多模态上下文 c
        cost_field = batch['dac_cost_field']      # 可行驶区域代价场

        # ── 计算三个损失分量 ───────────────────────────────
        L_DLM   = compute_dlm_loss(model, x0, context)
        L_SAP   = compute_autoedit_loss(model, x0, context)
        L_field = compute_field_loss(model, x0, context, cost_field)

        # ── 加权求和得到总监督损失 ─────────────────────────
        L_sup = L_DLM + lambda_SAP * L_SAP + lambda_field * L_field

        # ── 反向传播 & 更新参数 ────────────────────────────
        optimizer.zero_grad()
        L_sup.backward()
        optimizer.step()

        log(L_DLM=L_DLM, L_SAP=L_SAP, L_field=L_field, L_sup=L_sup)

    # SFT 结束后，进入 RL 微调阶段（见论文 Section 4.5）
    return model

# ============================================================
#  ReflectDrive-2 · RL 微调阶段伪代码
#  对应论文 Section 3.4 / 4.5
# ============================================================
#
#  核心思想：
#    SFT 之后，起草器和 AutoEdit 各自都有能力，但互不知道对方的存在。
#    RL 阶段把两者绑在一起：跑完"起草→AutoEdit"的完整 rollout，
#    只用最终轨迹的驾驶得分作为奖励，同时回传给两个阶段。
#    结果：起草器学会生成"可被 AutoEdit 改好"的草稿；
#          AutoEdit 学会做"真正提升驾驶分"的修改。
#
#  符号说明（与论文保持一致）
#    Ng       : 每帧采样的目标点数量（论文实验中 Ng=3）
#    I        : 每个目标点采样的草稿数（论文实验中 I=2）
#    G        : 总候选轨迹数 = Ng × I（论文实验中 G=6）
#    Sdraft   : 草稿阶段的 unmask 轮数
#    Sedit    : AutoEdit 阶段的修改轮数
#    S        : 总步数 = Sdraft + Sedit
#    ρg       : 第 g 条候选的完整 token 转换序列（长度 S+1）
#    τg       : 第 g 条候选的最终轨迹（连续坐标）
#    R(τg)    : 闭环规划奖励（即 PDMS 得分）
#    Ag       : 第 g 条候选的组相对优势
#    πθ       : 当前策略（即当前模型参数）
#    πθ_old   : rollout 时冻结的旧策略（用于计算重要性采样比）
#    πref     : SFT 之后冻结的参考策略（用于 KL 惩罚）
# ============================================================


import math
import copy


# ----------------------------------------------------------
# Step 1 · 单条候选的完整 Draft→AutoEdit Rollout
#          论文公式 (16)(17)
# ----------------------------------------------------------

def run_single_rollout(model, goal_token, context,
                       S_draft=3, S_edit=3):
    """
    对一个目标点执行完整的"起草→AutoEdit"rollout，
    并记录每一步的 token 转换序列（用于后续计算策略梯度）。

    Args:
        model      : 当前策略模型 πθ_old（rollout 时冻结）
        goal_token : 已选定的目标点 token（行为锚点）
        context    : 多模态上下文 c
        S_draft    : 草稿阶段 unmask 轮数
        S_edit     : AutoEdit 阶段修改轮数

    Returns:
        trajectory_history : 完整 token 状态序列
                             [x⁰, x¹, ..., x^S]  共 S+1 个快照
                             前 S_draft 步是草稿阶段，后 S_edit 步是 AutoEdit 阶段
        final_trajectory   : 最终连续坐标轨迹 τg（送入仿真器求奖励）
    """

    L = 16  # token 序列长度（8个路径点 × 2坐标）

    # ── 初始化：全部 assessments，目标点 token 固定 ────────────────
    x_current = [MASK_TOKEN_ID] * L
    x_current = set_goal_token(x_current, goal_token)  # 锁定目标点位置

    trajectory_history = [x_current.copy()]  # x⁰

    # ══════════════════════════════════════════════════════
    # 阶段 A · 草稿生成（Masked Diffusion Drafting）
    #   从全 assessments 出发，每轮并行 unmask 最有把握的 token
    # ══════════════════════════════════════════════════════
    for s in range(S_draft):

        # A-1. 模型预测每个位置的 token 概率分布
        #      输入：带 assessments 的不完整序列 + 多模态上下文
        logits = model(x_current, context)          # shape: [L, vocab_size]
        probs  = softmax(logits, dim=-1)            # πθ_old(· | x_current, c)

        # A-2. 对每个仍是 assessments 的位置，采样一个 token
        x_predicted = []
        confidence  = []
        for i in range(L):
            if x_current[i] == MASK_TOKEN_ID:
                token_id = sample_from(probs[i])    # 按概率采样（非贪心）
                conf     = probs[i][token_id]       # 该 token 的概率作为置信度
            else:
                token_id = x_current[i]             # 已确定的 token 保持不变
                conf     = 1.0
            x_predicted.append(token_id)
            confidence.append(conf)

        # A-3. 按置信度从高到低排序，本轮只 unmask 最有把握的那一批
        #      每轮 unmask 比例随步数递增（第 s 轮 unmask 约 s/S_draft 的 token）
        n_to_unmask = compute_unmask_count(s, S_draft, n_masked=count_masks(x_current))
        unmask_indices = top_k_indices(confidence, k=n_to_unmask,
                                       only_masked_positions=True)

        # A-4. 把选中位置的 assessments 替换为预测 token，其余保持 assessments
        x_next = x_current.copy()
        for idx in unmask_indices:
            x_next[idx] = x_predicted[idx]

        x_current = x_next
        trajectory_history.append(x_current.copy())   # 记录 x^(s+1)

    # 草稿完成，x_current 此时应没有 assessments（全部填满）
    x_draft = x_current.copy()   # x^(S_draft)

    # ══════════════════════════════════════════════════════
    # 阶段 B · AutoEdit 修改（Token-to-Token Rewriting）
    #   从完整草稿出发，直接 token→token 重写低置信度位置
    #   注意：不再引入 assessments，输入始终是完整的具体 token 序列
    # ══════════════════════════════════════════════════════
    for k in range(S_edit):

        # B-1. 模型对当前完整 token 序列预测替换 token
        #      输入：当前具体 token 序列（无 assessments）+ 多模态上下文
        logits      = model(x_current, context)     # shape: [L, vocab_size]
        probs       = softmax(logits, dim=-1)
        x_hat       = argmax(probs, dim=-1)         # 每个位置的最优替换 token
        confidence  = [probs[i][x_hat[i]] for i in range(L)]

        # B-2. 构造 commit mask：选出置信度低的非目标点 token
        #      论文公式 (10)：x^(k+1) = m^(k) ⊙ x̂^(k+1) + (1-m^(k)) ⊙ x^(k)
        #      目标点 token 永远不进入候选（行为锚点锁定）
        commit_mask = compute_commit_mask(
            confidence    = confidence,
            x_current     = x_current,
            goal_positions = get_goal_positions(),   # 目标点 token 的位置索引
            threshold      = CONFIDENCE_THRESHOLD    # 低于此值才会被替换
        )   # shape: [L]，1=替换，0=保留

        # B-3. 按 commit mask 执行原地替换
        x_next = [
            x_hat[i] if commit_mask[i] == 1 else x_current[i]
            for i in range(L)
        ]

        x_current = x_next
        trajectory_history.append(x_current.copy())   # 记录 x^(S_draft + k + 1)

    # 最终轨迹（连续坐标）送入闭环仿真器求奖励
    final_trajectory = detokenize(x_current)          # τg

    # trajectory_history 包含 S_draft + S_edit + 1 个快照
    # 论文公式 (16)：ρg = (x⁰g, x¹g, ..., x^(Sdraft+Sedit)g)
    return trajectory_history, final_trajectory


# ----------------------------------------------------------
# Step 2 · 对一个场景采样 G 条候选 rollout
#          论文公式 (16)(17)(18)
# ----------------------------------------------------------

def sample_group_rollouts(model, context, N_g=3, I=2,
                          S_draft=3, S_edit=3):
    """
    对当前场景：
      1. 采样 Ng 个目标点（top-k + NMS）
      2. 每个目标点生成 I 条草稿
      3. 每条草稿跑完完整的 AutoEdit rollout
      4. 把最终轨迹送入仿真器，得到 G 个驾驶奖励

    Args:
        model   : 旧策略模型 πθ_old（rollout 期间完全冻结）
        context : 多模态上下文 c
        N_g     : 目标点采样数（论文实验值 3）
        I       : 每个目标点的草稿数（论文实验值 2）

    Returns:
        group : 长度 G = Ng×I 的列表，每项包含：
                  - history       : 完整 token 转换序列 ρg
                  - final_traj    : 最终连续坐标轨迹 τg
                  - reward        : 闭环驾驶奖励 R(τg)
                  - advantage     : 组相对优势 Ag（统一计算后填入）
    """

    G = N_g * I
    group = []

    # ── 2-a. 预测目标点后验分布，top-k 采样 + NMS ────────────
    goal_logits = model.goal_head(context)        # 目标点概率分布
    goal_tokens = top_k_nms_sample(goal_logits,
                                   k=N_g,
                                   nms_threshold=1.2)   # NMS 阈值约 1.2m

    # ── 2-b. 对每个目标点生成 I 条 rollout ───────────────────
    for goal_token in goal_tokens:          # 外层：Ng 个目标点
        for _ in range(I):                  # 内层：每个目标点 I 条草稿

            history, final_traj = run_single_rollout(
                model, goal_token, context, S_draft, S_edit
            )

            # 送入闭环仿真器（NAVSIM）计算 PDMS 奖励
            # PDMS = 加权综合(NC, DAC, TTC, 舒适性, EP)
            reward = closed_loop_simulator.score(final_traj)   # R(τg)

            group.append({
                'history'    : history,      # ρg：[x⁰, x¹, ..., x^S]
                'final_traj' : final_traj,   # τg
                'reward'     : reward,       # R(τg)
                'goal_token' : goal_token,
            })

    # ── 2-c. 计算组相对优势（Group-Relative Advantage）────────
    #         论文公式 (18)：Ag = R(τg) - 1/G · Σj R(τj)
    #
    #         直觉：比组内平均分高 → 正优势（这条路走对了，强化它）
    #               比组内平均分低 → 负优势（这条路走错了，抑制它）
    mean_reward = sum(item['reward'] for item in group) / G

    for item in group:
        item['advantage'] = item['reward'] - mean_reward  # Ag

    return group   # 长度 G 的候选列表，每项含 ρg, τg, R(τg), Ag


# ----------------------------------------------------------
# Step 3 · 计算单条 rollout 的策略梯度损失
#          论文公式 (2)(19)
# ----------------------------------------------------------

def compute_pg_loss_single(model, model_old, item, epsilon=0.2):
    """
    对一条候选 rollout 计算 PPO-clip 风格的策略梯度损失。

    核心设计：
      · 只有"在这一步真正发生变化"的 token 才获得梯度信号
        （草稿阶段：从 assessments 变成具体 token 的位置）
        （AutoEdit 阶段：被 commit mask 选中并替换的位置）
      · 使用重要性采样比 r = πθ / πθ_old 做 off-policy 修正
      · clip 截断防止更新步幅过大（PPO 技巧）

    论文公式 (2) 中的指示函数：
      δ^s_{g,p} = 1{ x^(s+1)_{g,p} ≠ x^s_{g,p} }
      只有 token 在第 s 步发生了变化，才把这步的梯度纳入损失。

    Args:
        model     : 当前策略 πθ（需要更新）
        model_old : 旧策略 πθ_old（rollout 时冻结，用于计算比值）
        item      : 单条 rollout 信息（含 history, advantage）
        epsilon   : PPO clip 系数

    Returns:
        loss_pg : 该条 rollout 的策略梯度损失（标量）
    """

    history   = item['history']    # [x⁰, x¹, ..., x^S]  共 S+1 个快照
    Ag        = item['advantage']  # 组相对优势（标量）
    context   = item['context']
    S         = len(history) - 1   # 总步数 = S_draft + S_edit
    L         = len(history[0])    # token 序列长度 = 16

    loss_pg = 0.0
    n_credited = 0   # 记录实际获得梯度的 (步, 位置) 对数量

    for s in range(S):

        x_s      = history[s]       # 第 s 步的 token 状态（模型输入）
        x_s_next = history[s + 1]   # 第 s+1 步的 token 状态（目标）

        # 3-a. 用当前策略和旧策略分别计算每个位置的 token 概率
        logits_new = model(x_s, context)          # πθ 的输出
        logits_old = model_old(x_s, context)      # πθ_old 的输出（不求梯度）

        probs_new = softmax(logits_new, dim=-1)   # shape: [L, vocab_size]
        probs_old = softmax(logits_old, dim=-1)   # shape: [L, vocab_size]

        for p in range(L):

            # 3-b. 指示函数：只有这个位置在这一步发生了变化，才纳入梯度
            #      论文公式 (19)：δ^s_{g,p} = 1{ x^(s+1)_{g,p} ≠ x^s_{g,p} }
            #
            #      草稿阶段：assessments → 具体 token，发生变化 → δ=1
            #      AutoEdit：被替换的 token 位置，发生变化 → δ=1
            #      两个阶段没变化的位置：δ=0，跳过，不贡献梯度
            token_changed = (x_s_next[p] != x_s[p])
            if not token_changed:
                continue

            target_token = x_s_next[p]   # 这一步实际写入的 token

            # 3-c. 计算重要性采样比 r^s_{g,p} = πθ / πθ_old
            #      分子：当前策略对 target_token 的概率
            #      分母：旧策略对 target_token 的概率（采样时的概率）
            prob_new = probs_new[p][target_token]
            prob_old = probs_old[p][target_token]
            ratio    = prob_new / (prob_old + 1e-8)   # r^s_{g,p}

            # 3-d. PPO-clip：截断过大的更新步幅
            #      论文公式 (2)：min( r·Ag,  clip(r, 1-ε, 1+ε)·Ag )
            clipped_ratio = clip(ratio, 1 - epsilon, 1 + epsilon)
            objective     = min(ratio * Ag, clipped_ratio * Ag)

            loss_pg += -objective   # 最大化目标 = 最小化负目标
            n_credited += 1

    # 对所有 (步, 位置) 取平均
    if n_credited > 0:
        loss_pg /= n_credited

    return loss_pg


# ----------------------------------------------------------
# Step 4 · KL 散度惩罚项
#          论文公式 (2) 中的 λ_KL · DKL(πθ ‖ πref)
# ----------------------------------------------------------

def compute_kl_penalty(model, model_ref, x_masked, context):
    """
    防止 RL 微调跑偏太远，在 SFT 基础上保持合理的分布。

    πref：SFT 结束后冻结的参考策略（不随 RL 更新）
    πθ  ：当前正在更新的策略

    对一个随机掩码输入，计算两者输出分布的 KL 散度。

    Args:
        model     : 当前策略 πθ
        model_ref : SFT 后冻结的参考策略 πref
        x_masked  : 随机掩码后的 token 序列
        context   : 多模态上下文

    Returns:
        kl : KL 散度（标量）
    """

    probs_theta = softmax(model(x_masked, context),     dim=-1)  # πθ
    probs_ref   = softmax(model_ref(x_masked, context), dim=-1)  # πref（无梯度）

    # KL(πθ ‖ πref) = Σ πθ · log(πθ / πref)
    kl = 0.0
    L  = probs_theta.shape[0]

    for i in range(L):
        for v in range(VOCAB_SIZE):
            p = probs_theta[i][v]
            q = probs_ref[i][v]
            if p > 1e-9:
                kl += p * math.log(p / (q + 1e-9))

    return kl / L   # 对位置取平均


# ----------------------------------------------------------
# Step 5 · RL 微调主循环
#          论文公式 (2)
# ----------------------------------------------------------

def rl_finetuning(model_sft, dataloader,
                  N_g=3, I=2, S_draft=3, S_edit=3,
                  epsilon=0.2, lambda_kl=0.01,
                  n_epochs=1):
    """
    在 SFT 模型基础上做 RL 微调（Reinforcement Fine-Tuning, RFT）。

    Args:
        model_sft : SFT 阶段训练好的模型（作为起点和参考策略）
        dataloader: 驾驶场景数据集
        N_g       : 每场景目标点数
        I         : 每目标点草稿数
        S_draft   : 草稿阶段步数
        S_edit    : AutoEdit 阶段步数
        epsilon   : PPO clip 系数
        lambda_kl : KL 惩罚权重
        n_epochs  : RL 训练轮数
    """

    # ── 初始化 ──────────────────────────────────────────────
    model     = copy.deepcopy(model_sft)   # πθ：当前策略（持续更新）
    model_ref = copy.deepcopy(model_sft)   # πref：SFT 参考策略（永久冻结）
    freeze(model_ref)

    optimizer = AdamW(model.parameters(), lr=1e-6)

    for epoch in range(n_epochs):
        for scene in dataloader:

            context = scene['context']   # 多模态上下文 c

            # ══════════════════════════════════════════════
            # Phase 1 · Rollout（用旧策略采样，不求梯度）
            # ══════════════════════════════════════════════

            # 冻结当前策略作为 πθ_old，用于 rollout 和计算重要性采样比
            model_old = copy.deepcopy(model)
            freeze(model_old)

            # 采样 G = Ng × I 条完整 Draft→AutoEdit rollout
            # 并计算每条轨迹的驾驶奖励和组相对优势
            with no_grad():
                group = sample_group_rollouts(
                    model_old, context, N_g, I, S_draft, S_edit
                )
            # group 中每项已包含：history(ρg), final_traj(τg), reward(R), advantage(Ag)

            # ══════════════════════════════════════════════
            # Phase 2 · 计算总损失（需要梯度）
            # ══════════════════════════════════════════════

            total_loss = 0.0

            # 2-a. 对 G 条 rollout 分别计算策略梯度损失，取平均
            #      论文公式 (2) 的主项
            loss_pg_total = 0.0
            for item in group:
                item['context'] = context
                loss_pg_total  += compute_pg_loss_single(
                    model, model_old, item, epsilon
                )
            loss_pg = loss_pg_total / len(group)   # 对 G 条取平均

            # 2-b. KL 惩罚：防止策略偏离 SFT 太远
            x_sample = sample_masked_sequence(scene['trajectory_tokens'])
            loss_kl  = compute_kl_penalty(model, model_ref, x_sample, context)

            # 2-c. 合并
            #      L(θ) = L_PG + λ_KL · DKL(πθ ‖ πref)
            total_loss = loss_pg + lambda_kl * loss_kl

            # ══════════════════════════════════════════════
            # Phase 3 · 更新参数
            # ══════════════════════════════════════════════

            optimizer.zero_grad()
            total_loss.backward()
            clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            log(
                epoch       = epoch,
                loss_pg     = loss_pg,
                loss_kl     = loss_kl,
                total_loss  = total_loss,
                mean_reward = sum(item['reward'] for item in group) / len(group),
                autoedit_gain = compute_autoedit_gain(group),
                # autoedit_gain = 有 AutoEdit 的平均分 - 无 AutoEdit 的平均分
                # RL 训练成功的标志：这个值从 ~0.3 增长到 ~1.9
            )

    return model   # RL 微调完成，起草器和 AutoEdit 已联合优化

# ============================================================
#  Action-Expert FFN 实现伪代码
#  对应论文 Section 5（推理优化部分）
# ============================================================
#
#  核心思想：
#    轨迹 token 的语义空间极小（只是坐标格子编号），
#    不需要和语言/视觉 token 一样大的 FFN。
#    用一个更小的专用 FFN 处理轨迹 token，
#    在不损失（甚至提升）精度的前提下大幅降低计算量。
#
#  论文数值：
#    主 FFN   d_ffn = 4096（论文实验中约为 d_model 的 2 倍）
#    Action FFN d_ffn = 1024（压缩到 1/4）
#    加速比：2.6×
#    meanSADE：不降反升（说明大 FFN 对轨迹 token 反而是过拟合）
# ============================================================

import torch
import torch.nn as nn


# ----------------------------------------------------------
#  1. 标准 FFN（用于语言/视觉 token）
# ----------------------------------------------------------

class StandardFFN(nn.Module):
    """
    Transformer 主干的标准前馈网络。
    处理语言 token、视觉 token、导航指令 token、自车状态 token。

    维度：d_model → d_ffn(大) → d_model
    """

    def __init__(self, d_model=2048, d_ffn=4096, dropout=0.0):
        super().__init__()
        self.w1      = nn.Linear(d_model, d_ffn,   bias=False)
        self.w2      = nn.Linear(d_ffn,   d_model, bias=False)
        self.act     = nn.SiLU()   # 论文用 SwiGLU 变体，这里简化为 SiLU
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x: [batch, seq_len, d_model]
        return self.dropout(self.w2(self.act(self.w1(x))))


# ----------------------------------------------------------
#  2. Action-Expert FFN（专用于轨迹 token）
# ----------------------------------------------------------

class ActionExpertFFN(nn.Module):
    """
    轨迹 token 专用的小型前馈网络。
    只处理 16 个轨迹 token（8 个路径点 × 2 坐标）。

    维度：d_model → d_action_ffn(小) → d_model
    d_action_ffn = 1024（是标准 FFN 的 1/4）

    为什么更小也够用？
      轨迹 token 的 "语义" 极其简单：就是一个坐标格子的编号。
      它不需要理解"转弯"、"让行"这类高层语义（那是 Attention 的工作）。
      FFN 只需要做一个低维的坐标→隐空间的映射，1024 维绰绰有余。
    """

    def __init__(self, d_model=2048, d_action_ffn=1024, dropout=0.0):
        super().__init__()
        self.w1      = nn.Linear(d_model, d_action_ffn, bias=False)
        self.w2      = nn.Linear(d_action_ffn, d_model, bias=False)
        self.act     = nn.SiLU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x: [batch, n_action_tokens, d_model]  n_action_tokens ≤ 16
        return self.dropout(self.w2(self.act(self.w1(x))))


# ----------------------------------------------------------
#  3. 混合 FFN 层：按 token 类型路由
# ----------------------------------------------------------

class MixedExpertFFNLayer(nn.Module):
    """
    一个完整的 Transformer FFN 子层，内含两套 FFN：
      - StandardFFN    → 处理所有非轨迹 token
      - ActionExpertFFN→ 处理所有轨迹 token

    路由依据：token 位置索引（轨迹 token 始终在序列末尾固定位置）。

    序列结构（论文 Section 4.1）：
      [视觉 token (V)] [导航 token (N)] [自车状态 token (E)]
      [目标点 token (G)] [轨迹 token (A)]
                                         ↑
                               这 16 个位置走 Action FFN
    """

    def __init__(self,
                 d_model=2048,
                 d_ffn=4096,
                 d_action_ffn=1024,
                 n_action_tokens=16):
        super().__init__()

        self.standard_ffn     = StandardFFN(d_model, d_ffn)
        self.action_expert_ffn = ActionExpertFFN(d_model, d_action_ffn)
        self.n_action_tokens  = n_action_tokens

        # 轨迹 token 始终在序列末尾，位置固定
        # 这使得路由无需任何运行时判断，直接切片即可
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, x, action_token_mask=None):
        """
        Args:
            x                : [batch, seq_len, d_model]
                               完整序列（视觉+语言+目标点+轨迹 token 混在一起）
            action_token_mask: [batch, seq_len] bool tensor
                               True 表示该位置是轨迹 token
                               如果为 None，则默认末尾 n_action_tokens 个是轨迹 token

        Returns:
            out: [batch, seq_len, d_model]
        """

        batch, seq_len, d = x.shape
        out = x.clone()   # 输出在原位修改

        # ── 构造路由 mask ─────────────────────────────────────
        if action_token_mask is None:
            # 默认：序列末尾 n_action_tokens 个位置是轨迹 token
            action_token_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
            action_token_mask[:, -self.n_action_tokens:] = True

        context_mask = ~action_token_mask   # 非轨迹 token 的 mask

        # ── 路由 1：非轨迹 token → StandardFFN ───────────────
        # 把所有非轨迹位置的 token 聚合成一个 batch 送入大 FFN
        # 用 mask 索引实现，避免 for 循环
        if context_mask.any():
            x_context = x[context_mask]        # [n_context, d_model]
            x_context = self.layer_norm(x_context)
            out_context = self.standard_ffn(x_context)   # [n_context, d_model]
            out[context_mask] = out_context

        # ── 路由 2：轨迹 token → ActionExpertFFN ─────────────
        if action_token_mask.any():
            x_action = x[action_token_mask]    # [n_action, d_model]
                                               # n_action = batch × 16（或更少）
            x_action = self.layer_norm(x_action)
            out_action = self.action_expert_ffn(x_action)  # [n_action, d_model]
            out[action_token_mask] = out_action

        return out


# ----------------------------------------------------------
#  4. 将其嵌入完整 Transformer Block
# ----------------------------------------------------------

class TransformerBlockWithActionExpert(nn.Module):
    """
    完整的 Transformer 层，Attention 部分不变，
    FFN 部分替换为 MixedExpertFFNLayer。

    注意：Attention 对所有 token 统一计算（轨迹 token 需要
    "看"场景上下文才能生成正确坐标），只有 FFN 做了分路由。
    """

    def __init__(self, d_model=2048, n_heads=16,
                 d_ffn=4096, d_action_ffn=1024):
        super().__init__()

        # Attention 部分：所有 token 统一处理（双向注意力）
        self.attn     = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm1    = nn.LayerNorm(d_model)

        # FFN 部分：按 token 类型路由
        self.ffn      = MixedExpertFFNLayer(d_model, d_ffn, d_action_ffn)
        self.norm2    = nn.LayerNorm(d_model)

    def forward(self, x, attn_mask=None, action_token_mask=None):

        # ── Attention 子层（Pre-Norm + 残差）────────────────
        x_norm = self.norm1(x)
        attn_out, _ = self.attn(x_norm, x_norm, x_norm, attn_mask=attn_mask)
        x = x + attn_out                   # 残差连接

        # ── 混合 FFN 子层（Post-Norm + 残差）────────────────
        ffn_out = self.ffn(self.norm2(x), action_token_mask)
        x = x + ffn_out                    # 残差连接

        return x


# ----------------------------------------------------------
#  5. 参数量 & 计算量对比（以单层 FFN 为例）
# ----------------------------------------------------------

def compare_params_and_flops(d_model=2048, d_ffn=4096, d_action_ffn=1024,
                              n_context_tokens=512, n_action_tokens=16):
    """
    对比标准方案（所有 token 用大 FFN）
    vs   Action-Expert 方案（轨迹 token 用小 FFN）
    的参数量和 FLOPs 差异。
    """

    # ── 参数量 ────────────────────────────────────────────
    # 标准 FFN：W1(d_model→d_ffn) + W2(d_ffn→d_model)
    params_standard     = 2 * d_model * d_ffn          # 2 × 2048 × 4096

    # Action-Expert 方案：大 FFN + 小 Action FFN
    params_action_expert = (2 * d_model * d_ffn          # 大 FFN（不变）
                           + 2 * d_model * d_action_ffn)  # 小 Action FFN（新增）

    # 新增参数量很少：2 × 2048 × 1024 = 4M，相对于主干几乎可忽略

    # ── FLOPs（以 batch_size=1 为例）────────────────────────
    n_total = n_context_tokens + n_action_tokens

    # 标准方案：所有 token 都走大 FFN
    flops_standard = n_total * 2 * d_model * d_ffn * 2
    #                ↑token数  ↑两层线性  ↑每次乘加2op

    # Action-Expert 方案：
    flops_context = n_context_tokens * 2 * d_model * d_ffn   * 2
    flops_action  = n_action_tokens  * 2 * d_model * d_action_ffn * 2
    flops_expert  = flops_context + flops_action

    speedup = flops_standard / flops_expert

    print(f"标准方案 参数量：{params_standard/1e6:.1f}M")
    print(f"Expert 方案 参数量：{params_action_expert/1e6:.1f}M（+{(params_action_expert-params_standard)/1e6:.1f}M）")
    print()
    print(f"标准方案 FLOPs：{flops_standard/1e9:.2f}G")
    print(f"Expert 方案 FLOPs：{flops_expert/1e9:.2f}G")
    print(f"理论加速比：{speedup:.2f}×")
    print()
    print("直觉：轨迹 token 只占序列的 16/(512+16) ≈ 3%，")
    print("      但 Action FFN 是大 FFN 的 1/4，")
    print("      所以整体 FFN FLOPs 节省约 3% × 75% ≈ 2.25%。")
    print("      论文报告的 2.6× 是针对 Action FFN 模块本身的加速，")
    print("      不是全模型加速。")

    # 实际上 2.6× 指的是 Action FFN 这个操作本身的延迟下降：
    # 原来每个 token 都调用大 FFN kernel，现在轨迹 token 调用小 FFN kernel
    # kernel launch + 矩阵乘法的延迟同时下降 → 2.6×


# ----------------------------------------------------------
#  6. 为什么精度不降反升？
# ----------------------------------------------------------
#
#  论文里提到一个反直觉的现象：
#    Action FFN 从 d_ffn=4096 压缩到 d_action_ffn=1024 之后，
#    meanSADE（平均 SAD 误差，越小越好）反而变好了。
#
#  可能的解释：
#
#  ① 过参数化导致"记忆"而非"泛化"
#     大 FFN 有足够的容量把训练集里每条轨迹"背下来"，
#     小 FFN 容量受限，被迫学习更通用的坐标变换规律。
#
#  ② 正则化效应
#     小 FFN 相当于对轨迹 token 的表示施加了隐式的信息瓶颈，
#     逼着模型把高层语义（该往哪走）完全依赖 Attention 解决，
#     而 FFN 只做纯粹的坐标数值变换。
#     两者分工更清晰，反而更健康。
#
#  ③ 梯度干净
#     大 FFN 里，语言和轨迹 token 共享参数，梯度相互干扰。
#     分离之后，Action FFN 的梯度只来自轨迹相关的损失，
#     更新方向更纯粹。

# ============================================================
#  KV 缓存回退 + 合并重写 伪代码
#  对应论文 Section 5（推理优化部分）
# ============================================================
#
#  问题背景：
#    掩码扩散用双向注意力，每步有若干 token 从 assessments 变成具体值。
#    朴素做法：每步把全序列 K/V 从头算一遍 → 14.7ms/step
#    合并重写：只重算发生变化的位置 → 11.5ms/step，加速 1.28×
#
#  核心近似：
#    当步骤 s 中只有少数位置 Δs 发生变化时，
#    其他位置的 K/V 在 s 和 s+1 步之间变化很小，
#    可以直接复用上一步的缓存值。
# ============================================================

import torch


# ----------------------------------------------------------
#  数据结构：多层 KV 缓存
# ----------------------------------------------------------

class KVCache:
    """
    存储所有 Transformer 层的 K 和 V 张量。

    shape: [n_layers, batch, n_heads, seq_len, head_dim]

    同时维护一个"版本戳"，记录每个位置的 token 在哪一步最后更新，
    用于 AutoEdit 的回退操作。
    """

    def __init__(self, n_layers, batch, n_heads, seq_len, head_dim):
        self.K = torch.zeros(n_layers, batch, n_heads, seq_len, head_dim)
        self.V = torch.zeros(n_layers, batch, n_heads, seq_len, head_dim)

        # 版本戳：position_version[i] = 该位置的 K/V 在第几步被算出来
        self.position_version = torch.full((seq_len,), -1, dtype=torch.long)

        # 快照栈：用于 AutoEdit 阶段的回退（rollback）
        # 每次"保存快照"就往栈里压一帧，回退就弹出
        self._snapshots = []

    def update_positions(self, changed_positions, new_K, new_V, step):
        """
        合并重写：只更新发生变化的位置，其余保持不变。

        Args:
            changed_positions : List[int]，本步发生变化的 token 位置索引
            new_K             : [n_layers, batch, n_heads, len(changed), head_dim]
            new_V             : [n_layers, batch, n_heads, len(changed), head_dim]
            step              : 当前步骤编号
        """
        for k, pos in enumerate(changed_positions):
            self.K[:, :, :, pos, :] = new_K[:, :, :, k, :]
            self.V[:, :, :, pos, :] = new_V[:, :, :, k, :]
            self.position_version[pos] = step   # 记录版本

    def save_snapshot(self):
        """保存当前完整 KV 状态（用于 AutoEdit 回退）。"""
        self._snapshots.append({
            'K': self.K.clone(),
            'V': self.V.clone(),
            'version': self.position_version.clone()
        })

    def rollback(self):
        """
        回退到上一个快照。
        在 AutoEdit 阶段，如果当前修改使得奖励变差，
        可以撤销这一轮的 K/V 更新，重新选择要修改的位置。
        """
        assert len(self._snapshots) > 0, "没有可回退的快照"
        snap = self._snapshots.pop()
        self.K                 = snap['K']
        self.V                 = snap['V']
        self.position_version  = snap['version']


# ----------------------------------------------------------
#  核心函数：计算变化位置的新 K/V
# ----------------------------------------------------------

def recompute_kv_for_changed_positions(model, token_sequence,
                                       changed_positions, kv_cache):
    """
    只对 changed_positions 上的 token 重新计算 K 和 V，
    其余位置直接复用 kv_cache 中已有的值。

    这就是"合并重写"的实质：
      - 旧缓存（不变位置）：直接读取，零计算
      - 新计算（变化位置）：只算这几个 token，写回缓存

    Args:
        model             : Transformer 模型
        token_sequence    : 当前完整 token 序列 [batch, seq_len]
        changed_positions : 本步发生变化的位置索引列表
        kv_cache          : 当前 KV 缓存对象

    Returns:
        updated_kv_cache  : 更新后的 KV 缓存
    """

    if len(changed_positions) == 0:
        return kv_cache   # 没有变化，直接返回

    # ── Step 1：只取出变化位置的 token embedding ──────────
    #    不需要对整个序列做 embedding，只做几个 token
    changed_token_ids = token_sequence[:, changed_positions]  # [batch, n_changed]
    changed_embeddings = model.embed(changed_token_ids)        # [batch, n_changed, d_model]

    # ── Step 2：逐层计算这些位置的新 K 和 V ───────────────
    #
    #    关键近似在这里：
    #    每层的输入 h_i^(l-1) 理论上依赖上层所有 token 的注意力结果。
    #    近似处理：对变化位置，用当前层已有的完整 K/V 做注意力，
    #              得到这些位置的新隐状态，再算新的 K/V。
    #    对不变位置：直接复用缓存，不做任何计算。
    #
    new_K_all_layers = []
    new_V_all_layers = []

    h = changed_embeddings   # [batch, n_changed, d_model]  初始为 embedding

    for layer_idx, layer in enumerate(model.layers):

        # 2-a. 用当前层已缓存的完整 K/V 做 cross-attention：
        #      变化位置的 query 去 attend 整个序列（包括不变位置的缓存 K/V）
        #      这样变化位置的隐状态就能看到完整的上下文
        K_full = kv_cache.K[layer_idx]   # [batch, heads, seq_len, head_dim]
        V_full = kv_cache.V[layer_idx]   # [batch, heads, seq_len, head_dim]

        # 只对变化位置做 attention（query 来自变化位置，key/value 来自完整缓存）
        Q_changed = layer.W_Q(h)   # [batch, n_changed, d_model]
        h = layer.cross_attend(Q_changed, K_full, V_full)  # [batch, n_changed, d_model]
        h = layer.ffn(h)                                   # [batch, n_changed, d_model]

        # 2-b. 用更新后的隐状态计算这些位置的新 K 和 V
        new_K = layer.W_K(h)   # [batch, n_changed, d_model] → reshape to [batch, heads, n_changed, head_dim]
        new_V = layer.W_V(h)

        new_K_all_layers.append(new_K)
        new_V_all_layers.append(new_V)

    # ── Step 3：合并写回缓存（只写变化的位置）─────────────
    #    "合并"：新算的几个位置 + 原来缓存的其余位置 = 完整 KV 缓存
    new_K_tensor = torch.stack(new_K_all_layers)   # [n_layers, batch, heads, n_changed, head_dim]
    new_V_tensor = torch.stack(new_V_all_layers)

    kv_cache.update_positions(changed_positions, new_K_tensor, new_V_tensor,
                              step=current_step)

    return kv_cache


# ----------------------------------------------------------
#  草稿阶段：带合并重写的掩码扩散解码
# ----------------------------------------------------------

def draft_with_kv_merge_rewrite(model, context_tokens, S_draft=3):
    """
    草稿阶段的推理循环。
    每步只 unmask 少数 token，用合并重写更新 KV 缓存。

    Args:
        context_tokens : 场景上下文 token（视觉/导航/状态）[batch, n_ctx]
        S_draft        : 草稿阶段总步数

    Returns:
        draft_tokens   : 填满的轨迹 token 序列 [batch, L]
        kv_cache       : 最终 KV 缓存（供 AutoEdit 阶段复用）
    """

    L = 16   # 轨迹 token 数量
    batch = context_tokens.shape[0]

    # ── 初始化：轨迹部分全 assessments，和上下文拼在一起 ────────
    traj_tokens = torch.full((batch, L), MASK_TOKEN_ID)
    full_sequence = torch.cat([context_tokens, traj_tokens], dim=1)
    # full_sequence: [batch, n_ctx + L]

    # ── 计算上下文前缀的 KV 缓存（只算一次，整个推理过程复用）
    #    上下文 token 用因果注意力，所以这部分可以完美缓存
    kv_cache = compute_prefix_kv(model, context_tokens)   # 见论文"共享前缀 KV 缓存"

    # ── 初始化轨迹 token 的 KV（全 assessments 的 embedding）─────
    mask_embedding = model.embed(traj_tokens)   # [batch, L, d_model]
    initial_K, initial_V = model.compute_kv(mask_embedding)
    kv_cache.update_positions(list(range(L)), initial_K, initial_V, step=0)

    # ── 草稿循环 ─────────────────────────────────────────────
    unmasked_positions = set()   # 已经填好的位置

    for s in range(S_draft):

        # 用当前完整 KV 缓存做前向，得到所有位置的 logits
        # （只对仍是 assessments 的位置采样，已填好的位置直接保留）
        logits = model.forward_with_kv(full_sequence, kv_cache)
        # logits: [batch, n_ctx + L, vocab_size]

        traj_logits = logits[:, -L:, :]   # 只取轨迹部分的 logits
        probs = softmax(traj_logits, dim=-1)

        # 选出本步要填入的位置（最高置信度的那批 assessments）
        n_to_unmask = compute_unmask_count(s, S_draft, n_masked=(L - len(unmasked_positions)))
        still_masked = [i for i in range(L) if i not in unmasked_positions]

        confidences = {i: probs[0, i, :].max().item() for i in still_masked}
        newly_unmasked = sorted(still_masked,
                                key=lambda i: -confidences[i])[:n_to_unmask]

        # 采样并填入
        for pos in newly_unmasked:
            sampled_token = sample_from(probs[0, pos])
            full_sequence[0, -L + pos] = sampled_token
            unmasked_positions.add(pos)

        # ════════════════════════════════════════════════
        #  合并重写：只对新填入的位置重算 K/V
        #  之前已 unmask 的位置 + 仍是 assessments 的位置 → 直接复用缓存
        # ════════════════════════════════════════════════
        kv_cache = recompute_kv_for_changed_positions(
            model          = model,
            token_sequence = full_sequence,
            changed_positions = [n_ctx + pos for pos in newly_unmasked],
            #                    ↑ 转换为全序列索引（轨迹在末尾）
            kv_cache       = kv_cache
        )
        # 本步只重算了 n_to_unmask 个位置的 K/V，其余全部复用！

    draft_tokens = full_sequence[:, -L:]   # 提取轨迹 token
    return draft_tokens, kv_cache


# ----------------------------------------------------------
#  AutoEdit 阶段：带回退 + 合并重写的轨迹修正
# ----------------------------------------------------------

def autoedit_with_rollback(model, draft_tokens, context_tokens,
                           kv_cache, S_edit=3):
    """
    AutoEdit 阶段的推理循环。
    比草稿阶段多了一个"回退"能力：
    如果某次修改后置信度反而更低，可以撤销这次修改。

    Args:
        draft_tokens   : 草稿轨迹 token [batch, L]
        context_tokens : 场景上下文 token [batch, n_ctx]
        kv_cache       : 草稿阶段结束时的 KV 缓存（直接继承，无需重算上下文！）
        S_edit         : AutoEdit 迭代轮数

    Returns:
        edited_tokens  : 修正后的轨迹 token [batch, L]
    """

    L = 16
    full_sequence = torch.cat([context_tokens, draft_tokens], dim=1)

    for k in range(S_edit):

        # ── 保存当前快照，以备回退 ────────────────────────
        kv_cache.save_snapshot()
        prev_sequence = full_sequence.clone()

        # ── 用当前完整 KV 做前向，得到每个位置的替换建议 ──
        logits = model.forward_with_kv(full_sequence, kv_cache)
        traj_logits = logits[:, -L:, :]
        probs  = softmax(traj_logits, dim=-1)

        # 每个位置的"替换 token"和"置信度"
        x_hat       = probs.argmax(dim=-1)   # [batch, L]
        confidence  = probs.max(dim=-1).values   # [batch, L]

        # ── 构造 commit mask：选低置信度的非目标点位置 ────
        goal_positions = get_goal_token_positions()
        commit_mask = torch.zeros(L, dtype=torch.bool)

        for i in range(L):
            if i in goal_positions:
                commit_mask[i] = False   # 目标点永远不替换
            elif confidence[0, i] < CONFIDENCE_THRESHOLD:
                commit_mask[i] = True    # 置信度低 → 标记为替换

        changed_positions = commit_mask.nonzero().squeeze(-1).tolist()

        if len(changed_positions) == 0:
            # 没有低置信度的位置，不需要修改，丢弃快照直接退出
            kv_cache._snapshots.pop()
            break

        # ── 执行替换 ──────────────────────────────────────
        for pos in changed_positions:
            full_sequence[0, -L + pos] = x_hat[0, pos]

        # ── 合并重写：只更新被替换位置的 K/V ─────────────
        kv_cache = recompute_kv_for_changed_positions(
            model             = model,
            token_sequence    = full_sequence,
            changed_positions = [-L + pos for pos in changed_positions],
            kv_cache          = kv_cache
        )

        # ── 可选的质量检查：如果修改后置信度整体更低，回退 ─
        new_logits     = model.forward_with_kv(full_sequence, kv_cache)
        new_confidence = softmax(new_logits[:, -L:, :], dim=-1).max(dim=-1).values.mean()
        old_confidence = confidence.mean()

        if new_confidence < old_confidence * 0.95:
            # 修改让模型"更不确定"了，撤销这一轮
            kv_cache.rollback()
            full_sequence = prev_sequence
        else:
            # 修改有效，丢弃快照（不再需要回退到这个点）
            kv_cache._snapshots.pop()

    edited_tokens = full_sequence[:, -L:]
    return edited_tokens


# ----------------------------------------------------------
#  合并重写 vs 朴素重算：计算量对比
# ----------------------------------------------------------
#
#  设：
#    seq_len      = n_ctx + L = 512 + 16 = 528
#    每步变化位置  = n_changed（草稿阶段约 5~6，AutoEdit 阶段约 2~3）
#    n_layers     = 28
#
#  朴素重算（每步）：
#    对完整序列重算所有 K/V
#    FLOPs ∝ n_layers × seq_len × d_model²
#          = 28 × 528 × 2048²  ≈ 124G FLOPs
#
#  合并重写（每步）：
#    只对 n_changed 个位置重算 K/V
#    FLOPs ∝ n_layers × n_changed × d_model²
#          = 28 × 6 × 2048²  ≈ 1.4G FLOPs
#    （其余 522 个位置直接读缓存，零 FLOPs）
#
#  理论加速：124G / 1.4G ≈ 88×
#
#  实际加速 1.28×（14.7ms → 11.5ms）的原因：
#    · GPU kernel 启动开销、内存带宽瓶颈等 overhead 不随 FLOPs 线性缩减
#    · KV 缓存的读写本身也有延迟
#    · 合并重写引入了额外的 mask 判断和条件写操作
#    但在绝对延迟上节省了 3.2ms/帧，对 31.8ms 的总延迟已是可观的贡献