modules/octree_codec.py

"""fix3 tokenizer(已融合 numba 加速实现, 原 *_fast.py 并入本文件)。

encode_mesh / decode_tokens 走 numpy 向量化 + numba 枚举, 比纯 python 版快约 5x。
quantized_mesh 与 decode 输出和纯 python 版完全一致; token 的 face 块集合一致,
仅少数平局块顺序不同(本实现补了 (f2.y,f2.x) 决定性平局键, 比原版更可复现)。

原文件头:

fix3: 连接矩阵改为「octant 绝对原位 8x8 + 因子分解(按行)」编码，保留 F。

相对 fix2 的核心变化
--------------------
1. 不再按 F 排序、不再受 max_children 限制。每个连接矩阵(intra / inter)统一是
   8x8 的「octant 绝对原位」矩阵：行/列 = octant 0..7 的固定语义。
   → 天然 mc=8 无损，"超标 / 合并 / 退化分摊" 这套概念全部消失。
2. 8x8 矩阵不再编码成单个 token(那是 2^64 词表)，而是按行拆成 8 个「行子码」，
   每个行子码 = 该行 8 个 bit ∈ [0, 256)。intra / inter 各 8 个子码。
3. **逻辑序列长度不变**：每个 face 仍是 18 个「位置」(9 coord + 3 F + 3 intra + 3 inter)，
   只是 intra / inter 这 6 个位置每个携带 8 个子码(给因子分解头：同一位置 8 个并行 256-way 头)。
   coord / F 仍是单子码位置。tokenizer.encode 返回 List[List[int]]：每个内层 list = 一个位置的子码。
4. 保留 F(标识 active octant)。重建时用 F mask 掉非 active 行列(M &= F⊗F)，消除歧义。

词表(极小)：pad(1) + coord(resolution) + F(256) + intra_row(256) + inter_row(256)。
"""
import numpy as np
import trimesh
from tqdm import tqdm
import os
import sys
import glob
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from pathlib import Path
from utils.data_utils import undiscretize, discretize, normalize, clean_mesh
from utils.mesh_metric import compute_mesh_metrics
from modules.fast_octree_level import _triangles_csr
import itertools
from dataclasses import dataclass

NUM_OCTANTS = 8

# 预计算 24 个节点(3 槽位 x 8 octant)中选 3 个的所有组合，用于向量化找三角形
_TRI_INDICES = np.array(list(itertools.combinations(range(24), 3)))
_IDX_I = _TRI_INDICES[:, 0]
_IDX_J = _TRI_INDICES[:, 1]
_IDX_K = _TRI_INDICES[:, 2]


# ==========================================
# Tokenizer：octant 原位 + 按行子码
# ==========================================
class DynamicMeshTokenizer:
    """
    encode 输出格式(fine)：List[List[int]]，每个内层 list 是一个「逻辑位置」的子码。
      每 face 18 个位置： [v0x][v0y][v0z][v1x][v1y][v1z][v2x][v2y][v2z]  (9 个 coord, 各 1 子码)
                          [F0][F1][F2]                                   (3 个 F,     各 1 子码)
                          [M00 的 8 行][M11..][M22..]                    (3 个 intra, 各 8 子码)
                          [M01 的 8 行][M02..][M12..]                    (3 个 inter, 各 8 子码)
    coord / F 位置长度 1；intra / inter 位置长度 8(行子码)。逻辑序列长度 = 18/face，与 fix2 相同。
    """
    def __init__(self, resolution, max_children=NUM_OCTANTS):
        self.resolution = resolution
        self.max_children = NUM_OCTANTS  # octant 原位下固定为 8；保留字段仅为兼容 config

        self.pad_id = 0
        self.coord_offset = 1
        self.f_offset = self.coord_offset + self.resolution
        self.row_vocab = 1 << NUM_OCTANTS                       # 2^8 = 256
        self.f_vocab = 1 << NUM_OCTANTS                         # F 也是 8bit
        self.intra_row_offset = self.f_offset + self.f_vocab
        self.inter_row_offset = self.intra_row_offset + self.row_vocab
        self.vocab_size = self.inter_row_offset + self.row_vocab

    # ---- bit <-> int ----
    def _row_to_int(self, bits):
        v = 0
        for i, b in enumerate(bits):
            if b:
                v |= (1 << i)
        return v

    def _int_to_row(self, val):
        return [(val >> i) & 1 for i in range(NUM_OCTANTS)]

    def _mat_to_rows(self, M, offset):
        return [self._row_to_int(M[r]) + offset for r in range(NUM_OCTANTS)]

    def _rows_to_mat(self, rows, offset):
        return [self._int_to_row(rows[r] - offset) for r in range(NUM_OCTANTS)]

    # ------------------------------------
    # 细分网格 (Fine Mesh) —— 扁平 1D int 布局，每 face 60 个 token：
    #   9 coord + 3 F + 6 矩阵 x 8 行子码
    #   顺序: [coord(9)]  [F0 F1 F2]  [M00 M11 M22 各 8 行]  [M01 M02 M12 各 8 行]
    #   F: 每个 1 个 token；intra/inter: 每个矩阵按行拆成 8 个行子码(每行 8bit)
    #   所有真实 token 都带 offset(>=1)，0 仅作 trailing pad，所以 F 不需要复制避坑。
    # 与现有 dataset/collate(1D int, pad_id=0) 兼容；模型侧整理成 9 个 refine 位置/face
    #   (F 位置 1 子码、M 位置 8 行子码 sum 成一个 embedding token)，逻辑序列 18/face 不变。
    # ------------------------------------
    TOKENS_PER_FACE = 9 + 3 + 6 * NUM_OCTANTS   # 60
    REFINE_POS_PER_FACE = 9                       # 模型逻辑: 3 F + 3 intra + 3 inter

    def encode(self, faces_data, use_hourglass=True):
        seq = []
        for face in faces_data:
            v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12 = face
            for v in (v0, v1, v2):
                seq.extend([v[0] + self.coord_offset, v[1] + self.coord_offset, v[2] + self.coord_offset])
            for F in (F0, F1, F2):
                seq.append(self._row_to_int(F) + self.f_offset)        # F: 1 个 token
            for M in (M00, M11, M22):
                seq.extend(self._mat_to_rows(M, self.intra_row_offset))  # 各 8 行
            for M in (M01, M02, M12):
                seq.extend(self._mat_to_rows(M, self.inter_row_offset))  # 各 8 行
        return seq

    def decode(self, seq):
        faces_data = []
        B = self.TOKENS_PER_FACE  # 60
        clean_seq = [t for t in seq if t != self.pad_id]
        co = self.coord_offset
        for i in range(0, len(clean_seq) - B + 1, B):
            ch = clean_seq[i:i + B]
            cv = [ch[k] - co for k in range(9)]
            v0 = (cv[0], cv[1], cv[2]); v1 = (cv[3], cv[4], cv[5]); v2 = (cv[6], cv[7], cv[8])
            F0 = self._int_to_row(ch[9] - self.f_offset)
            F1 = self._int_to_row(ch[10] - self.f_offset)
            F2 = self._int_to_row(ch[11] - self.f_offset)
            M00 = self._rows_to_mat(ch[12:20], self.intra_row_offset)
            M11 = self._rows_to_mat(ch[20:28], self.intra_row_offset)
            M22 = self._rows_to_mat(ch[28:36], self.intra_row_offset)
            M01 = self._rows_to_mat(ch[36:44], self.inter_row_offset)
            M02 = self._rows_to_mat(ch[44:52], self.inter_row_offset)
            M12 = self._rows_to_mat(ch[52:60], self.inter_row_offset)
            faces_data.append((v0, v1, v2, F0, F1, F2, M00, M11, M22, M01, M02, M12))
        return faces_data

    # ------------------------------------
    # 粗糙网格 (Raw Mesh)：仅顶点坐标，每 face 9 个单子码位置
    # ------------------------------------
    def encode_raw(self, raw_faces_data, use_hourglass=True):
        seq = []
        for face in raw_faces_data:
            v0, v1, v2 = face
            for v in (v0, v1, v2):
                seq.append([v[0] + self.coord_offset])
                seq.append([v[1] + self.coord_offset])
                seq.append([v[2] + self.coord_offset])
        return seq

    def decode_raw(self, seq):
        faces_data = []
        clean_seq = [p for p in seq if not (len(p) == 1 and p[0] == self.pad_id)]
        for i in range(0, len(clean_seq) - 9 + 1, 9):
            chunk = clean_seq[i:i + 9]
            if len(chunk) < 9:
                break
            cv = [chunk[k][0] - self.coord_offset for k in range(9)]
            faces_data.append(((cv[0], cv[1], cv[2]), (cv[3], cv[4], cv[5]), (cv[6], cv[7], cv[8])))
        return faces_data


