"""舌の検出（DeepTIAS 2.7 の `Detection_Normal.pb` を ONNX に変換したもの）と枠の扱い.

DeepTIAS 2.7 は TF1 の物体検出 API の frozen graph を使う．入力 `image_tensor`（`(1, H, W, 3)`
の uint8 画像），出力 `detection_boxes`（`(1, N, 4)`．正規化 `[ymin, xmin, ymax, xmax]`）と
`detection_scores`（`(1, N)`）．スコアが 0.5 を超える中で最高の 1 つを採用し，その枠で切り出す．
ONNX への変換は `scripts/convert_detection_to_onnx.py`（tf2onnx は出力名に `:0` を付けるため，
名前は `:0` の有無を問わずに探す）．
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort


@dataclass(frozen=True)
class Rect:
    """画像上の矩形（画素．`x1`・`y1` は含まない）.

    Attributes:
        x0: 左端．
        y0: 上端．
        x1: 右端 + 1．
        y1: 下端 + 1．
    """

    x0: int
    y0: int
    x1: int
    y1: int

    @property
    def width(self) -> int:
        """幅."""
        return self.x1 - self.x0

    @property
    def height(self) -> int:
        """高さ."""
        return self.y1 - self.y0


def full_rect(height: int, width: int) -> Rect:
    """画像全体の矩形を返す.

    Args:
        height: 画像の高さ．
        width: 画像の幅．

    Returns:
        画像全体の矩形．
    """
    return Rect(0, 0, width, height)


def mask_bounding_rect(mask: np.ndarray) -> Rect | None:
    """マスクの前景の外接矩形を返す.

    Args:
        mask: `(H, W)` の bool．

    Returns:
        外接矩形．前景が無ければ None．
    """
    ys, xs = np.nonzero(mask)
    if ys.size == 0:
        return None
    return Rect(int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1)


def normalized_box_to_rect(box: Sequence[float], height: int, width: int) -> Rect | None:
    """正規化 `[ymin, xmin, ymax, xmax]` の枠を画素の矩形にする（画像の内側に収める）.

    画素への変換は切り捨て（C# の `(int)` のキャストと同じ）．

    Args:
        box: 正規化座標の枠（値域 0〜1）．
        height: 画像の高さ．
        width: 画像の幅．

    Returns:
        矩形．幅または高さが 0 になる場合は None．
    """
    ymin, xmin, ymax, xmax = (float(v) for v in box)
    x0 = min(max(int(xmin * width), 0), width)
    y0 = min(max(int(ymin * height), 0), height)
    x1 = min(max(int(xmax * width), 0), width)
    y1 = min(max(int(ymax * height), 0), height)
    if x1 <= x0 or y1 <= y0:
        return None
    return Rect(x0, y0, x1, y1)


def select_best_box(
    boxes: np.ndarray, scores: np.ndarray, threshold: float
) -> tuple[np.ndarray, float] | None:
    """スコアが `threshold` を超える中で最高の枠を選ぶ.

    Args:
        boxes: `(N, 4)` の float（正規化 `[ymin, xmin, ymax, xmax]`）．
        scores: `(N,)` の float．
        threshold: スコアの下限（これを超えるもののみ）．

    Returns:
        （枠 `(4,)`, スコア）．該当が無ければ None．
    """
    scores = np.asarray(scores, dtype=np.float64).reshape(-1)
    boxes = np.asarray(boxes, dtype=np.float64).reshape(-1, 4)
    if scores.size == 0:
        return None
    best = int(np.argmax(scores))
    if not scores[best] > threshold:
        return None
    return boxes[best], float(scores[best])


def _resolve_name(available: Sequence[str], wanted: str, what: str) -> str:
    for name in (wanted, f"{wanted}:0"):
        if name in available:
            return name
    raise ValueError(f"検出モデルに{what} {wanted!r} が無い（あるもの: {', '.join(available)}）")


@dataclass(frozen=True)
class Detection:
    """検出の結果.

    Attributes:
        rect: 採用した枠（画素）．検出できなければ None．
        score: 最高のスコア（検出できなくても記録する．出力が空なら None）．
    """

    rect: Rect | None
    score: float | None


class OnnxDetector:
    """検出モデルの ONNX を onnxruntime（CPU）で推論する薄いラッパ.

    Args:
        path: 検出モデルの ONNX のパス．
        score_threshold: 採用するスコアの下限．
        input_color: モデルへ渡す画像の色の順（`rgb`・`bgr`）．
        input_name: 入力名（`:0` の有無は問わない）．
        boxes_name: 枠の出力名．
        scores_name: スコアの出力名．

    Raises:
        ValueError: ファイルが無い，または入出力名が見つからない場合．
    """

    def __init__(
        self,
        path: Path,
        *,
        score_threshold: float,
        input_color: str,
        input_name: str,
        boxes_name: str,
        scores_name: str,
    ) -> None:
        if not path.is_file():
            raise ValueError(f"検出モデルの ONNX が無い: {path}")
        self._session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
        inputs = [i.name for i in self._session.get_inputs()]
        outputs = [o.name for o in self._session.get_outputs()]
        self._input = _resolve_name(inputs, input_name, "入力")
        self._boxes = _resolve_name(outputs, boxes_name, "出力")
        self._scores = _resolve_name(outputs, scores_name, "出力")
        self._threshold = score_threshold
        self._rgb = input_color == "rgb"

    def detect(self, bgr: np.ndarray) -> Detection:
        """撮影画像から舌の枠を 1 つ検出する.

        Args:
            bgr: BGR 画像（`(H, W, 3)`, uint8．OpenCV で読んだまま）．

        Returns:
            検出の結果．
        """
        image = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) if self._rgb else bgr
        boxes, scores = self._session.run(
            [self._boxes, self._scores], {self._input: np.ascontiguousarray(image)[None]}
        )
        scores = np.asarray(scores).reshape(-1)
        top = float(scores.max()) if scores.size else None
        picked = select_best_box(boxes, scores, self._threshold)
        if picked is None:
            return Detection(rect=None, score=top)
        height, width = bgr.shape[:2]
        return Detection(rect=normalized_box_to_rect(picked[0], height, width), score=picked[1])