from __future__ import print_function

import os
import os.path as osp
import pickle

import argparse

import numpy as np


def merge_models(smplh_fn, mano_left_fn, mano_right_fn, output_folder="output"):

    if smplh_fn.endswith(".pkl"):
        with open(smplh_fn, "rb") as body_file:
            body_data = pickle.load(body_file, encoding="latin1")
    elif smplh_fn.endswith(".npz"):
        body_data_np = np.load(smplh_fn)
        body_data = {}
        for key in body_data_np:
            body_data[key] = body_data_np[key]
    else:
        raise ValueError("The body model file should be either a .pkl or a .npz file.")

    with open(mano_left_fn, "rb") as lhand_file:
        lhand_data = pickle.load(lhand_file, encoding="latin1")

    with open(mano_right_fn, "rb") as rhand_file:
        rhand_data = pickle.load(rhand_file, encoding="latin1")

    out_fn = osp.split(smplh_fn)[1]
    if out_fn.endswith(".npz"):
        out_fn = out_fn.replace(".npz", ".pkl")

    output_data = body_data.copy()
    output_data["hands_componentsl"] = lhand_data["hands_components"]
    output_data["hands_componentsr"] = rhand_data["hands_components"]

    output_data["hands_coeffsl"] = lhand_data["hands_coeffs"]
    output_data["hands_coeffsr"] = rhand_data["hands_coeffs"]

    output_data["hands_meanl"] = lhand_data["hands_mean"]
    output_data["hands_meanr"] = rhand_data["hands_mean"]

    for key, data in output_data.items():
        if "chumpy" in str(type(data)):
            output_data[key] = np.array(data)
        else:
            output_data[key] = data

    out_path = osp.join(output_folder, out_fn)
    print(out_path)
    print("Saving to {}".format(out_path))
    with open(out_path, "wb") as output_file:
        pickle.dump(output_data, output_file)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--smplh-fn",
        dest="smplh_fn",
        required=True,
        type=str,
        help="The path to the SMPLH model",
    )
    parser.add_argument(
        "--mano-left-fn",
        dest="mano_left_fn",
        required=True,
        type=str,
        help="The path to the left hand MANO model",
    )
    parser.add_argument(
        "--mano-right-fn",
        dest="mano_right_fn",
        required=True,
        type=str,
        help="The path to the right hand MANO model",
    )
    parser.add_argument(
        "--output-folder",
        dest="output_folder",
        required=True,
        type=str,
        help="The path to the output folder",
    )

    args = parser.parse_args()

    smplh_fn = args.smplh_fn
    mano_left_fn = args.mano_left_fn
    mano_right_fn = args.mano_right_fn
    output_folder = args.output_folder

    if not osp.exists(output_folder):
        print("Creating directory: {}".format(output_folder))
        os.makedirs(output_folder)

    merge_models(smplh_fn, mano_left_fn, mano_right_fn, output_folder)
