class MotionTokenizer:
    def __init__(self, vocab_size, grid_range):
        self.vocab_size = vocab_size      # 例如 128x128 个网格类别 + 特殊 Token
        self.grid_range = grid_range      # 物理位移极限，例如 [-18m, 18m]
        self.ZERO_ACTION_TOKEN = 0        # 特殊单词：代表"保持匀速直线运动"

    def encode(self, continuous_traj):
        """
        输入: continuous_traj 形状[Agent数量, 时间步T, 2(x,y)]
        输出: discrete_tokens 形状 [Agent数量, 时间步T]
        """
        tokens = []
        # 计算每一步的位移变化 (Delta x, Delta y)
        displacements = compute_diff(continuous_traj)

        for t in range(T):
            if t > 0 and is_almost_equal(displacements[:, t], displacements[:, t-1]):
                # 核心魔法：如果当前速度/位移跟上一步一样，直接输出"零动作"单词
                token = self.ZERO_ACTION_TOKEN
            else:
                # 否则，将 映射到均匀划分的离散网格中，得到对应的类别ID
                token = self.quantize_to_grid(displacements[:, t], self.grid_range)
            tokens.append(token)

        return torch.stack(tokens)

class MotionLM(nn.Module):
    def __init__(self, num_agents, vocab_size, embed_dim):
        super().__init__()
        self.num_agents = num_agents

        # 1. 场景编码器 (借用 Wayformer 的 Early Fusion 架构)
        self.scene_encoder = WayformerEncoder(embed_dim)

        # 2. 各种 Embedding 层 (为"单词"赋予意义)
        self.token_embed = nn.Embedding(vocab_size, embed_dim)   # 动作值编码
        self.time_embed = nn.Embedding(MAX_TIME_STEPS, embed_dim) # 时间位置编码
        self.agent_embed = nn.Embedding(num_agents, embed_dim)    # 智能体身份编码

        # 3. 标准的 Transformer 解码器
        self.transformer_decoder = nn.TransformerDecoder(
            decoder_layer=nn.TransformerDecoderLayer(d_model=embed_dim, nhead=8),
            num_layers=6
        )

        # 4. 输出头：预测下一个单词的概率分布
        self.lm_head = nn.Linear(embed_dim, vocab_size)

    def forward(self, map_data, traffic_lights, history_traj, target_tokens):
        B, N, T = target_tokens.shape

        # 第一步：察言观色
        scene_memory = self.scene_encoder(map_data, traffic_lights, history_traj)

        # 第二步：准备要接龙的"单词"序列（三种 Embedding 逐元素相加）
        val_emb = self.token_embed(target_tokens)
        time_emb = self.time_embed(torch.arange(T))
        agent_emb = self.agent_embed(torch.arange(N))
        combined_emb = val_emb + time_emb + agent_emb

        # 第三步：降维打击 (Flattening) - 把所有智能体在所有时间的 Token 拉平
        flattened_sequence = combined_emb.view(B, N * T, -1)

        # 第四步：因果掩码（确保 t 时刻的预测只能看到 t-1 及以前的所有人的动作）
        causal_mask = generate_agent_time_causal_mask(N, T)

        # 第五步：Transformer 解码
        decoder_out = self.transformer_decoder(
            tgt=flattened_sequence,
            memory=scene_memory,
            tgt_mask=causal_mask
        )

        # 第六步：输出预测
        logits = self.lm_head(decoder_out)
        return logits

def train_step(model, batch_data, optimizer):
    model.train()
    optimizer.zero_grad()

    # 提取并分词真实未来的轨迹
    ground_truth_traj = batch_data['future_traj']
    tokenizer = MotionTokenizer()
    target_tokens = tokenizer.encode(ground_truth_traj)

    # 前向传播 (Teacher Forcing 模式)
    input_tokens = shift_right(target_tokens, pad_value='<BOS>')
    logits = model(batch_data['map'], batch_data['tl'], batch_data['hist'], input_tokens)

    # 【极致的极简主义】
    # 抛弃 Huber, 抛弃 L2 距离，抛弃一切复杂的轨迹 loss
    # 只有最经典的交叉熵损失
    loss_fn = nn.CrossEntropyLoss()
    loss = loss_fn(logits.view(-1, VOCAB_SIZE), target_tokens.view(-1))

    loss.backward()
    optimizer.step()

    return loss.item()

@torch.no_grad()
def inference(model, scene_data, num_rollouts=512, top_k=6):
    model.eval()

    # 1. 编码场景
    scene_memory = model.scene_encoder(...)
    scene_memory = scene_memory.repeat(num_rollouts, 1, 1)  # 复制 512 份

    # 2. 自回归解码
    generated_tokens = torch.full((num_rollouts, N, 1), '<BOS>')

    for t in range(MAX_TIME_STEPS):
        logits = model(scene_memory, generated_tokens)
        next_step_logits = logits[:, :, -1, :]  # [Rollouts, N, Vocab]

        # 按概率分布采样（保证多样性）
        probs = torch.softmax(next_step_logits, dim=-1)
        next_tokens = torch.multinomial(probs.view(-1, VOCAB_SIZE), 1)
        next_tokens = next_tokens.view(num_rollouts, N, 1)
        generated_tokens = torch.cat([generated_tokens, next_tokens], dim=2)

    # 3. 翻译回坐标系
    tokenizer = MotionTokenizer()
    continuous_rollouts = tokenizer.decode(generated_tokens)

    # 4. 后处理提炼
    final_trajs, final_probs = aggregate_rollouts(continuous_rollouts, top_k)
    return final_trajs, final_probs

def aggregate_rollouts(rollouts, k=6):
    """使用 NMS 和 K-Means 聚类"""
    distance_matrix = compute_pairwise_distances(rollouts)
    filtered_rollouts = apply_nms(rollouts, distance_matrix, threshold=2.0)
    kmeans = KMeans(n_clusters=k)
    cluster_centers = kmeans.fit(filtered_rollouts)
    cluster_probs = compute_cluster_probabilities(kmeans.labels_)
    return cluster_centers, cluster_probs