"""OpenPose_full (137-point) keypoint schema.

Layout: 25 BODY_25 body/foot points + 21 left-hand + 21 right-hand + 70 face
points, in that order, giving a flat 0..136 global keypoint index used
everywhere else in this plugin (mapping node attribute slots, draw override,
projector, rasterizer).

Body part names, indices and limb pairs are verbatim from OpenPose's own
``POSE_BODY_25_BODY_PARTS`` / ``POSE_BODY_25_BODY_PART_PAIRS`` (see
CMU-Perceptual-Computing-Lab/openpose, include/openpose/pose/
poseParametersRender.hpp). Hand and face keypoint layouts follow the
standard, widely-interoperable 21-point hand (wrist + 5x4-joint fingers,
identical ordering to OpenPose/MediaPipe hand models) and 70-point face
(iBUG 68-point facial landmark scheme + 2 pupils, OpenPose's own face model)
conventions.

Colors are generated from an HSV hue wheel per part rather than
hardcoded from upstream source: OpenPose's own hand/face color tables are
long generated gradients (not simple literals) that could not be reliably
transcribed byte-for-byte here, so we reproduce the same *visual* scheme
(rainbow limbs, distinct hue per finger, white face dots) programmatically,
which is verifiably correct instead of a guess.
"""

import colorsys


# --- Part sizes / offsets into the flat 0..136 keypoint index ---------------

BODY_COUNT = 25
HAND_COUNT = 21
FACE_COUNT = 70

BODY_START = 0
HAND_LEFT_START = BODY_START + BODY_COUNT            # 25
HAND_RIGHT_START = HAND_LEFT_START + HAND_COUNT       # 46
FACE_START = HAND_RIGHT_START + HAND_COUNT            # 67
TOTAL_KEYPOINTS = FACE_START + FACE_COUNT             # 137

PART_BODY = "body"
PART_HAND_LEFT = "hand_left"
PART_HAND_RIGHT = "hand_right"
PART_FACE = "face"

# (part id, start index, count) in draw order
PARTS = [
    (PART_BODY, BODY_START, BODY_COUNT),
    (PART_HAND_LEFT, HAND_LEFT_START, HAND_COUNT),
    (PART_HAND_RIGHT, HAND_RIGHT_START, HAND_COUNT),
    (PART_FACE, FACE_START, FACE_COUNT),
]


def _hsv255(hue, saturation=1.0, value=1.0):
    r, g, b = colorsys.hsv_to_rgb(hue % 1.0, saturation, value)
    return (int(round(r * 255)), int(round(g * 255)), int(round(b * 255)))


# --- Body: BODY_25 -----------------------------------------------------------

BODY_25_NAMES = [
    "Nose", "Neck", "RShoulder", "RElbow", "RWrist", "LShoulder", "LElbow",
    "LWrist", "MidHip", "RHip", "RKnee", "RAnkle", "LHip", "LKnee", "LAnkle",
    "REye", "LEye", "REar", "LEar", "LBigToe", "LSmallToe", "LHeel",
    "RBigToe", "RSmallToe", "RHeel",
]
assert len(BODY_25_NAMES) == BODY_COUNT

# Verbatim from OpenPose's POSE_BODY_25_BODY_PART_PAIRS_RENDER_GPU (local
# indices 0..24, i.e. relative to BODY_START).
BODY_25_PAIRS_LOCAL = [
    (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), (2, 17), (5, 18), (14, 19), (19, 20), (14, 21), (11, 22),
    (22, 23), (11, 24),
]

BODY_25_POINT_COLORS = [_hsv255(i / BODY_COUNT) for i in range(BODY_COUNT)]
BODY_25_LIMB_COLORS = [
    _hsv255(i / len(BODY_25_PAIRS_LOCAL)) for i in range(len(BODY_25_PAIRS_LOCAL))
]

BODY_POINT_RADIUS = 4.0
BODY_LIMB_THICKNESS = 4.0


# --- Body: COCO-18 (legacy OpenPose COCO model) ------------------------------
#
# COCO is *not* simply "BODY_25 with some points removed": it has its own
# limb connectivity (Neck connects directly to each hip -- there is no
# MidHip point at all -- and there are no foot points). Both the point names
# and the COCO section of the limb-pair table are verbatim from OpenPose's
# own POSE_COCO_BODY_PARTS / POSE_BODY_PART_PAIRS,
# src/openpose/pose/poseParameters.cpp.
#
# COCO mode does not need its own joint-mapping slots: it reuses the exact
# same BODY_25-indexed targetJoint mapping every character already has, and
# is purely a different "which of those points + which limbs to draw" view
# over it (COCO_TO_BODY25_INDEX below is that view).

