import numpy as np
import cv2


# draw pose
RADIUS = 5
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]])
SKELETON = [
    [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],
]
POSE_LIMB_COLOR = palette[[9, 9, 9, 9, 7, 7, 7,
                           0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]]
POSE_KPT_COLOR = palette[[16, 16, 16, 16,
                          16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]]


def display_frame_pose(image, bboxes, kpts):
    assert len(bboxes) == len(kpts), f"""Length of bboxes({len(bboxes)}) and 
                                        kpts({len(kpts)}) must be the same."""
    add_pose_display_meta(image, bboxes, kpts, draw_kpts=True, draw_limbs=True)


def add_pose_display_meta(image,
                          bboxes,
                          kpts,
                          draw_kpts=True,
                          draw_limbs=True,
                          KPT_THR=0.1):
    # Set radus = 5, thickness line = 4 by image size 1280
    ratio = image.shape[1]/1280
    RADIUS = int(5*ratio)
    THICKNESS = int(4*ratio)
    for bbox, kpt in zip(bboxes, kpts):

        # format kpt
        kpt = np.array(kpt).reshape(-1, 2)

        # keypoints
        if draw_kpts:
            # display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta)
            display_meta_num_circles = 1
            for j, pt in enumerate(kpt):
                # if sc < KPT_THR:
                #     continue
                r, g, b = POSE_KPT_COLOR[j]
                cnt = display_meta_num_circles-1
                # circle_params = display_meta_circle_params[cnt]
                circle_params_xc = int(pt[0])
                circle_params_yc = int(pt[1])
                circle_params_radius = RADIUS
                # circle_params_circle_color(r, g, b, 1.0)
                circle_params_has_bg_color = 1
                # circle_params_bg_color.(r, g, b, 1.0)
                image = cv2.circle(image, (circle_params_xc, circle_params_yc),
                                   circle_params_radius, (int(r), int(g), int(b)), -1)
                display_meta_num_circles += 1
                if display_meta_num_circles >= 16:
                    # display_meta = pyds.nvds_acquire_display_meta_from_pool(
                    #     batch_meta)
                    display_meta_num_circles = 1
                # import ipdb
                # ipdb.set_trace()
            # pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta)

        # limbs
        if draw_limbs:
            # display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta)
            display_meta_num_lines = 1
            for i, (s1, s2) in enumerate(SKELETON):

                pt1, pt2 = kpt[s1-1], kpt[s2-1]
                # sc1, sc2 = score[s1-1], score[s2-1]
                # if sc1 < KPT_THR or sc2 < KPT_THR:
                #     continue
                r, g, b = POSE_LIMB_COLOR[i]
                cnt = display_meta_num_lines-1
                # line_params = display_meta_line_params[cnt]
                line_params_x1 = int(pt1[0])
                line_params_y1 = int(pt1[1])
                line_params_x2 = int(pt2[0])
                line_params_y2 = int(pt2[1])
                line_params_line_width = 2
                # draw line
                cv2.line(image, (line_params_x1, line_params_y1),
                         (line_params_x2, line_params_y2), (int(r), int(g), int(b)), THICKNESS)
                # line_params_line_color.set(r, g, b, 1.0)
                display_meta_num_lines += 1

            #     if display_meta.num_lines >= 16:
            #         pyds.nvds_add_display_meta_to_frame(
            #             frame_meta, display_meta)
            #         display_meta = pyds.nvds_acquire_display_meta_from_pool(
            #             batch_meta)
            #         display_meta.num_lines = 1
            # pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta)