from PIL import Image, ImageDraw, ImageFont
import numpy as np
import utils


skeleton = np.array([
    [16, 14],
    [14, 12],
    [17, 15],
    [15, 13],
    [12, 13],
    [6, 12],
    [7, 13],
    [6, 7],
    [6, 8],
    [7, 9],
    [8, 10],
    [9, 11],
    [2, 3],
    [1, 2],
    [1, 3],
    [2, 4],
    [3, 5],
    [4, 6],
    [5, 7],
]) - 1


def draw_keypoints(img: Image.Image, keypoints: dict):
    draw = ImageDraw.Draw(img)

    font_size = max(round(sum(img.size) / 2 * 0.035), 12)
    font = ImageFont.truetype("DejaVuSerif.ttf", font_size)

    color = (0, 255, 0)
    lw = max(round(sum(img.size) / 2 * 0.005), 2)
    lw_poly = int(round(0.75*lw))
    r = max(round(sum(img.size) / 2 * 0.0132), 8)
    text_margin = r

    # draw skeleton
    keypoint_list = list(keypoints.values())
    for i, sk in enumerate(skeleton):
        if keypoint_list[sk[0]] and keypoint_list[sk[1]]:
            draw.line([keypoint_list[sk[0]], keypoint_list[sk[1]]], fill=color, width=lw)
            # draw rectangle polygons around keypoint_list
            poly1 = [keypoint_list[sk[0]][0] - r, keypoint_list[sk[0]][1] - r, keypoint_list[sk[0]][0] + r, keypoint_list[sk[0]][1] + r]
            poly1 = [(poly1[0], poly1[1]), (poly1[0], poly1[3]), (poly1[2], poly1[3]), (poly1[2], poly1[1])]
            draw.polygon(poly1, outline=color, width=lw_poly)
            poly2 = [keypoint_list[sk[1]][0] - r, keypoint_list[sk[1]][1] - r, keypoint_list[sk[1]][0] + r, keypoint_list[sk[1]][1] + r]
            poly2 = [(poly2[0], poly2[1]), (poly2[0], poly2[3]), (poly2[2], poly2[3]), (poly2[2], poly2[1])]
            draw.polygon(poly2, outline=color, width=lw_poly)

    # draw angles
    ear, eye = None, None
    is_left = None
    if keypoints["left_ear"] and keypoints["left_eye"]:
        ear = keypoints["left_ear"]
        eye = keypoints["left_eye"]
        is_left = True
    elif keypoints["right_ear"] and keypoints["right_eye"]:
        ear = keypoints["right_ear"]
        eye = keypoints["right_eye"]
        is_left = False

    # draw extended left and right eye lines
    if ear and eye:
        left_new_point = utils.extend_line(ear, eye, 3)
        l1 = [ear, left_new_point]
        draw.line(l1, fill='red', width=lw)
        # draw a horizontal line from ear forwards
        ear = np.array(ear)
        l1 = np.array(l1)
        l1_vector = l1[1] - l1[0]
        x_s = np.sign(l1_vector)[0]
        length_l1 = np.linalg.norm(l1_vector)
        p2 = ear + np.array([length_l1*x_s, 0])
        ear = tuple(ear.tolist())
        l = [ear, tuple(p2.tolist())]
        # draw.line(l, fill='gray', width=lw//2)
        # draw angle
        angle = utils.calculate_angle_to_horizontal(l1_vector)
        txt = f'{angle:.0f}°'
        txt_w, txt_h = font.getbbox(txt)[2:]
        if is_left:
            text_coords = (ear[0] - txt_w - text_margin, ear[1] - txt_h - text_margin)
        else:
            text_coords = (ear[0] + text_margin, ear[1] - txt_h - text_margin)
        draw.text(text_coords, txt, fill='red', font=font)

    # draw elbow angles
    left_elbow_angle, right_elbow_angle = utils.get_elbow_angles(keypoints)
    if left_elbow_angle and is_left is not False:
        txt = f'{left_elbow_angle:.0f}°'
        txt_w, txt_h = font.getbbox(txt)[2:]
        # adjust upwards and leftwards
        text_coords = (keypoints['left_elbow'][0] - txt_w - text_margin, keypoints['left_elbow'][1] - txt_h - text_margin)
        draw.text(text_coords, txt, fill='red', font=font)
        # draw polyline for left arm
        # draw.line([keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist']], fill='blue', width=lw)
    if right_elbow_angle and is_left is not True:
        txt = f'{right_elbow_angle:.0f}°'
        txt_w, txt_h = font.getbbox(txt)[2:]
        text_coords = (keypoints['right_elbow'][0] + text_margin, keypoints['right_elbow'][1] - txt_h - text_margin)
        draw.text(text_coords, txt, fill='red', font=font)
        # draw polyline for right arm
        # draw.line([keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist']], fill='blue', width=lw)

    return img