"""
Person detector — torchvision drop-in replacement for HSMR's detectron2/ViTDet.

Why: detectron2 + chumpy won't build from source on Apple Silicon (the Lane B blocker),
and this module's `import detectron2` was a load-time landmine for everything that touches
`lib.kits.hsmr_demo` (smoke_test.py, the future pose_engine.py). torchvision ships a
COCO-pretrained Faster R-CNN with ungated weights and no build step.

Same public contract as the original, so run_demo.py / smoke_test.py / pose_engine.py
import it unchanged:

    build_detector(...) -> callable(raw_imgs) -> (dets, downsample_ratios)

    dets[i]           : dict of CPU tensors (empty tensors if no person) —
                          pred_classes : all 0  (== hsmr_demo CLASS_HUMAN_ID)
                          scores       : float
                          pred_boxes   : (N, 4) left-upper-right-bottom pixels
    downsample_ratios : one float per image; boxes are in native-resolution pixels -> 1.0

Contract is checked by test_detector.py (runs today, no SKEL needed).

# ponytail: ~35 lines replaces the whole detectron2 stack. The ViT-H HMR forward pass is
# the runtime cost, not detection, so we don't chase detector throughput here.
"""
import numpy as np
import torch
from tqdm import tqdm
from torchvision.models.detection import (
    fasterrcnn_resnet50_fpn_v2,
    FasterRCNN_ResNet50_FPN_V2_Weights,
)

_COCO_PERSON = 1  # torchvision COCO label id for "person" (0 == background)


class _PersonDetector:
    """Callable over a list of RGB HxWx3 images (what HSMR's load_inputs yields).
    uint8 or float, [0,255] or [0,1] are all handled."""

    def __init__(self, device='cpu', score_thresh=0.25):
        # score_thresh 0.25 mirrors the old detectron2 test_score_thresh; hsmr_demo's
        # _img_det2patches re-filters at 0.5, so this only widens recall a little.
        weights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT
        self.model = fasterrcnn_resnet50_fpn_v2(
            weights=weights, box_score_thresh=score_thresh,
        ).eval().to(device)
        self.device = device

    @torch.no_grad()
    def __call__(self, raw_imgs):
        dets, ratios = [], []
        for img in tqdm(raw_imgs, desc='Detecting'):
            t = torch.as_tensor(np.ascontiguousarray(img), device=self.device).float()
            if float(t.max()) > 1.5:          # [0,255] -> [0,1]
                t = t / 255.0
            t = t.permute(2, 0, 1)            # HWC -> CHW
            out = self.model([t])[0]
            keep = out['labels'] == _COCO_PERSON
            dets.append({
                'pred_classes': torch.zeros(int(keep.sum()), dtype=torch.long),  # human -> 0
                'scores'      : out['scores'][keep].cpu(),
                'pred_boxes'  : out['boxes'][keep].cpu(),   # xyxy == left-upper-right-bottom
            })
            ratios.append(1.0)               # detected at native resolution
        return dets, ratios


def build_detector(batch_size=1, max_img_size=512, device='cpu'):
    """Signature-compatible with the old detectron2 builder. batch_size / max_img_size are
    accepted for parity but unused — detection runs per-frame at native resolution.

    # ponytail: detector pinned to CPU regardless of `device`. It's cheap next to the
    # ViT-H HMR pass, and CPU sidesteps the MPS coverage gaps in RoIAlign/NMS. Flip to
    # `device` only if per-frame detection ever dominates on long videos."""
    return _PersonDetector(device='cpu')