import os
from pathlib import Path

# These must be set before NumPy/Ultralytics load their native runtimes. Score's
# sandbox permits only 16 processes/threads, while the host libraries otherwise
# size pools from the underlying machine rather than the two-vCPU quota.
for thread_variable in (
    "OMP_NUM_THREADS",
    "MKL_NUM_THREADS",
    "OPENBLAS_NUM_THREADS",
    "NUMEXPR_NUM_THREADS",
    "VECLIB_MAXIMUM_THREADS",
    "BLIS_NUM_THREADS",
):
    os.environ[thread_variable] = "1"

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


class Miner:
    """SN44 PPE detector using a 320px SafetyVision v1 ONNX model."""

    # SafetyVision class ID -> active manifest object index:
    # [helmet, vest, gloves, eyewear]
    SOURCE_TO_MANIFEST = {1: 2, 2: 3, 3: 0, 12: 1}
    CONFIDENCE = 0.10
    NMS_IOU = 0.70
    IMAGE_SIZE = 320
    MAX_DETECTIONS = 300

    def __init__(self, path_hf_repo: Path) -> None:
        cv2.setNumThreads(1)
        options = ort.SessionOptions()
        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        options.intra_op_num_threads = 1
        options.inter_op_num_threads = 1
        options.add_session_config_entry("session.intra_op.allow_spinning", "0")
        self.session = ort.InferenceSession(
            str(path_hf_repo / "best.onnx"),
            sess_options=options,
            providers=["CPUExecutionProvider"],
        )
        self.input_name = self.session.get_inputs()[0].name

        # This graph was qualified on CPU. Keeping the same provider also
        # avoids provider-dependent output drift during the
        # validator's standardized two-vCPU compliance run.
        self.device = "cpu"

    def __repr__(self) -> str:
        return (
            "SafetyVision-v1-320 PPE ONNX detector "
            f"(confidence={self.CONFIDENCE}, device={self.device})"
        )

    @classmethod
    def _prepare_image(cls, image: ndarray) -> tuple[ndarray, float, tuple[int, int]]:
        height, width = image.shape[:2]
        scale = min(cls.IMAGE_SIZE / height, cls.IMAGE_SIZE / width)
        resized_width = round(width * scale)
        resized_height = round(height * scale)
        if (resized_width, resized_height) != (width, height):
            image = cv2.resize(
                image,
                (resized_width, resized_height),
                interpolation=cv2.INTER_LINEAR,
            )
        horizontal = (cls.IMAGE_SIZE - resized_width) / 2
        vertical = (cls.IMAGE_SIZE - resized_height) / 2
        left, right = round(horizontal - 0.1), round(horizontal + 0.1)
        top, bottom = round(vertical - 0.1), round(vertical + 0.1)
        image = cv2.copyMakeBorder(
            image,
            top,
            bottom,
            left,
            right,
            cv2.BORDER_CONSTANT,
            value=(114, 114, 114),
        )
        tensor = np.ascontiguousarray(image[:, :, ::-1].transpose(2, 0, 1))
        tensor = tensor.astype(np.float32) / 255.0
        return tensor[None], scale, (left, top)

    @staticmethod
    def _nms(boxes: ndarray, scores: ndarray, threshold: float) -> list[int]:
        if len(boxes) == 0:
            return []
        order = np.argsort(-scores)
        areas = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(
            0.0, boxes[:, 3] - boxes[:, 1]
        )
        keep: list[int] = []
        while order.size:
            current = int(order[0])
            keep.append(current)
            if order.size == 1:
                break
            remaining = order[1:]
            x1 = np.maximum(boxes[current, 0], boxes[remaining, 0])
            y1 = np.maximum(boxes[current, 1], boxes[remaining, 1])
            x2 = np.minimum(boxes[current, 2], boxes[remaining, 2])
            y2 = np.minimum(boxes[current, 3], boxes[remaining, 3])
            intersection = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
            union = areas[current] + areas[remaining] - intersection
            overlap = np.divide(
                intersection,
                union,
                out=np.zeros_like(intersection),
                where=union > 0,
            )
            order = remaining[overlap <= threshold]
        return keep

    def _detect(self, image: ndarray) -> list[BoundingBox]:
        height, width = image.shape[:2]
        tensor, scale, (pad_x, pad_y) = self._prepare_image(image)
        prediction = self.session.run(None, {self.input_name: tensor})[0][0].T
        source_boxes = prediction[:, :4]
        class_scores = prediction[:, 4:]
        source_classes = np.argmax(class_scores, axis=1)
        confidences = np.max(class_scores, axis=1)
        target_ids = np.fromiter(self.SOURCE_TO_MANIFEST, dtype=np.int64)
        selected = (confidences > self.CONFIDENCE) & np.isin(source_classes, target_ids)
        source_boxes = source_boxes[selected]
        source_classes = source_classes[selected]
        confidences = confidences[selected]
        if not len(source_boxes):
            return []

        boxes = np.empty_like(source_boxes)
        boxes[:, 0] = source_boxes[:, 0] - source_boxes[:, 2] / 2
        boxes[:, 1] = source_boxes[:, 1] - source_boxes[:, 3] / 2
        boxes[:, 2] = source_boxes[:, 0] + source_boxes[:, 2] / 2
        boxes[:, 3] = source_boxes[:, 1] + source_boxes[:, 3] / 2

        kept: list[int] = []
        for source_class in self.SOURCE_TO_MANIFEST:
            indices = np.flatnonzero(source_classes == source_class)
            local = self._nms(boxes[indices], confidences[indices], self.NMS_IOU)
            kept.extend(int(indices[index]) for index in local)
        kept.sort(key=lambda index: float(confidences[index]), reverse=True)

        output: list[BoundingBox] = []
        for index in kept[: self.MAX_DETECTIONS]:
            x1, y1, x2, y2 = boxes[index]
            x1 = np.clip((x1 - pad_x) / scale, 0, width)
            y1 = np.clip((y1 - pad_y) / scale, 0, height)
            x2 = np.clip((x2 - pad_x) / scale, 0, width)
            y2 = np.clip((y2 - pad_y) / scale, 0, height)
            output.append(
                BoundingBox(
                    x1=round(float(x1)),
                    y1=round(float(y1)),
                    x2=round(float(x2)),
                    y2=round(float(y2)),
                    cls_id=self.SOURCE_TO_MANIFEST[int(source_classes[index])],
                    conf=float(confidences[index]),
                )
            )
        return output

    def predict_batch(
        self,
        batch_images: list[ndarray],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        results: list[TVFrameResult] = []
        # The published ONNX graph has a fixed batch dimension of one. The
        # Chute contract may provide multiple frames, so execute them in order
        # rather than passing an invalid larger tensor to ONNX Runtime.
        for index, image in enumerate(batch_images):
            boxes = self._detect(image)
            results.append(
                TVFrameResult(
                    frame_id=offset + index,
                    boxes=boxes,
                    polygons=[],
                    keypoints=[(0, 0) for _ in range(n_keypoints)],
                )
            )
        return results