#!/usr/bin/env python
"""Detector smoke check — runs TODAY, no SKEL, no HMR forward pass.

Proves the torchvision person detector (lib/modeling/pipelines/vitdet) imports without
detectron2 and honors the contract lib.kits.hsmr_demo._img_det2patches depends on.

    cd ~/Documents/MOCAPGOD/.engine/HSMR && source .venv/bin/activate
    PYTHONPATH=. python test_detector.py
"""
from pathlib import Path
import cv2
import torch

import _headless  # noqa: F401  fake pyrender/OpenGL for headless macOS — MUST precede lib.* imports
from lib.modeling.pipelines.vitdet import build_detector  # must import without detectron2
from lib.kits.hsmr_demo import imgs_det2patches            # the real consumer of detector output

DEMO = Path('data_inputs/demo/example_imgs')
DET_THRESHOLD_SCORE = 0.5  # matches hsmr_demo._img_det2patches


def main():
    fns = sorted(p for p in DEMO.glob('*')
                 if p.suffix.lower() in {'.jpg', '.jpeg', '.png', '.webp'})
    assert fns, f'no demo images under {DEMO}'

    imgs = []
    for fn in fns:
        bgr = cv2.imread(str(fn))
        assert bgr is not None, f'cv2 failed to read {fn}'
        imgs.append(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))  # HSMR feeds RGB

    dets, ratios = build_detector(device='cpu')(imgs)
    assert len(dets) == len(imgs) == len(ratios), 'per-image list length mismatch'

    n_person_imgs = 0
    for fn, img, d, r in zip(fns, imgs, dets, ratios):
        H, W = img.shape[:2]
        assert set(d) == {'pred_classes', 'scores', 'pred_boxes'}, f'bad keys {set(d)}'
        assert d['pred_classes'].dtype == torch.long
        assert (d['pred_classes'] == 0).all(), 'human class must be 0 (CLASS_HUMAN_ID)'
        assert d['pred_boxes'].ndim == 2 and d['pred_boxes'].shape[1] == 4, 'boxes must be (N,4)'
        assert r == 1.0, 'native-resolution detection -> ratio 1.0'

        strong = d['scores'] > DET_THRESHOLD_SCORE
        if strong.any():
            n_person_imgs += 1
            x1, y1, x2, y2 = d['pred_boxes'][strong].unbind(1)
            assert (x2 > x1).all() and (y2 > y1).all(), 'boxes must be lurb (x2>x1, y2>y1)'
            assert (x1 >= -1).all() and (y1 >= -1).all() \
                and (x2 <= W + 1).all() and (y2 <= H + 1).all(), 'box out of frame'
            print(f'  ✓ {fn.name}: {int(strong.sum())} person(s), '
                  f'top score {float(d["scores"][strong].max()):.2f}')
        else:
            print(f'  · {fn.name}: no confident person')

    # Every demo image has an obvious human; allow one miss for robustness slack.
    assert n_person_imgs >= len(imgs) - 1, \
        f'expected a person in ~all demo imgs, got {n_person_imgs}/{len(imgs)}'
    print(f'DETECTOR OK: person found in {n_person_imgs}/{len(imgs)} demo images.')

    # Close the loop: feed detector output through HSMR's real cropper (SKEL-free) and
    # confirm it yields the 256x256 patches the HMR forward pass expects.
    patches, det_meta = imgs_det2patches(imgs, dets, ratios, max_instances_per_img=5)
    assert patches.ndim == 4 and patches.shape[1:] == (256, 256, 3), f'bad patch shape {patches.shape}'
    assert patches.shape[0] == sum(det_meta['n_patch_per_img']), 'patch count / meta mismatch'
    print(f'PATCHES OK: {patches.shape[0]} human patches @ 256x256 — ready for the HMR pass.')


if __name__ == '__main__':
    main()