"""
Pose engine: deteksi orang (YOLOX, dataset HumanArt) + estimasi pose whole-body
133 titik (RTMW / DWPose, keluarga RTMPose) dengan ONNX Runtime, lalu
rendering gaya OpenPose/ControlNet (badan 18 titik + tangan 2x21 + wajah 68).

Fokus utama: beberapa orang dalam satu foto yang saling bertumpuk / berhimpitan.
Strategi:
  1. Deteksi dengan ambang skor sangat rendah dan NMS longgar supaya kotak
     orang yang tertutup sebagian tidak ikut terhapus.
  2. Estimasi pose untuk sejumlah kandidat kotak teratas.
  3. Hapus duplikat di tingkat POSE (OKS), bukan hanya di tingkat kotak.
  4. Hapus keypoint "bajakan" (dua orang berbagi titik sendi yang persis sama).
  5. Pilih orang terbaik; bila jumlah yang lolos ambang ketat kurang dari target,
     ambil kandidat berikutnya (relaksasi otomatis).
"""
from __future__ import annotations

import colorsys
import math
import os
import shutil
import threading
import zipfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np

# --------------------------------------------------------------------------
# Registri model (mirror resmi rtmlib di Hugging Face: Tau-J/RTMPose)
# --------------------------------------------------------------------------
HF_REPO = "Tau-J/RTMPose"
MODEL_DIR = Path(os.environ.get("POSE_MODEL_DIR", Path.home() / ".cache" / "pose-skeleton-models"))

# nama tampilan -> (path zip di repo HF, ukuran input (lebar, tinggi))
DETECTORS: Dict[str, Tuple[str, Tuple[int, int]]] = {
    "YOLOX-l (akurasi tinggi, disarankan)": (
        "rtmposev1/onnx_sdk/yolox_l_8xb8-300e_humanart-ce1d7a62.zip", (640, 640)),
    "YOLOX-m (lebih cepat)": (
        "rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip", (640, 640)),
    "YOLOX-x (terbesar, paling lambat)": (
        "rtmposev1/onnx_sdk/yolox_x_8xb8-300e_humanart-a39d44ed.zip", (640, 640)),
}

# nama tampilan -> (path zip, ukuran input (lebar, tinggi), faktor skala skor)
# Skor SimCC bukan probabilitas; ekspor RTMW-l 384x288 menghasilkan skor ~0-10 (model lain ~0-1),
# sehingga diskalakan supaya ambang 0.3 bermakna sama untuk semua model.
POSE_MODELS: Dict[str, Tuple[str, Tuple[int, int], float]] = {
    "RTMW-l 384x288 (disarankan)": (
        "rtmw/onnx_sdk/rtmw-dw-x-l_simcc-cocktail14_270e-384x288_20231122.zip", (288, 384), 0.1),
    "RTMW-x 384x288 (terbesar)": (
        "rtmw/onnx_sdk/rtmw-x_simcc-cocktail13_pt-ucoco_270e-384x288-0949e3a9_20230925.zip", (288, 384), 1.0),
    "DWPose-l 384x288 (klasik, sama seperti ComfyUI)": (
        "rtmposev1/onnx_sdk/rtmpose-l_simcc-ucoco_dw-ucoco_270e-384x288-2438fd99_20230728.zip", (288, 384), 1.0),
    "RTMW-m 256x192 (paling cepat)": (
        "rtmw/onnx_sdk/rtmw-dw-l-m_simcc-cocktail14_270e-256x192_20231122.zip", (192, 256), 1.0),
}

DEFAULT_DETECTOR = "YOLOX-l (akurasi tinggi, disarankan)"
DEFAULT_POSE = "RTMW-l 384x288 (disarankan)"

# Ambang skor deteksi "paling longgar" yang dipakai internal untuk mengumpulkan kandidat
RELAXED_DET_THR = 0.05

# Sigma OKS COCO (17 sendi badan)
COCO_SIGMAS = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62,
                        1.07, 1.07, .87, .87, .89, .89], dtype=np.float32) / 10.0

_LOCK = threading.Lock()
_SESSIONS: Dict[str, "OnnxModel"] = {}


