"""fix3 编码格式的数据通路工具 —— 从训练脚本抽出, 保证推理与训练用同一份实现。

  split_fix3_token_list  : 扁平 60/face token -> (coord_list[9*nface], subcode_list[9*nface,8] local)
  subcodes_to_flat_token : (coord, 采样得到的 subcodes) -> fix3 扁平 60/face token(含 offset), 供 decode_tokens
"""
import torch

from models.generator import REFINE_POS_PER_FACE, NUM_OCTANTS

TOKENS_PER_FACE = 60          # fix3 扁平: 9 coord + 3 F + 6 矩阵 x 8 行

def _trim_trailing_pad(token, pad_id):
    pad_pos = (token == pad_id).nonzero(as_tuple=True)[0]
    return token[:pad_pos[0].item()] if pad_pos.numel() > 0 else token
def split_fix3_token_list(token_list, pad_id, device, codec_cfg):
    """扁平 60/face token -> (coord_list, subcode_list)。
       coord_list[i]   : [9*nface]    坐标(保留 coord_offset, 直接喂 raw_embedding)
       subcode_list[i] : [9*nface, 8] refine 子码(去 offset, local[0,256))
                         位置顺序 [F0 F1 F2 M00 M11 M22 M01 M02 M12]
                         F 位置只用 slot0(其余槽填 0, 模型按类型只读 slot0)；intra/inter 用 8 行
    """
    f_off = codec_cfg["f_offset"]; ia = codec_cfg["intra_row_offset"]; ie = codec_cfg["inter_row_offset"]
    coord_list, sub_list = [], []
    for tok in token_list:
        tok = _trim_trailing_pad(tok.to(device=device, dtype=torch.long), pad_id)
        if tok.numel() % TOKENS_PER_FACE != 0:
            raise ValueError(f"token 长度 {tok.numel()} 不是 {TOKENS_PER_FACE} 的倍数")
        tok = tok.view(-1, TOKENS_PER_FACE)                           # [nface, 60]
        nface = tok.size(0)
        coord = tok[:, :9].reshape(-1)                                # [9*nface]
        Fv = tok[:, 9:12] - f_off                                     # [nface, 3]
        intra = (tok[:, 12:36] - ia).view(nface, 3, NUM_OCTANTS)      # [nface,3,8]
        inter = (tok[:, 36:60] - ie).view(nface, 3, NUM_OCTANTS)      # [nface,3,8]
        sub = torch.zeros(nface, REFINE_POS_PER_FACE, NUM_OCTANTS, dtype=torch.long, device=device)
        sub[:, 0:3, 0] = Fv                                           # F -> slot0
        sub[:, 3:6, :] = intra
        sub[:, 6:9, :] = inter
        sub_list.append(sub.reshape(-1, NUM_OCTANTS))                 # [9*nface, 8]
        coord_list.append(coord)
    return coord_list, sub_list
def subcodes_to_flat_token(coord_seq, sub_seq, codec_cfg):
    """采样得到的 (coord[9*nface], subcodes[9*nface,8] local) 还原成 fix3 扁平 60/face token(含 offset)，
       供 codec.decode_tokens 使用。F 位置取 slot0。"""
    f_off = codec_cfg["f_offset"]; ia = codec_cfg["intra_row_offset"]; ie = codec_cfg["inter_row_offset"]
    nface = coord_seq.numel() // REFINE_POS_PER_FACE
    coord = coord_seq.view(nface, REFINE_POS_PER_FACE)                    # [nface,9]
    sub = sub_seq.view(nface, REFINE_POS_PER_FACE, NUM_OCTANTS)           # [nface,9,8]
    Fv = sub[:, 0:3, 0] + f_off                                          # [nface,3] 取 slot0
    intra = sub[:, 3:6, :] + ia                                          # [nface,3,8]
    inter = sub[:, 6:9, :] + ie
    flat = torch.cat([coord, Fv, intra.reshape(nface, 24), inter.reshape(nface, 24)], dim=1)  # [nface,60]
    return flat.reshape(-1).tolist()