#!/usr/bin/env python3
"""Face recall at a fixed confidence threshold on WIDER val (what fraction of GT faces get a box with IoU>=0.5)."""
import argparse
from pathlib import Path

import numpy as np
from ultralytics import YOLO


def iou(a, b):
    x1 = np.maximum(a[:, None, 0], b[None, :, 0]); y1 = np.maximum(a[:, None, 1], b[None, :, 1])
    x2 = np.minimum(a[:, None, 2], b[None, :, 2]); y2 = np.minimum(a[:, None, 3], b[None, :, 3])
    inter = np.clip(x2 - x1, 0, None) * np.clip(y2 - y1, 0, None)
    aa = (a[:, 2] - a[:, 0]) * (a[:, 3] - a[:, 1]); ab = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1])
    return inter / (aa[:, None] + ab[None, :] - inter + 1e-9)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="export/face_yolo11n_fp16.onnx")
    ap.add_argument("--conf", type=float, default=0.25)
    ap.add_argument("--iou", type=float, default=0.5)
    ap.add_argument("--val", default="dataset/images/val")
    ap.add_argument("--device", default="0", help="0 for GPU, cpu for CPU")
    ap.add_argument("--imgsz", type=int, default=640, help="inference size (.pt models only; ONNX exports are fixed)")
    a = ap.parse_args()

    model = YOLO(a.model, task="detect")
    files = sorted(Path(a.val).glob("*.jpg"))
    hits, sizes, n_pred, n_fp = [], [], 0, 0
    for i in range(0, len(files), 1):  # ONNX export is fixed batch=1
        batch = files[i:i + 1]
        for f, r in zip(batch, model.predict([str(x) for x in batch], conf=a.conf, imgsz=a.imgsz, device=a.device, verbose=False)):
            h, w = r.orig_shape
            gt = np.array([[float(v) for v in l.split()[1:]] for l in
                           Path(str(f).replace("images", "labels").replace(".jpg", ".txt")).read_text().split("\n") if l.strip()])
            g = np.stack([(gt[:, 0] - gt[:, 2] / 2) * w, (gt[:, 1] - gt[:, 3] / 2) * h,
                          (gt[:, 0] + gt[:, 2] / 2) * w, (gt[:, 1] + gt[:, 3] / 2) * h], 1)
            p = r.boxes.xyxy.cpu().numpy()
            n_pred += len(p)
            if len(p):
                m = iou(g, p)
                hits += list(m.max(1) >= a.iou)
                n_fp += int((m.max(0) < a.iou).sum())
            else:
                hits += [False] * len(g)
            sizes += list(gt[:, 3] * h)  # face height in px at the image's stored (<=640) scale
    hits, sizes = np.array(hits), np.array(sizes)
    print(f"RES faces={len(hits)} overall recall@conf{a.conf}/IoU{a.iou} = {hits.mean():.1%}  (false-positive boxes: {n_fp}/{n_pred})")
    for lo, hi in [(0, 10), (10, 20), (20, 40), (40, 80), (80, 1e9)]:
        m = (sizes >= lo) & (sizes < hi)
        print(f"RES   face height {lo:>3}-{int(min(hi, 999)):>3}px: {hits[m].mean():6.1%}  (n={m.sum()})")


if __name__ == "__main__":
    main()