"""
TurboVision miner for element `manak0/Detect-car-wash` — ONNX / CPU-safe.

Why ONNX-only: the latency-loop compliance checker (branch `latency-loop`) loads THIS
miner.py in a sandbox whose image has ONLY onnxruntime + cv2 + numpy + pydantic
(NO torch, NO ultralytics), forbids `.pt/.pth/.safetensors` files, forbids importing
socket/urllib/http/subprocess and calling open()/eval/exec, blocks the network during
inference, caps memory at 8 GiB, and requires the repo to contain a `.onnx` model.
It times `predict_batch` on CPU and needs p95 <= element.latency_p95_ms (target 100 ms),
and it re-checks that these outputs match your submitted responses at IoU >= 0.85.

So: pure onnxruntime, deterministic, small input size. Classes MUST be in manifest
order (cls_id == index): 0=broom 1=drainage gate 2=nozzle 3=track.
"""
from pathlib import Path
import os

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

CLASSES = ["broom", "drainage gate", "nozzle", "track"]

CONF = float(os.environ.get("CARWASH_CONF", "0.15"))  # global floor; per-class overrides below
# Per-class confidence floors (index == cls_id). Each object sits at its own map50/FP
# sweet spot. Override via CARWASH_CONF_PER_CLASS="0.30,0.45,0.20,0.35".
_pc = os.environ.get("CARWASH_CONF_PER_CLASS", "")
# Per-class conf floors — meta-informed prior copied from the current #1 (Alexei
# a36): broom/drain/track LOW to keep recall (map50 is 0.6 wt, dominates the FP
# cost), nozzle HIGH (0.45) because its low-conf boxes are the main FP source
# (challenge 69fa60042b). Re-tuned per-model at packaging via per_class_tune.
PER_CLASS_CONF = ([float(x) for x in _pc.split(",")] if _pc else [0.32, 0.17, 0.37, 0.41])
# Per-class rescue bonus: if a class has ZERO boxes after the floor, admit its
# top-1 candidate when conf >= (floor - bonus). Kept SMALL — the #1's sweep found
# aggressive rescue "admits more false positives" than it gains (matches the FP
# issue we saw). The global fallback below is our separate, safer empty-guard.
_bn = os.environ.get("CARWASH_BONUS_PER_CLASS", "")
BONUS_PER_CLASS = ([float(x) for x in _bn.split(",")] if _bn else [0.05, 0.10, 0.10, 0.10])
IOU_NMS = float(os.environ.get("CARWASH_IOU", "0.6"))
# Per-class same-class NMS IoU: suppress a box overlapping a kept same-class box by
# more than this. Nozzles cluster tightly and can double-fire on one spray head, so
# nozzle is STRICT (0.3). Also, any box whose CENTER sits inside a kept higher-conf
# same-class box is dropped (kills "same nozzle detected twice" regardless of IoU).
_pi = os.environ.get("CARWASH_IOU_PER_CLASS", "")
IOU_PER_CLASS = ([float(x) for x in _pi.split(",")] if _pi else [0.5, 0.6, 0.30, 0.6])
# Cross-class dedup IoU: suppress a physical object firing as >1 class (e.g. water
# spray fires both `nozzle` and `track`). 0 disables. This lifts the FP pillar (0.4 wt).
CROSS_IOU = float(os.environ.get("CARWASH_CROSS_IOU", "0.9"))
# Box sanity filter — loose because nozzle boxes are tiny (GT median ~290 px²).
MIN_SIDE = float(os.environ.get("CARWASH_MIN_SIDE", "3"))
MIN_AREA = float(os.environ.get("CARWASH_MIN_AREA", "16"))
MAX_AR = float(os.environ.get("CARWASH_MAX_AR", "12"))
# Per-class MIN area as a fraction of the frame. Kills a specific FP mode: a real
# car-wash `broom` is a large floor-to-ceiling rotating brush (>=1.3% of frame);
# a tiny distant "broom" box (<=0.6%) is almost always a misfire on a dark
# structure. Only broom is gated (nozzle/track/drain legitimately vary in size).
_ma = os.environ.get("CARWASH_MIN_AREA_FRAC", "")
MIN_AREA_FRAC = ([float(x) for x in _ma.split(",")] if _ma else [0.0, 0.0, 0.0, 0.0])
MAX_DET = int(os.environ.get("CARWASH_MAX_DET", "50"))
# Global fallback: when normal post-processing yields ZERO boxes, emit the single
# highest-probability raw detection (any class, ignoring the conf floor). On
# challenges where every miner returns empty, one plausible box can catch a missed
# object and score > 0 while others score 0. Set "0" to disable.
GLOBAL_FALLBACK = os.environ.get("CARWASH_GLOBAL_FALLBACK", "1") != "0"
MODEL_FILE = os.environ.get("CARWASH_MODEL", "carwash.onnx")


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 _letterbox(img: np.ndarray, new_shape: tuple[int, int]) -> tuple[np.ndarray, float, float, float]:
    """Resize+pad BGR image to new_shape (H,W), keep aspect. Return (img, ratio, pad_w, pad_h)."""
    h, w = img.shape[:2]
    nh, nw = new_shape
    r = min(nh / h, nw / w)
    uw, uh = int(round(w * r)), int(round(h * r))
    resized = cv2.resize(img, (uw, uh), interpolation=cv2.INTER_LINEAR)
    pad_w, pad_h = (nw - uw) / 2, (nh - uh) / 2
    top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
    left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
    out = cv2.copyMakeBorder(resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
    return out, r, left, top


def _nms(boxes: np.ndarray, scores: np.ndarray, iou_thr: float) -> list[int]:
    if len(boxes) == 0:
        return []
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(int(i))
        if order.size == 1:
            break
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
        iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-9)
        order = order[1:][iou <= iou_thr]
    return keep


