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

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

class Miner:

    def __init__(self, path_hf_repo: Path) -> None:
        cv2.setNumThreads(1)
        model_path = self._resolve_model_path(path_hf_repo)
        self.class_names = ['car', 'bus', 'truck', 'motorcycle']
        self._model_class_order = ['car', 'bus', 'truck', 'motorcycle']
        print('ORT version:', ort.__version__)
        print('ORT available providers:', ort.get_available_providers())
        self._model_path = str(model_path)
        self.session = self._open_session(self._model_path)
        print('ORT session providers:', self.session.get_providers())
        self.sink_name = 'tuktuk'
        self.sink_margin = 0.05
        self.sink_index = None
        self.vehicle_cols = None
        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}')
        if self.sink_name in model_class_order:
            self.sink_index = model_class_order.index(self.sink_name)
            self.vehicle_cols = np.array([model_class_order.index(n) for n in self.class_names], dtype=np.int32)
            print(f'cls order: drop {self.sink_name} col {self.sink_index} when score + {self.sink_margin} >= vehicle')
        self.cls_remap = np.full(len(model_class_order), -1, dtype=np.int32)
        for i, name in enumerate(model_class_order):
            if name in self.class_names:
                self.cls_remap[i] = self.class_names.index(name)
        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
        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}')
        self.input_height = self._safe_dim(self.input_shape[2], default=704)
        self.input_width = self._safe_dim(self.input_shape[3], default=704)
        self._iou_thres_array = np.array([0.37, 0.4, 0.38, 0.35], dtype=np.float32)
        self.cross_iou_thresh = 0.7
        self.max_det = 300
        self.pre_nms_topk = 1000
        self.use_tta = False
        self.use_cluster_boost = True
        self.use_soft_nms = False
        self.soft_nms_sigma = 0.5
        self._conf_thres_array = np.array([0.35, 0.2, 0.5, 0.45], dtype=np.float32)
        self._extra_conf_thres_array = np.array([0.25, 0.2, 0.2, 0.2], dtype=np.float32)
        self._bonus_array = np.array([0, 0, 0.15, 0], dtype=np.float32)
        self._min_box_area_array = np.array([64, 12, 12, 12], dtype=np.float32)
        self.min_side = 3
        self.max_aspect_ratio = 20
        self.max_area_frac = 0.95
        self.square_pad_frac = 0
        self.dual_label_pair = ('bus', 'truck')
        self.dual_label_below = 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()

    def _session_options(self) -> ort.SessionOptions:
        options = ort.SessionOptions()
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        options.intra_op_num_threads = 2
        options.inter_op_num_threads = 1
        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        return options

    def _open_session(self, model_path: str) -> ort.InferenceSession:
        # CPU is the default. onnxruntime-gpu session creation can abort the
        # process on a library mismatch, which leaves the chute unhealthy and
        # the scorer records a zero with no predictions.
        use_cuda = os.environ.get('MINER_USE_CUDA', '') == '1' and os.environ.get('MINER_FORCE_CPU', '') != '1'
        if use_cuda and 'CUDAExecutionProvider' in ort.get_available_providers():
            try:
                session = ort.InferenceSession(model_path, sess_options=self._session_options(), providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
                self._probe_session(session)
                print('ORT session: CUDA probe ok')
                return session
            except Exception as e:
                print(f'CUDA session unusable, using CPU: {e}')
        session = ort.InferenceSession(model_path, sess_options=self._session_options(), providers=['CPUExecutionProvider'])
        self._probe_session(session)
        print('ORT session: CPU probe ok')
        return session

    @staticmethod
    def _probe_session(session: ort.InferenceSession) -> None:
        inp = session.get_inputs()[0]
        dims = list(inp.shape)
        if len(dims) != 4:
            shape = [1, 3, 704, 704]
        else:
            shape = []
            for axis, dim in enumerate(dims):
                if isinstance(dim, int) and dim > 0:
                    shape.append(dim)
                elif axis == 1:
                    shape.append(3)
                elif axis == 0:
                    shape.append(1)
                else:
                    shape.append(704)
        dtype = np.float16 if 'float16' in inp.type else np.float32
        session.run(None, {inp.name: np.zeros(shape, dtype=dtype)})

    def _run(self, tensor: np.ndarray) -> list:
        try:
            return self.session.run(self.output_names, {self.input_name: tensor})
        except Exception as e:
            if self.session.get_providers() == ['CPUExecutionProvider']:
                raise
            print(f'inference on {self.session.get_providers()} failed, switching to CPU: {e}')
            self.session = self._open_session_cpu()
            return self.session.run(self.output_names, {self.input_name: tensor})

    def _open_session_cpu(self) -> ort.InferenceSession:
        session = ort.InferenceSession(self._model_path, sess_options=self._session_options(), providers=['CPUExecutionProvider'])
        self._probe_session(session)
        return session

    def _warmup(self, iters: int=1) -> None:
        dummy = np.zeros((1024, 1024, 3), dtype=np.uint8)
        for _ in range(max(1, iters)):
            self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
        print(f'warmup: {max(1, iters)} dummy predict_batch call(s) done')

    def __repr__(self) -> str:
        return f'PetrolMiner(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:
        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'

    def _read_model_class_order(self) -> list[str] | 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
        lower = [n.lower() for n in order]
        vehicles = [n for n in lower if n != self.sink_name]
        if sorted(vehicles) != sorted(self.class_names):
            print(f'cls order: ONNX names {order} != expected {self.class_names}; ignoring')
            return None
        return lower

    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)
        if padded.shape[0] != new_h or padded.shape[1] != new_w:
            padded = cv2.resize(padded, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        return (padded, ratio, (dw, dh))

    def _preprocess(self, image: ndarray) -> tuple[np.ndarray, dict]:
        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})

    @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:
        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-07)
            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]:
        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-07)
            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]:
        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:
        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:
        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-07)
        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:
        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]:
        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-06), bh / np.maximum(bw, 1e-06))
        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]:
        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-07, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
            iou = inter / (a_i + areas - inter + 1e-07)
            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 _apply_dual_label(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        if self.dual_label_below <= 0.0 or len(boxes) == 0:
            return (boxes, scores, cls_ids)
        a = self.class_names.index(self.dual_label_pair[0])
        b = self.class_names.index(self.dual_label_pair[1])
        sib = np.full(len(self.class_names), -1, dtype=np.int32)
        sib[a], sib[b] = (b, a)
        mask = (sib[cls_ids] >= 0) & (scores < self.dual_label_below)
        if not mask.any():
            return (boxes, scores, cls_ids)
        boxes = np.concatenate([boxes, boxes[mask]], axis=0)
        scores = np.concatenate([scores, scores[mask]], axis=0)
        cls_ids = np.concatenate([cls_ids, sib[cls_ids[mask]]], axis=0)
        return (boxes, scores, cls_ids)

    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]:
        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)

    def _unmap_and_finish(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, meta: dict) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        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))

    def _decode_final_dets(self, preds: np.ndarray, meta: dict) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        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)
        raw_cls = preds[:, 5].astype(np.int32)
        raw_cls = np.clip(raw_cls, 0, len(self.cls_remap) - 1)
        cls_ids = self.cls_remap[raw_cls]
        known = cls_ids >= 0
        boxes, scores, cls_ids = (boxes[known], scores[known], cls_ids[known])
        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]:
        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)
        elif self.vehicle_cols is not None and cls_part.shape[1] > int(self.vehicle_cols.max()):
            veh = cls_part[:, self.vehicle_cols]
            cls_ids = np.argmax(veh, axis=1).astype(np.int32)
            scores = veh[np.arange(len(veh)), cls_ids]
            if self.sink_index is not None and self.sink_index < cls_part.shape[1]:
                keep_veh = cls_part[:, self.sink_index] + self.sink_margin < scores
                boxes_xywh, scores, cls_ids = (boxes_xywh[keep_veh], scores[keep_veh], cls_ids[keep_veh])
        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[np.clip(cls_ids, 0, len(self.cls_remap) - 1)]
            keep_known = cls_ids >= 0
            boxes_xywh, scores, cls_ids = (boxes_xywh[keep_known], scores[keep_known], cls_ids[keep_known])
        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)

    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._run(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)
        boxes, scores, cls_ids = self._apply_dual_label(boxes, scores, cls_ids)
        return self._to_boxes(boxes, scores, cls_ids, meta['orig_size'])

    def _infer_tta(self, image: np.ndarray) -> list[BoundingBox]:
        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]]
        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)
        kb, boosted, kc = self._apply_dual_label(kb, boosted, kc)
        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):
            boxes = self._infer_tta(image) if self.use_tta else self._infer_single(image)
            keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
            results.append(TVFrameResult(frame_id=offset + idx, boxes=boxes, polygons=[], keypoints=keypoints))
        return results