import json
from typing import List
import numpy as np
import math
import cv2
from numpy import ndarray

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

coco_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]
]

yolo_coco_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]
]

yolo_coco_limbSeq = [
    [1, 2], [1, 5], [2, 3], [3, 4], 
    [5, 6], [6, 7], [1, 8], [8, 9], 
    [9, 10], [1, 11], [11, 12], [12, 13], 
    [1, 0], [0, 14], [14, 16], [0, 15], 
]

body_25_limbSeq = [
    [1, 8], [1, 2], [1, 5], [2, 3], 
    [3, 4], [5, 6], [6, 7], [8, 9], 
    [9, 10], [10, 11], [8, 12], [12, 13], 
    [13, 14], [1, 0], [0, 15], [15, 17], 
    [0, 16], [16, 18], [14, 19], [19, 20], 
    [14, 21], [11, 22], [22, 23], [11, 24]
]
body_25_colors = [
    [255, 0, 85], [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0],
    [0, 255, 0], [255, 0, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255],
    [0, 0, 255], [255, 0, 170], [170, 0, 255], [255, 0, 255], [85, 0, 255], [0, 0, 255], [0, 0, 255],
    [0, 0, 255], [0, 0, 255], [0, 255, 255], [0, 255, 255], [0, 255, 255]
]

body_25B_limbSeq = [
    [0, 1], [0, 2], [1, 3], [2, 4], [5, 7], [6, 8], [7, 9], [8, 10], 
    [5, 11], [6, 12], [11, 13], [12, 14], [13, 15], [14, 16], [15, 19], 
    [19, 20], [15, 21], [16, 22], [22, 23], [16, 24], [5, 17], [6, 17], 
    [17, 18], [11, 12]
]

body_25B_colors = [
    [255, 0, 85], [170, 0, 255], [255, 0, 170], [85, 0, 255], [255, 0, 255], 
    [170, 255, 0], [255, 85, 0], [85, 255, 0], [255, 170, 0], [0, 255, 0], 
    [255, 255, 0], [0, 170, 255], [0, 255, 85], [0, 85, 255], [0, 255, 170], 
    [0, 0, 255], [0, 255, 255], [255, 0, 0], [255, 0, 0], [0, 0, 255], 
    [0, 0, 255], [0, 0, 255], [0, 255, 255], [0, 255, 255], [0, 255, 255]
]

class Keypoint:
    def __init__(self, x: float, y: float, confidence: float = 1.0):
        """
        Initialize a Keypoint object.

        Args:
            x (float): The x-coordinate of the keypoint.
            y (float): The y-coordinate of the keypoint.
            confidence (float): The confidence score of the keypoint. Default is 1.0.
        """
        self.x = x
        self.y = y
        self.confidence = confidence


    def __repr__(self):
        return f"Keypoint(x={self.x}, y={self.y}, confidence={self.confidence})"

class KeypointEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, Keypoint):
            return {'x': obj.x, 'y': obj.y, 'confidence': obj.confidence}
        return super().default(obj)
    
class SkeletonEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, Skeleton):
            return {'keypoints': obj.keypoints}
        return super().default(obj)

class Encoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, Skeleton):
            return SkeletonEncoder().default(obj)
        elif isinstance(obj, Keypoint):
            return KeypointEncoder().default(obj)
        return super().default(obj)

class Skeleton:
    def __init__(self, keypoints: List[Keypoint]):
        self.keypoints = keypoints

    def __repr__(self):
        return f"Skeleton(keypoints={self.keypoints})"
    
    def is_healthy_skeleton(self):
        for keypoint in self.keypoints:
            if keypoint.confidence == 0.0:
                return False
        return True