# ==========================================
# Builder：直接构建 octant 原位 8x8 矩阵
# ==========================================
class OctreeFaceTokenBuilder:
    def __init__(self, tree, max_children=NUM_OCTANTS):
        self.tree = tree
        self.max_children = NUM_OCTANTS  # octant 原位固定 8

    def get_level_info(self, depth):
        nodes_at_depth = {}

        def collect_nodes(current, target_depth):
            if current.depth == target_depth:
                nodes_at_depth[current.coord] = current
                return
            for child in current.children:
                if child:
                    collect_nodes(child, target_depth)

        collect_nodes(self.tree.root, depth)

        coords = sorted(nodes_at_depth.keys(), key=lambda x: (x[2], x[1], x[0]))
        coord_to_idx = {c: i for i, c in enumerate(coords)}

        adj = {i: set() for i in range(len(coords))}
        for i, c in enumerate(coords):
            node = nodes_at_depth[c]
            all_neighbors = node.intra_neighbors | node.inter_neighbors
            for neighbor_coord in all_neighbors:
                if neighbor_coord in coord_to_idx:
                    j = coord_to_idx[neighbor_coord]
                    if i < j:
                        adj[i].add(j)

        raw_faces_set = set()
        covered_edges = set()

        # 1. 3-clique (三角形)
        for u in range(len(coords)):
            for v in adj[u]:
                common_neighbors = adj[u] & adj[v]
                for w in common_neighbors:
                    face_coords = tuple(sorted([coords[u], coords[v], coords[w]]))
                    raw_faces_set.add(face_coords)
                    covered_edges.add(tuple(sorted([face_coords[0], face_coords[1]])))
                    covered_edges.add(tuple(sorted([face_coords[1], face_coords[2]])))
                    covered_edges.add(tuple(sorted([face_coords[0], face_coords[2]])))

        # 2. 未被三角形覆盖的独立边 -> 退化面片 (A, B, B)
        for u in range(len(coords)):
            for v in adj[u]:
                edge = tuple(sorted([coords[u], coords[v]]))
                if edge not in covered_edges:
                    raw_faces_set.add((edge[0], edge[1], edge[1]))

        # 3. 孤立顶点 -> (A, A, A)
        for u in range(len(coords)):
            has_edges = len(adj[u]) > 0 or any(u in adj[i] for i in range(u))
            if not has_edges:
                raw_faces_set.add((coords[u], coords[u], coords[u]))

        raw_faces_at_depth = list(raw_faces_set)
        raw_faces_at_depth.sort(key=lambda f: (f[0][2], f[0][1], f[0][0],
                                               f[1][2], f[1][1], f[1][0],
                                               f[2][2]))
        return raw_faces_at_depth, nodes_at_depth

    def _active_octants(self, coord, n, nodes_map):
        """该节点哪些 octant 子节点存在(occupancy)。
        用 children 占据而非"有连接"标记 —— 因为「有连接 ⟹ 必有子节点」，occupancy 是
        "有连接"的超集，且正好等于 raw 的 get_level_info 收集的 depth+1 节点集；
        这样「有几何但无边」的孤立子节点也会被 F 标记，fine 重构能补回孤立点，与 raw 完全一致。
        M(连接矩阵)不变，故三角形/边重构不受影响。"""
        return set(i for i in range(NUM_OCTANTS) if n.children[i] is not None)

    def encode_level_tokens(self, depth):
        raw_faces, nodes_map = self.get_level_info(depth)
        empty = np.zeros((NUM_OCTANTS, NUM_OCTANTS), dtype=np.int8)

        # F：每个节点的 active octant 掩码 (8bit)
        node_F = {}
        for coord, n in nodes_map.items():
            active = self._active_octants(coord, n, nodes_map)
            node_F[coord] = [1 if i in active else 0 for i in range(NUM_OCTANTS)]

        intra_cache = {}

        def intra8(coord):
            if coord not in intra_cache:
                M = (nodes_map[coord].intra_matrix != 0).astype(np.int8)  # 8x8 octant 原位、对称
                intra_cache[coord] = M.tolist()
            return intra_cache[coord]

        def inter8(a, b):
            if a == b:
                return intra8(a)
            na = nodes_map[a]
            M = na.inter_matrices.get(b)
            if M is None:
                nb = nodes_map[b]
                Mt = nb.inter_matrices.get(a)
                M = Mt.T if Mt is not None else empty
            return (np.asarray(M) != 0).astype(np.int8).tolist()

        raw_tokens = []
        output_tokens = []
        for face in raw_faces:
            v0_c, v1_c, v2_c = face
            raw_tokens.append([v0_c, v1_c, v2_c])

            F0, F1, F2 = node_F[v0_c], node_F[v1_c], node_F[v2_c]
            M00, M11, M22 = intra8(v0_c), intra8(v1_c), intra8(v2_c)
            M01 = inter8(v0_c, v1_c)
            M02 = inter8(v0_c, v2_c)
            M12 = inter8(v1_c, v2_c)

            output_tokens.append((v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12))
        return raw_tokens, output_tokens


