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:
    def __init__(self,
        path_hf_repo: Path
    ) -> None:
        model_path = self._resolve_model_path(path_hf_repo)
        # car-wash element classes — cls_id order MUST match element `objects`
        # (0=broom, 1=drainage gate, 2=nozzle, 3=track). This is the canonical
        # order every downstream consumer (validator, BoundingBox.cls_id) sees.
        self.class_names = ["broom", "drainage gate", "nozzle", "track"]
        # FALLBACK model-emit order: the authoritative order is read from the
        # ONNX `names` metadata after the session is created (embedded by
        # Ultralytics at export, ships inside weights.onnx), so a retrained
        # model with a different class order is remapped correctly without
        # code changes. This list is used only when metadata is missing.
        self._model_class_order = ["broom", "drainage gate", "nozzle", "track"]
        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
        sess_options.intra_op_num_threads = 2
        sess_options.inter_op_num_threads = 1
        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL

        try:
            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 failed, falling back to CPU: {e}")
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CPUExecutionProvider"],
            )

        print("ORT session providers:", self.session.get_providers())

        # Build cls_remap: for each model-emit index i,
        #   cls_remap[i] = self.class_names.index(model_class_order[i])
        # The model-side order comes from the ONNX metadata when available,
        # else falls back to the static _model_class_order.
        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 = [output.name for output in self.session.get_outputs()]
        self.input_shape = self.session.get_inputs()[0].shape

        # Match the ONNX input dtype (this export is FP16 -> 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}")

        # ONNX is fixed-size 1408x1408 (v1 export); read actual shape to be safe.
        self.input_height = self._safe_dim(self.input_shape[2], default=1280)
        self.input_width = self._safe_dim(self.input_shape[3], default=1280)

        # Tuned for validator scoring (pillars: 0.6*map50 + 0.4*false_positive).
        # All values below are the measured optimum of a full grid sweep on
        # the validator-style val split (tune_miner.py, 241 1024x1024 crops,
        # composite 0.8002 -> 0.8103) -- re-run the sweep after any retrain.
        self.iou_thres = 0.5        # Per-class NMS IoU; lower = stricter dedup
        self.cross_iou_thresh = 0.9  # Cross-class dedup IoU (suppress same physical object firing multiple classes)
        self.max_det = 200
        # TTA = a 2nd (flipped) forward pass. Doubles latency; off for the
        # CPU latency gate. Re-enable only if the latency budget allows.
        self.use_tta = True