# --------------------------------------------------------------------------
# Unduh + ekstrak model
# --------------------------------------------------------------------------
def ensure_onnx(zip_rel_path: str) -> str:
    """Mengunduh zip SDK dari HF (sekali) dan mengekstrak file .onnx di dalamnya."""
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    stem = Path(zip_rel_path).name[:-4]  # tanpa .zip
    onnx_path = MODEL_DIR / f"{stem}.onnx"
    if onnx_path.exists() and onnx_path.stat().st_size > 0:
        return str(onnx_path)

    from huggingface_hub import hf_hub_download  # impor malas (lazy)

    zip_path = hf_hub_download(repo_id=HF_REPO, filename=zip_rel_path)
    tmp_path = MODEL_DIR / f"{stem}.onnx.part"
    with zipfile.ZipFile(zip_path) as zf:
        members = [m for m in zf.namelist() if m.lower().endswith(".onnx")]
        if not members:
            raise RuntimeError(f"Tidak ada file .onnx di dalam {zip_rel_path}")
        # utamakan end2end.onnx (format mmdeploy SDK)
        members.sort(key=lambda m: (not m.endswith("end2end.onnx"), len(m)))
        with zf.open(members[0]) as src, open(tmp_path, "wb") as dst:
            shutil.copyfileobj(src, dst, length=16 * 1024 * 1024)
    os.replace(tmp_path, onnx_path)
    return str(onnx_path)


# --------------------------------------------------------------------------
# Pembungkus ONNX Runtime
# --------------------------------------------------------------------------
class OnnxModel:
    def __init__(self, onnx_path: str):
        import onnxruntime as ort

        so = ort.SessionOptions()
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        providers = ["CPUExecutionProvider"]
        avail = ort.get_available_providers()
        if "CUDAExecutionProvider" in avail:
            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        self.session = ort.InferenceSession(onnx_path, sess_options=so, providers=providers)
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]

    def run(self, chw_float: np.ndarray) -> List[np.ndarray]:
        inp = np.ascontiguousarray(chw_float[None].astype(np.float32))
        return self.session.run(self.output_names, {self.input_name: inp})


def _get_model(key: str, zip_rel_path: str) -> OnnxModel:
    with _LOCK:
        m = _SESSIONS.get(key)
        if m is None:
            m = OnnxModel(ensure_onnx(zip_rel_path))
            _SESSIONS[key] = m
        return m


def get_detector(name: str) -> Tuple[OnnxModel, Tuple[int, int]]:
    zip_rel, size = DETECTORS[name]
    return _get_model("det:" + zip_rel, zip_rel), size


def get_pose_model(name: str) -> Tuple[OnnxModel, Tuple[int, int], float]:
    zip_rel, size, score_scale = POSE_MODELS[name]
    return _get_model("pose:" + zip_rel, zip_rel), size, score_scale


def preload(det_name: str = DEFAULT_DETECTOR, pose_name: str = DEFAULT_POSE) -> None:
    get_detector(det_name)
    get_pose_model(pose_name)


# --------------------------------------------------------------------------
# Deteksi orang (YOLOX, mmdet/mmdeploy, tanpa modul NMS di dalam ONNX)
# --------------------------------------------------------------------------
def _nms(boxes: np.ndarray, scores: np.ndarray, iou_thr: float) -> List[int]:
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = np.clip(x2 - x1, 0, None) * np.clip(y2 - y1, 0, None)
    order = scores.argsort()[::-1]
    keep: List[int] = []
    while order.size > 0:
        i = int(order[0])
        keep.append(i)
        if order.size == 1:
            break
        rest = order[1:]
        xx1 = np.maximum(x1[i], x1[rest])
        yy1 = np.maximum(y1[i], y1[rest])
        xx2 = np.minimum(x2[i], x2[rest])
        yy2 = np.minimum(y2[i], y2[rest])
        inter = np.clip(xx2 - xx1, 0, None) * np.clip(yy2 - yy1, 0, None)
        iou = inter / (areas[i] + areas[rest] - inter + 1e-9)
        order = rest[iou <= iou_thr]
    return keep


def box_iou(a: np.ndarray, b: np.ndarray) -> float:
    xx1, yy1 = max(a[0], b[0]), max(a[1], b[1])
    xx2, yy2 = min(a[2], b[2]), min(a[3], b[3])
    inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
    ua = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return float(inter / (ua + 1e-9))


