"""Standalone SCRFD face detector (ONNX Runtime, CPU).

A self-contained re-implementation of InsightFace's SCRFD post-processing so the
app can load a raw ``.onnx`` file directly — no ``insightface`` package, no model
pack directory. Handles the 6/9/10/15-output SCRFD variants; ``det_500m.onnx`` and
``det_2.5g.onnx`` are both 9-output (3 strides x {score, bbox, kps}) with 5-point
landmarks.

Reference: https://github.com/deepinsight/insightface (scrfd.py).
"""

import cv2
import numpy as np
import onnxruntime


def distance2bbox(points, distance):
    """Decode (left, top, right, bottom) distances from an anchor center to a box."""
    x1 = points[:, 0] - distance[:, 0]
    y1 = points[:, 1] - distance[:, 1]
    x2 = points[:, 0] + distance[:, 2]
    y2 = points[:, 1] + distance[:, 3]
    return np.stack([x1, y1, x2, y2], axis=-1)


def distance2kps(points, distance):
    """Decode per-keypoint (dx, dy) distances from an anchor center to landmarks."""
    preds = []
    for i in range(0, distance.shape[1], 2):
        px = points[:, i % 2] + distance[:, i]
        py = points[:, i % 2 + 1] + distance[:, i + 1]
        preds.append(px)
        preds.append(py)
    return np.stack(preds, axis=-1)


