import os
import pickle
from tqdm.auto import tqdm
from copy import deepcopy

            

if __name__ == "__main__":
    pred_path = \
        "/home/jovyan/users/bulat/workspace/3drec/Indoor/MaskClustering/data/prediction/arkit_gt_train"
    out_pkl_path = \
        "arkit_gt_train.pkl"
    gt_pkl_path = \
        "/home/jovyan/users/bulat/workspace/3drec/Indoor/OKNO/data/arkitscenes/arkitscenes_offline_infos_train.pkl"


    with open(gt_pkl_path, 'rb') as file:
        gt_data = pickle.load(file)

    new_data = {"metainfo": gt_data["metainfo"]}
    data_list = []

    picked_scenes = [scene[:-4] for scene in os.listdir(pred_path)]
    num = 0
    for scene in tqdm(gt_data['data_list']):
        scene_name = scene['lidar_points']['lidar_path'].split("_")[0]
        if scene_name not in picked_scenes:
            print(f"Scene {scene_name} not found in {pred_path}")
            continue
        num += 1
        tmp_scene = deepcopy(scene)
        data_list.append(tmp_scene)
    print(f"Number of scenes: {num}")
    new_data['data_list'] = data_list
    with open(out_pkl_path, 'wb') as f:
        pickle.dump(new_data, f)