# AUTOGENERATED! DO NOT EDIT! File to edit: 02_features.ipynb (unless otherwise specified).

__all__ = ['calculate_stats', 'stats_for_paths', 'jsonize']

# Cell
import json
import gzip
import joblib
import numpy as np
import llamass.core
from .tqdm import ProgressParallel


def calculate_stats(npz_path):
    data = np.load(npz_path)
    pose = np.concatenate([data['poses'], data['trans']], axis=1)
    return npz_path, (np.mean(pose, 0), np.std(pose, 0), pose.shape[0])

def stats_for_paths(unpacked_directory, n_jobs):
    paths = [p for p in llamass.core.npz_paths(unpacked_directory)]
    return ProgressParallel(n_jobs=n_jobs)(
        [joblib.delayed(calculate_stats)(npz_path) for npz_path in paths], total=len(paths)
    )

def jsonize(stats):
    return [(p, (list(m), list(s))) for p, (m, s) in stats]

#amass_loc = '/nobackup/gngdb/repos/amass/data'
#stats = stats_for_paths(amass_loc, 8)