"""PPE detector miner (SN44 public track, manak0/Detect-worksite-ppe-compliance-detection-probe-set).

Ensemble: every YOLO graph in weights*.onnx (FP32 or INT8 QDQ; one merged file may hold several, input i -> output i) runs on the frame, candidates (conf >= 0.05,
per-model NMS 0.7) are fused with weighted box fusion, then per-class thresholds + NMS are applied. The CPU provider keeps inference consistent between the local
compliance check and the chute output (IoU 0.95).
No file/network access here (sandbox rules): thresholds below are written by scripts/ppe/package_ens.py.
"""
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort
from numpy import ndarray
from pydantic import BaseModel

MANIFEST_CLASSES = ["safety helmet", "high-visibility safety vest", "protective gloves", "protective eyewear"]
CONF_THRESHOLD = 0.3
PER_CLASS_CONF: dict[str, float] = {"safety helmet": 0.3, "high-visibility safety vest": 0.3, "protective gloves": 0.35, "protective eyewear": 0.275}
NMS_IOU_THRESHOLD = 0.45
EYE_ALONE_CONF = 0.5  # 0 turns the rule off
SCENE_EYE_CONF = 0.7  # eyewear in a scene with no helmet and no vest must clear this; empty exams cost 0.10 per box; 0 turns it off
MAX_DETECTIONS = 100
CAND_CONF, CAND_NMS, CAND_MAX, WBF_IOU = 0.05, 0.7, 300, 0.55


class BoundingBox(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


class TVFrameResult(BaseModel):
    frame_id: int
    boxes: list[BoundingBox] | None = None
    polygons: list | None = None
    keypoints: list[tuple[int, int]] | None = None


def _letterbox(img: ndarray, size: int):
    h, w = img.shape[:2]
    r = min(size / h, size / w)
    nw, nh = int(round(w * r)), int(round(h * r))
    px, py = (size - nw) // 2, (size - nh) // 2
    canvas = np.full((size, size, 3), 114, dtype=np.uint8)
    canvas[py:py + nh, px:px + nw] = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
    return canvas, r, px, py


def _nms(b: ndarray, s: ndarray, thr: float) -> list[int]:
    area = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1])
    order, keep = s.argsort()[::-1], []
    while order.size:
        i = int(order[0]); keep.append(i); rest = order[1:]
        iw = np.clip(np.minimum(b[i, 2], b[rest, 2]) - np.maximum(b[i, 0], b[rest, 0]), 0, None)
        ih = np.clip(np.minimum(b[i, 3], b[rest, 3]) - np.maximum(b[i, 1], b[rest, 1]), 0, None)
        inter = iw * ih
        order = rest[inter / np.maximum(area[i] + area[rest] - inter, 1e-9) <= thr]
    return keep


