#! /usr/bin/python3


"""
Manipulate Openpose Data - Convert all JSON files Openpose extraction
 
Usage:
  extract-op-data.py convert <type> [--path <path>] [--out <out>] [--empty <val>]
  extract-op-data.py -h
  extract-op-data.py --version
   
Examples:

Arguments:

Options:
  -h, --help
  -p <path>, --path <path>              Json files directory
  -o <out>, --out <out>                 Output path and file name
  -d <detect>, --detect <detect>        Element's type to detect
  -e <val>, --empty <val>               keep empty values (1, true | 0, false) [default: 1]
  --version

"""
import time
import json
import os
import sys

from docopt import docopt
from collections import defaultdict

def main():
    opt = docopt(__doc__, sys.argv[1:])
    print(opt)
    out_file_name = opt["--out"]
    out_type      = opt["<type>"]
    
    list_files  = element_path_files(opt)
    points      = extract_op_data(list_files, "pose_keypoints_2d")
    
    # Delete '0' values
    if opt["--empty"] == '0' or opt["--empty"].lower() == 'false' :
        points = estimate_values(points)

    # Call dynamic function to export data
    getattr(sys.modules[__name__], "export_%s" % out_type)(points, out_file_name)

    
def element_path_files(opt):
    """
    Convert data from JSON files to list
    """
    path_dir = opt["--path"]
    if (path_dir):
        if os.path.exists(path_dir):
            print("chemin ok")

            dir_name = path_dir.split("/")[-1]
            # Suppression du "/" du chemin
            if dir_name == "":
                dir_name = path_dir.split("/")[-2] + "/"

            list_files = os.listdir(path_dir)
            list_files.sort()

            return {
                "list": list_files,
                "path": path_dir,
                "name": dir_name
            }
        else:
            print("L'arguments --path ne correspond pas à un repertoire")
            exit()

def extract_op_data(elements, *types):
    """
    Extract data from all JSON files
    """
    points = defaultdict(dict)
    
    for i,file_name in enumerate(elements["list"]):

        file_data = open(elements["path"] + file_name, "r").read()
        json_file = json.loads(file_data)
        
        # Get the file if there information in people key
        if json_file["people"] != []:

            for type_name in types:
                # print("1-")
                # print(json_file["people"][0][type_name])
                try:
                    points[type_name][i] = json_file["people"][0][type_name]
                except:
                    print("the type \"" + type_name + "\" not exist !")
                    exit()

    return points

def export_csv(points, out_file_name):
    """
    Export data into CSV file
    """

    for type_name in points:
        time_str   = time.strftime("%Y%m%d-%H%M%S")

        with open(out_file_name + "_" + type_name + "_" + time_str + ".csv", 'w') as text_file:
            for frame in points[type_name]:
                line = ""
                for i, val in enumerate(points[type_name][frame]):
                    if i % 3 != 2:
                        line += str(val) + ','
                        # text_file.write('\'' + str(val) + '\', ')
                # Remove the last comma
                line = line[0:-1]
                text_file.write(line + '\n')   
        print(type_name)
    
def estimate_values(points):
    """
    Replace each empty value by a new estimate value, between two values
    example : 
        15 -  0 -  0 -  0 - 19 --> Estimate empty values
        15 - 16 - 17 - 18 - 19
    """
    start = 0
    begin = 0.0

    for type_name in points:
        for i in range(len(points[type_name][0])):
            for frame in points[type_name]:
                if points[type_name][frame][i] == 0.0 and frame != 0:
                    if begin == 0.0:
                        begin = points[type_name][frame-1][i]
                        start = frame
                        # print(begin)
                elif start != 0:
                    ite  = frame - start
                    diff = abs(begin - points[type_name][frame][i]) / 3.0
                    
                    start = 0
                    begin = 0.0
                    for inc in range(ite):
                        points[type_name][frame-ite+inc][i] = round(points[type_name][frame-ite+inc-1][i] + diff, 3)

    return points



if __name__ == "__main__":
    main()