def _yolox_preprocess(img_bgr: np.ndarray, input_size: Tuple[int, int]) -> Tuple[np.ndarray, float]:
    """Letterbox ke (w, h) input model, padding 114 di kanan/bawah. Mengembalikan CHW float32 dan rasio."""
    in_w, in_h = input_size
    padded = np.full((in_h, in_w, 3), 114, dtype=np.uint8)
    r = min(in_h / img_bgr.shape[0], in_w / img_bgr.shape[1])
    new_w, new_h = max(1, int(img_bgr.shape[1] * r)), max(1, int(img_bgr.shape[0] * r))
    resized = cv2.resize(img_bgr, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    padded[:new_h, :new_w] = resized
    chw = np.ascontiguousarray(padded.transpose(2, 0, 1), dtype=np.float32)
    return chw, r


def _yolox_decode(outputs: np.ndarray, input_size: Tuple[int, int]) -> np.ndarray:
    """Dekode grid/stride YOLOX (keluaran mentah tanpa NMS) -> [cx, cy, w, h, obj, cls...]."""
    in_w, in_h = input_size
    grids, strides_all = [], []
    for stride in (8, 16, 32):
        hsize, wsize = in_h // stride, in_w // stride
        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
        grids.append(grid)
        strides_all.append(np.full((*grid.shape[:2], 1), stride))
    grids = np.concatenate(grids, 1).astype(np.float32)
    strides_all = np.concatenate(strides_all, 1).astype(np.float32)
    out = outputs.astype(np.float32).copy()
    out[..., :2] = (out[..., :2] + grids) * strides_all
    out[..., 2:4] = np.exp(out[..., 2:4]) * strides_all
    return out


def detect_people(img_bgr: np.ndarray, det_name: str, score_thr: float, nms_thr: float) -> np.ndarray:
    """Mengembalikan array (N, 5): x1, y1, x2, y2, skor — koordinat gambar asli, urut skor menurun."""
    model, input_size = get_detector(det_name)
    chw, ratio = _yolox_preprocess(img_bgr, input_size)
    outputs = model.run(chw)
    out0 = np.asarray(outputs[0])
    if out0.ndim == 2:
        out0 = out0[None]

    if out0.shape[-1] == 5:
        # Model dengan NMS bawaan: dets (1, N, 5) = x1,y1,x2,y2,skor pada koordinat input
        dets = out0[0]
        boxes = dets[:, :4] / ratio
        scores = dets[:, 4]
    else:
        pred = _yolox_decode(out0, input_size)[0]
        boxes_c = pred[:, :4]
        cls_scores = pred[:, 5:] if pred.shape[1] > 5 else np.ones((pred.shape[0], 1), np.float32)
        # kelas 0 = orang (model HumanArt hanya satu kelas)
        scores = pred[:, 4] * cls_scores[:, 0]
        boxes = np.empty_like(boxes_c)
        boxes[:, 0] = boxes_c[:, 0] - boxes_c[:, 2] / 2.0
        boxes[:, 1] = boxes_c[:, 1] - boxes_c[:, 3] / 2.0
        boxes[:, 2] = boxes_c[:, 0] + boxes_c[:, 2] / 2.0
        boxes[:, 3] = boxes_c[:, 1] + boxes_c[:, 3] / 2.0
        boxes = boxes / ratio

    valid = scores > score_thr
    boxes, scores = boxes[valid], scores[valid]
    if boxes.shape[0] == 0:
        return np.zeros((0, 5), np.float32)
    keep = _nms(boxes, scores, nms_thr)
    boxes, scores = boxes[keep], scores[keep]
    h, w = img_bgr.shape[:2]
    boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, w - 1)
    boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, h - 1)
    dets = np.concatenate([boxes, scores[:, None]], axis=1).astype(np.float32)
    ok = (dets[:, 2] - dets[:, 0] > 2) & (dets[:, 3] - dets[:, 1] > 2)
    dets = dets[ok]
    return dets[np.argsort(-dets[:, 4])]


# --------------------------------------------------------------------------
# Estimasi pose top-down (RTMPose / RTMW / DWPose, keluaran SimCC)
# --------------------------------------------------------------------------
MEAN = np.array([123.675, 116.28, 103.53], dtype=np.float32)
STD = np.array([58.395, 57.12, 57.375], dtype=np.float32)


def _bbox_xyxy2cs(bbox: np.ndarray, padding: float = 1.25) -> Tuple[np.ndarray, np.ndarray]:
    x1, y1, x2, y2 = [float(v) for v in bbox[:4]]
    center = np.array([(x1 + x2) * 0.5, (y1 + y2) * 0.5], dtype=np.float32)
    scale = np.array([(x2 - x1) * padding, (y2 - y1) * padding], dtype=np.float32)
    return center, scale


def _fix_aspect_ratio(scale: np.ndarray, aspect_ratio: float) -> np.ndarray:
    w, h = float(scale[0]), float(scale[1])
    if w > h * aspect_ratio:
        return np.array([w, w / aspect_ratio], dtype=np.float32)
    return np.array([h * aspect_ratio, h], dtype=np.float32)


