import math
import json
import numpy as np
import matplotlib
import cv2
from comfy.utils import ProgressBar
from typing import List, Dict

eps = 0.01

def extend_scalelist(scalelist_behavior, pose_json, hands_scale, body_scale, head_scale, overall_scale, match_scalelist_method, only_scale_pose_index) -> List[list]:
    if pose_json.startswith('{'):
        pose_json = '[{}]'.format(pose_json)
    poses = json.loads(pose_json)
    # initialize scale lists
    hands_scalelist, body_scalelist, head_scalelist, overall_scalelist = [], [], [], []
    num_imgs = 0
    num_poses = 0
    scale_values = [hands_scale, body_scale, head_scale, overall_scale]
    scale_lists = [hands_scalelist, body_scalelist, head_scalelist, overall_scalelist]
    for img in poses:
        default_scale = 0.0
        default_num_person = 1
        if 'people' in img:
            default_scale = 1.0
            default_num_person = len(img['people'])
            subscales = [default_scale]*default_num_person
            if scalelist_behavior == 'poses':
                for i, scales in enumerate(scale_values):
                    if isinstance(scales, (list, tuple)):
                        if len(scales) >= num_poses + default_num_person:
                            if only_scale_pose_index<default_num_person and only_scale_pose_index >= -default_num_person:
                                subscales[only_scale_pose_index] = scales[num_poses + only_scale_pose_index]
                            else:
                                subscales = scales[num_poses:num_poses + default_num_person]
                        else:
                            if match_scalelist_method == 'no extend':
                                subscales = [default_scale]*default_num_person
                            elif match_scalelist_method == 'loop extend':
                                extend_scaleslist = scales*math.ceil((num_poses+default_num_person) / len(scales))
                                if only_scale_pose_index<default_num_person and only_scale_pose_index >= -default_num_person:
                                    subscales[only_scale_pose_index] = extend_scaleslist[num_poses + only_scale_pose_index]
                                else:
                                    subscales = extend_scaleslist[num_poses:num_poses + default_num_person]
                            elif match_scalelist_method == 'clamp extend':
                                if only_scale_pose_index<default_num_person and only_scale_pose_index >= -default_num_person:
                                    subscales[only_scale_pose_index] = scales[-1]
                                else:
                                    subscales = [scales[-1]] * default_num_person
                    else:
                        if only_scale_pose_index<default_num_person and only_scale_pose_index >= -default_num_person:
                            subscales[only_scale_pose_index] = scales
                        else:
                            subscales = [scales] * default_num_person

                    scale_lists[i].append(subscales.copy())
            else:
                for i, scales in enumerate(scale_values):
                    if isinstance(scales, (list, tuple)):
                        if len(scales) >= num_imgs + 1:
                            if only_scale_pose_index<default_num_person and only_scale_pose_index >= -default_num_person:
                                subscales[only_scale_pose_index] = scales[num_imgs]
                            else:
                                subscales = scales[num_poses]*default_num_person
                        else:
                            if match_scalelist_method == 'no extend':
                                subscales = [default_scale]*default_num_person
                            elif match_scalelist_method == 'loop extend':
                                extend_scaleslist = scales*math.ceil((num_imgs+1) / len(scales))
                                if only_scale_pose_index<default_num_person and only_scale_pose_index >= -default_num_person:
                                    subscales[only_scale_pose_index] = extend_scaleslist[num_imgs]
                                else:
                                    subscales = extend_scaleslist[num_imgs]*default_num_person
                            elif match_scalelist_method == 'clamp extend':
                                if only_scale_pose_index<default_num_person and only_scale_pose_index >= -default_num_person:
                                    subscales[only_scale_pose_index] = scales[-1]
                                else:
                                    subscales = [scales[-1]] * default_num_person
                    else:
                        if only_scale_pose_index<default_num_person and only_scale_pose_index >= -default_num_person:
                            subscales[only_scale_pose_index] = scales
                        else:
                            subscales = [scales] * default_num_person

                    scale_lists[i].append(subscales.copy())

            num_poses += default_num_person
            num_imgs += 1
        else:
            # if no people in image
            for i in range(len(scale_values)):
                scale_lists[i].append([default_scale])

    return scale_lists

