import json
import os.path as osp
from LabanCalculate import xyz2laban
from tqdm import tqdm


def main():
    split = 'test'
    root_dir_path = osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__))))
    data_dir = osp.join(root_dir_path, 'data')
    annotations_dir = osp.join(data_dir, 'annotations')
    split_annotation_dir = osp.join(annotations_dir, split)
    data_file_name = 'InterHand2.6M_' + split + '_data.json'
    data_path = osp.join(split_annotation_dir, data_file_name)
    xyz_file_name = 'InterHand2.6M_' + split + '_joint_3d.json'
    xyz_path = osp.join(split_annotation_dir, xyz_file_name)
    laban_file_name = 'InterHand2.6M_' + split + '_laban.json'
    laban_path = osp.join(split_annotation_dir, laban_file_name)
    f = open(data_path, 'r')
    f_2 = open(xyz_path, 'r')
    data = json.load(f)
    xyz_data = json.load(f_2)
    laban_ans_dict = {}
    for i in tqdm(range(len(data['images']))):
        capture_id = data['images'][i]['capture']
        frame_id = data['images'][i]['frame_idx']
        capture_id_str = str(capture_id)
        frame_id_str = str(frame_id)
        joint_xyz = xyz_data[capture_id_str][frame_id_str]['world_coord']
        right_hand_joint = joint_xyz[:21].copy()
        left_hand_joint = joint_xyz[21:].copy()
        
        if capture_id_str not in laban_ans_dict:
            laban_ans_dict[capture_id_str] = {}
        if frame_id_str not in laban_ans_dict[capture_id_str]:
            laban_ans_dict[capture_id_str][frame_id_str] = {}
        
        right_laban_res = xyz2laban(right_hand_joint)
        left_laban_res = xyz2laban(left_hand_joint)

        laban_ans_dict[capture_id_str][frame_id_str]['right_hand_laban'] = right_laban_res
        laban_ans_dict[capture_id_str][frame_id_str]['left_hand_laban'] = left_laban_res
    with open(laban_path, 'w') as json_file:
        json.dump(laban_ans_dict, json_file)
    
if __name__ == '__main__':
    main()