### Visualization for advanced user
import math

import cv2
import numpy as np


def draw_points(
    image,
    keypoints,
    scores,
    pose_keypoint_color,
    keypoint_score_threshold,
    radius,
    show_keypoint_weight,
):
    if pose_keypoint_color is not None:
        assert len(pose_keypoint_color) == len(keypoints)
    for kid, (kpt, kpt_score) in enumerate(zip(keypoints, scores)):
        x_coord, y_coord = int(kpt[0]), int(kpt[1])
        if kpt_score > keypoint_score_threshold:
            color = tuple(int(c) for c in pose_keypoint_color[kid])
            if show_keypoint_weight:
                cv2.circle(image, (int(x_coord), int(y_coord)), radius, color, -1)
                transparency = max(0, min(1, kpt_score))
                cv2.addWeighted(
                    image, transparency, image, 1 - transparency, 0, dst=image
                )
            else:
                cv2.circle(image, (int(x_coord), int(y_coord)), radius, color, -1)


def draw_links(
    image,
    keypoints,
    scores,
    keypoint_edges,
    link_colors,
    keypoint_score_threshold,
    thickness,
    show_keypoint_weight,
    stick_width=2,
):
    height, width, _ = image.shape
    if keypoint_edges is not None and link_colors is not None:
        assert len(link_colors) == len(keypoint_edges)
        for sk_id, sk in enumerate(keypoint_edges):
            x1, y1, score1 = (
                int(keypoints[sk[0], 0]),
                int(keypoints[sk[0], 1]),
                scores[sk[0]],
            )
            x2, y2, score2 = (
                int(keypoints[sk[1], 0]),
                int(keypoints[sk[1], 1]),
                scores[sk[1]],
            )
            if (
                x1 > 0
                and x1 < width
                and y1 > 0
                and y1 < height
                and x2 > 0
                and x2 < width
                and y2 > 0
                and y2 < height
                and score1 > keypoint_score_threshold
                and score2 > keypoint_score_threshold
            ):
                color = tuple(int(c) for c in link_colors[sk_id])
                if show_keypoint_weight:
                    X = (x1, x2)
                    Y = (y1, y2)
                    mean_x = np.mean(X)
                    mean_y = np.mean(Y)
                    length = ((Y[0] - Y[1]) ** 2 + (X[0] - X[1]) ** 2) ** 0.5
                    angle = math.degrees(math.atan2(Y[0] - Y[1], X[0] - X[1]))
                    polygon = cv2.ellipse2Poly(
                        (int(mean_x), int(mean_y)),
                        (int(length / 2), int(stick_width)),
                        int(angle),
                        0,
                        360,
                        1,
                    )
                    cv2.fillConvexPoly(image, polygon, color)
                    transparency = max(
                        0, min(1, 0.5 * (keypoints[sk[0], 2] + keypoints[sk[1], 2]))
                    )
                    cv2.addWeighted(
                        image, transparency, image, 1 - transparency, 0, dst=image
                    )
                else:
                    cv2.line(image, (x1, y1), (x2, y2), color, thickness=thickness)


palette = np.array(
    [
        [255, 128, 0],
        [255, 153, 51],
        [255, 178, 102],
        [230, 230, 0],
        [255, 153, 255],
        [153, 204, 255],
        [255, 102, 255],
        [255, 51, 255],
        [102, 178, 255],
        [51, 153, 255],
        [255, 153, 153],
        [255, 102, 102],
        [255, 51, 51],
        [153, 255, 153],
        [102, 255, 102],
        [51, 255, 51],
        [0, 255, 0],
        [0, 0, 255],
        [255, 0, 0],
        [255, 255, 255],
    ]
)

link_colors = palette[[0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16]]
keypoint_colors = palette[
    [16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0] + [4] * (52 - 17)
]