import numpy as np
import h5py
import os
import json
import argparse
from links import joints_human36m, links_human36m


def main():
    # --------------------------------------------------------
    # Arguments
    # --------------------------------------------------------
    ap = argparse.ArgumentParser()

    ap.add_argument("-dataset", "--dataset_name",
                    help="Dataset name. Datasets must be inside hpe/images/. Please use -dataset human36m to optimize Human 3.6M.",
                    type=str, default='human36m')
    ap.add_argument("-s", "--section",
                    help="Human 3.6M section to optimize. Only works when the dataset is set to human36m",
                    type=str, default='S1')
    ap.add_argument("-a", "--action",
                    help="Human 3.6M action to optimize. Only works when the dataset is set to human36m",
                    type=str, default='Directions-1')

    args = vars(ap.parse_args())

    file_name = '../../images/' + args['dataset_name'] + '/' + "extra/una-dinosauria-data/h36m/" + args[
        'section'] + "/MyPoses/3D_positions/" + args['action'].replace("-", " ") + '.h5'

    gt = {}

    # Load some 3d poses
    with h5py.File(file_name) as h5f:
        poses3d = h5f["3D_positions"][:]
        print("File successfully read!")
        poses3d = poses3d / 1000
        # poses3d=poses3d[:,0]
        # print(poses3d.shape)
        # exit(0)

        # i=0

        for i in range(poses3d.shape[1]):
            gt[i] = {}
            pose = poses3d[:, i]
            joint_idx = 0
            for n in range(0, 96, 3):
                gt[i][joint_idx] = {}
                # print(n)
                gt[i][joint_idx]['x'] = pose[n]
                gt[i][joint_idx]['y'] = pose[n + 1]
                gt[i][joint_idx]['z'] = pose[n + 2]
                joint_idx += 1

    # print(gt)
    # print(pose.shape)
    # i+=1
    # print(i)
    # print(i)

    new_gt = {}

    # new_gt[joint_key]={}
    for frame in gt:
        new_gt[frame] = {}
        for joint_key in joints_human36m:
            # print(joint_key)
            idx = joints_human36m[joint_key]['idx']
            new_gt[frame][joint_key] = {}
            new_gt[frame][joint_key]['pose'] = {}
            # print(gt[frame][idx]['x'])
            new_gt[frame][joint_key]['pose']['x'] = gt[frame][idx]['x']
            new_gt[frame][joint_key]['pose']['y'] = gt[frame][idx]['y']
            new_gt[frame][joint_key]['pose']['z'] = gt[frame][idx]['z']

        # print(frame)
    # print(idx)

    # print(new_gt)
    output_path = '../../images/' + args['dataset_name'] + '/processed/'+ args[
        'section'] +'/' + args['action'] + "/ground_truth.txt"

    with open(output_path, "w") as fp:
        # Load the dictionary from the file
        json.dump(new_gt, fp)
        print("Ground truth file saved successfully!")


if __name__ == "__main__":
    main()