import os
import lzma
import pickle
from tqdm import tqdm

def load_pdm_dict(path):
    """The folder structure is like this:
    path/
        log_name/
        token/
            pdm_cache.pkl
        ...
    """
    pdm_dict = {}
    for root, _, files in tqdm(os.walk(path)):
        for file in files:
            if file.endswith('.pkl'):
                token = os.path.basename(root)
                with lzma.open(os.path.join(root, file), 'rb') as f:
                    data = pickle.load(f).trajectory
                pdm_dict[token] = data
    return pdm_dict



pdm_dict = load_pdm_dict('data/algengine_cache/pdm_8192_gt_cache_navtest2')
pickle.dump(pdm_dict, open('data/algengine_cache/pdm_8192_gt_cache_navtest2.pkl', 'wb'), protocol=pickle.HIGHEST_PROTOCOL)

pdm_dict = load_pdm_dict('data/algengine_cache/pdm_8192_gt_cache')
pickle.dump(pdm_dict, open('data/algengine_cache/pdm_8192_gt_cache.pkl', 'wb'), protocol=pickle.HIGHEST_PROTOCOL)