'''
@Author: Jiapeng Zhou
@Date: 2023-05-26 15:55:30
@Description: This is to get the smpl params "reconstructed_poses.hdf5" from videoavatar.
Usage: conda activate videoavatar && pythonn get_reconstructed_poses.py --root ~/Deexxx/xinyu_a --out videoavatars --gender xxx
'''
import argparse
import os, os.path as osp, sys
import pdb


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str)
    parser.add_argument('--out', type=str, default='videoavatars')
    parser.add_argument('--gender', type=str, choices=['male', 'female'],help=' "male" or "female" run_step1.sh needs to select smpl model of different gender')
    args = parser.parse_args()
    return args

def main():
    args = parse_args() 
    output_dir = osp.join(args.root, args.out)
    #os.makedirs(output_dir, exist_ok=True)
    # openpose => keypoints.hdf5
    output_pth =osp.join(output_dir, 'keypoints.hdf5') 
    if not osp.isfile(output_pth):
        os.system('python ./prepare_data/2djoints2hdf5.py {} {}'.\
                 format(osp.join(args.root, 'openpose'), output_pth) )

    # masks => masks.hdf5
    output_pth = osp.join(output_dir, 'masks.hdf5')
    if not osp.isfile(output_pth):
        os.system('python ./prepare_data/masks2hdf5.py {} {}'.\
                 format(osp.join(args.root, 'masks'), output_pth))

    # camera => camera.hdf5
    # TODO: the f and c is given by the camera intrinsics
    output_pth = osp.join(output_dir, 'camera.pkl')
    '''the compute formula:
    fx2 = fy, fy2 = fx, cx2 = 1080 - cy, cy2 = cx
    '''
    if not osp.isfile(output_pth):
#        os.system('python ./prepare_data/create_camera.py \
#                  {} 1080 1920 -f 914.87 915.067 -c 540.63 958.13 '\
#                  .format(output_pth))
        os.system('python ./prepare_data/create_camera.py \
                  {} 1080 1920  '\
                  .format(output_pth))

    # get reconstructed_poses.hdf5
    # TODO: the smpl model is needs to be selected according to gender
    os.system('bash run_step1.sh {} {} {}'.format(output_dir, output_dir, args.gender))
    return

if __name__ == '__main__':
    main()