"""Styled pose-skeleton renderer for AnimaPoseControl. Pure numpy + cv2 (no torch).
Single source of truth for the predefined skeleton looks.
Wholebody-133: body 0-16, feet 17-22, face 23-90, L-hand 91-111, R-hand 112-132."""
import numpy as np
import cv2

BODY = list(range(17)); FEET = list(range(17, 23)); FACE = list(range(23, 91)); HAND = list(range(91, 133))
COCO_LIMBS = [(5, 7), (7, 9), (6, 8), (8, 10), (11, 13), (13, 15), (12, 14), (14, 16),
              (5, 6), (11, 12), (5, 11), (6, 12), (0, 5), (0, 6)]
_PAL = [(255, 0, 0), (255, 128, 0), (255, 255, 0), (128, 255, 0), (0, 255, 0), (0, 255, 128), (0, 255, 255),
        (0, 128, 255), (0, 0, 255), (128, 0, 255), (255, 0, 255), (255, 0, 128), (180, 180, 180), (220, 220, 220)]
LIMB_COLOR = {l: _PAL[i % len(_PAL)] for i, l in enumerate(COCO_LIMBS)}
JOINT_COLOR = [_PAL[i % len(_PAL)] for i in range(17)]
STYLES = ("R0_thin", "R1_thick", "R2_puppet", "heatmap")
_W = {"R0_thin": {"stick": 2, "joint": 3}, "R1_thick": {"stick": 8, "joint": 6}, "R2_puppet": {"stick": 22, "joint": 11}, "heatmap": {"stick": 2, "joint": 3}}


def _toggle(sc, hands, face, feet):
    sc = np.asarray(sc, float).copy()
    if not hands: sc[HAND] = 0.0
    if not face: sc[FACE] = 0.0
    if not feet: sc[FEET] = 0.0
    return sc


def render(kp, sc, style, resolution, hands=True, face=True, feet=True, kpt_thr=0.3):
    if style not in STYLES:
        raise ValueError(f"unknown style {style!r}")
    H = W = int(resolution)
    kp = np.asarray(kp, float); sc = _toggle(sc, hands, face, feet)
    if style == "heatmap":
        return _heatmap(kp, sc, H, W, feet, kpt_thr)
    cfg = _W[style]; cv = np.zeros((H, W, 3), np.uint8)
    for a, b in COCO_LIMBS:
        if sc[a] >= kpt_thr and sc[b] >= kpt_thr:
            cv2.line(cv, (int(kp[a][0]), int(kp[a][1])), (int(kp[b][0]), int(kp[b][1])), LIMB_COLOR[(a, b)], cfg["stick"])
    for i in BODY:
        if sc[i] >= kpt_thr:
            cv2.circle(cv, (int(kp[i][0]), int(kp[i][1])), cfg["joint"], (255, 255, 255), -1)
    for i in FEET:
        if sc[i] >= kpt_thr:
            cv2.circle(cv, (int(kp[i][0]), int(kp[i][1])), max(1, cfg["joint"] // 2), (255, 255, 255), -1)
    if style == "R0_thin":              # only R0 draws hands + face (DWPose-style small visible dots)
        for i in HAND:
            if sc[i] >= kpt_thr:
                cv2.circle(cv, (int(kp[i][0]), int(kp[i][1])), 2, (0, 255, 255), -1)
        for i in FACE:
            if sc[i] >= kpt_thr:
                cv2.circle(cv, (int(kp[i][0]), int(kp[i][1])), 2, (255, 255, 255), -1)
    return cv


def _heatmap(kp, sc, H, W, feet, thr, sig=11):
    acc = np.zeros((H, W, 3), np.float32)
    for i in BODY + (FEET if feet else []):
        if sc[i] >= thr:
            x, y = int(round(kp[i][0])), int(round(kp[i][1]))
            if 0 <= x < W and 0 <= y < H:
                b = np.zeros((H, W), np.float32); b[y, x] = 1.0
                b = cv2.GaussianBlur(b, (0, 0), sig); b /= (b.max() + 1e-9)
                acc += b[..., None] * np.array(JOINT_COLOR[i % 17], np.float32)
    return np.clip(acc, 0, 255).astype(np.uint8)