class SCRFD:
    """SCRFD ONNX face detector.

    Parameters
    ----------
    model_file : str
        Path to the ``.onnx`` file.
    providers : list[str] | None
        ONNX Runtime execution providers. Defaults to CPU.
    """

    def __init__(self, model_file, providers=None):
        self.model_file = model_file
        providers = providers or ["CPUExecutionProvider"]
        self.session = onnxruntime.InferenceSession(model_file, providers=providers)
        self.center_cache = {}
        self.nms_thresh = 0.4
        self.input_mean = 127.5
        self.input_std = 128.0
        self._init_vars()

    def _init_vars(self):
        inp = self.session.get_inputs()[0]
        self.input_name = inp.name
        # Static input size if the model declares one (e.g. [1, 3, 640, 640]),
        # else fall back to 640x640 at detect() time.
        shape = inp.shape
        if isinstance(shape[2], int) and isinstance(shape[3], int):
            self.input_size = (shape[3], shape[2])  # (w, h)
        else:
            self.input_size = (640, 640)
        outputs = self.session.get_outputs()
        self.output_names = [o.name for o in outputs]
        self.use_kps = False
        self._num_anchors = 1
        n = len(outputs)
        if n == 6:
            self.fmc, self._feat_stride_fpn, self._num_anchors = 3, [8, 16, 32], 2
        elif n == 9:
            self.fmc, self._feat_stride_fpn, self._num_anchors = 3, [8, 16, 32], 2
            self.use_kps = True
        elif n == 10:
            self.fmc, self._feat_stride_fpn, self._num_anchors = 5, [8, 16, 32, 64, 128], 1
        elif n == 15:
            self.fmc, self._feat_stride_fpn, self._num_anchors = 5, [8, 16, 32, 64, 128], 1
            self.use_kps = True
        else:
            raise ValueError(f"Unexpected SCRFD output count: {n}")

    def forward(self, img, thresh):
        scores_list, bboxes_list, kpss_list = [], [], []
        blob = cv2.dnn.blobFromImage(
            img,
            1.0 / self.input_std,
            (img.shape[1], img.shape[0]),
            (self.input_mean, self.input_mean, self.input_mean),
            swapRB=True,
        )
        net_outs = self.session.run(self.output_names, {self.input_name: blob})

        input_height, input_width = blob.shape[2], blob.shape[3]
        fmc = self.fmc
        for idx, stride in enumerate(self._feat_stride_fpn):
            scores = net_outs[idx]
            bbox_preds = net_outs[idx + fmc] * stride
            if self.use_kps:
                kps_preds = net_outs[idx + fmc * 2] * stride

            height, width = input_height // stride, input_width // stride
            key = (height, width, stride)
            if key in self.center_cache:
                anchor_centers = self.center_cache[key]
            else:
                anchor_centers = np.stack(
                    np.mgrid[:height, :width][::-1], axis=-1
                ).astype(np.float32)
                anchor_centers = (anchor_centers * stride).reshape((-1, 2))
                if self._num_anchors > 1:
                    anchor_centers = np.stack(
                        [anchor_centers] * self._num_anchors, axis=1
                    ).reshape((-1, 2))
                if len(self.center_cache) < 100:
                    self.center_cache[key] = anchor_centers

            pos_inds = np.where(scores >= thresh)[0]
            bboxes = distance2bbox(anchor_centers, bbox_preds)
            scores_list.append(scores[pos_inds])
            bboxes_list.append(bboxes[pos_inds])
            if self.use_kps:
                kpss = distance2kps(anchor_centers, kps_preds)
                kpss = kpss.reshape((kpss.shape[0], -1, 2))
                kpss_list.append(kpss[pos_inds])
        return scores_list, bboxes_list, kpss_list

    def nms(self, dets):
        """Standard IoU NMS on ``[x1, y1, x2, y2, score]`` rows (score-sorted)."""
        x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = scores.argsort()[::-1]
        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(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:]])
            w = np.maximum(0.0, xx2 - xx1 + 1)
            h = np.maximum(0.0, yy2 - yy1 + 1)
            inter = w * h
            ovr = inter / (areas[i] + areas[order[1:]] - inter)
            inds = np.where(ovr <= self.nms_thresh)[0]
            order = order[inds + 1]
        return keep

    def detect(self, img, thresh=0.5, input_size=None, max_num=0, metric="default"):
        """Detect faces in a BGR ``uint8`` image.

        Returns ``(dets, kpss)`` where ``dets`` is ``[N, 5]`` (x1,y1,x2,y2,score)
        in original-image pixels and ``kpss`` is ``[N, 5, 2]`` landmarks (or None).
        """
        input_size = self.input_size if input_size is None else input_size

        # Letterbox: preserve aspect ratio, pad to the model's square input.
        im_ratio = float(img.shape[0]) / img.shape[1]
        model_ratio = float(input_size[1]) / input_size[0]
        if im_ratio > model_ratio:
            new_height = input_size[1]
            new_width = int(new_height / im_ratio)
        else:
            new_width = input_size[0]
            new_height = int(new_width * im_ratio)
        det_scale = float(new_height) / img.shape[0]
        resized = cv2.resize(img, (new_width, new_height))
        det_img = np.zeros((input_size[1], input_size[0], 3), dtype=np.uint8)
        det_img[:new_height, :new_width, :] = resized

        scores_list, bboxes_list, kpss_list = self.forward(det_img, thresh)
        scores = np.vstack(scores_list)
        order = scores.ravel().argsort()[::-1]
        bboxes = np.vstack(bboxes_list) / det_scale
        pre_det = np.hstack((bboxes, scores)).astype(np.float32, copy=False)
        pre_det = pre_det[order, :]
        keep = self.nms(pre_det)
        det = pre_det[keep, :]

        kpss = None
        if self.use_kps:
            kpss = np.vstack(kpss_list) / det_scale
            kpss = kpss[order, :, :][keep, :, :]

        if max_num > 0 and det.shape[0] > max_num:
            area = (det[:, 2] - det[:, 0]) * (det[:, 3] - det[:, 1])
            img_center = img.shape[0] // 2, img.shape[1] // 2
            offsets = np.vstack([
                (det[:, 0] + det[:, 2]) / 2 - img_center[1],
                (det[:, 1] + det[:, 3]) / 2 - img_center[0],
            ])
            offset_dist_squared = np.sum(np.power(offsets, 2.0), 0)
            values = area if metric == "max" else area - offset_dist_squared * 2.0
            bindex = np.argsort(values)[::-1][:max_num]
            det = det[bindex, :]
            if kpss is not None:
                kpss = kpss[bindex, :, :]
        return det, kpss