def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    direction = a - b
    return b + np.array([-direction[1], direction[0]], dtype=np.float32)


def _get_warp_matrix(center: np.ndarray, scale: np.ndarray, output_size: Tuple[int, int]) -> np.ndarray:
    src_w = float(scale[0])
    dst_w, dst_h = output_size
    src_dir = np.array([0.0, src_w * -0.5], dtype=np.float32)
    dst_dir = np.array([0.0, dst_w * -0.5], dtype=np.float32)
    src = np.zeros((3, 2), dtype=np.float32)
    src[0] = center
    src[1] = center + src_dir
    src[2] = _get_3rd_point(src[0], src[1])
    dst = np.zeros((3, 2), dtype=np.float32)
    dst[0] = [dst_w * 0.5, dst_h * 0.5]
    dst[1] = np.array([dst_w * 0.5, dst_h * 0.5], dtype=np.float32) + dst_dir
    dst[2] = _get_3rd_point(dst[0], dst[1])
    return cv2.getAffineTransform(src, dst)


def _simcc_decode(simcc_x: np.ndarray, simcc_y: np.ndarray, split_ratio: float = 2.0):
    n, k, _ = simcc_x.shape
    sx = simcc_x.reshape(n * k, -1)
    sy = simcc_y.reshape(n * k, -1)
    x_locs = sx.argmax(axis=1)
    y_locs = sy.argmax(axis=1)
    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
    vx, vy = sx.max(axis=1), sy.max(axis=1)
    vals = np.minimum(vx, vy)
    locs[vals <= 0.0] = -1
    return locs.reshape(n, k, 2) / split_ratio, vals.reshape(n, k)


def _pick_simcc(model: "OnnxModel", outputs: List[np.ndarray], w: int, h: int) -> Tuple[np.ndarray, np.ndarray]:
    """Memilih tensor simcc_x / simcc_y berdasarkan nama keluaran, atau bentuknya bila nama tidak baku."""
    names = [n.lower() for n in getattr(model, "output_names", [])]
    if "simcc_x" in names and "simcc_y" in names:
        return np.asarray(outputs[names.index("simcc_x")]), np.asarray(outputs[names.index("simcc_y")])
    a, b = np.asarray(outputs[0]), np.asarray(outputs[1])
    # simcc_x panjangnya ~ 2*lebar, simcc_y ~ 2*tinggi
    if a.shape[-1] != b.shape[-1] and abs(a.shape[-1] - 2 * w) > abs(b.shape[-1] - 2 * w):
        a, b = b, a
    return a, b


def estimate_poses(img_rgb: np.ndarray, boxes_xyxy: np.ndarray, pose_name: str) -> Tuple[np.ndarray, np.ndarray]:
    """Mengembalikan keypoints (N, 133, 2) pada koordinat gambar dan skor (N, 133), skor sudah diskalakan ~0-1."""
    model, input_size, score_scale = get_pose_model(pose_name)
    w, h = input_size
    kpts_all, scores_all = [], []
    for bbox in boxes_xyxy:
        center, scale = _bbox_xyxy2cs(np.asarray(bbox, dtype=np.float32), padding=1.25)
        scale = _fix_aspect_ratio(scale, w / h)
        warp = _get_warp_matrix(center, scale, (w, h))
        crop = cv2.warpAffine(img_rgb, warp, (int(w), int(h)), flags=cv2.INTER_LINEAR)
        crop = (crop.astype(np.float32) - MEAN) / STD
        outputs = model.run(crop.transpose(2, 0, 1))
        simcc_x, simcc_y = _pick_simcc(model, outputs, w, h)
        locs, vals = _simcc_decode(simcc_x, simcc_y)
        kp = locs[0] / np.array([w, h], dtype=np.float32) * scale + center - scale / 2.0
        kpts_all.append(kp.astype(np.float32))
        scores_all.append((vals[0] * score_scale).astype(np.float32))
    if not kpts_all:
        return np.zeros((0, 133, 2), np.float32), np.zeros((0, 133), np.float32)
    return np.stack(kpts_all), np.stack(scores_all)


