"""Executable Python script for running Python actions.

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
"""

'''
Some code is based on Ildoo Kim's code (https://github.com/ildoonet/tf-openpose) and https://gist.github.com/alesolano/b073d8ec9603246f766f9f15d002f4f4
and derived from the OpenPose Library (https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/LICENSE)
'''

from collections import defaultdict
from enum import Enum
import math
import numpy as np
import itertools
from scipy.ndimage.filters import maximum_filter
from PIL import Image, ImageDraw


class CocoPart(Enum):
    Nose = 0
    Neck = 1
    RShoulder = 2
    RElbow = 3
    RWrist = 4
    LShoulder = 5
    LElbow = 6
    LWrist = 7
    RHip = 8
    RKnee = 9
    RAnkle = 10
    LHip = 11
    LKnee = 12
    LAnkle = 13
    REye = 14
    LEye = 15
    REar = 16
    LEar = 17
    Background = 18


parts_dict = {'Nose': [0], 'Neck': [1], 'Shoulders': [2, 5], 'Elbows': [3, 6], 'Wrists': [
    4, 7], 'Hips': [8, 11], 'Knees': [9, 12], 'Ankles': [10, 13], 'Eyes': [14, 15], 'Ears': [16, 17]}
# parts_if_notfound_upper={'Eyes':'Ears','Ears':'Eyes','Nose':'Ears','Neck':'Nose','Shoulders':'Neck','Elbows':'Shoulders','Wrists':'Elbows','Hips':'Wrists#','Knees':'Hips'}
# parts_if_notfound_lower=#{'Ears':'Nose','Nose':'Neck','Neck':'Shoulders','Shoulders':'Elbows','Elbows':'Wrists','Wrists':'Hips','Hips':'Knees',
#                  'Knees':'Ankles','Ankles':'Knees'}
body_parts_dict = {0: 'Eyes', 1: 'Ears', 2: 'Nose', 3: 'Neck', 4: 'Shoulders',
                   5: 'Elbows', 6: 'Wrists', 7: 'Hips', 8: 'Knees', 9: 'Ankles'}

CocoPairs = [
    (1, 2), (1, 5), (2, 3), (3, 4), (5, 6), (6,
                                             7), (1, 8), (8, 9), (9, 10), (1, 11),
    (11, 12), (12, 13), (1, 0), (0, 14), (14,
                                          16), (0, 15), (15, 17), (2, 16), (5, 17)
]   # = 19
CocoPairsRender = CocoPairs[:-2]
CocoPairsNetwork = [
    (12, 13), (20, 21), (14, 15), (16, 17), (22,
                                             23), (24, 25), (0, 1), (2, 3), (4, 5),
    (6, 7), (8, 9), (10, 11), (28, 29), (30, 31), (34,
                                                   35), (32, 33), (36, 37), (18, 19), (26, 27)
]  # = 19

CocoColors = [(255, 0, 0), (255, 85, 0), (255, 170, 0), (255, 255, 0), (170, 255, 0), (85, 255, 0), (0, 255, 0),
              (0, 255, 85), (0, 255, 170), (0, 255, 255), (0,
                                                           170, 255), (0, 85, 255), (0, 0, 255), (85, 0, 255),
              (170, 0, 255), (255, 0, 255), (255, 0, 170), (255, 0, 85)]


NMS_Threshold = 0.1
InterMinAbove_Threshold = 6
Inter_Threashold = 0.1
Min_Subset_Cnt = 4
Min_Subset_Score = 0.8
Max_Human = 96


def human_conns_to_human_parts(human_conns, heatMat):
    human_parts = defaultdict(lambda: None)
    for conn in human_conns:
        human_parts[conn['partIdx'][0]] = (
            conn['partIdx'][0],  # part index
            (conn['coord_p1'][0] / heatMat.shape[2], conn['coord_p1']
             [1] / heatMat.shape[1]),  # relative coordinates
            heatMat[conn['partIdx'][0], conn['coord_p1']
                    [1], conn['coord_p1'][0]]  # score
        )
        human_parts[conn['partIdx'][1]] = (
            conn['partIdx'][1],
            (conn['coord_p2'][0] / heatMat.shape[2],
             conn['coord_p2'][1] / heatMat.shape[1]),
            heatMat[conn['partIdx'][1], conn['coord_p2']
                    [1], conn['coord_p2'][0]]
        )
    return human_parts


def non_max_suppression(heatmap, window_size=3, threshold=NMS_Threshold):
    heatmap[heatmap < threshold] = 0  # set low values to 0
    part_candidates = heatmap * \
        (heatmap == maximum_filter(heatmap, footprint=np.ones((window_size, window_size))))
    return part_candidates


