# 读取并且保存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)
#             if os.path.exists(path.replace('smplx_data', 'gmr_data')):
#                 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_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[:-20000], 'data/train_smplx_400k.pkl')
# joblib.dump(all_data[:100_000], 'data/train_smplx_100k.pkl')
joblib.dump(all_data[-20000:], 'data/test_smplx_20k.pkl')
# joblib.dump(all_data[-5000:], 'data/test_smplx_5k.pkl')
smplx_data = np.concatenate(all_data[::4], axis=0) 
Mean = smplx_data.mean(axis=0)
# Std = smplx_data.std(axis=0)
# joints_num = 30
# 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/gmr_mean.npy', Mean)
# np.save(f'data/gmr_std.npy', Std)