"""ScoreVision crime detector v27 — YOLOv11s trained on 4-voter consensus
(new vmodel0, hermes_sv, pmodel_9, SAM3-refined; >=2 agree, IoU>=0.5).
Per-class confidence + per-class rescue (bonus) logic, flip TTA, cross-class NMS."""
from pathlib import Path
import math

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


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]
    keypoints: list[tuple[int, int]]


class Miner:
    """YOLOv11s + per-class threshold + per-class rescue + flip TTA + cross-class dedup."""

    # validator output order
    class_names = ["balaclava", "hoodie", "glove", "bat", "spray paint", "graffiti"]
    input_size = 1280
    iou_thres = 0.30
    cross_iou_thresh = 0.80
    max_aspect_ratio = 10.0
    max_det = 150
    _conf_thres_array = np.array(
        [0.30, 0.70, 0.70, 0.50, 0.50, 0.40], dtype=np.float32
    )
    _bonus_array = np.array(
        [0.20, 0.25, 0.20, 0.30, 0.20, 0.15], dtype=np.float32
    )

    _glove_cls_id = 2
    _glove_spurious_max_conf = 0.04
    _glove_spurious_aspect = 2.5
    _glove_crowd_min_count = 5
    _glove_crowd_max_conf_trigger = 0.6
    _glove_crowd_keep_min_conf = 0.27
    _glove_split_low_max = 0.26
    _glove_split_low_min_count = 4
    _glove_split_high_min_count = 2
    _glove_pair_low_max = 0.1
    _glove_pair_high_min = 0.4
    _glove_pair_max_center_x_dist = 0.028
    _glove_reject_r_min = 100
    _glove_reject_r_max = 125
    _glove_reject_g_min = 88
    _glove_reject_g_max = 115
    _glove_reject_b_min = 68
    _glove_reject_b_max = 88
    _glove_reject_min_pixel_frac = 0.04
    _glove_reject_p90_percentile = 90.0

    _spray_paint_cls_id = 4
    _spray_orphan_max_conf = 0.3
    _spray_glove_min_overlap = 0.3

    def __init__(self, path_hf_repo: Path) -> None:
        model_path = path_hf_repo / "weights.onnx"
        print("ORT version:", ort.__version__)
        try:
            ort.preload_dlls()
            print("preload_dlls success")
        except Exception as e:
            print(f"preload_dlls failed: {e}")
        print("ORT providers:", ort.get_available_providers())

        opts = ort.SessionOptions()
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        try:
            self.session = ort.InferenceSession(
                str(model_path), sess_options=opts,
                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
            )
            print("CUDA session created")
        except Exception as e:
            print(f"CUDA failed, CPU fallback: {e}")
            self.session = ort.InferenceSession(
                str(model_path), sess_options=opts, providers=["CPUExecutionProvider"]
            )
        print("session providers:", self.session.get_providers())

        for inp in self.session.get_inputs():
            print("INPUT:", inp.name, inp.shape, inp.type)
        for out in self.session.get_outputs():
            print("OUTPUT:", out.name, out.shape, out.type)

        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]
        sh = self.session.get_inputs()[0].shape
        self.input_height = self._safe_dim(sh[2], self.input_size)
        self.input_width = self._safe_dim(sh[3], self.input_size)
        print(f"ONNX loaded: {model_path}")
        print(
            "per-class conf: "
            + ", ".join(f"{n}={t:.2f}" for n, t in zip(self.class_names, self._conf_thres_array.tolist()))
        )
        print(
            "per-class rescue bonus: "
            + ", ".join(f"{n}={t:.2f}" for n, t in zip(self.class_names, self._bonus_array.tolist()))
        )

    @staticmethod
    def _safe_dim(v, d):
        return v if isinstance(v, int) and v > 0 else d

    def _letterbox(self, image, new_shape, color=(114, 114, 114)):
        h, w = image.shape[:2]
        nw, nh = new_shape
        r = min(nw / w, nh / h)
        rw, rh = int(round(w * r)), int(round(h * r))
        if (rw, rh) != (w, h):
            interp = cv2.INTER_CUBIC if r > 1.0 else cv2.INTER_LINEAR
            image = cv2.resize(image, (rw, rh), interpolation=interp)
        dw, dh = (nw - rw) / 2.0, (nh - rh) / 2.0
        top, bot = int(round(dh - 0.1)), int(round(dh + 0.1))
        lt, rt = int(round(dw - 0.1)), int(round(dw + 0.1))
        padded = cv2.copyMakeBorder(image, top, bot, lt, rt, cv2.BORDER_CONSTANT, value=color)
        return padded, r, (dw, dh)

    def _preprocess(self, image):
        h, w = image.shape[:2]
        img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))[None, ...]
        return np.ascontiguousarray(img, dtype=np.float32), ratio, pad, (w, h)

    @staticmethod
    def _clip(boxes, size):
        w, h = size
        boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
        boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
        boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
        boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
        return boxes

    @staticmethod
    def _hard_nms(boxes, scores, iou_thr):
        n = len(boxes)
        if n == 0: return np.array([], dtype=np.intp)
        order = np.argsort(-scores)
        keep = []
        while len(order) > 0:
            i = int(order[0]); keep.append(i)
            if len(order) == 1: break
            rest = order[1:]
            xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
            yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
            xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
            yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
            inter = np.maximum(0.0, xx2-xx1)*np.maximum(0.0, yy2-yy1)
            a_i = max(0.0, boxes[i, 2]-boxes[i, 0])*max(0.0, boxes[i, 3]-boxes[i, 1])
            a_r = (np.maximum(0.0, boxes[rest, 2]-boxes[rest, 0])
                   * np.maximum(0.0, boxes[rest, 3]-boxes[rest, 1]))
            iou = inter / (a_i+a_r-inter+1e-7)
            order = rest[iou <= iou_thr]
        return np.array(keep, dtype=np.intp)

    def _per_class_hard_nms(self, boxes, scores, cls_ids, iou_thr):
        if len(boxes) == 0: return np.array([], dtype=np.intp)
        keep_all = []
        for c in np.unique(cls_ids):
            mask = cls_ids == c
            idx = np.where(mask)[0]
            keep = self._hard_nms(boxes[mask], scores[mask], iou_thr)
            keep_all.extend(idx[keep].tolist())
        keep_all.sort()
        return np.array(keep_all, dtype=np.intp)

    def _cross_class_dedup(self, boxes, scores, cls_ids, iou_thr):
        n = len(boxes)
        if n == 0: return np.array([], dtype=np.intp)
        order = np.argsort(-scores)
        keep = []; suppressed = np.zeros(n, dtype=bool)
        for i in order:
            if suppressed[i]: continue
            keep.append(int(i))
            ix1 = np.maximum(boxes[i, 0], boxes[:, 0])
            iy1 = np.maximum(boxes[i, 1], boxes[:, 1])
            ix2 = np.minimum(boxes[i, 2], boxes[:, 2])
            iy2 = np.minimum(boxes[i, 3], boxes[:, 3])
            inter = np.maximum(0.0, ix2-ix1)*np.maximum(0.0, iy2-iy1)
            a_i = max(0.0, boxes[i, 2]-boxes[i, 0])*max(0.0, boxes[i, 3]-boxes[i, 1])
            a_r = (np.maximum(0.0, boxes[:, 2]-boxes[:, 0])
                   * np.maximum(0.0, boxes[:, 3]-boxes[:, 1]))
            iou = inter / (a_i+a_r-inter+1e-7)
            iou[i] = 0.0
            suppressed |= iou >= iou_thr
        return np.array(keep, dtype=np.intp)

    def _filter_sane(self, boxes, scores, cls_ids, orig_size):
        if len(boxes) == 0: return boxes, scores, cls_ids
        w, h = orig_size
        area_img = float(w * h)
        bw = np.maximum(0.0, boxes[:, 2]-boxes[:, 0])
        bh = np.maximum(0.0, boxes[:, 3]-boxes[:, 1])
        area = bw * bh
        ar = np.where((bw > 0) & (bh > 0),
                      np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6)),
                      np.inf)
        keep = (area <= 0.95 * area_img) & (ar <= self.max_aspect_ratio)
        return boxes[keep], scores[keep], cls_ids[keep]

    def _conf_filter_with_rescue(self, boxes, scores, cls_ids):
        """Per-class threshold + per-class rescue (admit top-1 if score >= thr - bonus
        and no box of that class passed the normal threshold)."""
        if len(scores) == 0:
            return np.zeros(0, dtype=bool)
        thr = self._conf_thres_array[cls_ids]
        keep = scores >= thr
        for c in np.unique(cls_ids):
            bonus = float(self._bonus_array[c])
            if bonus <= 0.0: continue
            cm = cls_ids == c
            if keep[cm].any(): continue
            idx = np.where(cm)[0]
            top = int(idx[int(np.argmax(scores[idx]))])
            if scores[top] >= self._conf_thres_array[c] - bonus:
                keep[top] = True
        return keep

    def _decode(self, preds, ratio, pad, orig_size):
        if preds.ndim == 3 and preds.shape[0] == 1:
            preds = preds[0]
        if preds.ndim != 2 or preds.shape[1] < 6:
            return (np.empty((0, 4), dtype=np.float32),
                    np.empty((0,), dtype=np.float32),
                    np.empty((0,), dtype=np.int32))
        boxes = preds[:, :4].astype(np.float32)
        scores = preds[:, 4].astype(np.float32)
        cls_ids = preds[:, 5].astype(np.int32)
        keep = (scores > 0) & (cls_ids >= 0) & (cls_ids < len(self.class_names))
        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
        if len(boxes) == 0: return boxes, scores, cls_ids
        # per-class threshold WITH rescue
        keep = self._conf_filter_with_rescue(boxes, scores, cls_ids)
        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
        if len(boxes) == 0: return boxes, scores, cls_ids
        pad_w, pad_h = pad
        boxes[:, [0, 2]] -= pad_w
        boxes[:, [1, 3]] -= pad_h
        boxes /= ratio
        boxes = self._clip(boxes, orig_size)
        return boxes, scores, cls_ids

    @classmethod
    def _is_spurious_glove(cls, cls_id, x1, y1, x2, y2, conf):
        if cls_id != cls._glove_cls_id: return False
        w, h = x2 - x1, y2 - y1
        if w <= 0: return False
        return h > cls._glove_spurious_aspect * w and conf < cls._glove_spurious_max_conf

    @staticmethod
    def _box_intersection_area(a, b):
        ix1, iy1 = max(a.x1, b.x1), max(a.y1, b.y1)
        ix2, iy2 = min(a.x2, b.x2), min(a.y2, b.y2)
        return max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)

    @classmethod
    def _spray_overlaps_glove(cls, spray, glove):
        inter = cls._box_intersection_area(spray, glove)
        area = max(0.0, float(spray.x2 - spray.x1) * float(spray.y2 - spray.y1))
        if area <= 0.0: return False
        return inter / area >= cls._spray_glove_min_overlap

    @classmethod
    def _is_glove_rejection_color(cls, cls_id, image, x1, y1, x2, y2):
        if cls_id != cls._glove_cls_id or image is None: return False
        H, W = image.shape[:2]
        xi1 = max(0, int(math.floor(x1)))
        yi1 = max(0, int(math.floor(y1)))
        xi2 = min(W, int(math.ceil(x2)))
        yi2 = min(H, int(math.ceil(y2)))
        if xi2 <= xi1 or yi2 <= yi1: return False
        crop = image[yi1:yi2, xi1:xi2]
        if crop.size == 0: return False
        b = crop[:, :, 0].astype(np.float32)
        g = crop[:, :, 1].astype(np.float32)
        r = crop[:, :, 2].astype(np.float32)
        in_range = ((r >= cls._glove_reject_r_min) & (r <= cls._glove_reject_r_max)
                    & (g >= cls._glove_reject_g_min) & (g <= cls._glove_reject_g_max)
                    & (b >= cls._glove_reject_b_min) & (b <= cls._glove_reject_b_max))
        if float(in_range.mean()) >= cls._glove_reject_min_pixel_frac:
            return True
        p90_r = float(np.percentile(r, cls._glove_reject_p90_percentile))
        p90_g = float(np.percentile(g, cls._glove_reject_p90_percentile))
        p90_b = float(np.percentile(b, cls._glove_reject_p90_percentile))
        return (cls._glove_reject_r_min <= p90_r <= cls._glove_reject_r_max
                and cls._glove_reject_g_min <= p90_g <= cls._glove_reject_g_max
                and cls._glove_reject_b_min <= p90_b <= cls._glove_reject_b_max)

    @classmethod
    def _filter_glove_by_color(cls, results, image):
        if image is None: return results
        return [b for b in results
                if not cls._is_glove_rejection_color(b.cls_id, image, b.x1, b.y1, b.x2, b.y2)]

    @staticmethod
    def _glove_cx(b): return (b.x1 + b.x2) / 2.0

    @classmethod
    def _filter_paired_low_glove(cls, results, image_width):
        if image_width <= 0: return results
        gloves = [b for b in results if b.cls_id == cls._glove_cls_id]
        if len(gloves) < 2: return results
        remove = set()
        for i, a in enumerate(gloves):
            for b in gloves[i+1:]:
                if a.conf < cls._glove_pair_low_max and b.conf > cls._glove_pair_high_min:
                    low, high = a, b
                elif b.conf < cls._glove_pair_low_max and a.conf > cls._glove_pair_high_min:
                    low, high = b, a
                else:
                    continue
                cx_dist = abs(cls._glove_cx(low) - cls._glove_cx(high)) / float(image_width)
                if cx_dist < cls._glove_pair_max_center_x_dist:
                    remove.add(id(low))
        return [b for b in results if id(b) not in remove]

    @classmethod
    def _filter_weak_glove_with_strong_others(cls, results):
        others = [b for b in results if b.cls_id != cls._glove_cls_id]
        if not others: return results
        if not all(b.conf > cls._glove_pair_high_min for b in others): return results
        return [b for b in results
                if b.cls_id != cls._glove_cls_id or b.conf >= cls._glove_pair_low_max]

    @classmethod
    def _filter_split_glove_confidence(cls, results):
        gloves = [b for b in results if b.cls_id == cls._glove_cls_id]
        low = [b for b in gloves if b.conf < cls._glove_split_low_max]
        high = [b for b in gloves if b.conf > cls._glove_split_low_max]
        if len(low) <= cls._glove_split_low_min_count: return results
        if len(high) < cls._glove_split_high_min_count: return results
        best_low = max(low, key=lambda b: b.conf)
        others = [b for b in results if b.cls_id != cls._glove_cls_id]
        return others + high + [best_low]

    @classmethod
    def _filter_crowded_gloves(cls, results):
        gloves = [b for b in results if b.cls_id == cls._glove_cls_id]
        if len(gloves) <= cls._glove_crowd_min_count: return results
        if max(b.conf for b in gloves) <= cls._glove_crowd_max_conf_trigger: return results
        return [b for b in results
                if b.cls_id != cls._glove_cls_id or b.conf > cls._glove_crowd_keep_min_conf]

    @classmethod
    def _filter_orphan_low_conf_spray(cls, results):
        gloves = [b for b in results if b.cls_id == cls._glove_cls_id]
        kept = []
        for box in results:
            if box.cls_id != cls._spray_paint_cls_id or box.conf >= cls._spray_orphan_max_conf:
                kept.append(box); continue
            if any(cls._spray_overlaps_glove(box, g) for g in gloves):
                kept.append(box)
        return kept

    @classmethod
    def _box_iou_pair(cls, a, b):
        inter = cls._box_intersection_area(a, b)
        aw = max(0.0, float(a.x2 - a.x1)) * max(0.0, float(a.y2 - a.y1))
        bw = max(0.0, float(b.x2 - b.x1)) * max(0.0, float(b.y2 - b.y1))
        union = aw + bw - inter
        if union <= 0: return 0.0
        return inter / union

    @classmethod
    def _aggressive_dedup_balaclava_graffiti(cls, results, iou_thr=0.10):
        """For balaclava (cls 0) and graffiti (cls 5), drop lower-conf box
        whenever any two boxes of the same class overlap > iou_thr.
        Rationale: one face = one balaclava; one tag = one graffiti box."""
        target = {0, 5}
        others = [b for b in results if b.cls_id not in target]
        by_cls = {}
        for b in results:
            if b.cls_id in target:
                by_cls.setdefault(b.cls_id, []).append(b)
        kept = list(others)
        for cid, boxes in by_cls.items():
            boxes_sorted = sorted(boxes, key=lambda b: -b.conf)
            survivors = []
            for b in boxes_sorted:
                if all(cls._box_iou_pair(b, k) < iou_thr for k in survivors):
                    survivors.append(b)
            kept.extend(survivors)
        return kept

    @classmethod
    def _apply_rules(cls, results, image_width, image=None):
        results = cls._filter_glove_by_color(results, image)
        results = cls._filter_paired_low_glove(results, image_width)
        results = cls._filter_weak_glove_with_strong_others(results)
        results = cls._filter_split_glove_confidence(results)
        results = cls._filter_crowded_gloves(results)
        results = cls._filter_orphan_low_conf_spray(results)
        return cls._aggressive_dedup_balaclava_graffiti(results)

    def _predict_single(self, image):
        x, ratio, pad, orig_size = self._preprocess(image)
        out = self.session.run(self.output_names, {self.input_name: x})[0]
        return self._decode(out, ratio, pad, orig_size)

    def _predict_tta(self, image):
        b0, s0, c0 = self._predict_single(image)
        flipped = cv2.flip(image, 1)
        bf, sf, cf = self._predict_single(flipped)
        if len(bf):
            w = image.shape[1]
            bf2 = bf.copy()
            bf2[:, 0] = w - bf[:, 2]
            bf2[:, 2] = w - bf[:, 0]
            bf = bf2
        boxes = np.concatenate([b0, bf], axis=0) if len(b0) or len(bf) else b0
        scores = np.concatenate([s0, sf], axis=0) if len(b0) or len(bf) else s0
        cls_ids = np.concatenate([c0, cf], axis=0) if len(b0) or len(bf) else c0
        if len(boxes) == 0: return []
        boxes, scores, cls_ids = self._filter_sane(boxes, scores, cls_ids,
                                                    (image.shape[1], image.shape[0]))
        if len(boxes) == 0: return []
        keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
        if len(keep) == 0: return []
        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
        keep = self._cross_class_dedup(boxes, scores, cls_ids, self.cross_iou_thresh)
        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
        if len(scores) > self.max_det:
            top = np.argsort(-scores)[:self.max_det]
            boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
        results = []
        for b, s, c in zip(boxes, scores, cls_ids):
            if b[2] <= b[0] or b[3] <= b[1]: continue
            x1, y1, x2, y2 = float(b[0]), float(b[1]), float(b[2]), float(b[3])
            if self._is_spurious_glove(int(c), x1, y1, x2, y2, float(s)): continue
            results.append(BoundingBox(
                x1=int(math.floor(x1)), y1=int(math.floor(y1)),
                x2=int(math.ceil(x2)), y2=int(math.ceil(y2)),
                cls_id=int(c), conf=float(s),
            ))
        return self._apply_rules(results, image.shape[1], image=image)

    def predict_batch(self, batch_images, offset, n_keypoints):
        results = []
        for j, img in enumerate(batch_images):
            try:
                boxes = self._predict_tta(img)
            except Exception as e:
                print(f"Inference failed for frame {offset + j}: {e}")
                boxes = []
            results.append(
                TVFrameResult(
                    frame_id=offset + j, boxes=boxes,
                    keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
                )
            )
        return results