# --------------------------------------------------------------------------
# Pasca-proses multi-orang
# --------------------------------------------------------------------------
@dataclass(eq=False)
class Person:
    box: np.ndarray            # x1, y1, x2, y2
    det_score: float
    kpts: np.ndarray           # (133, 2)
    scores: np.ndarray         # (133,)
    relaxed: bool = False      # lolos hanya karena relaksasi ambang
    removed_joints: int = 0    # jumlah keypoint yang dihapus karena duplikat antar-orang

    @property
    def area(self) -> float:
        return float(max(1.0, (self.box[2] - self.box[0]) * (self.box[3] - self.box[1])))

    @property
    def center(self) -> np.ndarray:
        return np.array([(self.box[0] + self.box[2]) / 2.0, (self.box[1] + self.box[3]) / 2.0], np.float32)

    def n_visible_body(self, thr: float) -> int:
        return int((self.scores[:17] > thr).sum())

    def mean_body_score(self, thr: float) -> float:
        s = self.scores[:17]
        m = s > thr
        return float(s[m].mean()) if m.any() else 0.0


def oks_similarity(a: Person, b: Person, thr: float) -> Tuple[float, int]:
    """OKS atas sendi badan yang terlihat pada keduanya. Area = perkiraan area segmentasi (0.53 x area kotak)."""
    vis = (a.scores[:17] > thr) & (b.scores[:17] > thr)
    n = int(vis.sum())
    if n < 3:
        return 0.0, n
    d2 = ((a.kpts[:17] - b.kpts[:17]) ** 2).sum(-1)
    area = 0.53 * 0.5 * (a.area + b.area)
    var = (2.0 * COCO_SIGMAS) ** 2
    e = d2 / (2.0 * var * (area + 1e-6))
    return float(np.exp(-e[vis]).mean()), n


def suppress_duplicate_persons(persons: List[Person], kpt_thr: float, oks_thr: float) -> List[Person]:
    """NMS tingkat pose: instance dengan OKS tinggi terhadap instance berperingkat lebih tinggi = duplikat."""
    kept: List[Person] = []
    for cand in persons:  # sudah urut skor deteksi menurun
        dup = False
        for k in kept:
            o, _ = oks_similarity(cand, k, kpt_thr)
            if o > oks_thr:
                dup = True
                break
        if not dup:
            kept.append(cand)
    return kept


def remove_shared_joints(persons: List[Person], kpt_thr: float, rel_tol: float = 0.05) -> None:
    """Dua orang yang 'berbagi' titik sendi persis sama (indeks sama, jarak sangat dekat) =>
    titik itu dipertahankan untuk orang yang kotaknya paling 'memiliki' titik tersebut
    (jarak ternormalisasi ke pusat kotak lebih kecil); bila seimbang, skor keypoint yang menentukan."""
    for i in range(len(persons)):
        for j in range(i + 1, len(persons)):
            a, b = persons[i], persons[j]
            tol = max(3.0, rel_tol * math.sqrt(min(a.area, b.area)))
            d = np.linalg.norm(a.kpts - b.kpts, axis=-1)
            both = (a.scores > kpt_thr) & (b.scores > kpt_thr) & (d < tol)
            if not both.any():
                continue
            sa, sb = math.sqrt(a.area), math.sqrt(b.area)
            for jidx in np.flatnonzero(both):
                na = float(np.linalg.norm(a.kpts[jidx] - a.center)) / sa
                nb = float(np.linalg.norm(b.kpts[jidx] - b.center)) / sb
                if abs(na - nb) < 0.15:
                    loser = a if a.scores[jidx] < b.scores[jidx] else b
                else:
                    loser = a if na > nb else b
                loser.scores[jidx] = 0.0
                loser.removed_joints += 1


# pasangan kelompok kiri/kanan pada COCO-WholeBody: (kaki kiri, kaki kanan), (tangan kiri, tangan kanan)
_LR_GROUPS = [((17, 20), (20, 23)), ((91, 112), (112, 133))]


def remove_lr_duplicates(p: Person, kpt_thr: float, rel_tol: float = 0.025) -> None:
    """Artefak umum top-down: tangan (atau kaki) kiri dan kanan diprediksi di tempat yang persis sama
    karena salah satunya tersembunyi. Jika kedua kelompok hampir berimpit, kelompok dengan skor rata-rata
    lebih rendah dianggap halusinasi dan dimatikan."""
    tol = max(4.0, rel_tol * math.sqrt(p.area))
    for (l0, l1), (r0, r1) in _LR_GROUPS:
        ls, rs = p.scores[l0:l1], p.scores[r0:r1]
        if not (ls > kpt_thr).any() or not (rs > kpt_thr).any():
            continue
        d = np.linalg.norm(p.kpts[l0:l1] - p.kpts[r0:r1], axis=-1)
        if float(d.mean()) < tol:
            lo0, lo1 = (l0, l1) if float(ls.mean()) < float(rs.mean()) else (r0, r1)
            n_on = int((p.scores[lo0:lo1] > kpt_thr).sum())
            p.scores[lo0:lo1] = 0.0
            p.removed_joints += n_on