COCO_COUNT = 18

COCO_NAMES = [
    "Nose", "Neck", "RShoulder", "RElbow", "RWrist", "LShoulder", "LElbow",
    "LWrist", "RHip", "RKnee", "RAnkle", "LHip", "LKnee", "LAnkle", "REye",
    "LEye", "REar", "LEar",
]
assert len(COCO_NAMES) == COCO_COUNT

# COCO local index -> BODY_25 local index (== global index, since the body
# part starts at 0). BODY_25's MidHip (8) and the six foot points (19-24)
# have no COCO equivalent and are simply absent from this list.
COCO_TO_BODY25_INDEX = [0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]
assert len(COCO_TO_BODY25_INDEX) == COCO_COUNT
assert all(BODY_25_NAMES[COCO_TO_BODY25_INDEX[i]] == COCO_NAMES[i] for i in range(COCO_COUNT))

COCO_ACTIVE_BODY25_INDICES = sorted(COCO_TO_BODY25_INDEX)

# Verbatim (COCO section of POSE_BODY_PART_PAIRS), expressed in COCO-local
# 0..17 indices, then translated to BODY_25-local/global indices so it can
# be used directly against the same joint mapping as BODY_25 mode.
_COCO_PAIRS_COCO_SPACE = [
    (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),
    (2, 16), (5, 17),
]
COCO_PAIRS_LOCAL = [
    (COCO_TO_BODY25_INDEX[a], COCO_TO_BODY25_INDEX[b]) for (a, b) in _COCO_PAIRS_COCO_SPACE
]
COCO_LIMB_COLORS = [_hsv255(i / len(COCO_PAIRS_LOCAL)) for i in range(len(COCO_PAIRS_LOCAL))]

BODY_MODEL_BODY_25 = "BODY_25"
BODY_MODEL_COCO = "COCO"
BODY_MODELS = [BODY_MODEL_BODY_25, BODY_MODEL_COCO]


# --- Hands: 21 points each, standard wrist + 5 x 4-joint fingers ------------

_FINGER_NAMES = ["Thumb", "Index", "Middle", "Ring", "Pinky"]

HAND_NAMES = ["Wrist"] + [
    "{0}{1}".format(finger, joint)
    for finger in _FINGER_NAMES
    for joint in range(1, 5)
]
assert len(HAND_NAMES) == HAND_COUNT

# (0=wrist, 1-4=thumb, 5-8=index, 9-12=middle, 13-16=ring, 17-20=pinky)
HAND_PAIRS_LOCAL = []
for _finger_i in range(5):
    _base = 1 + _finger_i * 4
    HAND_PAIRS_LOCAL += [
        (0, _base), (_base, _base + 1), (_base + 1, _base + 2),
        (_base + 2, _base + 3),
    ]
assert len(HAND_PAIRS_LOCAL) == 20


def _hand_point_color(local_index):
    if local_index == 0:
        return (255, 255, 255)  # wrist
    finger = (local_index - 1) // 4
    joint = (local_index - 1) % 4  # 0 (nearest wrist) .. 3 (tip)
    hue = finger / 5.0
    value = 0.4 + 0.2 * joint  # brighter towards the fingertip
    return _hsv255(hue, 1.0, value)


HAND_POINT_COLORS = [_hand_point_color(i) for i in range(HAND_COUNT)]
HAND_LIMB_COLORS = [
    _hand_point_color(j) for (_, j) in HAND_PAIRS_LOCAL
]  # colored by the joint further from the wrist

HAND_POINT_RADIUS = 3.0
HAND_LIMB_THICKNESS = 2.0


# --- Face: 70 points (iBUG 68-point landmarks + 2 pupils) -------------------

FACE_NAMES = (
    ["Jaw{0}".format(i) for i in range(17)]
    + ["REyebrow{0}".format(i) for i in range(5)]
    + ["LEyebrow{0}".format(i) for i in range(5)]
    + ["Nose{0}".format(i) for i in range(9)]
    + ["REye{0}".format(i) for i in range(6)]
    + ["LEye{0}".format(i) for i in range(6)]
    + ["Mouth{0}".format(i) for i in range(20)]
    + ["RPupil", "LPupil"]
)
assert len(FACE_NAMES) == FACE_COUNT

# OpenPose's default renderer draws the face as keypoints only (no limb
# lines), all white.
FACE_PAIRS_LOCAL = []
FACE_POINT_COLORS = [(255, 255, 255)] * FACE_COUNT