#       conf thresholds:  broom=0.38  drainage gate=0.45  nozzle=0.30  track=0.60
        # Per-class confidence thresholds.
        # Indexed by class_names order: [broom, drainage gate, nozzle, track].
        # broom/nozzle sit low: under the validator metric the mAP gained
        # from the extra recall outweighs the FP-pillar cost (the previous
        # 0.5/0.5 silently discarded many valid detections); track is the
        # one class where false fires are common enough to need 0.38.
        self._conf_thres_array = np.array(
            [0.28, 0.38, 0.60, 0.45], dtype=np.float32
        )
        # Per-class rescue bonus: when a class has ZERO boxes passing the
        # threshold in a frame, its top-1 candidate is admitted when its score
        # is at least (per-class threshold - per-class bonus).
        # DISABLED (all zeros): the sweep showed rescue admits more false
        # positives than true positives under the validator's FP pillar.
        self._bonus_array = np.array(
            [0.05, 0.1, 0.25, 0.2], dtype=np.float32
        )

        # Box sanity filter — kept loose: car-wash `nozzle` boxes are tiny
        # (GT median ~290 px², smallest ~32 px²). Fire's 14x14/min_side 8
        # would delete valid nozzles, so thresholds are dropped here.
        self.min_box_area = 4 * 4  # 16 px²
        self.min_side = 3
        self.max_aspect_ratio = 12.0

        print(f"✅ ONNX model loaded from: {model_path}")
        print(f"✅ ONNX providers: {self.session.get_providers()}")
        print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")

    def __repr__(self) -> str:
        return (
            f"ONNXRuntime(session={type(self.session).__name__}, "
            f"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:
        """Locate the ONNX model in the repo dir.

        Prefers weights.onnx (FP16/FP32 export), then weights_int8.onnx (the
        training script's INT8-quantized export -- works as-is: quantization
        preserves the Ultralytics metadata and QDQ models take regular fp32
        input), then any other .onnx file. INT8 is the fallback when the FP16
        export exceeds the 30 MB deployment limit (e.g. yolo26m).
        """
        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 the model's class order from Ultralytics ONNX metadata.

        Returns the class names ordered by model-emit index, or None when
        metadata is missing/unparsable or doesn't match `class_names` as a
        set (in which case the static _model_class_order fallback is used).
        """
        try:
            import ast

            meta = self.session.get_modelmeta().custom_metadata_map
            names = ast.literal_eval(meta["names"])  # e.g. {0: 'broom', ...}
            if isinstance(names, dict):
                order = [str(names[i]) for i in sorted(names)]
            else:
                order = [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} do not match expected classes "
                f"{self.class_names}; ignoring metadata"
            )
            return None
        return order

    def _letterbox(
        self,
        image: ndarray,
        new_shape: tuple[int, int],
        color=(114, 114, 114),
    ) -> tuple[ndarray, float, tuple[float, float]]:
        """
        Resize with unchanged aspect ratio and pad to target shape.
        Returns:
            padded_image,
            ratio,
            (pad_w, pad_h)  # half-padding
        """
        h, w = image.shape[:2]
        new_w, new_h = new_shape

        ratio = min(new_w / w, new_h / h)
        resized_w = int(round(w * ratio))
        resized_h = int(round(h * ratio))

        if (resized_w, resized_h) != (w, h):
            interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
            image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)

        dw = new_w - resized_w
        dh = new_h - resized_h
        dw /= 2.0
        dh /= 2.0

        left = int(round(dw - 0.1))
        right = int(round(dw + 0.1))
        top = int(round(dh - 0.1))
        bottom = int(round(dh + 0.1))

        padded = cv2.copyMakeBorder(
            image,
            top,
            bottom,
            left,
            right,
            borderType=cv2.BORDER_CONSTANT,
            value=color,
        )
        return padded, ratio, (dw, dh)

    def _preprocess(
        self, image: ndarray
    ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
        """
        Preprocess for fixed-size ONNX export:
        - enhance image quality (CLAHE, denoise, sharpen)
        - letterbox to model input size
        - BGR -> RGB
        - normalize to [0,1]
        - HWC -> NCHW float32
        """
        orig_h, orig_w = image.shape[:2]

        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, pad, (orig_w, orig_h)

    @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

    def _soft_nms(
        self,
        boxes: np.ndarray,
        scores: np.ndarray,
        sigma: float = 0.5,
        score_thresh: float = 0.01,
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
        Returns (kept_original_indices, updated_scores).
        """
        N = len(boxes)
        if N == 0:
            return np.array([], dtype=np.intp), np.array([], dtype=np.float32)

        boxes = boxes.astype(np.float32, copy=True)
        scores = scores.astype(np.float32, copy=True)
        order = np.arange(N)

        for i in range(N):
            max_pos = i + int(np.argmax(scores[i:]))
            boxes[[i, max_pos]] = boxes[[max_pos, i]]
            scores[[i, max_pos]] = scores[[max_pos, i]]
            order[[i, max_pos]] = order[[max_pos, i]]

            if i + 1 >= N:
                break

            xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
            yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
            xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
            yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)

            area_i = max(0.0, float(
                (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
            ))
            areas_j = (
                np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
                * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
            )
            iou = inter / (area_i + areas_j - inter + 1e-7)
            scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)

        mask = scores > score_thresh
        return order[mask], scores[mask]

    @staticmethod
    def _hard_nms(
        boxes: np.ndarray,
        scores: np.ndarray,
        iou_thresh: float,
    ) -> np.ndarray:
        """
        Standard NMS: keep one box per overlapping cluster (the one with highest score).
        Returns indices of kept boxes (into the boxes/scores arrays).
        """
        N = len(boxes)
        if N == 0:
            return np.array([], dtype=np.intp)
        boxes = np.asarray(boxes, dtype=np.float32)
        scores = np.asarray(scores, dtype=np.float32)
        order = np.argsort(scores)[::-1]
        keep: list[int] = []
        suppressed = np.zeros(N, dtype=bool)
        for i in range(N):
            idx = order[i]
            if suppressed[idx]:
                continue
            keep.append(idx)
            bi = boxes[idx]
            for k in range(i + 1, N):
                jdx = order[k]
                if suppressed[jdx]:
                    continue
                bj = boxes[jdx]
                xx1 = max(bi[0], bj[0])
                yy1 = max(bi[1], bj[1])
                xx2 = min(bi[2], bj[2])
                yy2 = min(bi[3], bj[3])
                inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
                area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
                area_j = (bj[2] - bj[0]) * (bj[3] - bj[1])
                iou = inter / (area_i + area_j - inter + 1e-7)
                if iou > iou_thresh:
                    suppressed[jdx] = True
        return np.array(keep)

    def _per_class_hard_nms(
        self,
        boxes: np.ndarray,
        scores: np.ndarray,
        cls_ids: np.ndarray,
        iou_thresh: float,
    ) -> np.ndarray:
        """Hard NMS applied independently per class."""
        if len(boxes) == 0:
            return np.array([], dtype=np.intp)
        all_keep: list[int] = []
        for c in np.unique(cls_ids):
            mask = cls_ids == c
            indices = np.where(mask)[0]
            keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
            all_keep.extend(indices[keep].tolist())
        all_keep.sort()
        return np.array(all_keep, dtype=np.intp)

    def _per_class_soft_nms(
        self,
        boxes: np.ndarray,
        scores: np.ndarray,
        cls_ids: np.ndarray,
        sigma: float = 0.5,
        score_thresh: float = 0.01,
    ) -> tuple[np.ndarray, np.ndarray]:
        """Soft NMS applied independently per class."""
        if len(boxes) == 0:
            return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
        all_keep: list[int] = []
        all_scores: list[float] = []
        for c in np.unique(cls_ids):
            mask = cls_ids == c
            indices = np.where(mask)[0]
            keep, updated = self._soft_nms(boxes[mask], scores[mask], sigma, score_thresh)
            for k, s in zip(keep, updated):
                all_keep.append(int(indices[k]))
                all_scores.append(float(s))
        if not all_keep:
            return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
        return np.array(all_keep, dtype=np.intp), np.array(all_scores, dtype=np.float32)

    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]:
        """Filter out tiny, degenerate, or implausible boxes (common FP)."""
        if len(boxes) == 0:
            return boxes, scores, cls_ids
        orig_w, orig_h = orig_size
        image_area = float(orig_w * orig_h)
        keep = []
        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = box.tolist()
            bw = x2 - x1
            bh = y2 - y1
            if bw <= 0 or bh <= 0:
                continue
            if bw < self.min_side or bh < self.min_side:
                continue
            area = bw * bh
            if area < self.min_box_area:
                continue
            if area > 0.95 * image_area:
                continue
            ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
            if ar > self.max_aspect_ratio:
                continue
            keep.append(i)
        if not keep:
            return (
                np.empty((0, 4), dtype=np.float32),
                np.empty((0,), dtype=np.float32),
                np.empty((0,), dtype=np.int32),
            )
        k = np.array(keep, dtype=np.intp)
        return boxes[k], scores[k], cls_ids[k]

    @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:
        """For each kept (post-NMS) box, return the max score over the FULL
        candidate set among SAME-CLASS boxes with IoU >= iou_thresh.

        The previous version omitted the same-class constraint, which let a
        confident broom raise the score of a coincident nozzle (or vice
        versa) under TTA. That's a silent FP booster and is fixed here.
        """
        n = len(post_boxes)
        if n == 0:
            return np.empty(0, dtype=np.float32)
        full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
                      np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
        out = np.empty(n, dtype=np.float32)
        for i in range(n):
            bi = post_boxes[i]
            xx1 = np.maximum(bi[0], full_boxes[:, 0])
            yy1 = np.maximum(bi[1], full_boxes[:, 1])
            xx2 = np.minimum(bi[2], full_boxes[:, 2])
            yy2 = np.minimum(bi[3], full_boxes[:, 3])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
            a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
            iou = inter / (a_i + full_areas - inter + 1e-7)
            cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
            out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
        return out

    def _conf_filter_mask(
        self, scores: np.ndarray, cls_ids: np.ndarray
    ) -> np.ndarray:
        """Boolean keep-mask: score >= per-class threshold, with a per-class
        rescue -- if a class has zero boxes passing, admit its top-1 candidate
        when its score >= (per-class threshold - per-class bonus).
        """
        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):
            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] >= self._conf_thres_array[c] - b:
                keep[top] = True
        return 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]:
        """Remove near-duplicate boxes across classes.

        Order candidates by (score - per_class_threshold) margin, then by area;
        keep the highest, suppress every other box with IoU > iou_thresh. For
        car-wash this kills the common failure where water spray makes the
        model fire both `nozzle` and `track` on the same patch, or where a
        broom handle overlaps a drainage-gate detection.
        """
        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,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Per-view post-processing: per-class NMS -> cap -> cross-class dedup."""
        if len(boxes) > 1:
            keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
            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 _decode_final_dets(
        self,
        preds: np.ndarray,
        ratio: float,
        pad: tuple[float, float],
        orig_size: tuple[int, int],
        apply_optional_dedup: bool = False,
    ) -> list[BoundingBox]:
        """
        Primary path:
        expected output rows like [x1, y1, x2, y2, conf, cls_id]
        in letterboxed input coordinates.
        """
        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 ONNX final-det output shape: {preds.shape}")

        boxes = preds[:, :4].astype(np.float32)
        scores = preds[:, 4].astype(np.float32)
        cls_ids = preds[:, 5].astype(np.int32)
        cls_ids = self.cls_remap[cls_ids]

        # Per-class confidence filter with rescue (replaces scalar threshold)
        keep = self._conf_filter_mask(scores, cls_ids)
        boxes = boxes[keep]
        scores = scores[keep]
        cls_ids = cls_ids[keep]

        if len(boxes) == 0:
            return []

        pad_w, pad_h = pad
        orig_w, orig_h = orig_size

        # reverse letterbox
        boxes[:, [0, 2]] -= pad_w
        boxes[:, [1, 3]] -= pad_h
        boxes /= ratio
        boxes = self._clip_boxes(boxes, (orig_w, orig_h))

        # Box sanity filter (reduces FP)
        boxes, scores, cls_ids = self._filter_sane_boxes(
            boxes, scores, cls_ids, orig_size
        )
        if len(boxes) == 0:
            return []

        if apply_optional_dedup and len(boxes) > 1:
            # Soft-NMS path preserved as a tunable option; default below.
            keep_idx, scores = self._per_class_soft_nms(boxes, scores, cls_ids)
            boxes = boxes[keep_idx]
            cls_ids = cls_ids[keep_idx]
            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
                )
        else:
            # Default: per-class hard NMS -> cap -> cross-class dedup
            boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)

        results: list[BoundingBox] = []
        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=int(math.floor(x1)),
                    y1=int(math.floor(y1)),
                    x2=int(math.ceil(x2)),
                    y2=int(math.ceil(y2)),
                    cls_id=int(cls_id),
                    conf=float(conf),
                )
            )

        return results

    def _decode_raw_yolo(
        self,
        preds: np.ndarray,
        ratio: float,
        pad: tuple[float, float],
        orig_size: tuple[int, int],
    ) -> list[BoundingBox]:
        """
        Fallback path for raw YOLO predictions.
        Supports common layouts:
        - [1, C, N]
        - [1, N, C]
        """
        if preds.ndim != 3:
            raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")

        if preds.shape[0] != 1:
            raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")

        preds = preds[0]

        # Normalize to [N, C]
        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 output 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]

        # Per-class confidence filter with rescue (replaces scalar threshold)
        keep = self._conf_filter_mask(scores, cls_ids)
        boxes_xywh = boxes_xywh[keep]
        scores = scores[keep]
        cls_ids = cls_ids[keep]
        if len(boxes_xywh) == 0:
            return []

        boxes = self._xywh_to_xyxy(boxes_xywh)

        # Order matches fire001 / _decode_final_dets:
        # unscale -> clip -> sanity filter -> per-view pipeline (NMS, cap, cross-class dedup).
        pad_w, pad_h = pad
        orig_w, orig_h = orig_size
        boxes[:, [0, 2]] -= pad_w
        boxes[:, [1, 3]] -= pad_h
        boxes /= ratio
        boxes = self._clip_boxes(boxes, (orig_w, orig_h))

        boxes, scores, cls_ids = self._filter_sane_boxes(
            boxes, scores, cls_ids, (orig_w, orig_h)
        )
        if len(boxes) == 0:
            return []

        boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)

        results: list[BoundingBox] = []
        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=int(math.floor(x1)),
                    y1=int(math.floor(y1)),
                    x2=int(math.ceil(x2)),
                    y2=int(math.ceil(y2)),
                    cls_id=int(cls_id),
                    conf=float(conf),
                )
            )

        return results

    def _postprocess(
        self,
        output: np.ndarray,
        ratio: float,
        pad: tuple[float, float],
        orig_size: tuple[int, int],
    ) -> list[BoundingBox]:
        """
        Prefer final detections first.
        Fallback to raw decode only if needed.
        """
        # final detections: [N,6]
        if output.ndim == 2 and output.shape[1] >= 6:
            return self._decode_final_dets(output, ratio, pad, orig_size)

        # final detections: [1,N,6]
        if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
            return self._decode_final_dets(output, ratio, pad, orig_size)

        # fallback raw decode
        return self._decode_raw_yolo(output, ratio, pad, orig_size)

    def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
        if image is None:
            raise ValueError("Input image is None")
        if not isinstance(image, np.ndarray):
            raise TypeError(f"Input is not numpy array: {type(image)}")
        if image.ndim != 3:
            raise ValueError(f"Expected HWC image, got shape={image.shape}")
        if image.shape[0] <= 0 or image.shape[1] <= 0:
            raise ValueError(f"Invalid image shape={image.shape}")
        if image.shape[2] != 3:
            raise ValueError(f"Expected 3 channels, got shape={image.shape}")

        if image.dtype != np.uint8:
            image = image.astype(np.uint8)

        input_tensor, ratio, pad, orig_size = self._preprocess(image)

        expected_shape = (1, 3, self.input_height, self.input_width)
        if input_tensor.shape != expected_shape:
            raise ValueError(
                f"Bad input tensor shape={input_tensor.shape}, expected={expected_shape}"
            )

        outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
        det_output = outputs[0]
        return self._postprocess(det_output, ratio, pad, orig_size)

    def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
        """Horizontal-flip TTA.

        Strategy (ported from fire001):
          1. Predict on original and on flipped image.
          2. Map flipped boxes back to original coordinates.
          3. Per-class hard NMS on the union.
          4. For each kept box, compute the max SAME-CLASS score across the
             FULL union -- a high-confidence flipped detection raises a
             borderline original one, but never one of a different class.
          5. Cross-class dedup to suppress same-physical-object multi-class.
        """
        boxes_orig = self._predict_single(image)

        flipped = cv2.flip(image, 1)
        boxes_flip = self._predict_single(flipped)

        w = image.shape[1]
        boxes_flip = [
            BoundingBox(
                x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
                cls_id=b.cls_id, conf=b.conf,
            )
            for b in boxes_flip
        ]

        all_boxes = boxes_orig + boxes_flip
        if len(all_boxes) == 0:
            return []

        coords = np.array(
            [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
        )
        scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
        cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)

        hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
        if len(hard_keep) == 0:
            return []
        if len(hard_keep) > self.max_det:
            top = np.argsort(-scores[hard_keep])[: self.max_det]
            hard_keep = hard_keep[top]

        # Class-aware cluster-max score boost (fixes the silent cross-class
        # leak in the previous _max_score_per_cluster).
        boosted = self._max_score_per_cluster(
            coords[hard_keep], cls_ids[hard_keep],
            coords, scores, cls_ids, self.iou_thres,
        )

        kept_coords = coords[hard_keep]
        kept_cls = cls_ids[hard_keep]
        if len(kept_coords) > 1:
            kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
                kept_coords, boosted, kept_cls, self.cross_iou_thresh
            )

        return [
            BoundingBox(
                x1=int(math.floor(kept_coords[j, 0])),
                y1=int(math.floor(kept_coords[j, 1])),
                x2=int(math.ceil(kept_coords[j, 2])),
                y2=int(math.ceil(kept_coords[j, 3])),
                cls_id=int(kept_cls[j]),
                conf=float(boosted[j]),
            )
            for j in range(len(kept_coords))
        ]

    def predict_batch(
        self,
        batch_images: list[ndarray],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        results: list[TVFrameResult] = []

        for frame_number_in_batch, image in enumerate(batch_images):
            try:
                if self.use_tta:
                    boxes = self._predict_tta(image)
                else:
                    boxes = self._predict_single(image)
            except Exception as e:
                print(f"⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}")
                boxes = []

            results.append(
                TVFrameResult(
                    frame_id=offset + frame_number_in_batch,
                    boxes=boxes,
                    keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
                )
            )

        return results