import os
import shutil

def pack_scene_data(root, scene, output_dir):
    scene_dir = os.path.join(output_dir, scene)
    if not os.path.exists(scene_dir):
        os.makedirs(scene_dir)

    pts_dir = os.path.join(root, scene, "pts")
    if os.path.exists(pts_dir):
        shutil.move(pts_dir, os.path.join(scene_dir, "pts"))

    camera_dir = os.path.join(root, scene, "camera_params")
    if os.path.exists(camera_dir):
        shutil.move(camera_dir, os.path.join(scene_dir, "camera_params"))

    scene_info_file = os.path.join(root, scene, "scene_info.json")
    if os.path.exists(scene_info_file):
        shutil.move(scene_info_file, os.path.join(scene_dir, "scene_info.json"))

    label_dir = os.path.join(root, scene, "label")
    if os.path.exists(label_dir):
        shutil.move(label_dir, os.path.join(scene_dir, "label"))


def pack_all_scenes(root, scene_list, output_dir):
    for idx, scene in enumerate(scene_list):
        print(f"packing {scene} ({idx+1}/{len(scene_list)})")
        pack_scene_data(root, scene, output_dir)

if __name__ == "__main__":
    root = r"H:\AI\Datasets\nbv_rec_part2"
    output_dir = r"H:\AI\Datasets\upload_part2"
    scene_list = os.listdir(root)
    from_idx = 0
    to_idx = len(scene_list)
    print(f"packing {scene_list[from_idx:to_idx]}")

    pack_all_scenes(root, scene_list[from_idx:to_idx], output_dir)
    print("packing done")