import os
import glob

# ======================
# 配置部分（你要改这里）
# ======================

# 标签文件所在文件夹
LABEL_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/point2/2"

# 新的关键点顺序 (0-based)
# 例如：原顺序 1,2,3,4 → 新顺序 1,2,4,3
#new_order = [0, 1, 3, 2]
new_order = [0, 2, 3, 1]

# ======================
# 程序部分
# ======================

def reorder_keypoints(line, new_order):
    parts = line.strip().split()

    if len(parts) < 17:
        print(f"数据格式不对，字段数为 {len(parts)}，应 ≥ 17 ：", line)
        return line

    # 前 5 个保持不变
    head = parts[:5]

    # 后 12 个是 4 个关键点 (4 * 3)
    kps = parts[5:]

    # 分组成 [(x1,y1,v1),(x2,y2,v2)...]
    keypoints = [kps[i:i+3] for i in range(0, 12, 3)]

    # 按新顺序重排
    reordered = []
    for idx in new_order:
        reordered.extend(keypoints[idx])

    return " ".join(head + reordered)


def process_all_labels(label_dir, new_order):
    txt_files = glob.glob(os.path.join(label_dir, "*.txt"))

    print(f"找到 {len(txt_files)} 个标签文件\n")

    for txt_path in txt_files:
        lines_out = []

        with open(txt_path, "r") as f:
            lines = f.readlines()

        for line in lines:
            if line.strip():
                line = reorder_keypoints(line, new_order)
            lines_out.append(line)

        # 覆写文件
        with open(txt_path, "w") as f:
            f.write("\n".join(lines_out))

        print(f"✓ 已处理：{txt_path}")


if __name__ == "__main__":
    process_all_labels(LABEL_DIR, new_order)
    print("\n全部处理完成！")