def _same_class_nms(boxes: np.ndarray, scores: np.ndarray, iou_thr: float) -> list[int]:
    """NMS within one class, PLUS containment suppression: drop a box whose center
    lies inside a kept higher-conf box (or IoU > iou_thr). Kills duplicate detections
    of the same physical object (e.g. one nozzle boxed twice)."""
    n = len(boxes)
    if n == 0:
        return []
    order = scores.argsort()[::-1]
    cx = (boxes[:, 0] + boxes[:, 2]) / 2.0
    cy = (boxes[:, 1] + boxes[:, 3]) / 2.0
    areas = np.maximum(0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0, boxes[:, 3] - boxes[:, 1])
    keep = []
    for i in order:
        drop = False
        for j in keep:
            xx1 = max(boxes[i, 0], boxes[j, 0]); yy1 = max(boxes[i, 1], boxes[j, 1])
            xx2 = min(boxes[i, 2], boxes[j, 2]); yy2 = min(boxes[i, 3], boxes[j, 3])
            inter = max(0, xx2 - xx1) * max(0, yy2 - yy1)
            iou = inter / (areas[i] + areas[j] - inter + 1e-9)
            center_in = (boxes[j, 0] <= cx[i] <= boxes[j, 2]) and (boxes[j, 1] <= cy[i] <= boxes[j, 3])
            if iou > iou_thr or center_in:
                drop = True
                break
        if not drop:
            keep.append(int(i))
    return keep


def _sane_mask(xyxy: np.ndarray, img_area: float) -> np.ndarray:
    """Keep-mask dropping degenerate/implausible boxes (a common FP source)."""
    bw = xyxy[:, 2] - xyxy[:, 0]
    bh = xyxy[:, 3] - xyxy[:, 1]
    area = bw * bh
    ar = np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6))
    return ((bw >= MIN_SIDE) & (bh >= MIN_SIDE) & (area >= MIN_AREA)
            & (area <= 0.95 * img_area) & (ar <= MAX_AR))


def _cross_class_dedup(xyxy: np.ndarray, conf: np.ndarray, cls_id: np.ndarray,
                       margin: np.ndarray, iou_thr: float) -> list[int]:
    """Suppress near-duplicate boxes ACROSS classes: order by conf-margin then area,
    keep the best, drop any other-index box with IoU > iou_thr. Kills same-object
    multi-class fires (nozzle+track on one water patch)."""
    n = len(xyxy)
    if n <= 1:
        return list(range(n))
    areas = np.maximum(0, xyxy[:, 2] - xyxy[:, 0]) * np.maximum(0, xyxy[:, 3] - xyxy[:, 1])
    order = np.lexsort((-areas, -margin))
    suppressed = np.zeros(n, dtype=bool)
    keep = []
    for i in order:
        if suppressed[i]:
            continue
        keep.append(int(i))
        xx1 = np.maximum(xyxy[i, 0], xyxy[:, 0]); yy1 = np.maximum(xyxy[i, 1], xyxy[:, 1])
        xx2 = np.minimum(xyxy[i, 2], xyxy[:, 2]); yy2 = np.minimum(xyxy[i, 3], xyxy[:, 3])
        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
        iou = inter / (max(1e-7, areas[i]) + areas - inter + 1e-7)
        dup = iou > iou_thr
        dup[i] = False
        suppressed |= dup
    return keep


