"""RF-DETR official postprocess (transformer, NO NMS) for nvinferserver raw tensors.

Tensor layout (verified in-container, full-dims with batch prepended):
  detection  RFDETRMedium:        dets[1,Q,4]  labels[1,Q,91]
  pose       RFDETRKeypointPreview: dets[1,100,4] labels[1,100,2] keypoints[1,100,34,8]
  seg        RFDETRSegMedium:      dets[1,200,4] labels[1,200,91] masks[1,200,144,144]

Pose keypoints layout (official rfdetr/models/postprocess.py):
  [Q, C*maxK, 8] where C=2 classes, maxK=17 (COCO person).
  class c rows: c*17 .. c*17+16. Per keypoint:
    dim0 = x normalized, dim1 = y normalized, dim2 = visibility logit (sigmoid),
    dim3 = reserved/unused, dim4..6 = precision Cholesky (log_l11, l21, log_l22),
    dim7 = reserved.
  Class 0 (background) rows are all-zero.

Policy: per-class sigmoid on logits, index 0 = background (masked to -1),
topk by score, threshold filter. NO NMS (Hungarian 1-to-1 query matching).
"""
from __future__ import annotations

import numpy as np

COCO_KP_PER_PERSON = 17
KP_DIM = 8


def _sigmoid(a: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-np.clip(a.astype(np.float64), -500, 500)))


def squeeze_batch(a: np.ndarray, target_ndim: int) -> np.ndarray:
    """Drop a leading size-1 batch dim if present (full-dims tensors)."""
    while a.ndim > target_ndim and a.shape[0] == 1:
        a = a[0]
    return a


def parse_boxes_labels(dets: np.ndarray, labels: np.ndarray, threshold: float = 0.5):
    """Shared box+class parse for det/pose/seg.

    ``model_classes``는 background를 포함한 raw model index이고, ``classes``는
    background를 제외해 0부터 시작하는 canonical class_id다.
    """
    probs = _sigmoid(labels)
    if probs.shape[1] > 1:
        probs = probs.copy()
        probs[:, 0] = -1.0  # background
    scores = probs.max(axis=1)
    model_classes = probs.argmax(axis=1)
    classes = model_classes - 1
    keep = (scores > threshold) & (model_classes > 0)
    return {
        "dets": dets,
        "model_classes": model_classes,
        "classes": classes,
        "scores": scores,
        "keep": keep,
    }


def parse_pose_keypoints(kp: np.ndarray, classes: np.ndarray, num_keypoint_classes: int = 2,
                         max_keypoints: int = COCO_KP_PER_PERSON,
                         active_keypoints: tuple[int, ...] = (0, 17)) -> np.ndarray:
    """Decode official pose keypoints to [Q, maxK, 3] in NORMALIZED (x, y, conf).

    kp: [Q, C*maxK, 8] full layout. classes: [Q] predicted class per query.
    Non-keypoint classes (e.g. background, idx >= num_keypoint_classes) get all zeros.
    """
    q = kp.shape[0]
    out = np.zeros((q, max_keypoints, 3), dtype=np.float32)
    reshaped = kp.reshape(q, num_keypoint_classes, max_keypoints, kp.shape[-1])
    for c in range(num_keypoint_classes):
        n_active = active_keypoints[c] if c < len(active_keypoints) else 0
        if n_active <= 0:
            continue
        m = classes == c
        if not m.any():
            continue
        sel = reshaped[m, c, :n_active]          # [n, K, 8]
        out[m, :n_active, 0] = sel[..., 0]       # x norm
        out[m, :n_active, 1] = sel[..., 1]       # y norm
        out[m, :n_active, 2] = _sigmoid(sel[..., 2]).astype(np.float32)
    return out


def keypoints_to_pixels(kp_norm: np.ndarray, img_w: int, img_h: int) -> np.ndarray:
    """Normalized (x,y,conf) -> pixel (x,y,conf)."""
    out = kp_norm.copy()
    out[..., 0] = kp_norm[..., 0] * img_w
    out[..., 1] = kp_norm[..., 1] * img_h
    return out