def _wbf(lists: list[list[list[float]]], thr: float) -> list[list[float]]:
    """lists: one list of [x1, y1, x2, y2, cls, conf] per model; conf = summed conf / number of models."""
    allb = sorted((b for lst in lists for b in lst), key=lambda b: -b[5])
    clusters: list[list] = []  # [fused box, members]
    for b in allb:
        best, bv = None, thr
        for cl in clusters:
            f = cl[0]
            if f[4] == b[4]:
                iw = max(0.0, min(f[2], b[2]) - max(f[0], b[0])); ih = max(0.0, min(f[3], b[3]) - max(f[1], b[1]))
                inter = iw * ih; u = (f[2] - f[0]) * (f[3] - f[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
                v = inter / u if u > 0 else 0.0
                if v >= bv:
                    best, bv = cl, v
        if best is None:
            clusters.append([list(b), [b]])
        else:
            best[1].append(b)
            ws = np.array([m[5] for m in best[1]]); co = np.array([m[:4] for m in best[1]], float)
            best[0][:4] = list((co * ws[:, None]).sum(0) / ws.sum())
    return [[*cl[0][:4], cl[0][4], min(1.0, sum(m[5] for m in cl[1]) / len(lists))] for cl in clusters]


def _alone_eyewear(e: ndarray, helmets: list[ndarray], vests: list[ndarray]) -> bool:
    """No helmet above this eyewear box and no vest below it: ordinary glasses and sunglasses on people without PPE are mostly not labelled."""
    x, y = (e[0] + e[2]) / 2, (e[1] + e[3]) / 2
    for b in helmets:
        w, h = b[2] - b[0], b[3] - b[1]
        if b[0] - .3 * w <= x <= b[2] + .3 * w and b[1] - 5 <= y <= b[3] + 1.2 * h:
            return False
    for b in vests:
        w, h = b[2] - b[0], b[3] - b[1]
        if b[0] - .25 * w <= x <= b[2] + .25 * w and y <= b[1] + .3 * h and b[1] - y <= 1.2 * w:
            return False
    return True


class Miner:
    def __init__(self, path_hf_repo: Path) -> None:
        so = ort.SessionOptions()
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        so.intra_op_num_threads = 2
        so.inter_op_num_threads = 1
        so.add_session_config_entry("session.set_denormal_as_zero", "1")
        so.add_session_config_entry("session.intra_op.allow_spinning", "0")  # faster on 2 vCPUs, also with one merged session (p95 77 vs 80 ms)
        self.sessions = [ort.InferenceSession(str(p), so, providers=["CPUExecutionProvider"])
                         for p in sorted(Path(path_hf_repo).glob("weights*.onnx"))]
        # one branch per (session, input i, output i): a merged file (scripts/ppe/merge_onnx.py) has N inputs/outputs
        self.inputs = [(k, i.name, int(i.shape[2]), s.get_outputs()[j].name)
                       for k, s in enumerate(self.sessions) for j, i in enumerate(s.get_inputs())]
        self.cls_conf = np.array([PER_CLASS_CONF.get(n, CONF_THRESHOLD) for n in MANIFEST_CLASSES], np.float32)
        for _ in range(2):
            self._infer(np.zeros((640, 640, 3), np.uint8))

    def __repr__(self) -> str:
        return f"PPE YOLO ONNX ensemble inputs={[b[2] for b in self.inputs]}"

    def _candidates(self, img: ndarray) -> list[list[list[float]]]:
        feeds: list[dict] = [{} for _ in self.sessions]
        geo = []
        for k, name, size, _ in self.inputs:
            canvas, r, px, py = _letterbox(img, size)
            feeds[k][name] = np.ascontiguousarray(canvas[:, :, ::-1].transpose(2, 0, 1)[None], dtype=np.float32) / 255.0
            geo.append((r, px, py))
        outs = {}
        for k, s in enumerate(self.sessions):
            names = [b[3] for b in self.inputs if b[0] == k]
            outs.update(zip(((k, n) for n in names), s.run(names, feeds[k])))
        return [self._decode(outs[(k, oname)][0], img, *geo[b]) for b, (k, _, _, oname) in enumerate(self.inputs)]

    def _decode(self, out: ndarray, img: ndarray, r: float, px: int, py: int) -> list[list[float]]:
        if out.shape[-1] == 6:  # end-to-end head: x1,y1,x2,y2,conf,cls
            boxes, scores, cls = out[:, :4], out[:, 4], out[:, 5].astype(np.int32)
        else:  # raw head: (4+nc, N) cxcywh + class scores
            out = out.T
            c = out[:, 4:]
            cls = c.argmax(1).astype(np.int32)
            scores = c[np.arange(len(c)), cls]
            xy, wh = out[:, :2], out[:, 2:4] / 2
            boxes = np.concatenate([xy - wh, xy + wh], 1)
        ok = (cls >= 0) & (cls < len(MANIFEST_CLASSES)) & (scores >= CAND_CONF)
        boxes, scores, cls = boxes[ok].astype(np.float32), scores[ok], cls[ok]
        boxes[:, [0, 2]] = ((boxes[:, [0, 2]] - px) / r).clip(0, img.shape[1])
        boxes[:, [1, 3]] = ((boxes[:, [1, 3]] - py) / r).clip(0, img.shape[0])
        res = []
        for c_ in np.unique(cls):
            sel = np.where(cls == c_)[0]
            res += [[*map(float, boxes[sel[j]]), int(c_), float(scores[sel[j]])]
                    for j in _nms(boxes[sel], scores[sel], CAND_NMS)]
        res.sort(key=lambda b: -b[5])
        return res[:CAND_MAX]

    def _infer(self, img: ndarray) -> list[BoundingBox]:
        cands = self._candidates(img)
        # one model: plain per-class NMS below (WBF would also merge the model's own overlapping boxes)
        fused = cands[0] if len(cands) == 1 else _wbf(cands, WBF_IOU)
        fused = [b for b in fused if b[5] >= self.cls_conf[b[4]]]
        if not fused:
            return []
        boxes = np.array([b[:4] for b in fused], np.float32); scores = np.array([b[5] for b in fused], np.float32)
        cls = np.array([b[4] for b in fused], np.int32)
        kept = []
        for k in np.unique(cls):
            sel = np.where(cls == k)[0]
            kept += [int(sel[j]) for j in _nms(boxes[sel], scores[sel], NMS_IOU_THRESHOLD)]
        helmets = [boxes[i] for i in kept if cls[i] == 0]; vests = [boxes[i] for i in kept if cls[i] == 1]
        res = []
        for i in kept:
            x1, y1, x2, y2 = boxes[i]
            if cls[i] == 3 and scores[i] < EYE_ALONE_CONF and _alone_eyewear(boxes[i], helmets, vests):
                continue
            if cls[i] == 3 and scores[i] < SCENE_EYE_CONF and not helmets and not vests:
                continue
            if x2 - x1 >= 1 and y2 - y1 >= 1:
                res.append(BoundingBox(x1=int(round(float(x1))), y1=int(round(float(y1))), x2=int(round(float(x2))),
                                       y2=int(round(float(y2))), cls_id=int(cls[i]), conf=float(min(1.0, scores[i]))))
        res.sort(key=lambda b: -b.conf)
        return res[:MAX_DETECTIONS]

    def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]:
        kp = [(0, 0)] * max(0, int(n_keypoints))
        out = []
        for i, img in enumerate(batch_images):
            try:
                boxes = self._infer(img)
            except Exception as exc:
                print(f"inference failed on frame {offset + i}: {exc}")
                boxes = []
            out.append(TVFrameResult(frame_id=offset + i, boxes=boxes, polygons=[], keypoints=list(kp)))
        return out