"""OpenVINO-backed anime-face detector for anime_facesaver (iGPU path).

Used when --gpu igpu: replaces ultralytics' torch predict loop with a direct
OpenVINO inference request on Intel GPUs (Level Zero / ANV). Pre/post-processing
emulates ultralytics YOLO defaults: 1280 letterbox (stride 32), conf/iou gates,
class-agnostic-free single-class NMS.

On chibi (i3-N305): GPU fp32 measured 5.2 s/img vs CPU 14.5 s/img
(2026-10-20, models/yolov8x6_animeface.onnx, 25 real stills).
"""
import os

import numpy as np


def letterbox(img, new_shape=1280, stride=32):
    """Square ultralytics letterbox (auto=False): fit, then center-pad to
    new_shape x new_shape. The GPU plugin compiles one static shape this way
    and sidesteps dynamic-shape concat failures on non-square inputs
    (measured: (1,3,736,1280) fails inside the graph; 1280x1280 always works).

    Returns (blob CHW float32 /255, scale, (pad_left, pad_top)).
    Inverse map for boxes: x_orig = (x_lb - pad_left) / scale.
    """
    import cv2

    h, w = img.shape[:2]
    r = min(new_shape / h, new_shape / w)
    nw, nh = int(round(w * r)), int(round(h * r))
    interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
    resized = cv2.resize(img, (nw, nh), interpolation=interp)
    canvas = np.full((new_shape, new_shape, 3), 114, dtype=np.uint8)
    top, left = (new_shape - nh) // 2, (new_shape - nw) // 2
    canvas[top:top + nh, left:left + nw] = resized
    blob = canvas[..., ::-1].transpose(2, 0, 1)[None].astype(np.float32) / 255.0
    return blob, r, (left, top)


def numpy_ceil_div(a, b):
    return -(-a // b)


def nms_numpy(boxes, scores, iou_thres):
    """Plain greedy NMS; returns kept indices."""
    if len(boxes) == 0:
        return []
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size:
        i = order[0]
        keep.append(int(i))
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.clip(xx2 - xx1, 0, None) * np.clip(yy2 - yy1, 0, None)
        iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-9)
        order = order[1:][iou <= iou_thres]
    return keep


class OVFaceDetector:
    """model.predict(frame_bgr, conf=, iou=, imgsz=, device=) compatible shim.

    Returns a one-element list with .boxes being a Boxes-like object exposing
    .xyxy (numpy, already denormalized to original-frame pixels) and .conf.
    """

    class _Boxes:
        """numpy-backed stand-in for ultralytics Boxes (.xyxy.cpu().numpy())."""
        class _Lazy:
            def __init__(self, arr):
                self._arr = arr

            def cpu(self):
                return self

            def numpy(self):
                return self._arr

            def int(self):
                return OVFaceDetector._Boxes._Lazy(self._arr.astype(np.int32))

            def tolist(self):
                return self._arr.tolist()

        def __init__(self, xyxy, conf):
            self.xyxy = self._Lazy(xyxy)
            self.conf = conf

    class _Result:
        def __init__(self, boxes):
            self.boxes = boxes

    def __init__(self, model_path, device="GPU", imgsz=1280):
        import openvino as ov

        core = ov.Core()
        if device != "CPU" and device not in core.available_devices:
            raise RuntimeError(
                f"OpenVINO device {device!r} not available "
                f"(have {core.available_devices}). Is the render-group fix "
                f"active (sg render / re-login) and intel-level-zero-gpu "
                f"installed?")
        self._ov = ov
        self._req = core.compile_model(model_path, device).create_infer_request()
        self.imgsz = imgsz
        self.names = {0: "face"}

    def predict(self, frame, conf=0.25, iou=0.45, imgsz=None, device=None,
                verbose=False):
        imgsz = imgsz or self.imgsz
        blob, scale, (pad_l, pad_t) = letterbox(frame, imgsz)
        out = self._req.infer([self._ov.Tensor(blob)])
        pred = list(out.values())[0]  # [1, 5, N]: cx cy w h conf cls
        pred = pred[0].T  # [N, 5]
        scores = pred[:, 4]
        cand = pred[scores >= conf]
        if len(cand) == 0:
            return [self._Result(self._Boxes(
                np.zeros((0, 4), np.float32), np.zeros((0,), np.float32)))]
        cx, cy, w, h = cand[:, 0], cand[:, 1], cand[:, 2], cand[:, 3]
        boxes = np.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], 1)
        keep = nms_numpy(boxes, cand[:, 4], iou)
        boxes, scores = boxes[keep], cand[keep, 4]
        # map letterbox coords back to original frame pixels
        boxes[:, [0, 2]] -= pad_l
        boxes[:, [1, 3]] -= pad_t
        boxes = boxes / scale
        h0, w0 = frame.shape[:2]
        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w0)
        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h0)
        return [self._Result(self._Boxes(boxes.astype(np.float32),
                                         scores.astype(np.float32)))]