@dataclass
class PipelineResult:
    persons: List[Person]
    n_candidates: int
    n_after_dedup: int
    notes: List[str] = field(default_factory=list)


def run_pipeline(img_rgb: np.ndarray,
                 det_name: str = DEFAULT_DETECTOR,
                 pose_name: str = DEFAULT_POSE,
                 n_target: int = 2,
                 det_thr: float = 0.3,
                 nms_thr: float = 0.75,
                 kpt_thr: float = 0.3,
                 oks_thr: float = 0.75,
                 dedup_joints: bool = True,
                 min_body_visible: int = 3,
                 max_candidates: Optional[int] = None) -> PipelineResult:
    notes: List[str] = []
    h, w = img_rgb.shape[:2]
    img_bgr = np.ascontiguousarray(img_rgb[..., ::-1])

    # 1) kandidat kotak dengan ambang longgar (duplikat dihapus lagi di tingkat pose)
    dets = detect_people(img_bgr, det_name, score_thr=min(RELAXED_DET_THR, det_thr), nms_thr=nms_thr)
    if dets.shape[0] == 0:
        notes.append("Detektor tidak menemukan orang; seluruh foto dipakai sebagai satu kotak.")
        dets = np.array([[0, 0, w - 1, h - 1, 1.0]], dtype=np.float32)
    n_candidates = int(dets.shape[0])

    # pra-dedup tingkat kotak yang sangat murah (IoU > 0.9 jelas duplikat)
    pre: List[np.ndarray] = []
    for d in dets:
        if all(box_iou(d, k) < 0.9 for k in pre):
            pre.append(d)
    dets = np.stack(pre)

    if max_candidates is None:
        max_candidates = max(6, 2 * n_target + 4) if n_target > 0 else 12
    dets = dets[:max_candidates]

    # 2) estimasi pose untuk semua kandidat
    kpts, scores = estimate_poses(img_rgb, dets[:, :4], pose_name)
    persons = [Person(box=dets[i, :4].copy(), det_score=float(dets[i, 4]),
                      kpts=kpts[i].copy(), scores=scores[i].copy()) for i in range(len(dets))]

    # 3) hapus duplikat tingkat pose (OKS)
    persons = suppress_duplicate_persons(persons, kpt_thr, oks_thr)
    n_after_dedup = len(persons)

    # 4) hapus keypoint halusinasi: tangan/kaki kiri-kanan yang berimpit pada orang yang sama,
    #    lalu keypoint yang dibagi dua orang berbeda
    if dedup_joints:
        for p in persons:
            remove_lr_duplicates(p, kpt_thr)
        if len(persons) > 1:
            remove_shared_joints(persons, kpt_thr)

    # 5) saring kualitas minimum
    persons = [p for p in persons if p.n_visible_body(kpt_thr) >= min_body_visible]

    # 6) seleksi akhir + relaksasi otomatis
    strong = [p for p in persons if p.det_score >= det_thr]
    if n_target > 0:
        final = strong[:n_target]
        if len(final) < n_target:
            extra = [p for p in persons if p not in final][: n_target - len(final)]
            for p in extra:
                p.relaxed = True
            if extra:
                notes.append(
                    f"{len(extra)} orang ditambahkan lewat relaksasi ambang (skor deteksi < {det_thr:.2f}) "
                    f"agar mencapai target {n_target} orang.")
            final += extra
        if len(final) < n_target:
            notes.append(f"Hanya {len(final)} orang yang dapat ditemukan (target {n_target}).")
    else:
        final = strong

    # urutkan kiri -> kanan supaya penomoran konsisten
    final.sort(key=lambda p: float(p.center[0]))
    return PipelineResult(persons=final, n_candidates=n_candidates, n_after_dedup=n_after_dedup, notes=notes)


