#!/usr/bin/env python3

import os
import json
import argparse
import numpy as np


def flatten(l):
    return [item for sublist in l for item in sublist]


def get_score(pose):
    total = len(pose)
    score = 0
    for keypoint in pose:
        score += keypoint[2]

    score = score / total
    # print(score)
    return score


def main():
    # --------------------------------------------------------
    # Arguments
    # --------------------------------------------------------

    ap = argparse.ArgumentParser()
    ap.add_argument("-dataset_name", "--dataset_name",
                    help="Dataset name. Datasets must be inside hpe/images/.",
                    type=str, default='sim_moving')
    ap.add_argument("-cam", "--camera",
                    help="Choose camera detection. This camera must be present in the dataset.",
                    type=str, default='camera_2')
    args = vars(ap.parse_args())


    dataset_folder = '../../images/' + args['dataset_name'] + '/'

    if not os.path.exists(dataset_folder):
        raise Exception("The dataset folder does not exist!")

    if not os.path.exists(dataset_folder + args['camera']):
        raise Exception("The given sensor " + args['camera'] + " does not exist in the dataset!")

    openpose_file = dataset_folder + 'cache/' + args['camera'] + '/2d_pose.npy'
    poses_np = np.load(openpose_file, allow_pickle=True)
    # print(openpose_file)

    alphapose_json = []

    i = 0
    for pose_np in poses_np.tolist():
        temp_dict = {}
        temp_dict['image_id'] = str(1) + '.png'
        temp_dict['category_id'] = 1
        temp_dict['keypoint_list'] = flatten(pose_np)
        temp_dict['score'] = get_score(pose_np)
        # print(file)
        # print("POSE")
        # print(pose_np)
        i += 1
        alphapose_json.append(temp_dict)

        # print(alphapose_json)
    alphapose_file_path=dataset_folder + 'cache/' + args['camera'] + "/alphapose.json"
    with open(alphapose_file_path, "w") as fp:
        json.dump(alphapose_json, fp)


if __name__ == "__main__":
    main()

# [
#  // for person_1 in image_1
#  {
#     "image_id" : string, image_1_name,
#     "category_id" : int, 1 for person
#     "keypoint_list" : [x1,y1,c1,...,xk,yk,ck],
#     "score" : float,
#  },
#  // for person_2 in image_1
#  {
#     "image_id" : string, image_1_name,
#     "category_id" : int, 1 for person
#     "keypoint_list" : [x1,y1,c1,...,xk,yk,ck],
#     "score" : float,
#  },
#  ...
#  // for persons in image_2
# {
#     "image_id" : string, image_2_name,
#     "category_id" : int, 1 for person
#     "keypoint_list" : [x1,y1,c1,...,xk,yk,ck],
#     "score" : float,
#  },
#  ...
# ]