from enum import IntEnum, unique

from typing import List

import cv2

import numpy as np

@unique

class CocoPart(IntEnum):

    """Body part locations in the 'coordinates' list."""

    Nose = 0

    LEye = 1

    REye = 2

    LEar = 3

    REar = 4

    LShoulder = 5

    RShoulder = 6

    LElbow = 7

    RElbow = 8

    LWrist = 9

    RWrist = 10

    LHip = 11

    RHip = 12

    LKnee = 13

    RKnee = 14

    LAnkle = 15

    RAnkle = 16

SKELETON_CONNECTIONS_COCO = [(0, 1, (210, 182, 247)), (0, 2, (127, 127, 127)), (1, 2, (194, 119, 227)),

                             (1, 3, (199, 199, 199)), (2, 4, (34, 189, 188)), (3, 5, (141, 219, 219)),

                             (4, 6, (207, 190, 23)), (5, 6, (150, 152, 255)), (5, 7, (189, 103, 148)),

                             (5, 11, (138, 223, 152)), (6, 8, (213, 176, 197)), (6, 12, (40, 39, 214)),

                             (7, 9, (75, 86, 140)), (8, 10, (148, 156, 196)), (11, 12, (44, 160, 44)),

                             (11, 13, (232, 199, 174)), (12, 14,

                                                         (120, 187, 255)), (13, 15, (180, 119, 31)),

                             (14, 16, (14, 127, 255))]

SKELETON_CONNECTIONS_5P = [('H', 'N', (210, 182, 247)), ('N', 'B', (210, 182, 247)), ('B', 'KL', (210, 182, 247)),

                           ('B', 'KR', (210, 182, 247)), ('KL', 'KR', (210, 182, 247))]

COLOR_ARRAY = [(210, 182, 247), (127, 127, 127), (194, 119, 227), (199, 199, 199), (34, 189, 188),

               (141, 219, 219), (207, 190, 23), (150, 152, 255), (189, 103, 148), (138, 223, 152)]

UNMATCHED_COLOR = (180, 119, 31)

# activity_dict = {

#     1.0: "Falling forward using hands",

#     2.0: "Falling forward using knees",

#     3: "Falling backwards",

#     4: "Falling sideward",

#     5: "Falling",

#     6: "Walking",

#     7: "Standing",

#     8: "Sitting",

#     9: "Picking up an object",

#     10: "Jumping",

#     11: "Laying",

#     12: "False Fall",

#     20: "None"

# }

activity_dict = {

    1.0: "Falling forward using hands",

    2.0: "Falling forward using knees",

    3: "Falling backwards",

    4: "Falling sideward",

    5: "FALL",

    6: "Normal",

    7: "Normal",

    8: "Normal",

    9: "Normal",

    10: "Normal",

    11: "Normal",

    12: "FALL Warning",

    20: "None"

}

def write_on_image(img: np.ndarray, text: str, color: List) -> np.ndarray:

    """Write text at the top of the image."""

    # Add a white border to top of image for writing text

    img = cv2.copyMakeBorder(src=img,

                             top=int(0.1 * img.shape[0]),

                             bottom=0,

                             left=0,

                             right=0,

                             borderType=cv2.BORDER_CONSTANT,

                             dst=None,

                             value=[255, 255, 255])

    for i, line in enumerate(text.split('\n')):

        y = 30 + i * 30

        cv2.putText(img=img,

                    text=line,

                    org=(0, y),

                    fontFace=cv2.FONT_HERSHEY_SIMPLEX,

                    fontScale=0.7,

                    color=color,

                    thickness=2)

    return img

def visualise(img: np.ndarray, keypoint_sets: List, width: int, height: int, vis_keypoints: bool = False,

              vis_skeleton: bool = False, CocoPointsOn: bool = False) -> np.ndarray:

    """Draw keypoints/skeleton on the output video frame."""

    if CocoPointsOn:

        SKELETON_CONNECTIONS = SKELETON_CONNECTIONS_COCO

    else:

        SKELETON_CONNECTIONS = SKELETON_CONNECTIONS_5P

    if vis_keypoints or vis_skeleton:

        for keypoints in keypoint_sets:

            if not CocoPointsOn:

                keypoints = keypoints["keypoints"]

            if vis_skeleton:

                for p1i, p2i, color in SKELETON_CONNECTIONS:

                    if keypoints[p1i] is None or keypoints[p2i] is None:

                        continue

                    p1 = (int(keypoints[p1i][0] * width), int(keypoints[p1i][1] * height))

                    p2 = (int(keypoints[p2i][0] * width), int(keypoints[p2i][1] * height))

                    if p1 == (0, 0) or p2 == (0, 0):

                        continue

                    cv2.line(img=img, pt1=p1, pt2=p2, color=color, thickness=3)

    return img

def visualise_tracking(img: np.ndarray, keypoint_sets: List, width: int, height: int, num_matched: int, vis_keypoints: bool = False,

                       vis_skeleton: bool = False, CocoPointsOn: bool = False) -> np.ndarray:

    """Draw keypoints/skeleton on the output video frame."""

    if CocoPointsOn:

        SKELETON_CONNECTIONS = SKELETON_CONNECTIONS_COCO

    else:

        SKELETON_CONNECTIONS = SKELETON_CONNECTIONS_5P

    if vis_keypoints or vis_skeleton:

        for i, keypoints in enumerate(keypoint_sets):

            if keypoints is None:

                continue

            if not CocoPointsOn:

                keypoints = keypoints["keypoints"]

            if vis_skeleton:

                for p1i, p2i, color in SKELETON_CONNECTIONS:

                    if keypoints[p1i] is None or keypoints[p2i] is None:

                        continue

                    p1 = (int(keypoints[p1i][0] * width), int(keypoints[p1i][1] * height))

                    p2 = (int(keypoints[p2i][0] * width), int(keypoints[p2i][1] * height))

                    if p1 == (0, 0) or p2 == (0, 0):

                        continue

                    if i < num_matched:

                        color = COLOR_ARRAY[i % 10]

                    else:

                        color = UNMATCHED_COLOR

                    cv2.line(img=img, pt1=p1, pt2=p2, color=color, thickness=3)

    return img