class OctreeNode:
    def __init__(self, x, y, z, depth):
        self.coord = (x, y, z)
        self.depth = depth
        self.children = [None] * 8
        self.intra_neighbors = set()
        self.inter_neighbors = set()
        self.intra_matrix = np.zeros((8, 8), dtype=np.int8)
        self.inter_matrices = {}


class MatrixOctree:
    def __init__(self, max_depth=7):
        self.max_depth = max_depth
        self.root = None

    def _get_parent_coord(self, coord):
        return (coord[0] >> 1, coord[1] >> 1, coord[2] >> 1)

    def _get_local_index(self, coord):
        return ((coord[2] & 1) << 2) | ((coord[1] & 1) << 1) | (coord[0] & 1)

    def build(self, vertices, faces):
        resolution = 2 ** self.max_depth
        quantized = discretize(vertices, num_discrete=resolution)

        leaf_nodes_map = {}
        for i in range(len(vertices)):
            coord = tuple(quantized[i])
            if coord not in leaf_nodes_map:
                leaf_nodes_map[coord] = OctreeNode(coord[0], coord[1], coord[2], self.max_depth)

        for face in faces:
            coords = [tuple(quantized[i]) for i in face]
            for i in range(3):
                u, v = coords[i], coords[(i + 1) % 3]
                if u == v:
                    continue
                parent_u = self._get_parent_coord(u)
                parent_v = self._get_parent_coord(v)
                node_u = leaf_nodes_map[u]
                node_v = leaf_nodes_map[v]
                if parent_u == parent_v:
                    node_u.intra_neighbors.add(v)
                    node_v.intra_neighbors.add(u)
                else:
                    node_u.inter_neighbors.add(v)
                    node_v.inter_neighbors.add(u)

        current_layer_nodes = list(leaf_nodes_map.values())

        for d in range(self.max_depth, 0, -1):
            parent_layer_map = {}
            for child in current_layer_nodes:
                parent_coord = self._get_parent_coord(child.coord)
                if parent_coord not in parent_layer_map:
                    parent_layer_map[parent_coord] = OctreeNode(*parent_coord, d - 1)
                parent_node = parent_layer_map[parent_coord]
                idx = self._get_local_index(child.coord)
                parent_node.children[idx] = child

            for child in current_layer_nodes:
                parent_coord = self._get_parent_coord(child.coord)
                parent_node = parent_layer_map[parent_coord]
                child_idx = self._get_local_index(child.coord)

                for sibling_coord in child.intra_neighbors:
                    sibling_idx = self._get_local_index(sibling_coord)
                    parent_node.intra_matrix[child_idx, sibling_idx] = 1

                for neighbor_coord in child.inter_neighbors:
                    neighbor_parent_coord = self._get_parent_coord(neighbor_coord)
                    if neighbor_parent_coord not in parent_node.inter_matrices:
                        parent_node.inter_matrices[neighbor_parent_coord] = np.zeros((8, 8), dtype=np.int8)
                    neighbor_child_idx = self._get_local_index(neighbor_coord)
                    parent_node.inter_matrices[neighbor_parent_coord][child_idx, neighbor_child_idx] = 1

                    grandparent_p = self._get_parent_coord(parent_coord)
                    grandparent_n = self._get_parent_coord(neighbor_parent_coord)
                    if grandparent_p == grandparent_n:
                        parent_node.intra_neighbors.add(neighbor_parent_coord)
                    else:
                        parent_node.inter_neighbors.add(neighbor_parent_coord)

            current_layer_nodes = list(parent_layer_map.values())
        if current_layer_nodes:
            self.root = current_layer_nodes[0]


