import os
import math
from pathlib import Path

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


# =============================================================================
# Person miner (manak0/Detect-Person) — port of the proven carwash001 miner,
# same structure as VehicleDetection/miner/my_miner.
#
# Model: yolo26n end2end ONNX @704 (person_0904_704 / powerful_selected soup),
# single element-canonical class: 0=person.
#
# Carried over (all validated on the carwash/vehicle elements):
#   * robust model loading, FP16 auto-detect, metadata class order, ORT thread
#     pinning, raw-YOLO decode fallback, NO CLAHE (train/test match)
#   * vectorized hard NMS + cluster-max boost, per-class NMS IoU / min-area,
#     pre_nms_topk, cross-class dedup (a no-op guard with one class),
#     square-domain hook
#   * soft-NMS with post-decay re-filter (use_soft_nms), rescue bonus — both
#     SWEEPABLE via tune_miner.py; which setting wins is element-specific
#     (carwash: soft on / cluster on / rescue off; vehicle: the opposite).
#
# The vehicle miner's dual-label bus/truck lever is intentionally absent:
# person has no sibling class to hedge against.
#
# Run tune_miner.py after ANY retrain/re-export and paste its printout below.
# A failing conformity tuple (latency/format) removes the miner from scoring
# entirely — bench latency with bench_latency.py (2 pinned CPUs) too.
# =============================================================================


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:
    def __init__(self, path_hf_repo: Path) -> None:
        model_path = self._resolve_model_path(path_hf_repo)
        # Canonical class order every downstream consumer (validator,
        # BoundingBox.cls_id) sees — MUST equal element.objects order:
        self.class_names = ["person"]
        # Fallback model-emit order (used only when ONNX metadata is missing).
        self._model_class_order = ["person"]
        print("ORT version:", ort.__version__)

        try:
            ort.preload_dlls()
            print("onnxruntime.preload_dlls() success")
        except Exception as e:
            print(f"preload_dlls failed: {e}")

        print("ORT available providers BEFORE session:", ort.get_available_providers())

        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        # Pin threads for the CPU latency gate.
        sess_options.intra_op_num_threads = 2
        sess_options.inter_op_num_threads = 1
        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL

        # MINER_FORCE_CPU=1 skips the CUDA attempt — used by bench_latency.py to
        # measure the conformity-gate path (this box has onnxruntime-gpu too, so
        # an unforced "CPU" benchmark silently runs on the GPU).
        force_cpu = os.environ.get("MINER_FORCE_CPU", "") == "1"
        try:
            if force_cpu:
                raise RuntimeError("MINER_FORCE_CPU=1")
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
            )
            print("Created ORT session with preferred CUDA provider list")
        except Exception as e:
            print(f"CUDA session creation skipped/failed, using CPU: {e}")
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CPUExecutionProvider"],
            )
        print("ORT session providers:", self.session.get_providers())

        # cls_remap[i] = self.class_names.index(model_class_order[i]); order comes
        # from ONNX metadata when present, else the static fallback.
        model_class_order = self._read_model_class_order()
        if model_class_order is None:
            model_class_order = list(self._model_class_order)
            print(f"cls order: no usable ONNX metadata, FALLBACK {model_class_order}")
        else:
            print(f"cls order: from ONNX metadata {model_class_order}")
        self.cls_remap = np.array(
            [self.class_names.index(n) for n in model_class_order], dtype=np.int32
        )

        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()]
        self.input_shape = self.session.get_inputs()[0].shape

        # Match the ONNX input dtype (FP16 export needs float16 input).
        input_type = self.session.get_inputs()[0].type
        self.np_dtype = np.float16 if "float16" in input_type else np.float32
        print(f"ONNX input dtype: {input_type} -> numpy {self.np_dtype}")

        # The miner MUST run at the ONNX's baked input size (704). Dropping to
        # 640 is a latency lever but an EXPORT decision (re-export + re-eval),
        # not a miner-side change.
        self.input_height = self._safe_dim(self.input_shape[2], default=704)
        self.input_width = self._safe_dim(self.input_shape[3], default=704)

        # ── Post-processing config (tune on person1-powerful/val) ─────────────
        # Per-class NMS IoU ([person]).
        self._iou_thres_array = np.array([0.45], dtype=np.float32)
        self.cross_iou_thresh = 0.9   # near-duplicate guard (no-op across classes
                                      # with one class; catches post-NMS dupes)
        self.max_det = 300            # person frames can be busy
        self.pre_nms_topk = 1000      # cap candidates before NMS (crowded-frame speed guard)
        self.use_tta = False          # 2nd flipped pass; doubles latency -> off for CPU gate

        # Single-view cluster boost: raise each survivor's conf to its
        # IoU-cluster max. OFF per the 2026-09-04 person sweep (carwash: on,
        # vehicle: off -- element-specific, re-sweep on new weights).
        self.use_cluster_boost = False

        # Soft-NMS: Gaussian decay of overlapping scores instead of hard
        # deletion, then RE-FILTER decayed scores against the confidence
        # threshold (keeps decayed junk away from the FP pillar).
        # OFF per the 2026-09-04 person sweep (off 0.5670 vs on 0.5662) --
        # element- and weights-specific, re-sweep after every retrain.
        self.use_soft_nms = True
        self.soft_nms_sigma = 0.4

        # Confidence threshold ([person]). `_conf_thres_array` is the SQUARE /
        # validator-eval set (every scored frame is 1024x1024 squished, so this
        # is the one that matters); `_extra` covers non-square inputs (warmup /
        # any 16:9 frame). Tuned 2026-09-04 (tune_miner.py, 1853-img
        # person1-powerful/val, adaptive-IoU matcher): comp 0.5385 -> 0.5670.
        # NOTE: the val is dense (29 GT/frame avg) and recall is the binding
        # constraint. A low-edge extension sweep found a plateau below the
        # floor (0.05-0.09 -> comp 0.5713, +0.004) but on REAL challenge
        # frames the 0.05-0.12 band adds ~25 boxes/frame -- too FP-risky to
        # ship on dense-val evidence alone. Decide with test_miner.py + SAM3
        # GT on real frames; retune after every retrain (especially once the
        # sparse regenerated data lands).
        self._conf_thres_array = np.array([0.35], dtype=np.float32)
        self._extra_conf_thres_array = np.array([0.35], dtype=np.float32)

        # Rescue bonus: if NO box passes the threshold, admit the top-1 when
        # score >= (threshold - bonus). ON per the 2026-09-04 sweep.
        self._bonus_array = np.full(1, 0.03, dtype=np.float32)

        # Min box area (px^2), loose: distant people are tiny in 1024x1024
        # frames and tiny-GT matches only need IoU 0.30 under the adaptive
        # matcher — do not silently drop them.
        self._min_box_area_array = np.array([12.0], dtype=np.float32)
        self.min_side = 3
        self.max_aspect_ratio = 12.0
        self.max_area_frac = 0.95

        # Square-domain de-squish: widen a square (validator) frame by this
        # fraction before letterboxing, then drop pad-center boxes and un-pad.
        # 0.0 = off — the model was TRAINED on squished 1024 frames, so the
        # domains already match; enabling changes aspect -> validate first.
        self.square_pad_frac = 0.0

        self._avg_iou = float(np.mean(self._iou_thres_array))
        print(f"ONNX model loaded from: {model_path}")
        print(f"ONNX input: name={self.input_name}, shape={self.input_shape}")
        self._warmup()

    # ── Setup helpers ────────────────────────────────────────────────────────

    def _warmup(self, iters: int = 3) -> None:
        try:
            dummy = np.zeros((1024, 1024, 3), dtype=np.uint8)  # validator domain
            for _ in range(max(1, iters)):
                self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
            print(f"warmup: {iters} dummy predict_batch call(s) done")
        except Exception as e:
            print(f"warmup skipped: {e}")

    def __repr__(self) -> str:
        return f"PersonMiner(classes={len(self.class_names)}, providers={self.session.get_providers()})"

    @staticmethod
    def _safe_dim(value, default: int) -> int:
        return value if isinstance(value, int) and value > 0 else default

    @staticmethod
    def _resolve_model_path(repo: Path) -> Path:
        """Prefer weights.onnx, then weights_int8.onnx, then any .onnx."""
        for name in ("weights.onnx", "weights_int8.onnx"):
            p = repo / name
            if p.exists():
                if name != "weights.onnx":
                    print(f"model: weights.onnx not found, using {name}")
                return p
        candidates = sorted(repo.glob("*.onnx"))
        if candidates:
            print(f"model: using {candidates[0].name}")
            return candidates[0]
        return repo / "weights.onnx"  # let session creation raise the error

    def _read_model_class_order(self) -> list[str] | None:
        """Read class order from Ultralytics ONNX `names` metadata, or None."""
        try:
            import ast
            meta = self.session.get_modelmeta().custom_metadata_map
            names = ast.literal_eval(meta["names"])
            order = ([str(names[i]) for i in sorted(names)] if isinstance(names, dict)
                     else [str(n) for n in names])
        except Exception as e:
            print(f"cls order: could not read ONNX names metadata ({e})")
            return None
        if sorted(order) != sorted(self.class_names):
            print(f"cls order: ONNX names {order} != expected {self.class_names}; ignoring")
            return None
        return order

    # ── Preprocessing ────────────────────────────────────────────────────────

    def _letterbox(self, image: ndarray, new_shape: tuple[int, int],
                   color=(114, 114, 114)) -> tuple[ndarray, float, tuple[float, float]]:
        h, w = image.shape[:2]
        new_w, new_h = new_shape
        ratio = min(new_w / w, new_h / h)
        rw, rh = int(round(w * ratio)), int(round(h * ratio))
        if (rw, rh) != (w, h):
            interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
            image = cv2.resize(image, (rw, rh), interpolation=interp)
        dw = (new_w - rw) / 2.0
        dh = (new_h - rh) / 2.0
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        padded = cv2.copyMakeBorder(image, top, bottom, left, right,
                                    cv2.BORDER_CONSTANT, value=color)
        return padded, ratio, (dw, dh)

    def _preprocess(self, image: ndarray) -> tuple[np.ndarray, dict]:
        """Letterbox to the ONNX input size; NO CLAHE/denoise/sharpen (any
        enhancement the model wasn't trained on is a train/test mismatch that
        HURTS accuracy and also risks the CPU latency gate).

        Square-domain de-squish: when the frame is square (the validator's
        1024x1024 squished eval image) and square_pad_frac > 0, widen it with
        gray bars first -- objects come out horizontally compressed in those
        frames, and mild widening partly restores their aspect.
        """
        orig_h, orig_w = image.shape[:2]
        extra_left = extra_right = 0
        if self.square_pad_frac > 0.0 and orig_w == orig_h:
            target_w = int(orig_w * (1.0 + self.square_pad_frac))
            if target_w > orig_w:
                total = target_w - orig_w
                extra_left = total // 2
                extra_right = total - extra_left
                image = cv2.copyMakeBorder(image, 0, 0, extra_left, extra_right,
                                           cv2.BORDER_CONSTANT, value=(114, 114, 114))

        img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))[None, ...]
        img = np.ascontiguousarray(img, dtype=self.np_dtype)
        return img, {
            "ratio": ratio, "pad": pad, "orig_size": (orig_w, orig_h),
            "extra_left": extra_left, "extra_right": extra_right,
        }

    # ── Vectorized box ops ───────────────────────────────────────────────────

    @staticmethod
    def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
        w, h = image_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 _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
        out = np.empty_like(boxes)
        out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
        out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
        out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
        out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
        return out

    @staticmethod
    def _hard_nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float) -> np.ndarray:
        """Vectorized greedy NMS: areas precomputed once, each step is a single
        numpy IoU vector -- no inner Python loop."""
        n = len(boxes)
        if n == 0:
            return np.array([], dtype=np.intp)
        x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
        areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
        order = np.argsort(-scores)
        keep = []
        while order.size > 0:
            i = int(order[0])
            keep.append(i)
            if order.size == 1:
                break
            rest = order[1:]
            xx1 = np.maximum(x1[i], x1[rest])
            yy1 = np.maximum(y1[i], y1[rest])
            xx2 = np.minimum(x2[i], x2[rest])
            yy2 = np.minimum(y2[i], y2[rest])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
            iou = inter / (areas[i] + areas[rest] - inter + 1e-7)
            order = rest[iou <= iou_thresh]
        return np.array(keep, dtype=np.intp)

    @staticmethod
    def _soft_nms(boxes: np.ndarray, scores: np.ndarray, sigma: float
                  ) -> tuple[np.ndarray, np.ndarray]:
        """Gaussian soft-NMS: decay overlapping scores instead of deleting.
        Returns (original indices in decay order, decayed scores)."""
        n = len(boxes)
        if n == 0:
            return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
        b = boxes.astype(np.float32, copy=True)
        sc = scores.astype(np.float32, copy=True)
        order = np.arange(n)
        for i in range(n):
            mp = i + int(np.argmax(sc[i:]))
            b[[i, mp]] = b[[mp, i]]; sc[[i, mp]] = sc[[mp, i]]
            order[[i, mp]] = order[[mp, i]]
            if i + 1 >= n:
                break
            xx1 = np.maximum(b[i, 0], b[i+1:, 0]); yy1 = np.maximum(b[i, 1], b[i+1:, 1])
            xx2 = np.minimum(b[i, 2], b[i+1:, 2]); yy2 = np.minimum(b[i, 3], b[i+1:, 3])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
            ai = max(0.0, float((b[i, 2]-b[i, 0]) * (b[i, 3]-b[i, 1])))
            aj = np.maximum(0.0, b[i+1:, 2]-b[i+1:, 0]) * np.maximum(0.0, b[i+1:, 3]-b[i+1:, 1])
            iou = inter / (ai + aj - inter + 1e-7)
            sc[i+1:] *= np.exp(-(iou ** 2) / sigma)
        return order, sc

    def _per_class_soft_nms(self, boxes: np.ndarray, scores: np.ndarray,
                            cls_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Soft-NMS per class, then re-filter decayed scores against the
        per-class confidence thresholds so no decayed box reaches the FP pillar
        below its class bar."""
        keep_idx: list[int] = []
        keep_sc: list[float] = []
        for c in np.unique(cls_ids):
            idx = np.where(cls_ids == c)[0]
            order, dec = self._soft_nms(boxes[idx], scores[idx], self.soft_nms_sigma)
            thr = float(self._conf_thres_array[c])
            for k, s_dec in zip(order, dec):
                if s_dec >= thr:
                    keep_idx.append(int(idx[k])); keep_sc.append(float(s_dec))
        return np.array(keep_idx, dtype=np.intp), np.array(keep_sc, dtype=np.float32)

    def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray,
                            cls_ids: np.ndarray) -> np.ndarray:
        """Per-class NMS using the per-class IoU thresholds."""
        if len(boxes) == 0:
            return np.array([], dtype=np.intp)
        all_keep: list[int] = []
        for c in np.unique(cls_ids):
            idx = np.where(cls_ids == c)[0]
            cls_iou = float(self._iou_thres_array[c])
            keep = self._hard_nms(boxes[idx], scores[idx], cls_iou)
            all_keep.extend(idx[keep].tolist())
        all_keep.sort()
        return np.array(all_keep, dtype=np.intp)

    @staticmethod
    def _max_score_per_cluster(post_boxes: np.ndarray, post_cls: np.ndarray,
                               full_boxes: np.ndarray, full_scores: np.ndarray,
                               full_cls: np.ndarray, iou_thresh: float) -> np.ndarray:
        """Each survivor's confidence -> max score in its SAME-CLASS IoU cluster.
        Vectorized single (n_post x n_full) IoU matrix."""
        n = len(post_boxes)
        if n == 0:
            return np.empty(0, dtype=np.float32)
        m = len(full_boxes)
        if m == 0:
            return np.zeros(n, dtype=np.float32)
        pa = (np.maximum(0.0, post_boxes[:, 2] - post_boxes[:, 0]) *
              np.maximum(0.0, post_boxes[:, 3] - post_boxes[:, 1]))
        fa = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
              np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
        xx1 = np.maximum(post_boxes[:, 0][:, None], full_boxes[:, 0][None, :])
        yy1 = np.maximum(post_boxes[:, 1][:, None], full_boxes[:, 1][None, :])
        xx2 = np.minimum(post_boxes[:, 2][:, None], full_boxes[:, 2][None, :])
        yy2 = np.minimum(post_boxes[:, 3][:, None], full_boxes[:, 3][None, :])
        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
        iou = inter / (pa[:, None] + fa[None, :] - inter + 1e-7)
        mask = (iou >= iou_thresh) & (post_cls[:, None] == full_cls[None, :])
        tiled = np.where(mask, full_scores[None, :], -np.inf)
        out = tiled.max(axis=1)
        out[~np.isfinite(out)] = 0.0
        return out.astype(np.float32)

    def _conf_filter_mask(self, scores: np.ndarray, cls_ids: np.ndarray,
                          extra_left: int) -> np.ndarray:
        """Per-class threshold (square vs non-square set) + per-class rescue."""
        if len(scores) == 0:
            return np.zeros(0, dtype=bool)
        thr_arr = self._extra_conf_thres_array if extra_left > 0 else self._conf_thres_array
        keep = scores >= thr_arr[cls_ids]
        for c in np.unique(cls_ids):
            b = float(self._bonus_array[c])
            if b <= 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] >= float(self._conf_thres_array[c]) - b:
                keep[top] = True
        return keep

    def _filter_sane_boxes(self, boxes: np.ndarray, scores: np.ndarray,
                           cls_ids: np.ndarray, orig_size: tuple[int, int]
                           ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Vectorized per-class min-area / max-area / min-side / aspect filter."""
        if len(boxes) == 0:
            return boxes, scores, cls_ids
        orig_w, orig_h = orig_size
        image_area = float(orig_w * orig_h)
        bw = np.maximum(0.0, boxes[:, 2] - boxes[:, 0])
        bh = np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
        area = bw * bh
        ar = np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6))
        keep = (
            (bw >= self.min_side) & (bh >= self.min_side) &
            (area >= self._min_box_area_array[cls_ids]) &
            (area <= self.max_area_frac * image_area) &
            (ar <= self.max_aspect_ratio)
        )
        return boxes[keep], scores[keep], cls_ids[keep]

    def _cross_class_dedup_op(self, boxes: np.ndarray, scores: np.ndarray,
                              cls_ids: np.ndarray, iou_thresh: float
                              ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Suppress near-duplicate boxes at IoU>thresh regardless of class.
        With a single class this only catches rare post-NMS duplicates."""
        n = len(boxes)
        if n <= 1:
            return boxes, scores, cls_ids
        boxes = np.asarray(boxes, dtype=np.float32)
        scores = np.asarray(scores, dtype=np.float32)
        cls_ids = np.asarray(cls_ids, dtype=np.int32)
        areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
                 np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
        margins = scores - self._conf_thres_array[cls_ids]
        order = np.lexsort((-areas, -margins))
        suppressed = np.zeros(n, dtype=bool)
        keep: list[int] = []
        for i in order:
            if suppressed[i]:
                continue
            keep.append(int(i))
            bi = boxes[i]
            xx1 = np.maximum(bi[0], boxes[:, 0])
            yy1 = np.maximum(bi[1], boxes[:, 1])
            xx2 = np.minimum(bi[2], boxes[:, 2])
            yy2 = np.minimum(bi[3], boxes[:, 3])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
            a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
            iou = inter / (a_i + areas - inter + 1e-7)
            dup = iou > iou_thresh
            dup[i] = False
            suppressed |= dup
        keep_idx = np.array(keep, dtype=np.intp)
        return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]

    def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray,
                           cls_ids: np.ndarray, orig_size: tuple[int, int]
                           ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """sane filter -> top-k cap -> per-class NMS -> max_det cap -> dedup."""
        boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
        if len(boxes) == 0:
            return boxes, scores, cls_ids
        if len(scores) > self.pre_nms_topk:
            top = np.argpartition(-scores, self.pre_nms_topk)[: self.pre_nms_topk]
            boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
        if len(boxes) > 1:
            if self.use_soft_nms:
                keep, scores = self._per_class_soft_nms(boxes, scores, cls_ids)
                boxes, cls_ids = boxes[keep], cls_ids[keep]
            else:
                keep = self._per_class_hard_nms(boxes, scores, cls_ids)
                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]
        if len(boxes) > 1:
            boxes, scores, cls_ids = self._cross_class_dedup_op(
                boxes, scores, cls_ids, self.cross_iou_thresh)
        return boxes, scores, cls_ids

    # ── Coordinate un-mapping (shared by both decode paths) ──────────────────

    def _unmap_and_finish(self, boxes: np.ndarray, scores: np.ndarray,
                          cls_ids: np.ndarray, meta: dict
                          ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Undo letterbox + square-pad, drop pad-center boxes, clip, post-process."""
        pad_w, pad_h = meta["pad"]
        ratio = meta["ratio"]
        orig_w, orig_h = meta["orig_size"]
        extra_left, extra_right = meta["extra_left"], meta["extra_right"]

        boxes[:, [0, 2]] -= pad_w
        boxes[:, [1, 3]] -= pad_h
        boxes /= ratio
        if extra_left:
            boxes[:, [0, 2]] -= extra_left
        if extra_left or extra_right:
            cx = (boxes[:, 0] + boxes[:, 2]) * 0.5
            inside = (cx >= 0) & (cx <= orig_w)
            boxes, scores, cls_ids = boxes[inside], scores[inside], cls_ids[inside]
            if len(boxes) == 0:
                return boxes, scores, cls_ids
        boxes = self._clip_boxes(boxes, (orig_w, orig_h))
        return self._per_view_pipeline(boxes, scores, cls_ids, (orig_w, orig_h))

    # ── Decoding ─────────────────────────────────────────────────────────────

    def _decode_final_dets(self, preds: np.ndarray, meta: dict
                           ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """End2end output rows [x1, y1, x2, y2, conf, cls_id] (NMS in graph)."""
        empty = (np.empty((0, 4), np.float32), np.empty(0, np.float32), np.empty(0, np.int32))
        if preds.ndim == 3 and preds.shape[0] == 1:
            preds = preds[0]
        if preds.ndim != 2 or preds.shape[1] < 6:
            raise ValueError(f"Unexpected final-det output shape: {preds.shape}")
        boxes = preds[:, :4].astype(np.float32)
        scores = preds[:, 4].astype(np.float32)
        cls_ids = self.cls_remap[preds[:, 5].astype(np.int32)]
        keep = self._conf_filter_mask(scores, cls_ids, meta["extra_left"])
        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
        if len(boxes) == 0:
            return empty
        return self._unmap_and_finish(boxes, scores, cls_ids, meta)

    def _decode_raw_yolo(self, preds: np.ndarray, meta: dict
                         ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Raw YOLO output [1,C,N] or [1,N,C]; xywh + per-class scores."""
        empty = (np.empty((0, 4), np.float32), np.empty(0, np.float32), np.empty(0, np.int32))
        if preds.ndim != 3 or preds.shape[0] != 1:
            raise ValueError(f"Unexpected raw output shape: {preds.shape}")
        preds = preds[0]
        if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
            preds = preds.T
        if preds.ndim != 2 or preds.shape[1] < 5:
            raise ValueError(f"Unexpected normalized raw shape: {preds.shape}")
        boxes_xywh = preds[:, :4].astype(np.float32)
        cls_part = preds[:, 4:].astype(np.float32)
        if cls_part.shape[1] == 1:
            scores = cls_part[:, 0]
            cls_ids = np.zeros(len(scores), dtype=np.int32)
        else:
            cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
            scores = cls_part[np.arange(len(cls_part)), cls_ids]
        cls_ids = self.cls_remap[cls_ids]
        keep = self._conf_filter_mask(scores, cls_ids, meta["extra_left"])
        boxes_xywh, scores, cls_ids = boxes_xywh[keep], scores[keep], cls_ids[keep]
        if len(boxes_xywh) == 0:
            return empty
        boxes = self._xywh_to_xyxy(boxes_xywh)
        return self._unmap_and_finish(boxes, scores, cls_ids, meta)

    def _decode(self, output: np.ndarray, meta: dict
                ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        if output.ndim == 2 and output.shape[1] >= 6:
            return self._decode_final_dets(output, meta)
        if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
            return self._decode_final_dets(output, meta)
        return self._decode_raw_yolo(output, meta)

    # ── Inference ────────────────────────────────────────────────────────────

    def _predict_single_arrays(self, image: np.ndarray
                               ) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict]:
        if not isinstance(image, np.ndarray) or image.ndim != 3 or image.shape[2] != 3:
            raise ValueError(f"Expected HWC BGR image, got {type(image)} {getattr(image,'shape',None)}")
        if image.dtype != np.uint8:
            image = image.astype(np.uint8)
        inp, meta = self._preprocess(image)
        outputs = self.session.run(self.output_names, {self.input_name: inp})
        boxes, scores, cls_ids = self._decode(outputs[0], meta)
        return boxes, scores, cls_ids, meta

    @staticmethod
    def _to_boxes(boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray,
                  orig_size: tuple[int, int]) -> list[BoundingBox]:
        results: list[BoundingBox] = []
        orig_w, orig_h = orig_size
        for box, conf, cls_id in zip(boxes, scores, cls_ids):
            x1, y1, x2, y2 = box.tolist()
            if x2 <= x1 or y2 <= y1:
                continue
            results.append(BoundingBox(
                x1=max(0, min(orig_w, int(math.floor(x1)))),
                y1=max(0, min(orig_h, int(math.floor(y1)))),
                x2=max(0, min(orig_w, int(math.ceil(x2)))),
                y2=max(0, min(orig_h, int(math.ceil(y2)))),
                cls_id=int(cls_id),
                conf=float(max(0.0, min(1.0, conf))),
            ))
        return results

    def _infer_single(self, image: np.ndarray) -> list[BoundingBox]:
        boxes, scores, cls_ids, meta = self._predict_single_arrays(image)
        if len(boxes) == 0:
            return []
        if self.use_cluster_boost and len(boxes) > 1:
            scores = self._max_score_per_cluster(
                boxes, cls_ids, boxes, scores, cls_ids, self._avg_iou)
        return self._to_boxes(boxes, scores, cls_ids, meta["orig_size"])

    def _infer_tta(self, image: np.ndarray) -> list[BoundingBox]:
        """Horizontal-flip TTA: union of original + flipped, per-class NMS,
        same-class cluster-max boost, dedup."""
        b0, s0, c0, meta = self._predict_single_arrays(image)
        w = image.shape[1]
        bf, sf, cf, _ = self._predict_single_arrays(cv2.flip(image, 1))
        if len(bf):
            bf = bf.copy()
            bf[:, [0, 2]] = w - bf[:, [2, 0]]  # mirror x back to original coords
        coords = np.concatenate([b0, bf], axis=0) if len(bf) else b0
        scores = np.concatenate([s0, sf], axis=0) if len(sf) else s0
        cls_ids = np.concatenate([c0, cf], axis=0) if len(cf) else c0
        if len(coords) == 0:
            return []
        keep = self._per_class_hard_nms(coords, scores, cls_ids)
        if len(keep) == 0:
            return []
        if len(keep) > self.max_det:
            keep = keep[np.argsort(-scores[keep])[: self.max_det]]
        boosted = self._max_score_per_cluster(
            coords[keep], cls_ids[keep], coords, scores, cls_ids, self._avg_iou)
        kb, kc = coords[keep], cls_ids[keep]
        if len(kb) > 1:
            kb, boosted, kc = self._cross_class_dedup_op(kb, boosted, kc, self.cross_iou_thresh)
        return self._to_boxes(kb, boosted, kc, meta["orig_size"])

    def predict_batch(self, batch_images: list[ndarray], offset: int,
                      n_keypoints: int) -> list[TVFrameResult]:
        results: list[TVFrameResult] = []
        for idx, image in enumerate(batch_images):
            try:
                boxes = self._infer_tta(image) if self.use_tta else self._infer_single(image)
            except Exception as e:
                print(f"Inference failed for frame {offset + idx}: {e}")
                boxes = []
            keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
            results.append(TVFrameResult(frame_id=offset + idx, boxes=boxes, keypoints=keypoints))
        return results