# 读取并且保存smplx路径
# import os
# import joblib
# import random
# from tqdm import tqdm
# data_root = "data/smplx_data"
# all_data = []
# for root, dirs, files in tqdm(os.walk(data_root)):
#     for file in files:
#         if file.endswith('.npy'):
#             path = os.path.join(root, file)
#             # import ipdb;  ipdb.set_trace()  
#             if os.path.exists(path.replace('smplx_data', '2_gmr_retarget_rl_0_120')) or \
#                 os.path.exists(path.replace('smplx_data', '2_gmr_retarget_rl_200_300')) or \
#                 os.path.exists(path.replace('smplx_data', '2_gmr_retarget_rl_400_500')):
#                 continue
#             all_data.append(path)
#         if len(all_data) % 10000 == 0:
#             print(len(all_data))
# joblib.dump(all_data, 'data/smplx_path.pkl')


import joblib
import numpy as np
from joblib import Parallel, delayed
from tqdm import tqdm
import random
data_paths = joblib.load('data/smplx_path.pkl')
all_data = Parallel(n_jobs=24)(delayed(np.load)(file_path) for file_path in tqdm(data_paths))
random.shuffle(all_data)
import ipdb; ipdb.set_trace()
joblib.dump(all_data[:-10000], 'data/train_smplx.pkl')
joblib.dump(all_data[-10000:], 'data/test_smplx.pkl')

smplx_data = np.concatenate(all_data, axis=0) 
Mean = smplx_data.mean(axis=0)
Std = smplx_data.std(axis=0)
joints_num = 22
Std[:2] = Std[:2].mean() / 1.0
Std[2:8] = Std[2:8].mean() / 1.0
Std[8:8+3*joints_num] = Std[8:8+3*joints_num].mean() / 1.0
Std[8+3*joints_num:8+6*joints_num] = Std[8+3*joints_num:8+6*joints_num].mean() / 1.0
Std[8+6*joints_num:] = Std[8+6*joints_num:].mean() / 1.0
np.save(f'data/smplx_mean.npy', Mean)
np.save(f'data/smplx_std.npy', Std)