import json
from pprint import pprint
import cv2
import numpy as np
from scipy.spatial import distance
import os
import matplotlib.pyplot as plt

path_root_file = "D:\Work\ComputerV[261458]\Contest\Examples\\"
path_out = "D:\Work\ComputerV[261458]\Contest\Frame_example\\"
path_list_file = "D:\Work\ComputerV[261458]\Contest\list.txt"
path_list_file_test = "D:\Work\ComputerV[261458]\Contest\\test.txt"
path_root_json = "D:\Work\ComputerV[261458]\Contest\Json\\"
path_pose_out = "D:\Work\ComputerV[261458]\Contest\Pose_out\\"
path_save_plot = "D:\Work\ComputerV[261458]\Contest\Graph\\"

file_ans_r = open("D:\Work\ComputerV[261458]\Contest\\ans_r.txt","w")
file_ans_l = open("D:\Work\ComputerV[261458]\Contest\\ans_l.txt","w")
file_out = open("D:\Work\ComputerV[261458]\Contest\\result.txt",'w')
data_window_size = 7
l = []
r = []
def main():
    try:
        os.stat(path_out)
    except:
        os.mkdir(path_out)
    try:
        os.stat(path_root_json)
    except:
        os.mkdir(path_root_json)
    # f = open(path_list_file_test,"r")
    f = open(path_list_file,"r")
    lines = f.readlines()
    # for i in lines:
    #     name = i.split(".")[0]
    #     path_file = path_root_file + name + ".mp4"
    #     print (i.split("\n")[0])
    #     frame = sampling_frame(path_file)
    #     c = 1
    #     for i in frame:
    #         if i.shape[0] > i.shape[1]:
    #             i = cv2.resize(i, (360,640))
    #         else:
    #              i = cv2.resize(i, (640,360))
    #         cv2.imwrite(path_out + name + "_" + str(c) + ".jpg",i)
    #         c += 1
    # script_generate()

    for line in lines:
        ans_l = ans_r = 0
        print (line.strip('\n'))
        point_neck,point_l_shoulder_l,point_l_shoulder_r,point_r_shoulder_l,point_r_shoulder_r,point_l_elbow,point_r_elbow,point_l_wrist,point_r_wrist,conf_l,conf_r,size_l,size_r = read_json(path_root_json,line)

        left_top_arm_dist = cal_distance(point_l_shoulder_l,point_l_elbow)
        right_top_arm_dist = cal_distance(point_r_shoulder_r,point_r_elbow)
        left_bt_arm_dist = cal_distance(point_l_wrist,point_l_elbow)
        right_bt_arm_dist = cal_distance(point_r_wrist,point_r_elbow)

        ratio_shoulder_l = cal_distance(point_l_shoulder_l,point_r_shoulder_l)
        ratio_shoulder_r = cal_distance(point_l_shoulder_r,point_r_shoulder_r)
        if size_l > 20:
            punch_l_dist,max_l,min_l,del_conf_l = count_punch_distance(left_top_arm_dist,ratio_shoulder_l,left_bt_arm_dist,conf_l,line.strip('\n') + '_left_dist')
            punch_l_conf = count_punch_conf(conf_l,del_conf_l,line.strip('\n') + '_left_conf')
        if size_r > 20:
            punch_r_dist,max_r,min_r,del_conf_r = count_punch_distance(right_top_arm_dist,ratio_shoulder_r,right_bt_arm_dist,conf_r,line.strip('\n') + '_right_dist')
            punch_r_conf = count_punch_conf(conf_r,del_conf_r,line.strip('\n') + '_right_conf')
        ans_l = punch_l_dist + punch_l_conf
        ans_r = punch_r_dist + punch_r_conf
        l.append(ans_l)
        r.append(ans_r)
        file_ans_l.write(str(ans_l) + "\n")
        file_ans_r.write(str(ans_r) + "\n")

    for i in range(len(l)):
        file_out.write(line.strip('\n') + ":" + str(l[i]) + " " + str(r[i]) + "\n")
    print (ans_l,ans_r)


def sampling_frame(path):
    cap = cv2.VideoCapture(path)
    video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 1
    # print (video_length)
    frames = []
    frame_ids = []
    if cap.isOpened() and video_length > 0:
        if video_length >= 4:
            for i in range(1,video_length):
                if video_length <= 150 :
                    frame_ids.append(i)
                else:
                    if i % 2 == 0:
                        frame_ids.append(i)
        count = 0
        success, image = cap.read()
        while success:
            if count in frame_ids:
                frames.append(image)
            success, image = cap.read()
            count += 1
    return frames