class Miner:
    def __init__(self, path_hf_repo: Path) -> None:
        model_path = str(Path(path_hf_repo) / MODEL_FILE)
        providers = os.environ.get("CARWASH_PROVIDERS", "CPUExecutionProvider").split(",")
        avail = ort.get_available_providers()
        providers = [p for p in providers if p in avail] or ["CPUExecutionProvider"]
        so = ort.SessionOptions()
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        so.intra_op_num_threads = int(os.environ.get("CARWASH_THREADS", "0"))  # 0 = ORT default
        self.sess = ort.InferenceSession(model_path, sess_options=so, providers=providers)
        self.inp = self.sess.get_inputs()[0]
        # match the model's input dtype (FP16 exports need float16 input)
        self.np_dtype = np.float16 if "float16" in (self.inp.type or "") else np.float32
        shape = self.inp.shape  # [1,3,H,W]; may contain strings if dynamic
        self.H = int(shape[2]) if isinstance(shape[2], int) else 640
        self.W = int(shape[3]) if isinstance(shape[3], int) else 640
        self.nc = len(CLASSES)
        # warmup so first real call isn't a cold-start outlier in p95
        dummy = np.zeros((1, 3, self.H, self.W), dtype=self.np_dtype)
        self.sess.run(None, {self.inp.name: dummy})
        print(f"✅ Car-wash ONNX loaded {MODEL_FILE} input={self.H}x{self.W} providers={providers} conf={CONF}")

    def __repr__(self) -> str:
        return f"CarWash ONNX ({MODEL_FILE}) {self.H}x{self.W} classes={CLASSES} conf={CONF}"

    def _preprocess(self, img_bgr: np.ndarray):
        lb, r, pad_w, pad_h = _letterbox(img_bgr, (self.H, self.W))
        rgb = lb[:, :, ::-1].astype(np.float32) / 255.0  # BGR->RGB, 0-1 (manifest norm rgb-01)
        chw = np.transpose(rgb, (2, 0, 1))
        return chw, r, pad_w, pad_h

    def _postprocess(self, out: np.ndarray, r: float, pad_w: float, pad_h: float,
                     orig_w: int, orig_h: int) -> list[BoundingBox]:
        # YOLOv8/11 detect ONNX head: (1, 4+nc, N) -> (N, 4+nc), xywh in input pixels
        pred = out[0]
        # Two supported ONNX heads (both in letterboxed input-pixel coords):
        #   end2end [N,6] = (x1,y1,x2,y2,conf,cls), already NMS'd (e.g. yolo11 nms=True)
        #   raw     [4+nc,N] or [N,4+nc] = xywh + per-class scores (yolo11 nms=False)
        if pred.ndim == 2 and pred.shape[1] == 6:
            pred = pred[pred[:, 4] > 1e-3]  # drop end2end padding rows
            boxes_in_all = pred[:, :4].astype(np.float32)
            conf_all = pred[:, 4].astype(np.float32)
            cls_id_all = pred[:, 5].astype(np.int32)
            _is_xywh = False
        else:
            if pred.shape[0] == (4 + self.nc):
                pred = pred.transpose(1, 0)
            boxes_in_all = pred[:, :4].astype(np.float32)
            cls_scores = pred[:, 4:4 + self.nc]
            cls_id_all = cls_scores.argmax(1).astype(np.int32)
            conf_all = cls_scores.max(1).astype(np.float32)
            _is_xywh = True

        def _to_xyxy(b: np.ndarray) -> np.ndarray:
            if _is_xywh:
                cx, cy, w, h = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
                xy = np.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], 1)
            else:
                xy = b.astype(np.float32, copy=True)
            xy[:, [0, 2]] = (xy[:, [0, 2]] - pad_w) / r
            xy[:, [1, 3]] = (xy[:, [1, 3]] - pad_h) / r
            xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, orig_w)
            xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, orig_h)
            return xy

        def _fallback() -> list[BoundingBox]:
            # emit the single highest-prob raw detection so we're never empty-handed
            if not GLOBAL_FALLBACK or len(conf_all) == 0:
                return []
            g = int(conf_all.argmax())
            xy = _to_xyxy(boxes_in_all[g:g + 1])[0]
            if xy[2] - xy[0] < 1 or xy[3] - xy[1] < 1:
                return []
            return [BoundingBox(x1=int(xy[0]), y1=int(xy[1]), x2=int(xy[2]), y2=int(xy[3]),
                                cls_id=int(cls_id_all[g]), conf=float(conf_all[g]))]

        # per-class confidence floor + rescue bonus
        floor = np.array(PER_CLASS_CONF, dtype=np.float32)
        bonus = np.array(BONUS_PER_CLASS, dtype=np.float32)
        m = conf_all >= floor[cls_id_all]
        for c in range(self.nc):
            if bonus[c] <= 0:
                continue
            cmask = cls_id_all == c
            if not cmask.any() or m[cmask].any():
                continue  # class absent, or already kept a box
            idx = np.where(cmask)[0]
            top = idx[int(conf_all[idx].argmax())]
            if conf_all[top] >= floor[c] - bonus[c]:
                m[top] = True  # rescue the top-1 candidate
        if not m.any():
            return _fallback()
        cls_id, conf = cls_id_all[m], conf_all[m]
        xyxy = _to_xyxy(boxes_in_all[m])
        # box sanity filter (drops degenerate/implausible FPs)
        img_area = float(orig_w * orig_h)
        sm = _sane_mask(xyxy, img_area)
        if not sm.any():
            return _fallback()
        xyxy, cls_id, conf = xyxy[sm], cls_id[sm], conf[sm]
        # per-class minimum-area gate (kills tiny-broom FPs on distant structures)
        mafrac = np.array(MIN_AREA_FRAC, dtype=np.float32)
        if mafrac.any():
            bw = xyxy[:, 2] - xyxy[:, 0]; bh = xyxy[:, 3] - xyxy[:, 1]
            am = (bw * bh) >= (mafrac[cls_id] * img_area)
            if not am.any():
                return _fallback()
            xyxy, cls_id, conf = xyxy[am], cls_id[am], conf[am]
        # cap candidates before the O(n^2) dedup so pathological frames stay fast
        if len(conf) > 150:
            top = np.argsort(-conf)[:150]
            xyxy, cls_id, conf = xyxy[top], cls_id[top], conf[top]
        # per-class NMS + containment dedup (nozzle strict) -> collect survivors
        keep_idx = []
        for c in np.unique(cls_id):
            idx = np.where(cls_id == c)[0]
            iou_c = IOU_PER_CLASS[c] if 0 <= c < len(IOU_PER_CLASS) else IOU_NMS
            for k in _same_class_nms(xyxy[idx], conf[idx], iou_c):
                keep_idx.append(int(idx[k]))
        keep_idx = np.array(keep_idx, dtype=np.intp)
        xyxy, cls_id, conf = xyxy[keep_idx], cls_id[keep_idx], conf[keep_idx]
        # cross-class dedup: suppress same physical object firing as multiple classes
        if CROSS_IOU > 0 and len(xyxy) > 1:
            margin = conf - floor[cls_id]
            cd = _cross_class_dedup(xyxy, conf, cls_id, margin, CROSS_IOU)
            xyxy, cls_id, conf = xyxy[cd], cls_id[cd], conf[cd]
        out_boxes = [BoundingBox(
            x1=int(xyxy[j, 0]), y1=int(xyxy[j, 1]), x2=int(xyxy[j, 2]), y2=int(xyxy[j, 3]),
            cls_id=int(cls_id[j]), conf=float(conf[j]),
        ) for j in range(len(xyxy))]
        if not out_boxes:
            return _fallback()
        out_boxes.sort(key=lambda b: b.conf, reverse=True)
        return out_boxes[:MAX_DET]

    def predict_batch(self, batch_images, offset: int, n_keypoints: int) -> list[TVFrameResult]:
        # Run one frame at a time: the exported ONNX has a fixed batch dim of 1, and
        # per-challenge latency is what the checker measures, so keep each run minimal.
        results: list[TVFrameResult] = []
        for i, img in enumerate(batch_images):
            chw, r, pw, ph = self._preprocess(img)
            inp = np.ascontiguousarray(chw[None], dtype=self.np_dtype)
            out = self.sess.run(None, {self.inp.name: inp})[0]  # (1, 4+nc, N)
            boxes = self._postprocess(out, r, pw, ph, img.shape[1], img.shape[0])
            results.append(TVFrameResult(frame_id=offset + i, boxes=boxes, polygons=[], keypoints=[]))
        return results