def estimate_pose(heatMat, pafMat):
    if heatMat.shape[2] == 19:
        # transform from [height, width, n_parts] to [n_parts, height, width]
        heatMat = np.rollaxis(heatMat, 2, 0)
    if pafMat.shape[2] == 38:
        # transform from [height, width, 2*n_pairs] to [2*n_pairs, height, width]
        pafMat = np.rollaxis(pafMat, 2, 0)

    # reliability issue.
    heatMat = heatMat - heatMat.min(axis=1).min(axis=1).reshape(19, 1, 1)
    heatMat = heatMat - heatMat.min(axis=2).reshape(19, heatMat.shape[1], 1)

    _NMS_Threshold = max(np.average(heatMat) * 4.0, NMS_Threshold)
    _NMS_Threshold = min(_NMS_Threshold, 0.3)

    coords = []  # for each part index, it stores coordinates of candidates
    for heatmap in heatMat[:-1]:  # remove background
        part_candidates = non_max_suppression(heatmap, 5, _NMS_Threshold)
        coords.append(np.where(part_candidates >= _NMS_Threshold))

    # all connections detected. no information about what humans they belong to
    connection_all = []
    for (idx1, idx2), (paf_x_idx, paf_y_idx) in zip(CocoPairs, CocoPairsNetwork):
        connection = estimate_pose_pair(
            coords, idx1, idx2, pafMat[paf_x_idx], pafMat[paf_y_idx])
        connection_all.extend(connection)

    conns_by_human = dict()
    for idx, c in enumerate(connection_all):
        # at first, all connections belong to different humans
        conns_by_human['human_%d' % idx] = [c]

    no_merge_cache = defaultdict(list)
    empty_set = set()
    while True:
        is_merged = False
        for h1, h2 in itertools.combinations(conns_by_human.keys(), 2):
            if h1 == h2:
                continue
            if h2 in no_merge_cache[h1]:
                continue
            for c1, c2 in itertools.product(conns_by_human[h1], conns_by_human[h2]):
                # if two humans share a part (same part idx and coordinates), merge those humans
                if set(c1['uPartIdx']) & set(c2['uPartIdx']) != empty_set:
                    is_merged = True
                    # extend human1 connectios with human2 connections
                    conns_by_human[h1].extend(conns_by_human[h2])
                    conns_by_human.pop(h2)  # delete human2
                    break
            if is_merged:
                no_merge_cache.pop(h1, None)
                break
            else:
                no_merge_cache[h1].append(h2)

        if not is_merged:  # if no more mergings are possible, then break
            break

    conns_by_human = {h: conns for (
        h, conns) in conns_by_human.items() if len(conns) >= Min_Subset_Cnt}
    conns_by_human = {h: conns for (h, conns) in conns_by_human.items() if max(
        [conn['score'] for conn in conns]) >= Min_Subset_Score}

    humans = [human_conns_to_human_parts(
        human_conns, heatMat) for human_conns in conns_by_human.values()]
    return humans


def estimate_pose_pair(coords, partIdx1, partIdx2, pafMatX, pafMatY):
    connection_temp = []  # all possible connections
    peak_coord1, peak_coord2 = coords[partIdx1], coords[partIdx2]

    for idx1, (y1, x1) in enumerate(zip(peak_coord1[0], peak_coord1[1])):
        for idx2, (y2, x2) in enumerate(zip(peak_coord2[0], peak_coord2[1])):
            score, count = get_score(x1, y1, x2, y2, pafMatX, pafMatY)
            if (partIdx1, partIdx2) in [(2, 3), (3, 4), (5, 6), (6, 7)]:  # arms
                if count < InterMinAbove_Threshold // 2 or score <= 0.0:
                    continue
            elif count < InterMinAbove_Threshold or score <= 0.0:
                continue
            connection_temp.append({
                'score': score,
                'coord_p1': (x1, y1),
                'coord_p2': (x2, y2),
                'idx': (idx1, idx2),  # connection candidate identifier
                'partIdx': (partIdx1, partIdx2),
                'uPartIdx': ('{}-{}-{}'.format(x1, y1, partIdx1), '{}-{}-{}'.format(x2, y2, partIdx2))
            })

    connection = []
    used_idx1, used_idx2 = [], []
    for conn_candidate in sorted(connection_temp, key=lambda x: x['score'], reverse=True):
        if conn_candidate['idx'][0] in used_idx1 or conn_candidate['idx'][1] in used_idx2:
            continue
        connection.append(conn_candidate)
        used_idx1.append(conn_candidate['idx'][0])
        used_idx2.append(conn_candidate['idx'][1])

    return connection