def script_generate():
    cmd = 'bin\OpenPoseDemo.exe --image_dir ' + path_out + ' -model_pose COCO -number_people_max 1 --write_json ' + path_root_json + ' --display 0 --render_pose 0'
    # cmd = 'bin\OpenPoseDemo.exe --image_dir ' + path_out + ' -model_pose COCO -number_people_max 1 --write_json ' + path_root_json + ' --write_images ' + path_pose_out + ' --display 0'
    # cmd = 'bin\OpenPoseDemo.exe --image_dir ' + path_out + ' -model_pose COCO -number_people_max 1 --write_images ' + path_pose_out + ' --display 0 '
    os.system(cmd)

def read_json(path,line):
    point_neck = []
    point_r_shoulder_r = []
    point_r_shoulder_l = []
    point_l_shoulder_r = []
    point_l_shoulder_l = []
    point_r_elbow = []
    point_l_elbow = []
    point_r_wrist = []
    point_l_wrist = []
    conf_l = []
    conf_r = []
    n_frame = 1
    size_l = 0
    size_r = 0
    for file in os.listdir(path):
        l = file.split("_")
        # print (l[0],line.strip('\n').strip('.mp4'))
        if l[0] == line.split('.')[0]:
            path_json = path + l[0] + '_' + str(n_frame) + '_keypoints.json'
            # print (path_json)
            is_exists = os.path.isfile(path_json)
            if is_exists:
                with open(path_json) as f:
                    data = json.load(f)
                # pprint(data)
                position = []
                confident = []
                count_3 = 1
                pair = []
                if data[u'people']:
                    for i in data[u'people'][0][u'pose_keypoints_2d']:
                        if count_3 % 3 == 0:
                            position.append(pair)
                            confident.append(i)
                            pair = []
                        elif count_3 % 3 == 1:
                            pair.append(i)
                        elif count_3 % 3 == 2:
                            pair.append(i)
                        count_3 += 1
                    check_left = True
                    check_right = True
                    check_right_list = [1,2,3,5]
                    check_left_list = [1,2,5,6]
                    for i in check_left_list:
                        if position[i][0] == 0:
                            check_left = False
                    for i in check_right_list:
                        if position[i][0] == 0:
                            check_right = False
                    if check_right:
                        size_r += 1
                        point_neck.append(position[1])
                        point_r_shoulder_r.append(position[2])
                        point_r_elbow.append(position[3])
                        point_l_shoulder_r.append(position[5])
                        point_r_wrist.append(position[4])
                        conf_r.append((confident[2] + confident[3])/2)
                    if check_left:
                        size_l += 1
                        point_l_shoulder_l.append(position[5])
                        point_r_shoulder_l.append(position[2])
                        point_l_elbow.append(position[6])
                        point_l_wrist.append(position[7])
                        conf_l.append((confident[5] + confident[6])/2)
            else :
                n_frame = 0
        n_frame += 1
    return point_neck,point_l_shoulder_l,point_l_shoulder_r,point_r_shoulder_l,point_r_shoulder_r,point_l_elbow,point_r_elbow,point_l_wrist,point_r_wrist,conf_l,conf_r,size_l,size_r

def find_bd_box(point_1,point_2):
    min_x = 50000
    min_y = 50000
    max_x = -1
    max_y = -1
    for i in point_1:
        min_x = min(min_x,i[0])
        min_y = min(min_y,i[1])
        max_x = max(max_x,i[0])
        max_y = max(max_y,i[1])
    for i in point_2:
        min_x = min(min_x,i[0])
        min_y = min(min_y,i[1])
        max_x = max(max_x,i[0])
        max_y = max(max_y,i[1])
    return min_x,min_y,max_x-min_x,max_y-min_y

def cal_distance(point_list_1,point_list_2):
    dist_list = []
    for i in range(len(point_list_1)):
        if point_list_1[i] == [0,0] or point_list_2[i] == [0,0]:
            dist_list.append(0)
        else :
            dist_list.append(distance.euclidean(point_list_1[i],point_list_2[i]))
    return dist_list

def data_process (data,window_filter_size):
    data_filter = filter_data(data,window_filter_size,1)
    max_data,min_data,avg_data,per_avg_data = math_process_data(data_filter,window_filter_size)
    return data_filter,max_data,min_data,avg_data,per_avg_data

