"""Petrol-Gauntlet v8 — the absolute solution. Single ONNX, max preprocessing,
multi-scale + h-flip TTA, WBF fusion (no aggressive filtering), class-aware NMS.

Architecture choice:
  Honest benchmark proved havengo's weights have the highest m50 (0.991) and
  fp (0.981) on the trusted eval. Every "smart filter" we tried (agreement
  gates, post-WBF thresholds) HURT the score by dropping real boxes. So v8
  layers only the techniques that ADD information (preprocessing, TTA) and
  removes any technique that could DROP information.

Pipeline per image:
  1. CLAHE on L-channel — rescues low-contrast / dusk frames
  2. 4 inference passes:
       a. original @ 1280
       b. h-flip   @ 1280
       c. original @ 1408
       d. h-flip   @ 1408
  3. WBF cluster across all 4 outputs — fused score is conf-weighted mean
     (NO agreement gate; singletons survive at their decayed score)
  4. Class-aware confidence threshold (pumps need lower threshold than canopy)
  5. Aspect-ratio sanity filter (drop 0px or full-image junk boxes)

Repo size: 22.7 MB (≤ 30 MB rule).
"""
from __future__ import annotations
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Tuple

import cv2
import numpy as np
from numpy import ndarray


@dataclass
class BoundingBox:
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


@dataclass
class FrameResult:
    frame_id: int
    boxes: List[BoundingBox] = field(default_factory=list)
    keypoints: List[Tuple[int, int]] = field(default_factory=list)


CLASS_NAMES = ["petrol hose", "petrol pump", "price board", "roof canopy"]

# Per-class confidence thresholds — pumps tend to dominate, give them a
# lower bar to surface; canopy false positives are common, raise it.
CONF_PER_CLASS = {0: 0.40, 1: 0.30, 2: 0.40, 3: 0.50}
# Default threshold during raw decode (lower than per-class, then refine)
DECODE_CONF = 0.25
WBF_IOU = 0.55
NMS_IOU = 0.50
MAX_DET = 300

# TTA OFF — single-pass to fit the harness latency budget (v8 timed out at 4
# passes/frame on the A4000). v8b retains the information-adding techniques
# (CLAHE, mean-color pad, class-aware conf, sanity filter) without TTA.
TTA_SCALES = [1280]
USE_HFLIP = False