def get_score(x1, y1, x2, y2, pafMatX, pafMatY):
    num_inter = 10
    dx, dy = x2 - x1, y2 - y1
    normVec = math.sqrt(dx ** 2 + dy ** 2)

    if normVec < 1e-4:
        return 0.0, 0

    vx, vy = dx / normVec, dy / normVec

    xs = np.arange(
        x1, x2, dx / num_inter) if x1 != x2 else np.full((num_inter, ), x1)
    ys = np.arange(
        y1, y2, dy / num_inter) if y1 != y2 else np.full((num_inter, ), y1)
    xs = (xs + 0.5).astype(np.int8)
    ys = (ys + 0.5).astype(np.int8)

    # without vectorization
    pafXs = np.zeros(num_inter)
    pafYs = np.zeros(num_inter)
    for idx, (mx, my) in enumerate(zip(xs, ys)):
        pafXs[idx] = pafMatX[my][mx]
        pafYs[idx] = pafMatY[my][mx]

    local_scores = pafXs * vx + pafYs * vy
    thidxs = local_scores > Inter_Threashold

    return sum(local_scores * thidxs), sum(thidxs)


def draw_humans(img1_raw, human_list):
    img = np.asarray(img1_raw)
    img_copied = np.copy(img)
    image_h, image_w = img_copied.shape[:2]
    centers = {}
    c = 10
    for human in human_list:
        part_idxs = human.keys()

        # draw point
        draw = ImageDraw.Draw(img1_raw)
        for i in range(CocoPart.Background.value):
            if i not in part_idxs:
                continue
            part_coord = human[i][1]
            center = (int(part_coord[0] * image_w + 0.5),
                      int(part_coord[1] * image_h + 0.5))
            centers[i] = center
            bbox = (center[0] - c, center[1] - c, center[0] + c, center[1] + c)
            draw.ellipse(bbox, fill=CocoColors[i])

        # draw line
        ctr = 1
        for pair_order, pair in enumerate(CocoPairsRender):
            if pair[0] not in part_idxs or pair[1] not in part_idxs:
                continue
            draw.line((centers[pair[0]][0], centers[pair[0]][1], centers[pair[1]]
                       [0], centers[pair[1]][1]), fill=CocoColors[pair_order], width=5)
    img1_raw = np.asarray(img1_raw)
    del draw
    return img1_raw


def crop_image(img, humans_list, upper_body, lower_body):
    upper_coord = 0.0
    upper_coord_x = 0.0
    lower_coord = 0.0
    lower_coord_x = 0.0

    img = np.asarray(img)
    image_h, image_w = img.shape[:2]

    if upper_body == 'Ankles' or lower_body == 'Eyes':
        raise NameError('Body parts not consistent')

    for human in humans_list:
        parts = human.keys()
        inte = parts_dict[upper_body]  # could be [1] or [2,3]

        if upper_body == 'Nose' or upper_body == 'Neck':
            upper_coord = human[inte[0]][1][1]  # interested only in heights.
            upper_coord_x = human[inte[0]][1][0]
        else:
            upper_coord = (human[inte[0]][1][1] + human[inte[1]][1][1])/2
            upper_coord_x = (human[inte[0]][1][0] + human[inte[1]][1][0])/2

        inte = parts_dict[lower_body]
        if lower_body == 'Nose' or lower_body == 'Neck':
            lower_coord = human[inte[0]][1][1]  # interested only in heights.
            lower_coord_x = human[inte[0]][1][0]
        else:
            lower_coord = (human[inte[0]][1][1] + human[inte[1]][1][1])/2
            lower_coord_x = (human[inte[0]][1][0] + human[inte[1]][1][0])/2

    image_h_u = int(upper_coord * image_h)
    image_h_l = int(lower_coord * image_h)

    image_w_left = int(upper_coord_x * image_w)
    image_w_right = int(lower_coord_x * image_w)
    aspect_ratio = image_h / image_w
    image_w = int((image_w_left + image_w_right)/2)

    img = img[image_h_u:image_h_l]
    wid = int((img.shape[0]/aspect_ratio)/2)
    img = img.transpose(1, 0, 2)
    img = img[image_w-2*wid:image_w+2*wid]
    img = img.transpose(1, 0, 2)

    crop_position = (image_w-2*wid, image_h_u)
    crop_size = (img.shape[1], img.shape[0])

    return img, crop_position, crop_size