def reconstruct_next_level_optimized(output_tokens, resolution, max_children=NUM_OCTANTS):
    """octant 绝对原位重建：M 已是 8x8、行列=octant；用 F mask 掉非 active 行列。"""
    indices = np.arange(8)
    offsets = np.stack([indices & 1, (indices >> 1) & 1, (indices >> 2) & 1], axis=1).astype(np.int32)

    tri_blocks = []          # 逐 face 向量化收集的三角形 (N,3,3)，最后统一去重
    edge_keys = set()        # 全局所有边: canonical (key_a, key_b)
    tri_edge_keys = set()    # 被三角形覆盖的边(全局)
    active_keys = set()      # 全局 active 节点 key
    key_to_coord = {}        # key(int 三元组) -> np.int32 坐标(3,)

    for f_token in output_tokens:
        v0_c, v1_c, v2_c, F0, F1, F2, M00, M11, M22, M01, M02, M12 = f_token
        fa = [np.asarray(F0, dtype=bool), np.asarray(F1, dtype=bool), np.asarray(F2, dtype=bool)]

        def intra(F, M):
            return np.asarray(M, dtype=bool) & np.outer(F, F)

        def inter(Fa, Fb, M):
            return np.asarray(M, dtype=bool) & np.outer(Fa, Fb)

        in0 = intra(fa[0], M00); in1 = intra(fa[1], M11); in2 = intra(fa[2], M22)
        m01 = inter(fa[0], fa[1], M01); m02 = inter(fa[0], fa[2], M02); m12 = inter(fa[1], fa[2], M12)

        A = np.zeros((24, 24), dtype=bool)
        A[0:8, 0:8] = in0 | in0.T
        A[8:16, 8:16] = in1 | in1.T
        A[16:24, 16:24] = in2 | in2.T
        A[0:8, 8:16] = m01;  A[8:16, 0:8] = m01.T
        A[0:8, 16:24] = m02; A[16:24, 0:8] = m02.T
        A[8:16, 16:24] = m12; A[16:24, 8:16] = m12.T

        # 24 个节点(slot p, octant c)的世界坐标 = v_parents[p]*2 + offset[c]
        v_parents_arr = np.stack([np.array(v0_c[:3], dtype=np.int32),
                                  np.array(v1_c[:3], dtype=np.int32),
                                  np.array(v2_c[:3], dtype=np.int32)])
        node_coords = v_parents_arr[np.arange(24) // 8] * 2 + offsets[np.arange(24) % 8]   # (24,3)
        keys = [tuple(int(x) for x in node_coords[i]) for i in range(24)]
        for i in range(24):
            if keys[i] not in key_to_coord:
                key_to_coord[keys[i]] = node_coords[i]

        # 三角形(逐 face 向量化)；同时把三角形的边记入全局 covered 集合
        valid_mask = A[_IDX_I, _IDX_J] & A[_IDX_J, _IDX_K] & A[_IDX_I, _IDX_K]
        tris = _TRI_INDICES[valid_mask]
        if len(tris) > 0:
            p_idx = tris // 8
            c_idx = tris % 8
            tri_blocks.append(v_parents_arr[p_idx] * 2 + offsets[c_idx])
            for i, j, k in tris:
                ka, kb, kc = sorted((keys[int(i)], keys[int(j)], keys[int(k)]))
                tri_edge_keys.add((ka, kb)); tri_edge_keys.add((kb, kc)); tri_edge_keys.add((ka, kc))

        # 全局收集所有边
        ei, ej = np.where(np.triu(A, 1))
        for i, j in zip(ei.tolist(), ej.tolist()):
            ki, kj = keys[i], keys[j]
            edge_keys.add((ki, kj) if ki <= kj else (kj, ki))

        # 全局收集 active 节点
        for i in np.where(np.concatenate(fa))[0].tolist():
            active_keys.add(keys[i])

    all_blocks = list(tri_blocks)

    # 未被任何三角形覆盖的边 -> 退化面 (A, B, B)   [全局判断，与 raw 的 get_level_info 一致]
    for a, b in (edge_keys - tri_edge_keys):
        all_blocks.append(np.stack([key_to_coord[a], key_to_coord[b], key_to_coord[b]])[None])

    # 孤立的 active 节点(全局无任何边) -> 点 (A, A, A)
    edge_nodes = set()
    for a, b in edge_keys:
        edge_nodes.add(a); edge_nodes.add(b)
    for a in (active_keys - edge_nodes):
        c = key_to_coord[a]
        all_blocks.append(np.stack([c, c, c])[None])

    if not all_blocks:
        return np.array([]), np.array([])

    flat_v = np.concatenate(all_blocks, axis=0).reshape(-1, 3)
    V, inv = np.unique(flat_v, axis=0, return_inverse=True)
    F_raw = inv.reshape(-1, 3)
    F_sorted = np.sort(F_raw, axis=1)
    _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True)
    F = F_raw[unique_face_indices]
    V = undiscretize(V, num_discrete=resolution)
    return V, F


def reconstruct_base_mesh(output_tokens, resolution):
    all_face_coords = []
    for f_token in output_tokens:
        v0_c, v1_c, v2_c = f_token
        c0 = np.array(v0_c[:3], dtype=np.int32)
        c1 = np.array(v1_c[:3], dtype=np.int32)
        c2 = np.array(v2_c[:3], dtype=np.int32)
        all_face_coords.append(np.stack([c0, c1, c2]))

    if not all_face_coords:
        return np.array([]), np.array([])

    flat_v = np.concatenate(all_face_coords, axis=0).reshape(-1, 3)
    V, inv = np.unique(flat_v, axis=0, return_inverse=True)
    F_raw = inv.reshape(-1, 3)
    F_sorted = np.sort(F_raw, axis=1)
    _, unique_face_indices = np.unique(F_sorted, axis=0, return_index=True)
    F = F_raw[unique_face_indices]
    V = undiscretize(V, num_discrete=resolution)
    return V, F


class MeshCodec:
    def __init__(self, args):
        self.max_depth = args.max_depth
        self.max_resolution = args.max_resolution
        self.use_hourglass = args.use_hourglass
        self.max_children = NUM_OCTANTS  # octant 原位固定 8(忽略 config 的 max_children)

        self.tokenizer = DynamicMeshTokenizer(resolution=self.max_resolution)

    def encode_mesh(self, mesh: trimesh.Trimesh, target_depth: int, corrupt=None) -> tuple:
        if isinstance(mesh, trimesh.Scene):
            mesh = mesh.dump(concatenate=True)
        V = np.asarray(mesh.vertices, dtype=np.float64)

        res = 2 ** (target_depth + 1)
        v_restored = undiscretize(discretize(V, num_discrete=res), num_discrete=res)
        quantized_mesh = trimesh.Trimesh(vertices=v_restored, faces=mesh.faces, process=False)

        # corrupt(仅训练时由 dataset 传入)= 2a 暴露偏差腐蚀; None 时与原版逐位一致(推理不受影响)
        raw_seq, seq = _fast_fix3_level_tokens(V, mesh.faces, target_depth, self.tokenizer, corrupt=corrupt)
        return raw_seq, seq, quantized_mesh


    def decode_tokens(self, token_sequence: list, target_depth: int):
        if os.environ.get("FIX3_FAST_DECODE", "1") != "1":
            return super().decode_tokens(token_sequence, target_depth)
        if token_sequence is None or len(token_sequence) == 0:
            return None
        return _fast_decode_tokens(token_sequence, target_depth, self.tokenizer)


    def decode_raw_tokens(self, raw_token_sequence: list, target_depth: int) -> trimesh.Trimesh:
        if not raw_token_sequence:
            return None
        resolution = 2 ** (target_depth)
        raw_t = self.tokenizer.decode_raw(raw_token_sequence)
        v, f = reconstruct_base_mesh(raw_t, resolution)
        if len(v) == 0 or len(f) == 0:
            return None
        return trimesh.Trimesh(vertices=v, faces=f, process=False)

    def get_model_config(self) -> dict:
        t = self.tokenizer
        return {
            "total_vocab_size": t.vocab_size,
            "coord_vocab_size": t.resolution,
            "f_vocab_size": t.f_vocab,
            "row_vocab_size": t.row_vocab,            # intra/inter 行子码各 256
            "coord_offset": t.coord_offset,
            "f_offset": t.f_offset,
            "intra_row_offset": t.intra_row_offset,
            "inter_row_offset": t.inter_row_offset,
            "pad_token_id": t.pad_id,
            "num_octants": NUM_OCTANTS,
            "subcodes_per_position": NUM_OCTANTS,     # 每个 refine 位置 8 个子码槽(F 复制；intra/inter 为 8 行)
            "tokens_per_face": DynamicMeshTokenizer.TOKENS_PER_FACE,   # 60 扁平 token/face (9+3+6*8)
            "coord_per_face": 9,
            "refine_positions_per_face": DynamicMeshTokenizer.REFINE_POS_PER_FACE,  # 9 (3F+3intra+3inter)
            "input_num_categories": t.resolution + 1,
            # 输出端：F / intra_row / inter_row 三类，每类 256；coord 为输入端条件
            "output_num_categories": 1 + t.f_vocab + 2 * t.row_vocab,
        }