class Skeleton_Seqence:
    def __init__(self, skeletons: List[Skeleton]):
        self.skeletons_frame = skeletons

    def __repr__(self):
        return f"Skeleton_Seqence(Skeleton_frames={self.skeletons_frame})"
    
    def get_frame(self, frame_index: int) -> Skeleton:
        return self.skeletons_frame[frame_index]
    
    def add_frame(self, skeleton: Skeleton):
        self.skeletons_frame.append(skeleton)

    def is_healthy_seqence(self):
        for skeleton in self.skeletons_frame:
            if not skeleton.is_healthy_skeleton():
                return False
        return True
    
def fix_keypoints(keypoints, last_keypoints, next_keypoints):
    if not keypoints or not last_keypoints or not next_keypoints:
        return keypoints
    for i, keypoint in enumerate(keypoints):
        if keypoint.confidence == 0.0 and last_keypoints and next_keypoints:
            last_keypoint = last_keypoints[i]
            next_keypoint = next_keypoints[i]
            if last_keypoint.confidence > 0 and next_keypoint.confidence > 0:
                keypoint.x = (last_keypoint.x + next_keypoint.x) / 2
                keypoint.y = (last_keypoint.y + next_keypoint.y) / 2
                keypoint.confidence = (last_keypoint.confidence + next_keypoint.confidence) / 2
    return keypoints

def get_time_slice_for_Skeleton_Seqences(skeleton_seqences: List[Skeleton_Seqence], frame_index: int) -> List[Skeleton]:
    return [skeleton_seq.get_frame(frame_index) for skeleton_seq in skeleton_seqences]


def is_normalized(keypoints: List[Keypoint]) -> bool:
    for keypoint in keypoints:
        if not (0 <= keypoint.x <= 1 and 0 <= keypoint.y <= 1):
            return False
    return True


def draw_bodypose(canvas: ndarray, keypoints: List[Keypoint], limbSeq, colors, xinsr_stick_scaling: bool = False) -> np.ndarray:
    """
    Draw keypoints and limbs representing body pose on a given canvas.

    Args:
        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the body pose.
        keypoints (List[Keypoint]): A list of Keypoint objects representing the body keypoints to be drawn.
        xinsr_stick_scaling (bool): Whether or not scaling stick width for xinsr ControlNet

    Returns:
        np.ndarray: A 3D numpy array representing the modified canvas with the drawn body pose.

    Note:
        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
    """
    if not is_normalized(keypoints):
        H, W = 1.0, 1.0
    else:
        H, W, _ = canvas.shape

    CH, CW, _ = canvas.shape
    stickwidth = 2

    # Ref: https://huggingface.co/xinsir/controlnet-openpose-sdxl-1.0
    max_side = max(CW, CH)
    if xinsr_stick_scaling:
        stick_scale = 1 if max_side < 500 else min(2 + (max_side // 1000), 7)
    else :
        stick_scale = 1
    
    if keypoints is None or len(keypoints) == 0:
        return canvas

    for (k1_index, k2_index), color in zip(limbSeq, colors):
        keypoint1 = keypoints[k1_index]
        keypoint2 = keypoints[k2_index]

        if keypoint1 is None or keypoint2 is None or keypoint1.confidence == 0 or keypoint2.confidence == 0 or keypoint1.x <= 0 or keypoint1.y <= 0 or keypoint2.x <= 0 or keypoint2.y <= 0:
            # if keypoint1 is None or keypoint1.confidence == 0:
            #     print(f"keypoint failed: {k1_index}")
            # if keypoint2 is None or keypoint2.confidence == 0:
            #     print(f"keypoint failed: {k2_index}")
            continue

        Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
        X = np.array([keypoint1.y, keypoint2.y]) * 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), stickwidth*stick_scale), int(angle), 0, 360, 1)
        cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])

    for keypoint, color in zip(keypoints, colors):
        if keypoint is None or keypoint.confidence == 0 or keypoint.x <= 0 or keypoint.y <= 0:
            continue

        x, y = keypoint.x, keypoint.y
        x = int(x * W)
        y = int(y * H)
        cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)

    return canvas