# merges the output of the main transfer_model script

import torch
from pathlib import Path
import pickle
from scipy.spatial.transform import Rotation as R
import numpy as np
KEYS = [
"transl",
"betas",
"full_pose",
]


def aggregate_rotmats(x):
    x = np.concatenate(x, axis=0)
    s = x.shape[:-2]
    try:
        x = R.from_matrix(x.reshape(-1, 3, 3)).as_rotvec()
    except:
        pass
    x = x.reshape(s[0], -1)
    return x

aggregate_function = {k: lambda x: np.concatenate(x, axis=0) for k in KEYS}
aggregate_function["betas"] = lambda x: np.concatenate(x, axis=0).mean(0)

for k in ["global_orient", "body_pose", "left_hand_pose", "right_hand_pose", "jaw_pose", "full_pose"]:
    aggregate_function[k] = aggregate_rotmats

def merge(output_dir, gender):
    output_dir = Path(output_dir)
    assert output_dir.exists()
    assert output_dir.is_dir()

    # get list of all pkl files in output_dir with fixed length numeral names
    pkl_files = [f for f in output_dir.glob("*.pkl") if f.stem != "merged"]
    pkl_files = [f for f in sorted(pkl_files, key=lambda x: int(x.stem))]
    assert "merged.pkl" not in [f.name for f in pkl_files]

    merged = {}
    # iterate over keys and put all values in lists
    keys = set(KEYS) 
    for k in keys:
        merged[k] = []
    for pkl_file in pkl_files:
        with open(pkl_file, "rb") as f:
            data = pickle.load(f)
        for k in keys:
            if k in data:
                merged[k].append(data[k])
    b = np.concatenate(merged["betas"], axis=0)
    print("betas:")
    for mu, sigma in zip(b.mean(0), b.std(0)):
        print("  {:.3f} +/- {:.3f}".format(mu, sigma))

    # aggregate all values
    for k in keys:
        merged[k] = aggregate_function[k](merged[k])

    # add gender

    poses = merged["full_pose"]
    trans = merged["transl"]
    if gender == "female":
        gender = np.zeros([poses.shape[0], 1])
    elif gender == "male":
        gender = np.ones([poses.shape[0], 1])
    else:
        gender = np.ones([poses.shape[0], 1]) * 2
    
    merged = np.concatenate([poses, trans, gender], axis=1)
    
    return merged


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Merge output of transfer_model script')
    parser.add_argument('output_dir', type=str, help='output directory of transfer_model script')
    parser.add_argument('--gender', type=str, choices=['male', 'female', 'neutral'], help='gender of actor in motion sequence')
    args = parser.parse_args()
    merge(args.output_dir, args.gender)