import os
import numpy as np
from tqdm import tqdm
import joblib
import random
from joblib import Parallel, delayed

def func(smplx_path):
    gmr_path = smplx_path.replace('smplx_data', 'gmr_data')
    assert os.path.exists(smplx_path) and os.path.exists(gmr_path)
    return [np.load(smplx_path), np.load(gmr_path), smplx_path.replace('data/smplx_data/', '')]

data_paths = joblib.load('data/smplx_data/smplx_path.pkl')
all_data = Parallel(n_jobs=48)(delayed(func)(file_path) for file_path in tqdm(data_paths))
random.shuffle(all_data)
smplx_data = [item[0] for item in all_data]
gmr_data = [item[1] for item in all_data]

joblib.dump(all_data[:-20_000], 'data/train_400k.pkl')
joblib.dump(all_data[:200_000], 'data/train_200k.pkl')
joblib.dump(all_data[:100_000], 'data/train_100k.pkl')

joblib.dump(all_data[-20_000:], 'data/test_20k.pkl')
joblib.dump(all_data[-5_000:], 'data/test_5k.pkl')


joblib.dump(smplx_data[:-20_000], 'data/train_smplx_400k.pkl')
joblib.dump(smplx_data[:200_000], 'data/train_smplx_200k.pkl')
joblib.dump(smplx_data[:100_000], 'data/train_smplx_100k.pkl')

joblib.dump(smplx_data[-20_000:], 'data/test_smplx_20k.pkl')
joblib.dump(smplx_data[-5_000:], 'data/test_smplx_5k.pkl')


joblib.dump(gmr_data[:-20_000], 'data/train_gmr_400k.pkl')
joblib.dump(gmr_data[:200_000], 'data/train_gmr_200k.pkl')
joblib.dump(gmr_data[:100_000], 'data/train_gmr_100k.pkl')

joblib.dump(gmr_data[-20_000:], 'data/test_gmr_20k.pkl')
joblib.dump(gmr_data[-5_000:], 'data/test_gmr_5k.pkl')