class Miner:
    WEIGHTS_FILENAME = "detector.onnx"

    def __init__(self, project_root: Path):
        self.project_root = Path(project_root)
        self.class_names = CLASS_NAMES
        weights_path = self._find_weights()

        import onnxruntime as ort
        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        try:
            self.session = ort.InferenceSession(
                str(weights_path), sess_options=sess_options,
                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
            )
        except Exception:
            self.session = ort.InferenceSession(
                str(weights_path), sess_options=sess_options,
                providers=["CPUExecutionProvider"],
            )
        self.input_name = self.session.get_inputs()[0].name
        in_shape = self.session.get_inputs()[0].shape
        # ONNX input is fixed at training resolution (1280); we'll resize
        # the input image to each TTA scale via letterbox before feeding.
        self.native_input = in_shape[2] if isinstance(in_shape[2], int) else 1280

    def __repr__(self) -> str:
        return f"PetrolV8b(native={self.native_input}, TTA_scales={TTA_SCALES}, hflip={USE_HFLIP}, CLAHE=on)"

    def _find_weights(self) -> Path:
        for c in (
            self.project_root / "weights" / self.WEIGHTS_FILENAME,
            self.project_root / self.WEIGHTS_FILENAME,
        ):
            if c.is_file():
                return c
        for d in (self.project_root / "weights", self.project_root):
            if d.is_dir():
                for f in sorted(d.iterdir()):
                    if f.suffix == ".onnx" and f.is_file():
                        return f
        raise FileNotFoundError("detector.onnx not found")

    # ── preprocessing ────────────────────────────────────────────────
    @staticmethod
    def _clahe_lab(img_bgr: ndarray) -> ndarray:
        """Apply CLAHE on the L channel of LAB color space."""
        lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        l = clahe.apply(l)
        return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2BGR)

    @staticmethod
    def _mean_color(img: ndarray) -> Tuple[int, int, int]:
        m = img.reshape(-1, 3).mean(axis=0)
        return int(m[0]), int(m[1]), int(m[2])

    def _letterbox(self, img: ndarray, target_size: int) -> Tuple[ndarray, float, Tuple[float, float]]:
        h, w = img.shape[:2]
        ratio = min(target_size / w, target_size / 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
            img = cv2.resize(img, (rw, rh), interpolation=interp)
        dw = (target_size - rw) / 2.0
        dh = (target_size - rh) / 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))
        pad_color = self._mean_color(img)
        img = cv2.copyMakeBorder(img, top, bottom, left, right,
                                 cv2.BORDER_CONSTANT, value=pad_color)
        return img, ratio, (dw, dh)

    def _preprocess(self, image: ndarray, scale: int, flip: bool):
        orig_h, orig_w = image.shape[:2]
        img = cv2.flip(image, 1) if flip else image
        img, ratio, pad = self._letterbox(img, scale)
        # If model expects native_input (1280) and we letterboxed to 1408,
        # resize down to native_input for the ONNX call.
        if scale != self.native_input:
            img = cv2.resize(img, (self.native_input, self.native_input),
                             interpolation=cv2.INTER_CUBIC if scale < self.native_input
                             else cv2.INTER_AREA)
            ratio *= self.native_input / scale
            pad = (pad[0] * self.native_input / scale, pad[1] * self.native_input / scale)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))[None, ...]
        return np.ascontiguousarray(img, dtype=np.float32), ratio, pad, (orig_w, orig_h)

    # ── decoding ─────────────────────────────────────────────────────
    @staticmethod
    def _xywh_to_xyxy(boxes: ndarray) -> 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 _clip_boxes(boxes: ndarray, image_size: Tuple[int, int]) -> 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

    def _decode(self, preds: ndarray, ratio: float, pad: Tuple[float, float],
                orig_size: Tuple[int, int], flip: bool) -> Tuple[ndarray, ndarray, ndarray]:
        preds = preds[0]
        nc = len(self.class_names)
        if preds.shape[0] == 4 + nc:
            preds = preds.T
        boxes_xywh = preds[:, :4].astype(np.float32)
        cls_probs = preds[:, 4:].astype(np.float32)
        cls_ids = np.argmax(cls_probs, axis=1).astype(np.int32)
        scores = cls_probs[np.arange(len(cls_probs)), cls_ids]
        keep = scores >= DECODE_CONF
        if not keep.any():
            return np.empty((0, 4)), np.empty(0), np.empty(0, dtype=np.int32)
        boxes_xywh = boxes_xywh[keep]; scores = scores[keep]; cls_ids = cls_ids[keep]
        boxes = self._xywh_to_xyxy(boxes_xywh)
        # Undo letterbox
        pad_w, pad_h = pad
        boxes[:, [0, 2]] -= pad_w
        boxes[:, [1, 3]] -= pad_h
        boxes /= ratio
        # Undo flip if needed
        if flip:
            orig_w = orig_size[0]
            x1 = orig_w - 1 - boxes[:, 2]
            x2 = orig_w - 1 - boxes[:, 0]
            boxes[:, 0] = x1; boxes[:, 2] = x2
        boxes = self._clip_boxes(boxes, orig_size)
        return boxes, scores, cls_ids

    # ── WBF (no agreement gate) ──────────────────────────────────────
    @staticmethod
    def _wbf(all_boxes: ndarray, all_scores: ndarray, all_cls: ndarray,
             iou_thresh: float) -> Tuple[ndarray, ndarray, ndarray]:
        """Per-class WBF without agreement gate. Singletons keep their score."""
        if len(all_boxes) == 0:
            return np.empty((0, 4)), np.empty(0), np.empty(0, dtype=np.int32)
        fused_b, fused_s, fused_c = [], [], []
        for cls in np.unique(all_cls):
            mask = all_cls == cls
            B = all_boxes[mask]; S = all_scores[mask]
            order = np.argsort(S)[::-1]
            B = B[order]; S = S[order]
            used = np.zeros(len(B), dtype=bool)
            for i in range(len(B)):
                if used[i]: continue
                cluster = [i]; used[i] = True
                xi1, yi1, xi2, yi2 = B[i]
                area_i = max(0.0, xi2 - xi1) * max(0.0, yi2 - yi1)
                for j in range(i + 1, len(B)):
                    if used[j]: continue
                    xj1, yj1, xj2, yj2 = B[j]
                    xx1, yy1 = max(xi1, xj1), max(yi1, yj1)
                    xx2, yy2 = min(xi2, xj2), min(yi2, yj2)
                    inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
                    area_j = max(0.0, xj2 - xj1) * max(0.0, yj2 - yj1)
                    iou = inter / (area_i + area_j - inter + 1e-9)
                    if iou >= iou_thresh:
                        cluster.append(j); used[j] = True
                cb = B[cluster]; cs = S[cluster]
                w = cs / (cs.sum() + 1e-9)
                fb = (cb * w[:, None]).sum(axis=0)
                fs = float(cs.max())   # use MAX score (not mean) to avoid signal loss
                fused_b.append(fb); fused_s.append(fs); fused_c.append(int(cls))
        if not fused_b:
            return np.empty((0, 4)), np.empty(0), np.empty(0, dtype=np.int32)
        return np.array(fused_b), np.array(fused_s), np.array(fused_c, dtype=np.int32)

    # ── geometric sanity filter ──────────────────────────────────────
    @staticmethod
    def _sanity_filter(boxes: ndarray, scores: ndarray, cls_ids: ndarray,
                       img_size: Tuple[int, int]) -> Tuple[ndarray, ndarray, ndarray]:
        """Drop boxes that are 0 area, ≥95% of image, or extreme aspect ratios."""
        if len(boxes) == 0:
            return boxes, scores, cls_ids
        W, H = img_size
        w = boxes[:, 2] - boxes[:, 0]
        h = boxes[:, 3] - boxes[:, 1]
        area = w * h
        img_area = W * H
        # Keep: positive area, area < 95% of image, aspect ratio reasonable
        ar = np.maximum(w, h) / np.maximum(np.minimum(w, h), 1.0)
        keep = (w > 2) & (h > 2) & (area < 0.95 * img_area) & (ar < 25.0)
        return boxes[keep], scores[keep], cls_ids[keep]

    # ── inference ────────────────────────────────────────────────────
    def _predict_single(self, image: ndarray) -> List[BoundingBox]:
        if image.dtype != np.uint8:
            image = image.astype(np.uint8)
        # 1) CLAHE preprocessing
        image_clahe = self._clahe_lab(image)
        orig_size = (image.shape[1], image.shape[0])

        # 2) Inference passes (TTA off in v8b — single pass to fit budget)
        all_boxes = []
        all_scores = []
        all_cls = []
        flip_modes = (False, True) if USE_HFLIP else (False,)
        for scale in TTA_SCALES:
            for flip in flip_modes:
                blob, ratio, pad, _ = self._preprocess(image_clahe, scale, flip)
                out = self.session.run(None, {self.input_name: blob})[0]
                b, s, c = self._decode(out, ratio, pad, orig_size, flip)
                if len(b) > 0:
                    all_boxes.append(b); all_scores.append(s); all_cls.append(c)
        if not all_boxes:
            return []
        all_boxes = np.concatenate(all_boxes, axis=0)
        all_scores = np.concatenate(all_scores, axis=0)
        all_cls = np.concatenate(all_cls, axis=0)

        # 3) WBF fuse (no agreement gate — singletons survive)
        b, s, c = self._wbf(all_boxes, all_scores, all_cls, WBF_IOU)

        # 4) Class-aware conf threshold
        keep = np.ones(len(s), dtype=bool)
        for i, (sc, cls) in enumerate(zip(s, c)):
            keep[i] = sc >= CONF_PER_CLASS.get(int(cls), 0.40)
        b, s, c = b[keep], s[keep], c[keep]

        # 5) Geometric sanity filter
        b, s, c = self._sanity_filter(b, s, c, orig_size)

        # Cap
        if len(b) > MAX_DET:
            top = np.argsort(s)[::-1][:MAX_DET]
            b, s, c = b[top], s[top], c[top]

        return [
            BoundingBox(
                x1=int(math.floor(box[0])), y1=int(math.floor(box[1])),
                x2=int(math.ceil(box[2])),  y2=int(math.ceil(box[3])),
                cls_id=int(cid), conf=float(sc),
            )
            for box, sc, cid in zip(b, s, c)
            if box[2] > box[0] and box[3] > box[1]
        ]

    def predict_batch(self, batch_images: List[ndarray], offset: int,
                      n_keypoints: int) -> List[FrameResult]:
        results: List[FrameResult] = []
        for i, img in enumerate(batch_images):
            try:
                boxes = self._predict_single(img)
            except Exception:
                boxes = []
            results.append(FrameResult(frame_id=offset + i, boxes=boxes))
        return results


Detector = Miner
