cot_generate/rl_data/select_single_pattern.py

"""
从 full_delta_frame_mapping.json 中为每个视频随机选取一个 mask pattern，
输出精简版 mapping，用于加速 oracle 推理。

随机性保证:
  - 使用 os.urandom 生成种子（系统级熵源）
  - 每条 entry 独立随机选取
  - 打印种子方便复现

使用方式:
  python select_single_pattern.py \
    --input full_delta_frame_mapping.json \
    --output single_delta_frame_mapping.json
"""

import argparse
import json
import os
import random


def main():
    parser = argparse.ArgumentParser(
        description="为每个视频随机选取一个 mask pattern"
    )
    parser.add_argument("--input", required=True, help="完整 mapping JSON")
    parser.add_argument("--output", required=True, help="输出精简 mapping JSON")
    parser.add_argument("--seed", type=int, default=None,
                        help="随机种子（不指定则用 os.urandom 自动生成）")
    args = parser.parse_args()

    # 随机种子：优先用户指定，否则用系统熵源
    if args.seed is not None:
        seed = args.seed
    else:
        seed = int.from_bytes(os.urandom(8), "big")
    random.seed(seed)
    print(f"Random seed: {seed}  (可用 --seed {seed} 复现)")

    with open(args.input) as f:
        data = json.load(f)

    mappings = data["mappings"]
    print(f"输入: {len(mappings)} entries")

    total_before = 0
    total_after = 0
    no_pattern = 0

    for entry in mappings:
        patterns = entry.get("mask_patterns", [])
        total_before += len(patterns)
        if patterns:
            selected = random.choice(patterns)
            entry["mask_patterns"] = [selected]
            total_after += 1
        else:
            no_pattern += 1

    data["metadata"]["selection"] = "single_random_pattern"
    data["metadata"]["random_seed"] = seed

    with open(args.output, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    print(f"输出: {len(mappings)} entries")
    print(f"Patterns: {total_before} → {total_after} ({total_before - total_after} 减少)")
    print(f"无 pattern 的 entry: {no_pattern}")
    print(f"Saved to {args.output}")


if __name__ == "__main__":
    main()