import os
import pickle

traj_root = os.getenv('NAVSIM_TRAJPDM_ROOT')

if __name__ == '__main__':
    out_dir = 'vocab_expanded_8192_navtest'
    os.makedirs(f'{traj_root}/{out_dir}', exist_ok=True)

    ins = [f'navtest_sub{i}.pkl' for i in range(1, 6)]
    out = 'navtest.pkl'

    result = {}
    for in_pkl in ins:
        postfix = in_pkl.split('.')[0]
        sub = postfix.split('_')[1]

        curr_pickle = pickle.load(open(f'{traj_root}/{out_dir}_{sub}/{in_pkl}', 'rb'))
        print(f'{traj_root}/{out_dir}_{sub}/{in_pkl}', len(curr_pickle))
        for k, v in curr_pickle.items():
            result[k] = v
    print(f'Length: {len(result)}')
    pickle.dump(result, open(f'{traj_root}/{out_dir}/{out}', 'wb'))