# ==========================================
# 自测：encode -> decode -> reconstruct 往返
# ==========================================
if __name__ == "__main__":
    # ==========================================
    # 编解码 + 重构导出流程 (参照 fix2 的 __main__)
    #   每个 mesh: 导出 GT；对每个 depth 做
    #     encode -> decode_tokens(细分/Fine) / decode_raw_tokens(粗糙/Coarse) / 量化网格
    #   并打印计时、token 数(60/face)与 token 级往返一致性。
    # ==========================================
    import time

    # 1. 输入/输出文件夹 (按需修改)
    INPUT_DIRECTORY = "/root/junkai_lin/mul_res_diffusion/demo/complex_test_data/files"
    OUTPUT_DIRECTORY = "./output/fix3_octant_factorized2"
    os.makedirs(OUTPUT_DIRECTORY, exist_ok=True)

    @dataclass
    class MeshCodecArgs:
        max_depth: int = 10
        max_resolution: int = 1024
        use_hourglass: bool = True
        max_children: int = 8      # fix3 octant 原位固定 8 (此字段实际被忽略)

    # 2. 初始化 Codec
    args = MeshCodecArgs()
    codec = MeshCodec(args=args)
    print(f"[*] Initialized fix3 Codec. Total Vocab Size: {codec.tokenizer.vocab_size} "
          f"(tokens/face={DynamicMeshTokenizer.TOKENS_PER_FACE})")

    # 3. 收集所有 .obj / .glb
    input_dir_path = Path(INPUT_DIRECTORY)
    mesh_files = sorted(input_dir_path.rglob("*.obj")) + sorted(input_dir_path.rglob("*.glb"))
    if not mesh_files:
        print(f"[!] No .obj or .glb files found in {INPUT_DIRECTORY}")

    DEPTHS = [0, 1, 2, 3, 4, 5, 6]

    # 4. 逐文件处理
    for idx, file_path in enumerate(mesh_files, 1):
        obj_name = file_path.stem
        print(f"\n{'=' * 50}")
        print(f"[{idx}/{len(mesh_files)}] Processing: {file_path.name}")
        print(f"{'=' * 50}")
        file_start_time = time.perf_counter()

        try:
            raw_mesh = trimesh.load(str(file_path), force="mesh")
            normalize(raw_mesh)
        except Exception as e:
            print(f"[!] load/normalize failed: {e}")
            continue

        # 导出 GT
        gt_path = os.path.join(OUTPUT_DIRECTORY, f"{obj_name}_gt.obj")
        raw_mesh.export(gt_path)
        print(f"[*] Saved GT -> {gt_path}")

        # 不同深度的编解码
        for depth in DEPTHS:
            print(f"\n--- Depth {depth} | res(fine)={2 ** (depth + 1)} | octant 原位 ---")

            # Encoding
            t0 = time.perf_counter()
            raw_tokens, tokens, quantized_mesh = codec.encode_mesh(raw_mesh, target_depth=depth)
            encode_time = time.perf_counter() - t0
            nface = len(tokens) // DynamicMeshTokenizer.TOKENS_PER_FACE
            print(f"[Timer] Encoding: {encode_time:.4f}s | flat_tokens={len(tokens)} "
                  f"(nface={nface}, 模型逻辑位置={18 * nface})")

            # token 级往返一致性 (encode->decode->re-encode)
            same = (tokens == codec.tokenizer.encode(codec.tokenizer.decode(tokens)))
            print(f"[Check] token 往返一致: {same}")

            # Decoding (Fine)
            if len(tokens) > 0:
                t1 = time.perf_counter()
                recon_mesh = codec.decode_tokens(tokens, target_depth=depth)
                decode_fine_time = time.perf_counter() - t1
                if recon_mesh is not None:
                    out_path = os.path.join(
                        OUTPUT_DIRECTORY, f"{obj_name}_recon_d{depth}_res{2 ** (depth + 1)}_fix3.obj")
                    recon_mesh.export(out_path)
                    print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s "
                          f"-> V={len(recon_mesh.vertices)} F={len(recon_mesh.faces)}")
                else:
                    print(f"[Timer] Decoding (Fine): {decode_fine_time:.4f}s -> empty")

            # Decoding (Coarse / Raw)
            if len(raw_tokens) > 0:
                t2 = time.perf_counter()
                raw_recon_mesh = codec.decode_raw_tokens(raw_tokens, target_depth=depth)
                decode_coarse_time = time.perf_counter() - t2
                if raw_recon_mesh is not None:
                    raw_out_path = os.path.join(
                        OUTPUT_DIRECTORY, f"{obj_name}_raw_recon_d{depth}_res{2 ** depth}.obj")
                    raw_recon_mesh.export(raw_out_path)
                    print(f"[Timer] Decoding (Coarse): {decode_coarse_time:.4f}s "
                          f"-> V={len(raw_recon_mesh.vertices)} F={len(raw_recon_mesh.faces)}")

            # 量化网格 (输入离散化后的 GT，用于对照重构质量上限)
            if quantized_mesh is not None:
                quant_out_path = os.path.join(
                    OUTPUT_DIRECTORY, f"{obj_name}_dis_d{depth + 1}_res{2 ** (depth + 1)}.obj")
                quantized_mesh.export(quant_out_path)

        file_total_time = time.perf_counter() - file_start_time
        print(f"\n[*] Finished {file_path.name} in {file_total_time:.2f}s")


# ==========================================================================
# encode_mesh / decode_tokens 的 numba 加速实现细节(原 *_fast.py 并入)
# ==========================================================================