def filter_data(data,window_filter_size,scale_rate):
    data_filtered = []
    window = []
    for i in range(window_filter_size):
        if i < window_filter_size/2:
            window.append(0)
        else:
            window.append(data[i])
    for i in range(len(data)):
        s = sum(window) + (window[int(window_filter_size/2)+1]*scale_rate) #increase mean Doble mid
        data_filtered.append((s/window_filter_size)/2)

        for j in range(window_filter_size):
            if j+1 < window_filter_size:
                window[j] = window[j+1]
            else:
                if i+int(window_filter_size/2)+1 < len(data):
                    window[j] = data[i+int(window_filter_size/2)+1]
                else:
                    window[j] = 0
    return data_filtered

def math_process_data (data_filter,window_filter_size):
    cut_zero_data = []
    for i in range(len(data_filter)):
        if int(window_filter_size/2) < i < len(data_filter) - int(window_filter_size/2):
            cut_zero_data.append(data_filter[i])
    max_data = max(cut_zero_data)
    min_data = min(cut_zero_data)
    avg_data = sum(cut_zero_data)/len(cut_zero_data)
    per_avg_data = (avg_data - min_data)/(max_data - min_data)*100
    return max_data,min_data,avg_data,per_avg_data

def count_punch_distance(top_arm_distance,ratio_shoulder,bt_arm_distance,conf,label):
    upper_rate_shoulder_top = []
    lower_rate_shoulder_top = []
    del_conf = []
    top_arm_dist_filtered,max_data,min_data,avg_data,per_avg_data = data_process(top_arm_distance,data_window_size)
    bt_arm_dist_filtered = filter_data(bt_arm_distance,data_window_size,0)
    for i in range(len(ratio_shoulder)):
        if ratio_shoulder[i] != 0:
            upper_rate_shoulder_top.append(ratio_shoulder[i]*0.80)
            lower_rate_shoulder_top.append(ratio_shoulder[i]*0.74)
    upper_rate_shoulder_top = filter_data(upper_rate_shoulder_top,15,1)
    lower_rate_shoulder_top = filter_data(lower_rate_shoulder_top,15,1)
    # plt.xlabel(label)
    # plt.ylim(0, 120)
    # plt.plot(np.arange(len(top_arm_dist_filtered)), top_arm_dist_filtered, color="red")
    # plt.plot(np.arange(len(upper_rate_shoulder_top)), upper_rate_shoulder_top, color="black")
    # plt.plot(np.arange(len(lower_rate_shoulder_top)), lower_rate_shoulder_top, color="gray")
    # plt.savefig(path_save_plot + label + '.png')
    # plt.close()
    if len(top_arm_dist_filtered) != 0:
        if per_avg_data > 40:
            if max_data - min_data > 10:
                punch = 0
                punching = True
                index = 0
                for i in range(len(top_arm_dist_filtered)):
                    if top_arm_dist_filtered[i] > upper_rate_shoulder_top[i]:
                        punching = False
                    if punching:
                        del_conf.append(i)
                    if top_arm_dist_filtered[i] < lower_rate_shoulder_top[i] and not punching and conf[i] > 0.2:
                        punch += 1
                        punching = True
                        del_conf.append(i)
                return punch,max_data,min_data,del_conf
            return 0,max_data,min_data,del_conf
        else:
            return 0,max_data,min_data,del_conf
    else :
        # print("D")
        return 0,0,0,del_conf

def count_punch_conf(conf,del_conf,label):
    for i in del_conf:
        conf.pop(i)
        for j in range(len(del_conf)):
            if del_conf[j] > i :
                del_conf[j] -= 1

    conf_filtered = filter_data(conf,3,0)
    upper_ratio = ((max(conf_filtered) - min(conf_filtered))*0.79) + min(conf_filtered)
    lower_ratio = ((max(conf_filtered) - min(conf_filtered))*0.38)  + min(conf_filtered)
    x = []
    y = []
    for i in range(len(conf_filtered)):
        x.append(upper_ratio)
        y.append(lower_ratio)
    # plt.ylim(0, 1)
    # plt.xlabel(label)
    # plt.plot(np.arange(len(conf_filtered)),conf_filtered, color="green")
    # plt.plot(np.arange(len(conf_filtered)),x, color="black")
    # plt.plot(np.arange(len(conf_filtered)),y, color="gray")
    # plt.show()
    # plt.savefig(path_save_plot + label + '.png')
    # print (max(conf) - min(conf))
    if max(conf) - min(conf) > 0.1:
        punching = True
        punch = 0
        for i in range(1,len(conf_filtered)-1):
            # print (conf_filtered[i],lower_ratio)
            if conf_filtered[i] < lower_ratio and not punching:
                punch += 1
                punching = True
            if conf_filtered[i] > upper_ratio:
                punching = False
        return punch
    return 0

main()