# --------------------------------------------------------------------------
# Konversi ke format OpenPose (BODY_18 + wajah 68/70 + tangan 2x21)
# --------------------------------------------------------------------------
# urutan OpenPose BODY_18: 0 hidung, 1 leher, 2 bahu ka, 3 siku ka, 4 pergelangan ka,
# 5 bahu ki, 6 siku ki, 7 pergelangan ki, 8 pinggul ka, 9 lutut ka, 10 mata kaki ka,
# 11 pinggul ki, 12 lutut ki, 13 mata kaki ki, 14 mata ka, 15 mata ki, 16 telinga ka, 17 telinga ki
_OP_FROM_COCO = {0: 0, 2: 6, 3: 8, 4: 10, 5: 5, 6: 7, 7: 9, 8: 12, 9: 14, 10: 16,
                 11: 11, 12: 13, 13: 15, 14: 2, 15: 1, 16: 4, 17: 3}


@dataclass(eq=False)
class OpenPosePerson:
    body: np.ndarray       # (18, 2)
    body_s: np.ndarray     # (18,)
    face: np.ndarray       # (68, 2)
    face_s: np.ndarray
    lhand: np.ndarray      # (21, 2)
    lhand_s: np.ndarray
    rhand: np.ndarray      # (21, 2)
    rhand_s: np.ndarray
    feet: np.ndarray       # (6, 2)
    feet_s: np.ndarray
    box: np.ndarray
    det_score: float
    relaxed: bool
    removed_joints: int


def to_openpose(p: Person, kpt_thr: float, scale: float = 1.0) -> OpenPosePerson:
    k = p.kpts * scale
    s = p.scores
    body = np.zeros((18, 2), np.float32)
    body_s = np.zeros(18, np.float32)
    for op_i, coco_i in _OP_FROM_COCO.items():
        body[op_i] = k[coco_i]
        body_s[op_i] = s[coco_i]
    # leher = titik tengah kedua bahu, terlihat hanya jika kedua bahu terlihat
    if s[5] > kpt_thr and s[6] > kpt_thr:
        body[1] = (k[5] + k[6]) / 2.0
        body_s[1] = min(s[5], s[6])
    return OpenPosePerson(
        body=body, body_s=body_s,
        face=k[23:91].copy(), face_s=s[23:91].copy(),
        lhand=k[91:112].copy(), lhand_s=s[91:112].copy(),
        rhand=k[112:133].copy(), rhand_s=s[112:133].copy(),
        feet=k[17:23].copy(), feet_s=s[17:23].copy(),
        box=p.box * scale, det_score=p.det_score, relaxed=p.relaxed,
        removed_joints=p.removed_joints,
    )


def _flat(kp: np.ndarray, sc: np.ndarray, thr: float) -> List[float]:
    out: List[float] = []
    for (x, y), c in zip(kp, sc):
        if c > thr:
            out += [round(float(x), 2), round(float(y), 2), round(float(c), 3)]
        else:
            out += [0.0, 0.0, 0.0]
    return out


def openpose_json(people: List[OpenPosePerson], width: int, height: int, kpt_thr: float) -> List[dict]:
    """Format JSON OpenPose (kompatibel dengan POSE_KEYPOINT ComfyUI / OpenPose Editor): list of frames."""
    plist = []
    for idx, p in enumerate(people):
        # wajah 70 titik: 68 landmark + 2 pupil (mata kanan, mata kiri dari badan)
        face70 = np.concatenate([p.face, p.body[[14, 15]]], axis=0)
        face70_s = np.concatenate([p.face_s, p.body_s[[14, 15]]], axis=0)
        plist.append({
            "person_id": [idx],
            "pose_keypoints_2d": _flat(p.body, p.body_s, kpt_thr),
            "face_keypoints_2d": _flat(face70, face70_s, kpt_thr),
            "hand_left_keypoints_2d": _flat(p.lhand, p.lhand_s, kpt_thr),
            "hand_right_keypoints_2d": _flat(p.rhand, p.rhand_s, kpt_thr),
            "foot_keypoints_2d": _flat(p.feet, p.feet_s, kpt_thr),
            "bbox_xyxy": [round(float(v), 1) for v in p.box],
            "detection_score": round(float(p.det_score), 3),
        })
    return [{"version": 1.3, "people": plist, "canvas_width": int(width), "canvas_height": int(height)}]


# --------------------------------------------------------------------------
# Rendering gaya OpenPose / ControlNet (warna RGB, identik dengan controlnet_aux)
# --------------------------------------------------------------------------
LIMB_SEQ = [[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]]
BODY_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]]
HAND_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]]


def _pt(v: np.ndarray) -> Tuple[int, int]:
    return int(round(float(v[0]))), int(round(float(v[1])))


