"""ResUNet の前処理・後処理（DeepTIAS 2.7 の手順の再現）と onnxruntime の薄いラッパ.

DeepTIAS 2.7（C#）の逆コンパイルと ResUNet の学習コードで確認した手順:

1. 切り出した画像（OpenCV で読んだ **BGR のまま**）を `cv2.resize` の既定（線形補間）で
   256×256 にし，/255 して CHW に並べる（学習時も `cv2.imread` の BGR を /255 している）
2. ResUNet（入力 `(1, 3, 256, 256)`・出力 `(1, 2, 256, 256)`）で推論する
3. 出力の**チャンネル 0** にシグモイドを掛けて ×255 し uint8 にする
4. 大津の方法で 2 値化する
5. 輪郭（Canny → findContours）が 0 個または 2 個を超えるときだけ小領域除去（前景・背景とも
   面積 `min_area` 未満の領域を除く．元は C++ の DLL なので連結成分で再実装した）
6. 切り出した枠のサイズに戻す

前処理・後処理は純関数とし，推論だけを `ResUNetRunner` に閉じ込める．
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort

# Canny のしきい値．入力は 0/255 の 2 値画像なので，境界の勾配は常にこれを大きく超え，値に
# よって輪郭の数は変わらない（DeepTIAS の値は不明のため OpenCV の例でよく使う値にした）．
_CANNY_LOW = 100
_CANNY_HIGH = 200
# 連結成分の近傍（前景・背景とも 8 近傍）．
_CONNECTIVITY = 8


def preprocess_resunet(bgr: np.ndarray, input_size: int) -> np.ndarray:
    """切り出した画像を ResUNet の入力にする（BGR のまま /255・CHW）.

    Args:
        bgr: BGR 画像（`(H, W, 3)`, uint8．OpenCV で読んだまま）．
        input_size: 入力の一辺 S．

    Returns:
        `(1, 3, S, S)` の float32（値域 0.0〜1.0．チャンネルは B・G・R の順）．

    Raises:
        ValueError: 形状・dtype が不正，または画像が空の場合．
    """
    arr = np.asarray(bgr)
    if arr.ndim != 3 or arr.shape[-1] != 3 or arr.dtype != np.uint8:
        raise ValueError(
            f"画像は (H, W, 3) の uint8 でなければならない（{arr.shape}, {arr.dtype}）"
        )
    if arr.shape[0] == 0 or arr.shape[1] == 0:
        raise ValueError("画像が空である（幅または高さが 0）")
    resized = cv2.resize(arr, (input_size, input_size))  # 既定は線形補間（DeepTIAS と同じ）
    return (resized.astype(np.float32) / 255.0).transpose(2, 0, 1)[None]


def probability_map_u8(output: np.ndarray, channel: int) -> np.ndarray:
    """ResUNet の出力の 1 チャンネルにシグモイドを掛けて 0〜255 の uint8 にする.

    Args:
        output: ResUNet の出力（`(1, 2, S, S)`, float．ロジット）．
        channel: 使うチャンネル（DeepTIAS は 0）．

    Returns:
        `(S, S)` の uint8（シグモイド × 255 の小数部を切り捨て）．
    """
    logits = np.asarray(output, dtype=np.float64)[0, channel]
    # ロジットが約 -709 未満だと exp が inf に桁あふれするが，1 / (1 + inf) = 0 で値は正しい．
    # DeepTIAS と同じ式のまま（切り捨ての結果を変えないため）警告だけを抑える．
    with np.errstate(over="ignore"):
        prob = 1.0 / (1.0 + np.exp(-logits))
    return (prob * 255.0).astype(np.uint8)


def binarize_otsu(prob_u8: np.ndarray) -> np.ndarray:
    """大津の方法で 2 値化する.

    Args:
        prob_u8: `(S, S)` の uint8．

    Returns:
        `(S, S)` の uint8（前景 255・背景 0）．
    """
    _, binary = cv2.threshold(prob_u8, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    return binary


def count_contours(binary: np.ndarray) -> int:
    """2 値画像の輪郭の数を数える（Canny でエッジを取ってから外側の輪郭を数える）.

    DeepTIAS の方式の再現であり，連結成分の数とは一致しない．Canny のエッジにならない数画素の
    小さな島は数えないことがあり，外側の輪郭だけを数えるため領域の内側の穴（とその中の島）も
    数えない．全面が前景・全面が背景ならエッジが無く 0 になる．

    Args:
        binary: `(H, W)` の uint8（0/255）．

    Returns:
        輪郭の数．
    """
    edges = cv2.Canny(binary, _CANNY_LOW, _CANNY_HIGH)
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    return len(contours)


def needs_small_region_removal(n_contours: int) -> bool:
    """輪郭の数から小領域除去が必要かを返す（0 個または 2 個を超えるとき．DeepTIAS と同じ条件）.

    Args:
        n_contours: `count_contours` の戻り値．

    Returns:
        小領域除去をするなら True．
    """
    return n_contours == 0 or n_contours > 2


def remove_small_regions(binary: np.ndarray, min_area: int) -> np.ndarray:
    """前景・背景とも面積 `min_area` 未満の連結成分を取り除く.

    先に前景の小さな島を背景にし，次に背景の小さな穴を前景で埋める．

    Args:
        binary: `(H, W)` の uint8（0/255）．
        min_area: 取り除く面積の上限（未満を除く．画素数）．

    Returns:
        `(H, W)` の uint8（0/255）．
    """
    out = np.where(binary > 0, 255, 0).astype(np.uint8)
    for value, fill in ((255, 0), (0, 255)):
        region = (out == value).astype(np.uint8)
        n, labels, stats, _ = cv2.connectedComponentsWithStats(region, connectivity=_CONNECTIVITY)
        small = [i for i in range(1, n) if stats[i, cv2.CC_STAT_AREA] < min_area]
        if small:
            out[np.isin(labels, small)] = fill
    return out


def resize_mask(mask: np.ndarray, width: int, height: int, threshold: int) -> np.ndarray:
    """ResUNet の出力サイズのマスクを枠のサイズに線形補間で戻し，2 値にする.

    Args:
        mask: `(S, S)` の uint8（0/255）．
        width: 枠の幅．
        height: 枠の高さ．
        threshold: この値より大きい画素を前景とする．

    Returns:
        `(height, width)` の bool．
    """
    resized = cv2.resize(mask, (width, height))
    return resized > threshold


@dataclass(frozen=True)
class CropMasks:
    """切り出した画像 1 枚のセグメンテーションの結果（枠のサイズ）.

    Attributes:
        post: `(H, W)` の bool．前景チャンネル・後処理あり（DeepTIAS と同じ手順）．
        no_post: `(H, W)` の bool．前景チャンネル・後処理（小領域除去）なし．
        other_channel: `(H, W)` の bool．もう一方のチャンネル・後処理あり
            （出力チャンネルの取り違えに気づくため）．
        postprocess_applied: 前景チャンネルで小領域除去をしたか（輪郭が 0 個または 2 個超）．
    """

    post: np.ndarray
    no_post: np.ndarray
    other_channel: np.ndarray
    postprocess_applied: bool


def output_to_masks(
    output: np.ndarray,
    width: int,
    height: int,
    *,
    foreground_channel: int,
    min_area: int,
    mask_threshold: int,
) -> CropMasks:
    """ResUNet の出力から，後処理の有無・チャンネルの違いによるマスクを作る.

    Args:
        output: ResUNet の出力（`(1, 2, S, S)`, float）．
        width: 切り出した枠の幅．
        height: 切り出した枠の高さ．
        foreground_channel: 前景として使うチャンネル．
        min_area: 小領域除去の面積の上限．
        mask_threshold: 枠のサイズに戻したマスクの 2 値化のしきい値．

    Returns:
        マスク（枠のサイズ）．
    """

    def binary_of(channel: int) -> tuple[np.ndarray, np.ndarray, bool]:
        binary = binarize_otsu(probability_map_u8(output, channel))
        applied = needs_small_region_removal(count_contours(binary))
        post = remove_small_regions(binary, min_area) if applied else binary
        return binary, post, applied

    raw, post, applied = binary_of(foreground_channel)
    _, other_post, _ = binary_of(1 - foreground_channel)
    return CropMasks(
        post=resize_mask(post, width, height, mask_threshold),
        no_post=resize_mask(raw, width, height, mask_threshold),
        other_channel=resize_mask(other_post, width, height, mask_threshold),
        postprocess_applied=applied,
    )


class ResUNetRunner:
    """ResUNet の ONNX を onnxruntime（CPU）で推論する薄いラッパ.

    Args:
        path: ResUNet の ONNX のパス．
        input_size: 入力の一辺 S（モデルの入力形状と照合する）．

    Raises:
        ValueError: ファイルが無い，入力・出力の数や形状が想定（`(1, 3, S, S)` → `(1, 2, S, S)`）
            と違う場合．
    """

    def __init__(self, path: Path, input_size: int) -> None:
        if not path.is_file():
            raise ValueError(f"ResUNet の ONNX が無い: {path}")
        self._session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
        inputs = self._session.get_inputs()
        outputs = self._session.get_outputs()
        if len(inputs) != 1 or len(outputs) < 1:
            raise ValueError(
                f"ResUNet の入力は 1 つ・出力は 1 つ以上のはずである（入力 {len(inputs)}・"
                f"出力 {len(outputs)}）"
            )
        shape = inputs[0].shape
        static = [d for d in shape[2:] if isinstance(d, int)]
        if len(shape) != 4 or any(d != input_size for d in static):
            raise ValueError(
                f"ResUNet の入力形状 {shape} が設定の input_size {input_size} と合わない"
            )
        self._input_name = inputs[0].name
        self._output_name = outputs[0].name
        self.input_size = input_size

    def run(self, bgr: np.ndarray) -> np.ndarray:
        """切り出した画像 1 枚を推論する.

        Args:
            bgr: BGR 画像（`(H, W, 3)`, uint8）．

        Returns:
            `(1, 2, S, S)` の float32 の出力（ロジット）．

        Raises:
            ValueError: 画像の形状・dtype が不正，または出力の形状が想定と違う場合．
        """
        x = preprocess_resunet(bgr, self.input_size)
        (output,) = self._session.run([self._output_name], {self._input_name: x})
        output = np.asarray(output)
        if output.ndim != 4 or output.shape[:2] != (1, 2):
            raise ValueError(f"ResUNet の出力形状 {output.shape} が (1, 2, S, S) でない")
        return output