"""SN44 TurboVision public-track miner for a raw YOLO head (the owner's element_trainer baselines: YOLO11n, detect or
seg, output [1, 4 + NC (+ mask coefficients), N]), decoded exactly as the owner's own baseline miner decodes it.

The validator hands predict_batch OpenCV BGR uint8 frames. Each frame is converted to RGB, resized (plain stretch, no
letterbox) to the model's input size, scaled to 0..1, and run on the CPU with 2 threads. Per anchor: the arg-max of the
NC class scores (mask coefficients ignored), kept when it reaches CONF; one class-agnostic greedy NMS at IOU over all
kept boxes; boxes mapped back by the stretch factors, floored / ceiled to whole pixels and clipped to the frame.
CLASS_MAP maps the model's classes to the element's `objects` indices (-1 drops a class; empty: the same order).
"""

import math
from pathlib import Path

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

MODEL_FILE = "model.onnx"
NC = 5
CONF = 0.25
IOU = 0.45
THREADS = 2
CLASS_MAP = ()
WARMUP_RUNS = 3


class BoundingBox(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


class Polygon(BaseModel):
    cls_id: int
    conf: float
    points: list[tuple[int, int]]


class TVFrameResult(BaseModel):
    frame_id: int
    boxes: list[BoundingBox] | None = None
    polygons: list[Polygon] | None = None
    keypoints: list[tuple[int, int]] | None = None


def _as_bgr(img) -> ndarray:
    if not isinstance(img, np.ndarray) or img.ndim not in (2, 3) or img.size == 0:
        raise ValueError(f"unsupported frame: {type(img).__name__} {getattr(img, 'shape', None)}")
    if img.dtype != np.uint8:
        img = np.clip(np.nan_to_num(img), 0, 255).astype(np.uint8)
    if img.ndim == 2 or img.shape[2] == 1:
        return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    if img.shape[2] == 4:
        return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
    return img


def _nms(boxes: ndarray, scores: ndarray, iou_thr: float) -> list[int]:
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(int(i))
        rest = order[1:]
        xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
        yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
        xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
        yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
        area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
        area_r = (boxes[rest, 2] - boxes[rest, 0]) * (boxes[rest, 3] - boxes[rest, 1])
        iou = inter / np.maximum(area_i + area_r - inter, 1e-6)
        order = rest[np.where(iou <= iou_thr)[0]]
    return keep


class Miner:
    def __init__(self, path_hf_repo: Path) -> None:
        so = ort.SessionOptions()
        so.intra_op_num_threads = THREADS
        so.inter_op_num_threads = 1
        so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        so.add_session_config_entry("session.intra_op.allow_spinning", "0")
        cv2.setNumThreads(1)
        self.session = ort.InferenceSession(str(Path(path_hf_repo) / MODEL_FILE), sess_options=so,
                                            providers=["CPUExecutionProvider"])
        inp = self.session.get_inputs()[0]
        self.input_name = inp.name
        self.in_h, self.in_w = int(inp.shape[2]), int(inp.shape[3])
        dummy = np.zeros((1024, 1024, 3), np.uint8)
        for _ in range(WARMUP_RUNS):
            self._frame_boxes(dummy)

    def __repr__(self) -> str:
        return f"{MODEL_FILE} @ {self.in_w}x{self.in_h} stretch, CPU {THREADS} threads, conf {CONF}, agnostic NMS {IOU}"

    def _frame_boxes(self, img) -> list[BoundingBox]:
        img = _as_bgr(img)
        h, w = img.shape[:2]
        rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        x = cv2.resize(rgb, (self.in_w, self.in_h)).astype(np.float32) / 255.0
        x = np.ascontiguousarray(np.transpose(x, (2, 0, 1))[None])
        pred = self.session.run(None, {self.input_name: x})[0][0]
        if pred.shape[0] < pred.shape[1]:
            pred = pred.T
        scores = pred[:, 4:4 + NC]
        cls = scores.argmax(1)
        conf = scores.max(1)
        m = np.isfinite(pred[:, :4]).all(1) & (conf >= CONF)
        b, cls, conf = pred[m, :4], cls[m], conf[m]
        if b.shape[0] == 0:
            return []
        sx, sy = w / float(self.in_w), h / float(self.in_h)
        xyxy = np.stack([(b[:, 0] - b[:, 2] / 2) * sx, (b[:, 1] - b[:, 3] / 2) * sy,
                         (b[:, 0] + b[:, 2] / 2) * sx, (b[:, 1] + b[:, 3] / 2) * sy], 1)
        boxes = []
        for i in _nms(xyxy, conf, IOU):
            c = int(cls[i])
            if CLASS_MAP:
                if c >= len(CLASS_MAP) or CLASS_MAP[c] < 0:
                    continue
                c = CLASS_MAP[c]
            x1, y1, x2, y2 = xyxy[i]
            ix1, iy1 = max(0, min(w, math.floor(x1))), max(0, min(h, math.floor(y1)))
            ix2, iy2 = max(0, min(w, math.ceil(x2))), max(0, min(h, math.ceil(y2)))
            if ix2 - ix1 < 1 or iy2 - iy1 < 1:
                continue
            boxes.append(BoundingBox(x1=ix1, y1=iy1, x2=ix2, y2=iy2, cls_id=c, conf=float(max(0.0, min(1.0, conf[i])))))
        return boxes

    def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]:
        results: list[TVFrameResult] = []
        for i, img in enumerate(batch_images):
            try:
                boxes = self._frame_boxes(img)
            except Exception as exc:  # one bad frame must not fail the whole challenge
                print(f"inference_failure frame {offset + i}: {type(exc).__name__}: {exc}")
                boxes = []
            results.append(TVFrameResult(frame_id=offset + i, boxes=boxes, polygons=[],
                                         keypoints=[(0, 0)] * int(n_keypoints) if n_keypoints else []))
        return results