import joblib
import random
import numpy as np
from tqdm import tqdm

data_root = 'data/final_data'
all_data = []
all_g1_motion = []
all_smplx_motion = []
with open(f'{data_root}/all.txt', 'r') as f:
    data_paths = f.readlines()
    for data_path in tqdm(data_paths):
        data = joblib.load(data_path.strip())
        all_data.append(data)
        all_g1_motion.append(data['g1_motion'])
        all_smplx_motion.append(data['smplx_motion'])

g1_data = np.concatenate(all_g1_motion, axis=0) 
Mean = g1_data.mean(axis=0)
Std = g1_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_root}/g1_mean.npy', Mean)
np.save(f'{data_root}/g1_std.npy', Std)

smplx_data = np.concatenate(all_smplx_motion, 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_root}/smplx_mean.npy', Mean)
np.save(f'{data_root}/smplx_std.npy', Std)

random.shuffle(all_data)
random.shuffle(all_data)
random.shuffle(all_data)
random.shuffle(all_data)
random.shuffle(all_data)
joblib.dump(all_data[:12500], f'{data_root}/train.pkl')
joblib.dump(all_data[12500:], f'{data_root}/test.pkl')