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

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})"

def is_normalized(keypoints: List[Keypoint]) -> bool:
    """
    Check if the keypoints are normalized between 0 and 1.

    Args:
        keypoints (List[Keypoint]): A list of Keypoint objects.

    Returns:
        bool: True if all keypoints are normalized, False otherwise.
    """
    for keypoint in keypoints:
        if not (0 <= keypoint.x <= 1 and 0 <= keypoint.y <= 1):
            return False
    return True

def draw_bodypose(canvas: np.ndarray, keypoints: List[Keypoint], 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 = 4

    # 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

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

    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 (k1_index, k2_index), color in zip(limbSeq, colors):
        keypoint1 = keypoints[k1_index - 1]
        keypoint2 = keypoints[k2_index - 1]

        if keypoint1 is None or keypoint2 is None:
            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:
            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

def openpose_json_to_keypoints(json_file: str) -> List[Keypoint]:
    """
    Convert OpenPose JSON file to a list of Keypoint objects.

    Args:
        json_file (str): The path to the OpenPose JSON file.

    Returns:
        List[Keypoint]: A list of Keypoint objects representing the keypoints extracted from the JSON file.
    """
    with open(json_file, 'r') as file:
        data = json.load(file)
        keypoints = data[0]['people'][0]['pose_keypoints_2d']
        keypoints = [Keypoint(keypoints[i], keypoints[i + 1]) for i in range(0, len(keypoints), 3)]
    return keypoints

def main():
    # Create a blank canvas
    canvas = np.zeros((768, 768, 3), dtype=np.uint8)

    # Load the keypoints from the OpenPose JSON file
    keypoints = openpose_json_to_keypoints('openpose_output.json')

    # Draw the body pose on the canvas
    canvas = draw_bodypose(canvas, keypoints)

    # Display the result
    cv2.imshow('Body Pose', canvas)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

# Here is an example of input json file, create a List of Keypoint objects from the json file and draw the body pose on a canvas using the draw_bodypose function.
# 
# [
#   {
#     "people": [
#       {
#         "pose_keypoints_2d": [423.575046753433, 113.537678172812, 1, 378.126093103822, 185.044031914867, 1, 416.303214169495, 169.288394649668, 1, 470.841958549029, 122.021482854073, 1, 519.320842441947, 48.0911849173717, 1, 339.948972038149, 200.799669180065, 1, 322.981362675627, 289.273632284642, 1, 316.921502189012, 391.079288459771, 1, 439.330684018632, 363.203930221343, 1, 412.667297877527, 541.363828527819, 1, 375.096162860515, 702.556117471773, 1, 375.096162860515, 364.415902318666, 1, 359.340525595316, 530.456079651912, 1, 292.682060242553, 674.680759233345, 1, 411.455325780204, 105.053873491551, 1, 418.727158364141, 102.629929296905, 1, 383.579967541775, 120.80951075675, 1, 371.460246568546, 122.021482854073, 1]
#       }
#     ],
#     "canvas_height": 768,
#     "canvas_width": 768
#   }
# ]