_OCT_W = np.array([1, 2, 4], dtype=np.int64)   # octant 局部索引 = x + 2y + 4z

def _corrupt_fid_nodes(fid, nodes, und, nN, target_depth, corrupt):
    """2a 暴露偏差训练腐蚀(仅训练): 模拟级联实测的两类粗输入误差。
      (a) 伪面: 在已存在的 GT 节点间加不该有的面(过连接)。下游打包按 fid 查表 ->
          非边对的 inter 行自动全 0(empty),教模型"看 PC 把伪面 refine 成空"。
      (b) 节点坐标抖动 ±jitter_max(位移): 只改 coord_part 的输入坐标,目标 F/intra/inter
          仍按真节点身份查表不变,教模型"输入坐标偏了也 refine 对"。
    目标无需重新导出 —— 全部沿用 fid 索引,与干净路径同一套向量化打包。返回 (fid_aug, nodes_coord)。
    """
    sp_rate = float(corrupt.get("spurious_rate", 0.0))
    jp = float(corrupt.get("jitter_prob", 0.0))

    if sp_rate > 0.0 and nN >= 3 and und.shape[0] > 0:
        radius = int(corrupt.get("spatial_radius", 2))
        n_add = int(sp_rate * fid.shape[0])
        if n_add > 0:
            coord2idx = {(int(c[0]), int(c[1]), int(c[2])): i for i, c in enumerate(nodes)}
            existing = set(map(tuple, np.sort(fid, axis=1).tolist()))
            ne = und.shape[0]
            spurious, tries, maxtries = [], 0, n_add * 6
            while len(spurious) < n_add and tries < maxtries:
                tries += 1
                ei = int(np.random.randint(ne))
                ai, bi = int(und[ei, 0]), int(und[ei, 1])
                base = nodes[ai] if np.random.random() < 0.5 else nodes[bi]
                off = np.random.randint(-radius, radius + 1, size=3)
                if not off.any():
                    continue
                ck = (int(base[0] + off[0]), int(base[1] + off[1]), int(base[2] + off[2]))
                ci = coord2idx.get(ck)
                if ci is None or ci == ai or ci == bi:
                    continue
                tri = tuple(sorted((ai, bi, ci)))
                if tri in existing:
                    continue
                existing.add(tri)
                spurious.append(tri)
            if spurious:
                fid = np.concatenate([fid, np.asarray(spurious, dtype=fid.dtype)], axis=0)

    nodes_coord = nodes
    if jp > 0.0:
        jmax = int(corrupt.get("jitter_max", 1))
        res_c = 2 ** target_depth
        nodes_coord = nodes.copy()
        mask = np.random.random(nN) < jp
        k = int(mask.sum())
        if k > 0:
            noise = np.random.randint(-jmax, jmax + 1, size=(k, 3))
            nodes_coord[mask] = np.clip(nodes_coord[mask] + noise, 0, res_c - 1)
    return fid, nodes_coord