def draw_bodies(canvas: np.ndarray, people: List[OpenPosePerson], kpt_thr: float, s: float) -> np.ndarray:
    stickwidth = max(1, int(round(4 * s)))
    radius = max(1, int(round(4 * s)))
    for p in people:
        for i, (a, b) in enumerate(LIMB_SEQ):
            ia, ib = a - 1, b - 1
            if p.body_s[ia] <= kpt_thr or p.body_s[ib] <= kpt_thr:
                continue
            (x1, y1), (x2, y2) = p.body[ia], p.body[ib]
            mx, my = (x1 + x2) / 2.0, (y1 + y2) / 2.0
            length = math.hypot(x1 - x2, y1 - y2)
            angle = math.degrees(math.atan2(y1 - y2, x1 - x2))
            poly = cv2.ellipse2Poly((int(mx), int(my)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
            cv2.fillConvexPoly(canvas, poly, BODY_COLORS[i])
    canvas = (canvas.astype(np.float32) * 0.6).astype(np.uint8)
    for p in people:
        for i in range(18):
            if p.body_s[i] > kpt_thr:
                cv2.circle(canvas, _pt(p.body[i]), radius, BODY_COLORS[i], thickness=-1)
    return canvas


def draw_hands(canvas: np.ndarray, people: List[OpenPosePerson], kpt_thr: float, s: float) -> np.ndarray:
    thickness = max(1, int(round(2 * s)))
    radius = max(1, int(round(3 * s)))
    n_edges = len(HAND_EDGES)
    for p in people:
        for hand, hand_s in ((p.lhand, p.lhand_s), (p.rhand, p.rhand_s)):
            for ie, (a, b) in enumerate(HAND_EDGES):
                if hand_s[a] <= kpt_thr or hand_s[b] <= kpt_thr:
                    continue
                r, g, bl = colorsys.hsv_to_rgb(ie / float(n_edges), 1.0, 1.0)
                color = (int(r * 255), int(g * 255), int(bl * 255))
                cv2.line(canvas, _pt(hand[a]), _pt(hand[b]), color, thickness=thickness)
            for i in range(21):
                if hand_s[i] > kpt_thr:
                    cv2.circle(canvas, _pt(hand[i]), radius, (0, 0, 255), thickness=-1)
    return canvas


def draw_faces(canvas: np.ndarray, people: List[OpenPosePerson], kpt_thr: float, s: float) -> np.ndarray:
    radius = max(1, int(round(2 * s)))
    for p in people:
        for i in range(p.face.shape[0]):
            if p.face_s[i] > kpt_thr:
                cv2.circle(canvas, _pt(p.face[i]), radius, (255, 255, 255), thickness=-1)
    return canvas


def render_pose_map(people: List[OpenPosePerson], width: int, height: int, kpt_thr: float,
                    line_scale: float = 1.0, draw_body: bool = True, draw_hand: bool = True,
                    draw_face: bool = True, background: Optional[np.ndarray] = None) -> np.ndarray:
    canvas = np.zeros((height, width, 3), np.uint8) if background is None else background.copy()
    s = line_scale * max(1.0, max(width, height) / 512.0)
    if draw_body:
        canvas = draw_bodies(canvas, people, kpt_thr, s)
    if draw_hand:
        canvas = draw_hands(canvas, people, kpt_thr, s)
    if draw_face:
        canvas = draw_faces(canvas, people, kpt_thr, s)
    return canvas


def render_overlay(img_rgb: np.ndarray, pose_map: np.ndarray, people: List[OpenPosePerson],
                   alpha_bg: float = 0.45) -> np.ndarray:
    """Foto digelapkan + skeleton ditimpa + kotak & nomor orang untuk pengecekan."""
    base = (img_rgb.astype(np.float32) * alpha_bg).astype(np.uint8)
    mask = pose_map.any(axis=-1)
    base[mask] = pose_map[mask]
    s = max(1.0, max(base.shape[:2]) / 512.0)
    for idx, p in enumerate(people):
        x1, y1, x2, y2 = [int(round(float(v))) for v in p.box]
        color = (255, 255, 255) if not p.relaxed else (255, 200, 0)
        cv2.rectangle(base, (x1, y1), (x2, y2), color, max(1, int(round(s))))
        label = f"#{idx + 1} {p.det_score:.2f}" + (" (relaksasi)" if p.relaxed else "")
        cv2.putText(base, label, (x1 + 3, max(12, y1 - 4)), cv2.FONT_HERSHEY_SIMPLEX,
                    0.45 * s, color, max(1, int(round(s))), cv2.LINE_AA)
    return base