FACE_POINT_RADIUS = 2.0
FACE_LIMB_THICKNESS = 0.0


# --- Flattened 0..136 global schema ------------------------------------------

KEYPOINT_NAMES = (
    ["Body_{0}".format(n) for n in BODY_25_NAMES]
    + ["HandLeft_{0}".format(n) for n in HAND_NAMES]
    + ["HandRight_{0}".format(n) for n in HAND_NAMES]
    + ["Face_{0}".format(n) for n in FACE_NAMES]
)
assert len(KEYPOINT_NAMES) == TOTAL_KEYPOINTS

NAME_TO_INDEX = {name: i for i, name in enumerate(KEYPOINT_NAMES)}

KEYPOINT_PART = (
    [PART_BODY] * BODY_COUNT
    + [PART_HAND_LEFT] * HAND_COUNT
    + [PART_HAND_RIGHT] * HAND_COUNT
    + [PART_FACE] * FACE_COUNT
)

KEYPOINT_POINT_COLORS = (
    BODY_25_POINT_COLORS + HAND_POINT_COLORS + HAND_POINT_COLORS + FACE_POINT_COLORS
)

KEYPOINT_POINT_RADIUS = (
    [BODY_POINT_RADIUS] * BODY_COUNT
    + [HAND_POINT_RADIUS] * HAND_COUNT
    + [HAND_POINT_RADIUS] * HAND_COUNT
    + [FACE_POINT_RADIUS] * FACE_COUNT
)

def part_of(global_index):
    """Return the part id (PART_BODY / PART_HAND_LEFT / ...) for a keypoint index."""
    for part_id, start, count in PARTS:
        if start <= global_index < start + count:
            return part_id
    raise IndexError(global_index)


# Hand limb pairs, offset into the flat 0..136 range -- unaffected by body
# model choice, so precomputed once.
_HAND_LEFT_PAIRS_GLOBAL = [(a + HAND_LEFT_START, b + HAND_LEFT_START) for (a, b) in HAND_PAIRS_LOCAL]
_HAND_RIGHT_PAIRS_GLOBAL = [(a + HAND_RIGHT_START, b + HAND_RIGHT_START) for (a, b) in HAND_PAIRS_LOCAL]


def _body_limb_data(body_model):
    """(pairs, colors, thickness) for just the body part, global-indexed."""
    if body_model == BODY_MODEL_COCO:
        pairs = COCO_PAIRS_LOCAL  # already translated to BODY_25-local/global indices
        colors = COCO_LIMB_COLORS
    else:
        pairs = [(a + BODY_START, b + BODY_START) for (a, b) in BODY_25_PAIRS_LOCAL]
        colors = BODY_25_LIMB_COLORS
    thickness = [BODY_LIMB_THICKNESS] * len(pairs)
    return pairs, colors, thickness


def full_limb_data(body_model=BODY_MODEL_BODY_25):
    """(pairs, colors, thickness) for the whole rig: the chosen body model
    plus both full hands, global-indexed into the flat 0..136 range. Faces
    render as dots only and contribute no limbs.
    """
    body_pairs, body_colors, body_thickness = _body_limb_data(body_model)
    pairs = list(body_pairs) + _HAND_LEFT_PAIRS_GLOBAL + _HAND_RIGHT_PAIRS_GLOBAL
    colors = list(body_colors) + HAND_LIMB_COLORS + HAND_LIMB_COLORS
    thickness = (
        list(body_thickness)
        + [HAND_LIMB_THICKNESS] * len(HAND_PAIRS_LOCAL)
        + [HAND_LIMB_THICKNESS] * len(HAND_PAIRS_LOCAL)
    )
    assert len(pairs) == len(colors) == len(thickness)
    return pairs, colors, thickness


def active_point_indices(body_model=BODY_MODEL_BODY_25):
    """Global keypoint indices to draw as points for the given body model.

    Hands and face are always fully active; the body part is either all 25
    BODY_25 points or just the 18 COCO points (no MidHip, no feet).
    """
    if body_model == BODY_MODEL_COCO:
        body_indices = list(COCO_ACTIVE_BODY25_INDICES)
    else:
        body_indices = list(range(BODY_START, BODY_START + BODY_COUNT))
    return (
        body_indices
        + list(range(HAND_LEFT_START, HAND_LEFT_START + HAND_COUNT))
        + list(range(HAND_RIGHT_START, HAND_RIGHT_START + HAND_COUNT))
        + list(range(FACE_START, FACE_START + FACE_COUNT))
    )