def _fast_fix3_level_tokens(vertices, faces, target_depth, tok, corrupt=None):
    """等价替换 MatrixOctree(max_depth=td+1).build + encode_level_tokens(td) + tokenizer.encode/encode_raw。
    返回 (raw_token_sequence, token_sequence)，与原版逐位一致。"""
    res = 2 ** (target_depth + 1)
    Q = discretize(np.asarray(vertices, dtype=np.float64), num_discrete=res).astype(np.int64)
    faces = np.asarray(faces, dtype=np.int64).reshape(-1, 3)

    # ---- 节点(td) = 所有顶点叶子坐标的 parent；占据 F = 所有顶点的 octant ----
    P_all = Q >> 1
    i_all = (Q & 1) @ _OCT_W

    # ---- 叶子边表(每面 3 条，跳过零长边 u==v，与原版一致) ----
    if faces.size:
        e = np.concatenate([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0)
        u, v = Q[e[:, 0]], Q[e[:, 1]]
        keep = np.any(u != v, axis=1)
        u, v = u[keep], v[keep]
    else:
        u = v = np.zeros((0, 3), dtype=np.int64)
    pu, pv = u >> 1, v >> 1
    iu = (u & 1) @ _OCT_W
    iv = (v & 1) @ _OCT_W

    # 节点表(np.unique 行字典序 = 坐标 (x,y,z) 升序 = 原版 face 内顶点 sorted() 同序)
    nodes, inv = np.unique(np.concatenate([P_all, pu, pv], axis=0), axis=0, return_inverse=True)
    nV, M = P_all.shape[0], pu.shape[0]
    occ_idx = inv[:nV]
    a, b = inv[nV:nV + M], inv[nV + M:]
    nN = nodes.shape[0]
    intra = a == b

    # ---- F 行子码(8bit int)：所有顶点的 octant 占据 ----
    Fbits = np.zeros(nN, dtype=np.int64)
    np.bitwise_or.at(Fbits, occ_idx, np.int64(1) << i_all)

    # ---- intra 行子码 [nN,8]：绝对原位、对称双向、无对角填充(与原版一致) ----
    intra_rows = np.zeros((nN, 8), dtype=np.int64)
    ai, iui, ivi = a[intra], iu[intra], iv[intra]
    np.bitwise_or.at(intra_rows, (ai, iui), np.int64(1) << ivi)
    np.bitwise_or.at(intra_rows, (ai, ivi), np.int64(1) << iui)

    # ---- inter 行子码：有向 (a,b) 分组(双向都存，与原版 build 一致) ----
    im = ~intra
    ea, eb, eiu, eiv = a[im], b[im], iu[im], iv[im]
    dkey = np.concatenate([ea * nN + eb, eb * nN + ea])
    drow = np.concatenate([eiu, eiv])
    dcol = np.concatenate([eiv, eiu])
    ukeys, kinv = np.unique(dkey, return_inverse=True)
    inter_rows = np.zeros((max(len(ukeys), 1), 8), dtype=np.int64)
    if dkey.size:
        np.bitwise_or.at(inter_rows, (kinv, drow), np.int64(1) << dcol)

    # ---- 无向节点邻接 -> CSR -> numba 三角形枚举 ----
    if ea.size:
        und = np.unique(np.stack([np.minimum(ea, eb), np.maximum(ea, eb)], axis=1), axis=0)
    else:
        und = np.zeros((0, 2), dtype=np.int64)
    deg = np.zeros(nN, dtype=np.int64)
    if und.size:
        np.add.at(deg, und[:, 0], 1)
        np.add.at(deg, und[:, 1], 1)
        order = np.lexsort((und[:, 1], und[:, 0]))
        fu, fv = und[order, 0], und[order, 1]
        indptr = np.zeros(nN + 1, dtype=np.int64)
        np.add.at(indptr, fu + 1, 1)
        indptr = np.cumsum(indptr)
        tu, tv, tw = _triangles_csr(indptr, fv.astype(np.int64))
    else:
        tu = tv = tw = np.zeros(0, dtype=np.int64)

    # ---- faces_ids：三角形 + 未覆盖边退化 (A,B,B) + 孤立点 (A,A,A)，集合与原版一致 ----
    tris = np.stack([tu, tv, tw], axis=1) if tu.size else np.zeros((0, 3), dtype=np.int64)
    covered = set()
    for x, y, z in tris:
        x, y, z = int(x), int(y), int(z)
        covered.add((x, y)); covered.add((y, z)); covered.add((x, z))
    extra = []
    for x, y in und:
        x, y = int(x), int(y)
        if (x, y) not in covered:
            extra.append((x, y, y))
    for x in np.flatnonzero(deg == 0):
        x = int(x)
        extra.append((x, x, x))
    fid = np.concatenate([tris, np.asarray(extra, dtype=np.int64).reshape(-1, 3)], axis=0)
    if fid.shape[0] == 0:
        return [], []

    # ---- 原版排序键 + 决定性平局键 ----
    c0, c1, c2 = nodes[fid[:, 0]], nodes[fid[:, 1]], nodes[fid[:, 2]]
    order = np.lexsort((c2[:, 0], c2[:, 1], c2[:, 2],
                        c1[:, 0], c1[:, 1], c1[:, 2],
                        c0[:, 0], c0[:, 1], c0[:, 2]))   # 主键=c0.z(最后) ... 与原版 7 键一致+补齐平局
    fid = fid[order]
    nodes_coord = nodes
    if corrupt is not None:
        fid, nodes_coord = _corrupt_fid_nodes(fid, nodes, und, nN, target_depth, corrupt)
    c0, c1, c2 = nodes_coord[fid[:, 0]], nodes_coord[fid[:, 1]], nodes_coord[fid[:, 2]]

    # ---- 60/face 扁平 token 向量化打包 ----
    # 注: F_part/intra_part/inter_part 下面仍用 fid(节点身份)查表 -> 腐蚀的伪面自动得正确目标
    # (非边对 inter 全 0),坐标抖动只影响上面 coord_part,目标不变。
    co, fo = tok.coord_offset, tok.f_offset
    io, eo = tok.intra_row_offset, tok.inter_row_offset
    nF = fid.shape[0]

    coord_part = np.concatenate([c0, c1, c2], axis=1) + co                      # (nF,9)
    F_part = np.stack([Fbits[fid[:, 0]], Fbits[fid[:, 1]], Fbits[fid[:, 2]]], axis=1) + fo
    intra_part = np.concatenate([intra_rows[fid[:, 0]],
                                 intra_rows[fid[:, 1]],
                                 intra_rows[fid[:, 2]]], axis=1) + io           # (nF,24)

    inter_blocks = []
    for x, y in ((0, 1), (0, 2), (1, 2)):
        A_, B_ = fid[:, x], fid[:, y]
        rows = np.zeros((nF, 8), dtype=np.int64)
        same = A_ == B_
        rows[same] = intra_rows[A_[same]]                  # inter8(a,a)=intra8(a)，与原版一致
        ns = np.flatnonzero(~same)
        if ns.size:
            dk = A_[ns] * nN + B_[ns]
            pos = np.searchsorted(ukeys, dk)
            pos_c = np.minimum(pos, max(len(ukeys) - 1, 0))
            found = (ukeys[pos_c] == dk) if len(ukeys) else np.zeros(ns.size, dtype=bool)
            rows[ns[found]] = inter_rows[pos_c[found]]     # 未找到 -> 保持全 0(原版 empty 矩阵)
        inter_blocks.append(rows)
    inter_part = np.concatenate(inter_blocks, axis=1) + eo                      # (nF,24)

    token = np.concatenate([coord_part, F_part, intra_part, inter_part], axis=1)  # (nF,60)
    token_sequence = token.ravel().tolist()
    raw_token_sequence = coord_part.reshape(-1, 1).tolist()                       # 9 coord/face, 各 [t]
    return raw_token_sequence, token_sequence

_BIT_SHIFTS = np.arange(NUM_OCTANTS, dtype=np.int64)                 # _int_to_row: bit i = (v>>i)&1

_NODE_PARENT = np.arange(24) // NUM_OCTANTS                          # 24 个节点 -> 哪个父顶点 slot

_NODE_OCTANT = np.arange(24) % NUM_OCTANTS                           # 24 个节点 -> 哪个 octant

_OCT_OFFSETS = np.stack([np.arange(8) & 1, (np.arange(8) >> 1) & 1, (np.arange(8) >> 2) & 1],
                        axis=1).astype(np.int32)                     # octant -> (dx,dy,dz)

_INTER_ROW_SLOT = np.array([0, 0, 1])

_INTER_COL_SLOT = np.array([1, 2, 2])

def _unpack_rows(values):
    """int 行子码 -> bool 位。values (...,) -> (..., 8), 与 _int_to_row 同序(bit i = (v>>i)&1)。"""
    return ((values[..., None].astype(np.int64) >> _BIT_SHIFTS) & 1).astype(bool)

def _pack_keys(coords, resolution):
    """(...,3) int 坐标 -> (...) int64 唯一 key。坐标范围 [0, resolution)。"""
    c = coords.astype(np.int64)
    return (c[..., 0] * resolution + c[..., 1]) * resolution + c[..., 2]

def _fast_reconstruct_next_level(token_array, resolution, tokenizer):
    """token_array: (nface, TOKENS_PER_FACE) int。返回 (V, F), 语义同 reconstruct_next_level_optimized。"""
    from modules.octree_codec import _TRI_INDICES

    nface = token_array.shape[0]
    parents = (token_array[:, :9] - tokenizer.coord_offset).reshape(nface, 3, 3).astype(np.int32)
    occupancy = _unpack_rows(token_array[:, 9:12] - tokenizer.f_offset)          # (nface,3,8) bool

    def rows_to_mat(lo, hi, offset):
        return _unpack_rows(token_array[:, lo:hi] - offset)                      # (nface,8,8) bool, [行,列]

    intra = [rows_to_mat(12 + 8 * i, 20 + 8 * i, tokenizer.intra_row_offset) for i in range(3)]
    inter = [rows_to_mat(36 + 8 * i, 44 + 8 * i, tokenizer.inter_row_offset) for i in range(3)]

    # ---- 24x24 邻接: intra 块按 F&F 掩码后对称化, inter 块按 Fa&Fb 掩码 ----
    adjacency = np.zeros((nface, 24, 24), dtype=bool)
    for slot in range(3):
        f = occupancy[:, slot]                                                   # (nface,8)
        block = intra[slot] & (f[:, :, None] & f[:, None, :])
        block = block | np.swapaxes(block, 1, 2)
        adjacency[:, slot * 8:(slot + 1) * 8, slot * 8:(slot + 1) * 8] = block
    for b in range(3):
        ra, ca = _INTER_ROW_SLOT[b], _INTER_COL_SLOT[b]
        block = inter[b] & (occupancy[:, ra][:, :, None] & occupancy[:, ca][:, None, :])
        adjacency[:, ra * 8:(ra + 1) * 8, ca * 8:(ca + 1) * 8] = block
        adjacency[:, ca * 8:(ca + 1) * 8, ra * 8:(ra + 1) * 8] = np.swapaxes(block, 1, 2)

    # ---- 24 个节点的世界坐标与 key ----
    node_coords = parents[:, _NODE_PARENT] * 2 + _OCT_OFFSETS[_NODE_OCTANT]      # (nface,24,3)
    node_keys = _pack_keys(node_coords, resolution)                              # (nface,24)

    # ---- 三角形: 三条边都在邻接里 ----
    i_idx, j_idx, k_idx = _TRI_INDICES[:, 0], _TRI_INDICES[:, 1], _TRI_INDICES[:, 2]
    valid = (adjacency[:, i_idx, j_idx] & adjacency[:, j_idx, k_idx] & adjacency[:, i_idx, k_idx])
    face_of_tri, tri_of_tri = np.nonzero(valid)                                  # 各 (ntri,)
    tri_nodes = _TRI_INDICES[tri_of_tri]                                         # (ntri,3) 节点编号
    tri_coords = node_coords[face_of_tri[:, None], tri_nodes]                    # (ntri,3,3)
    tri_keys = node_keys[face_of_tri[:, None], tri_nodes]                        # (ntri,3)

    # ---- 全部边(上三角) ----
    upper = np.triu(np.ones((24, 24), dtype=bool), 1)
    face_of_edge, node_i, node_j = np.nonzero(adjacency & upper)
    key_i = node_keys[face_of_edge, node_i]
    key_j = node_keys[face_of_edge, node_j]
    edge_pairs = np.unique(np.stack([np.minimum(key_i, key_j), np.maximum(key_i, key_j)], axis=1), axis=0)

    # ---- 被三角形覆盖的边 ----
    if len(tri_keys):
        tri_sorted = np.sort(tri_keys, axis=1)
        covered = np.concatenate([tri_sorted[:, [0, 1]], tri_sorted[:, [1, 2]], tri_sorted[:, [0, 2]]], axis=0)
        covered = np.unique(covered, axis=0)
    else:
        covered = np.zeros((0, 2), dtype=np.int64)

    blocks = [tri_coords] if len(tri_coords) else []

    # ---- 未被任何三角形覆盖的边 -> 退化面 (A,B,B) ----
    if len(edge_pairs):
        edge_flat = edge_pairs[:, 0] * (resolution ** 3) + edge_pairs[:, 1]
        covered_flat = (covered[:, 0] * (resolution ** 3) + covered[:, 1]) if len(covered) \
            else np.zeros(0, dtype=np.int64)
        open_edges = edge_pairs[~np.isin(edge_flat, covered_flat)]
        if len(open_edges):
            key_to_coord = _build_key_lookup(node_keys, node_coords)
            ca = key_to_coord[open_edges[:, 0]]
            cb = key_to_coord[open_edges[:, 1]]
            blocks.append(np.stack([ca, cb, cb], axis=1))

    # ---- 孤立的 active 节点(无任何边) -> 点 (A,A,A) ----
    active_keys = np.unique(node_keys[occupancy.reshape(nface, 24)])
    if len(active_keys):
        edge_nodes = np.unique(edge_pairs.reshape(-1)) if len(edge_pairs) else np.zeros(0, dtype=np.int64)
        isolated = active_keys[~np.isin(active_keys, edge_nodes)]
        if len(isolated):
            key_to_coord = _build_key_lookup(node_keys, node_coords)
            c = key_to_coord[isolated]
            blocks.append(np.stack([c, c, c], axis=1))

    if not blocks:
        return np.array([]), np.array([])

    flat_v = np.concatenate(blocks, axis=0).reshape(-1, 3)
    V, inverse = np.unique(flat_v, axis=0, return_inverse=True)
    faces_raw = inverse.reshape(-1, 3)
    _, unique_face_indices = np.unique(np.sort(faces_raw, axis=1), axis=0, return_index=True)
    F = faces_raw[unique_face_indices]
    return undiscretize(V, num_discrete=resolution), F

def _build_key_lookup(node_keys, node_coords):
    """key -> 坐标 的稠密查表(用 np.unique 的第一次出现), 供退化面/孤立点还原坐标。"""
    flat_keys = node_keys.reshape(-1)
    flat_coords = node_coords.reshape(-1, 3)
    uniq, first = np.unique(flat_keys, return_index=True)
    lookup = np.zeros((int(uniq.max()) + 1, 3), dtype=np.int32)
    lookup[uniq] = flat_coords[first]
    return lookup

def _fast_decode_tokens(token_sequence, target_depth, tokenizer):
    """decode_tokens 的向量化实现。返回 trimesh.Trimesh 或 None。"""
    arr = np.asarray(token_sequence).reshape(-1)
    arr = arr[arr != tokenizer.pad_id]                                # 与原版 clean_seq 一致
    per_face = tokenizer.TOKENS_PER_FACE
    nface = arr.shape[0] // per_face
    if nface == 0:
        return None
    token_array = arr[:nface * per_face].reshape(nface, per_face).astype(np.int64)
    V, F = _fast_reconstruct_next_level(token_array, 2 ** (target_depth + 1), tokenizer)
    if len(V) == 0 or len(F) == 0:
        return None
    return trimesh.Trimesh(vertices=V, faces=F, process=False)


FastMeshCodec = MeshCodec   # 兼容别名: 加速实现已并入 MeshCodec