def pose_normalized(pose_json):
    if pose_json.startswith('{'):
        pose_json = '[{}]'.format(pose_json)
    images = json.loads(pose_json)
    for image in images:
        if 'people' not in image:
            continue
        figures = image['people']
        H = image['canvas_height']
        W = image['canvas_width']
        normalized = 0.0
        for figure in figures:
            if 'pose_keypoints_2d' in figure:
                body = figure['pose_keypoints_2d']
                if body:
                    normalized = max(body)
                    if normalized > 2.0:
                        break
            if 'face_keypoints_2d' in figure:
                face = figure['face_keypoints_2d']
                if face:
                    normalized = max(face)
                    if normalized > 2.0:
                        break
            if 'hand_left_keypoints_2d' in figure:
                lhand = figure['hand_left_keypoints_2d']
                if lhand:
                    normalized = max(lhand)
                    if normalized > 2.0:
                        break
            if 'hand_right_keypoints_2d' in figure:
                rhand = figure['hand_right_keypoints_2d']
                if rhand:
                    normalized = max(rhand)
                    if normalized > 2.0:
                        break
        if normalized > 2.0:
            for figure in figures:
                if 'pose_keypoints_2d' in figure:
                    body = figure['pose_keypoints_2d']
                    if body is not None:
                        for i in range(0, len(body), 3):
                            body[i] = body[i] / float(W)
                            body[i+1] = body[i+1] / float(H)
                if 'face_keypoints_2d' in figure:
                    face = figure['face_keypoints_2d']
                    if face is not None:
                        for i in range(0, len(face), 3):
                            face[i] = face[i] / float(W)
                            face[i+1] = face[i+1] / float(H)
                if 'hand_left_keypoints_2d' in figure:
                    lhand = figure['hand_left_keypoints_2d']
                    if lhand is not None:
                        for i in range(0, len(lhand), 3):
                            lhand[i] = lhand[i] / float(W)
                            lhand[i+1] = lhand[i+1] / float(H)
                if 'hand_right_keypoints_2d' in figure:
                    rhand = figure['hand_right_keypoints_2d']
                    if rhand is not None:
                        for i in range(0, len(rhand), 3):
                            rhand[i] = rhand[i] / float(W)
                            rhand[i+1] = rhand[i+1] / float(H)
    return json.dumps(images)

def scale(point, scale_factor, pivot):
    return [(point[i] - pivot[i])*scale_factor + pivot[i] for i in range(len(point))]

