"""miner.py — uploaded to nexu02/ScoreVision HF repo (R17 ONNX migration).

Migrated from .pt → ONNX FP16 to comply with subnet requirement
(.onnx-only models). Same R17 weights (mAP50 0.928, mAP50-95 0.764) +
identical inference recipe to keep the #1 dashboard standing.

Inference (same as R17 .pt version):
  - imgsz=1280, conf=0.50, iou=0.45
  - hflip TTA (manual: run twice, merge with per-class NMS)
  - cross-class NMS at IoU 0.6

Runtime: onnxruntime-gpu (CUDAExecutionProvider) with CPU fallback.
FP16 input/weights to fit under 30 MB HF cap (19.3 MB total).
"""
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_NAMES = ["cup", "bottle", "can"]


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


def _iou_xyxy(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Vectorised IoU between one box (a) and array of boxes (b)."""
    xx1 = np.maximum(a[0], b[:, 0])
    yy1 = np.maximum(a[1], b[:, 1])
    xx2 = np.minimum(a[2], b[:, 2])
    yy2 = np.minimum(a[3], b[:, 3])
    inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
    a_area = max(0.0, (a[2] - a[0]) * (a[3] - a[1]))
    b_area = np.maximum(0.0, (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1]))
    return inter / (a_area + b_area - inter + 1e-7)


def _hard_nms(boxes: np.ndarray, scores: np.ndarray, iou_thr: float) -> np.ndarray:
    """Per-class hard NMS — assumes boxes already filtered to one class."""
    n = len(boxes)
    if n == 0:
        return np.array([], dtype=np.intp)
    order = np.argsort(-scores)
    keep = []
    while len(order) > 0:
        i = int(order[0])
        keep.append(i)
        if len(order) == 1:
            break
        rest = order[1:]
        iou = _iou_xyxy(boxes[i], boxes[rest])
        order = rest[iou <= iou_thr]
    return np.array(keep, dtype=np.intp)


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


def _cross_class_nms(boxes, scores, cls_ids, iou_thr):
    """Cross-class NMS — drop overlapping boxes regardless of class."""
    if len(boxes) <= 1:
        return np.arange(len(boxes))
    order = np.argsort(-scores)
    keep = []
    suppressed = np.zeros(len(boxes), dtype=bool)
    for i in order:
        if suppressed[i]:
            continue
        keep.append(int(i))
        iou = _iou_xyxy(boxes[i], boxes)
        dup = iou > iou_thr
        dup[i] = False
        suppressed |= dup
    return np.array(sorted(keep), dtype=np.intp)


class Miner:
    """R17 ONNX miner. Same recipe as .pt version: 1280 + flip TTA + cross-class NMS."""

    INPUT_SIZE = 1280
    CONF_THR = 0.50
    IOU_THR = 0.45
    CROSS_CLASS_IOU = 0.6

    def __init__(self, path_hf_repo: Path) -> None:
        model_path = path_hf_repo / "best.onnx"
        if not model_path.exists():
            raise FileNotFoundError(f"missing weights at {model_path}")

        print(f"ORT version: {ort.__version__}")
        try:
            ort.preload_dlls()
        except Exception:
            pass

        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        try:
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
            )
        except Exception as e:
            print(f"CUDA session failed, fallback CPU: {e}")
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CPUExecutionProvider"],
            )
        print(f"ORT providers: {self.session.get_providers()}")
        for inp in self.session.get_inputs():
            print(f"INPUT  {inp.name} shape={inp.shape} dtype={inp.type}")
        for out in self.session.get_outputs():
            print(f"OUTPUT {out.name} shape={out.shape} dtype={out.type}")

        self.input_name = self.session.get_inputs()[0].name
        # FP16 model expects float16 inputs
        in_type = self.session.get_inputs()[0].type
        self.input_dtype = np.float16 if "float16" in in_type else np.float32
        print(f"✅ R17 ONNX loaded, input dtype={self.input_dtype.__name__}")

    def __repr__(self) -> str:
        return f"R17_ONNX(imgsz={self.INPUT_SIZE}, conf={self.CONF_THR}, iou={self.IOU_THR})"

    def _letterbox(self, img: np.ndarray, size: int):
        h, w = img.shape[:2]
        r = min(size / w, size / h)
        new_w, new_h = int(round(w * r)), int(round(h * r))
        if (new_w, new_h) != (w, h):
            interp = cv2.INTER_LINEAR
            img = cv2.resize(img, (new_w, new_h), interpolation=interp)
        dw, dh = (size - new_w) / 2.0, (size - new_h) / 2.0
        top = int(round(dh - 0.1)); bottom = int(round(dh + 0.1))
        left = int(round(dw - 0.1)); right = int(round(dw + 0.1))
        padded = cv2.copyMakeBorder(img, top, bottom, left, right,
                                    borderType=cv2.BORDER_CONSTANT, value=(114, 114, 114))
        return padded, r, (dw, dh)

    def _preprocess(self, img_bgr: np.ndarray):
        h, w = img_bgr.shape[:2]
        padded, r, pad = self._letterbox(img_bgr, self.INPUT_SIZE)
        rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB)
        x = rgb.astype(self.input_dtype) / 255.0
        x = np.transpose(x, (2, 0, 1))[None, ...]
        return np.ascontiguousarray(x, dtype=self.input_dtype), r, pad, (w, h)

    def _decode_raw(self, raw: np.ndarray, r: float, pad, orig_size):
        """Decode YOLO11 raw output (1, 7, N) → boxes + scores + class.
        Output shape: 4 box (xywh) + 3 class scores.
        """
        if raw.ndim == 3:
            raw = raw[0]
        if raw.shape[0] < raw.shape[1]:
            raw = raw.T  # → (N, 7)
        boxes_xywh = raw[:, :4].astype(np.float32)
        cls_scores = raw[:, 4:].astype(np.float32)
        cls_ids = np.argmax(cls_scores, axis=1)
        scores = cls_scores[np.arange(len(cls_scores)), cls_ids]

        keep = scores >= self.CONF_THR
        if not keep.any():
            return (np.empty((0, 4)), np.empty(0), np.empty(0, dtype=int))
        boxes_xywh, scores, cls_ids = boxes_xywh[keep], scores[keep], cls_ids[keep]

        # xywh → xyxy
        boxes = np.empty_like(boxes_xywh)
        boxes[:, 0] = boxes_xywh[:, 0] - boxes_xywh[:, 2] / 2
        boxes[:, 1] = boxes_xywh[:, 1] - boxes_xywh[:, 3] / 2
        boxes[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2] / 2
        boxes[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3] / 2

        # Undo letterbox padding/scale
        pad_w, pad_h = pad
        boxes[:, [0, 2]] -= pad_w
        boxes[:, [1, 3]] -= pad_h
        boxes /= r

        # Clip to original image
        w, h = orig_size
        boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, w - 1)
        boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, h - 1)

        return boxes, scores, cls_ids

    def _predict_single(self, img_bgr: np.ndarray):
        x, r, pad, orig = self._preprocess(img_bgr)
        out = self.session.run(None, {self.input_name: x})[0]
        return self._decode_raw(out, r, pad, orig)

    def _predict_with_tta(self, img_bgr: np.ndarray):
        """Predict + horizontal flip TTA, merge with per-class NMS."""
        boxes1, scores1, cls1 = self._predict_single(img_bgr)
        flipped = cv2.flip(img_bgr, 1)
        boxes2, scores2, cls2 = self._predict_single(flipped)
        if len(boxes2):
            w = img_bgr.shape[1]
            new = boxes2.copy()
            new[:, 0] = w - boxes2[:, 2]
            new[:, 2] = w - boxes2[:, 0]
            boxes2 = new
        if not len(boxes1) and not len(boxes2):
            return np.empty((0, 4)), np.empty(0), np.empty(0, dtype=int)
        boxes = np.concatenate([boxes1, boxes2]) if len(boxes1) and len(boxes2) else (boxes1 if len(boxes1) else boxes2)
        scores = np.concatenate([scores1, scores2]) if len(boxes1) and len(boxes2) else (scores1 if len(scores1) else scores2)
        cls_ids = np.concatenate([cls1, cls2]) if len(boxes1) and len(boxes2) else (cls1 if len(cls1) else cls2)
        keep = _per_class_nms(boxes, scores, cls_ids, self.IOU_THR)
        return boxes[keep], scores[keep], cls_ids[keep]

    def predict_batch(self, batch_images: list[ndarray], offset: int,
                       n_keypoints: int) -> list[TVFrameResult]:
        out: list[TVFrameResult] = []
        kp_zeros = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
        for i, image in enumerate(batch_images):
            frame_id = offset + i
            try:
                if image is None or image.ndim != 3 or image.shape[2] != 3:
                    out.append(TVFrameResult(frame_id=frame_id, boxes=[], keypoints=kp_zeros))
                    continue
                if image.dtype != np.uint8:
                    image = image.astype(np.uint8)

                boxes, scores, cls_ids = self._predict_with_tta(image)
                if len(boxes):
                    # Cross-class NMS (validator counts cross-class overlap as FP)
                    keep = _cross_class_nms(boxes, scores, cls_ids, self.CROSS_CLASS_IOU)
                    boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]

                results = []
                for b, s, c in zip(boxes, scores, cls_ids):
                    x1, y1, x2, y2 = b
                    if x2 <= x1 or y2 <= y1:
                        continue
                    c_int = int(c)
                    if c_int < 0 or c_int >= len(CLASS_NAMES):
                        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=c_int, conf=float(s),
                    ))
                out.append(TVFrameResult(frame_id=frame_id, boxes=results, keypoints=kp_zeros))
            except Exception as e:
                print(f"Inference err for frame {frame_id}: {e}")
                out.append(TVFrameResult(frame_id=frame_id, boxes=[], keypoints=kp_zeros))
        return out