"""画像 1 枚のセグメンテーション（枠で切り出す → ResUNet → 元の画像の位置へ戻す）.

枠の決め方は呼び出し側が選ぶ（`detection.py`）:

- 検出なし: 画像全体（`full_rect`）を ResUNet に掛ける．切り出し済み画像はこれ
- 正解マスクの外接矩形（`mask_bounding_rect`）: 検出が完璧な場合の上限を見るため
- 検出モデル（`OnnxDetector`）: DeepTIAS 2.7 と同じ
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from tias_judge_model.segmentation.detection import Rect
from tias_judge_model.segmentation.resunet import CropMasks, ResUNetRunner, output_to_masks


@dataclass(frozen=True)
class SegmentationParams:
    """後処理の設定（`SegEvalConfig` の該当部分）.

    Attributes:
        foreground_channel: 前景として使う ResUNet の出力チャンネル．
        min_area: 小領域除去の面積の上限．
        mask_threshold: 枠のサイズに戻したマスクの 2 値化のしきい値．
    """

    foreground_channel: int
    min_area: int
    mask_threshold: int


def place_masks(crop: CropMasks, rect: Rect, height: int, width: int) -> CropMasks:
    """枠のサイズのマスクを元の画像サイズのマスクに配置する（枠の外は背景）.

    Args:
        crop: 枠のサイズのマスク．
        rect: 枠．
        height: 元の画像の高さ．
        width: 元の画像の幅．

    Returns:
        `(height, width)` の bool のマスク．
    """

    def place(mask: np.ndarray) -> np.ndarray:
        out = np.zeros((height, width), dtype=bool)
        out[rect.y0 : rect.y1, rect.x0 : rect.x1] = mask
        return out

    return CropMasks(
        post=place(crop.post),
        no_post=place(crop.no_post),
        other_channel=place(crop.other_channel),
        postprocess_applied=crop.postprocess_applied,
    )


def empty_masks(height: int, width: int) -> CropMasks:
    """前景の無いマスク（検出できなかった場合）を返す.

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

    Returns:
        すべて背景のマスク．
    """
    empty = np.zeros((height, width), dtype=bool)
    return CropMasks(post=empty, no_post=empty, other_channel=empty, postprocess_applied=False)


def segment_in_rect(
    bgr: np.ndarray, rect: Rect, runner: ResUNetRunner, params: SegmentationParams
) -> CropMasks:
    """枠で切り出して ResUNet に掛け，元の画像サイズのマスクを返す.

    Args:
        bgr: BGR 画像（`(H, W, 3)`, uint8）．
        rect: 切り出す枠（画像の内側．幅・高さは 1 以上）．
        runner: ResUNet の推論器．
        params: 後処理の設定．

    Returns:
        `(H, W)` の bool のマスク（後処理あり・なし・もう一方のチャンネル）．
    """
    height, width = bgr.shape[:2]
    crop = bgr[rect.y0 : rect.y1, rect.x0 : rect.x1]
    output = runner.run(crop)
    masks = output_to_masks(
        output,
        rect.width,
        rect.height,
        foreground_channel=params.foreground_channel,
        min_area=params.min_area,
        mask_threshold=params.mask_threshold,
    )
    return place_masks(masks, rect, height, width)