def draw_pose_json(pose_json, resolution_x, show_body, show_face, show_hands, pose_marker_size, face_marker_size, hand_marker_size, hands_scalelist, body_scalelist, head_scalelist, overall_scalelist):
    pose_imgs = []
    pose_scaled = []

    if pose_json:
        if pose_json.startswith('{'):
            pose_json = '[{}]'.format(pose_json)
        images = json.loads(pose_json)
        pbar = ProgressBar(len(images))
        for img_idx, image in enumerate(images):
            if 'people' not in image:
                pbar.update(len(images))
                return pose_imgs
            figures = image['people']
            H = image['canvas_height']
            W = image['canvas_width']

            bodies = []
            candidate = []
            subset = [[]]
            faces = []
            hands = []

            openpose_json = []
            for pose_idx, figure in enumerate(figures):
                body_scale = body_scalelist[img_idx][pose_idx]
                hands_scale = hands_scalelist[img_idx][pose_idx]
                head_scale = head_scalelist[img_idx][pose_idx]
                overall_scale = overall_scalelist[img_idx][pose_idx]
                body = []
                face = []
                lhand = []
                rhand = []
                if 'pose_keypoints_2d' in figure:
                    body = figure['pose_keypoints_2d']
                    if body is None:
                        body = []
                if 'face_keypoints_2d' in figure:
                    face = figure['face_keypoints_2d']
                    if face is None:
                        face = []
                if 'hand_left_keypoints_2d' in figure:
                    lhand = figure['hand_left_keypoints_2d']
                    if lhand is None:
                        lhand = []
                if 'hand_right_keypoints_2d' in figure:
                    rhand = figure['hand_right_keypoints_2d']
                    if rhand is None:
                        rhand = []

                body_scaled = body.copy()
                face_scaled = face.copy()
                lhand_scaled = lhand.copy()
                rhand_scaled = rhand.copy()

                face_offset = [0, 0]
                lhand_offset = [0, 0]
                rhand_offset = [0, 0]

                overall_pivot = [0.5, 0.5]
                lhand_pivot = [0.25, 0.5]
                rhand_pivot = [0.75, 0.5]
                face_pivot = [0.5, 0.5]

                if body:
                    candidate_start_idx = len(candidate)

                    for i in range(0,len(body),3):
                        p_scaled = scale(body[i:i+2], body_scale, overall_pivot)
                        p_scaled = scale(p_scaled, overall_scale, overall_pivot)
                        body_scaled[i:i+2] = p_scaled
                        candidate.append(p_scaled)

                    figure_head_idx = candidate_start_idx
                    if figure_head_idx < len(candidate):
                        factor = 0.8
                        face_offset = [(candidate[figure_head_idx][0] - body[0])*factor, (candidate[figure_head_idx][1] - body[1])*factor]
                        face_pivot = candidate[figure_head_idx]

                    wrist_left_idx = candidate_start_idx + 7
                    wrist_right_idx = candidate_start_idx + 4

                    if wrist_left_idx < len(candidate) and len(body) > 22:
                        lhand_offset = [candidate[wrist_left_idx][0] - body[21], candidate[wrist_left_idx][1] - body[22]]
                        lhand_pivot = candidate[wrist_left_idx]

                    if wrist_right_idx < len(candidate) and len(body) > 13:
                        rhand_offset = [candidate[wrist_right_idx][0] - body[12], candidate[wrist_right_idx][1] - body[13]]
                        rhand_pivot = candidate[wrist_right_idx]

                    if not subset[0]:
                        subset[0].extend([candidate_start_idx+(i//3) if body[i+2]>0 else -1 for i in range(0,len(body),3)])
                    else:
                        new_subset = [candidate_start_idx+(i//3) if body[i+2]>0 else -1 for i in range(0,len(body),3)]
                        subset.append(new_subset)

                if face:
                    f = []
                    for i in range(0,len(face),3):
                        p = face[i:i+2]
                        p_offset = [p[0] + face_offset[0], p[1] + face_offset[1]]
                        p_scaled = scale(p_offset, head_scale, face_pivot)
                        p_scaled = scale(p_scaled, overall_scale, overall_pivot)
                        face_scaled[i:i+2] = p_scaled
                        f.append(p_scaled)
                    faces.append(f)

                if lhand:
                    lh = []
                    for i in range(0, len(lhand), 3):
                        p = lhand[i:i+2]
                        p_offset = [p[0] + lhand_offset[0], p[1] + lhand_offset[1]]
                        p_scaled = scale(p_offset, hands_scale, lhand_pivot)
                        p_scaled = scale(p_scaled, overall_scale, overall_pivot)
                        lhand_scaled[i:i+2] = p_scaled
                        lh.append(p_scaled)
                    hands.append(lh)

                if rhand:
                    rh = []
                    for i in range(0, len(rhand), 3):
                        p = rhand[i:i+2]
                        p_offset = [p[0] + rhand_offset[0], p[1] + rhand_offset[1]]
                        p_scaled = scale(p_offset, hands_scale, rhand_pivot)
                        p_scaled = scale(p_scaled, overall_scale, overall_pivot)
                        rhand_scaled[i:i+2] = p_scaled
                        rh.append(p_scaled)
                    hands.append(rh)
                openpose_json.append(dict(pose_keypoints_2d=body_scaled, face_keypoints_2d=face_scaled, hand_left_keypoints_2d=lhand_scaled, hand_right_keypoints_2d=rhand_scaled))

            bodies = dict(candidate=candidate, subset=subset)
            pose = dict(bodies=bodies, faces=faces, hands=hands)
            pose = dict(bodies=bodies if show_body else {'candidate':[], 'subset':[]}, faces=faces if show_face else [], hands=hands if show_hands else [])

            W_scaled = resolution_x
            if resolution_x < 64:
                W_scaled = W
            H_scaled = int(H*(W_scaled*1.0/W))
            openpose_json = {
                            'people': openpose_json,
                            'canvas_height': H_scaled,
                            'canvas_width': W_scaled,
                            }

            pose_img = draw_pose(pose, H_scaled, W_scaled, pose_marker_size, face_marker_size, hand_marker_size)
            pose_imgs.append(pose_img)
            pose_scaled.append(openpose_json)
            pbar.update(1)

    return pose_imgs, pose_scaled

def draw_pose(pose, H, W, pose_marker_size, face_marker_size, hand_marker_size):
    bodies = pose['bodies']
    faces = pose['faces']
    hands = pose['hands']
    candidate = bodies['candidate']
    subset = bodies['subset']
    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)

    if len(candidate) > 0:
        canvas = draw_bodypose(canvas, candidate, subset, pose_marker_size)

    if len(hands) > 0:
        canvas = draw_handpose(canvas, hands, hand_marker_size)

    if len(faces) > 0:
        canvas = draw_facepose(canvas, faces, face_marker_size)

    return canvas

def draw_bodypose(canvas, candidate, subset, pose_marker_size):
    H, W, C = canvas.shape
    candidate = np.array(candidate)
    subset = np.array(subset)

    # stickwidth = 4

    limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
               [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
               [1, 16], [16, 18], [3, 17], [6, 18]]

    colors = [[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]]

    for i in range(17):
        for n in range(len(subset)):
            index = subset[n][np.array(limbSeq[i]) - 1]
            if -1 in index:
                continue
            Y = candidate[index.astype(int), 0] * float(W)
            X = candidate[index.astype(int), 1] * float(H)
            mX = np.mean(X)
            mY = np.mean(Y)
            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), pose_marker_size), int(angle), 0, 360, 1)
            cv2.fillConvexPoly(canvas, polygon, colors[i])

    canvas = (canvas * 0.6).astype(np.uint8)

    for i in range(18):
        for n in range(len(subset)):
            index = int(subset[n][i])
            if index == -1:
                continue
            x, y = candidate[index][0:2]
            x = int(x * W)
            y = int(y * H)
            cv2.circle(canvas, (int(x), int(y)), pose_marker_size, colors[i], thickness=-1)

    return canvas


def draw_handpose(canvas, all_hand_peaks, hand_marker_size):
    H, W, C = canvas.shape

    edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
             [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]

    for peaks in all_hand_peaks:
        peaks = np.array(peaks)

        for ie, e in enumerate(edges):
            x1, y1 = peaks[e[0]]
            x2, y2 = peaks[e[1]]
            x1 = int(x1 * W)
            y1 = int(y1 * H)
            x2 = int(x2 * W)
            y2 = int(y2 * H)
            if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
                cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255, thickness=1 if hand_marker_size == 0 else hand_marker_size)

        joint_size=0
        if hand_marker_size < 2:
            joint_size = hand_marker_size + 1
        else:
            joint_size = hand_marker_size + 2
        for i, keyponit in enumerate(peaks):
            x, y = keyponit
            x = int(x * W)
            y = int(y * H)
            if x > eps and y > eps:
                cv2.circle(canvas, (x, y), joint_size, (0, 0, 255), thickness=-1)
    return canvas


def draw_facepose(canvas, all_lmks, face_marker_size):
    H, W, C = canvas.shape
    for lmks in all_lmks:
        lmks = np.array(lmks)
        for lmk in lmks:
            x, y = lmk
            x = int(x * W)
            y = int(y * H)
            if x > eps and y > eps:
                cv2.circle(canvas, (x, y), face_marker_size, (255, 255, 255), thickness=-1)
    return canvas