"""画像 1 枚のセグメンテーション（`segmentation/pipeline.py`）のテスト.

仕様（実装サマリーを仕様として扱う）:

- 枠で切り出して ResUNet に掛け，枠のサイズに戻したマスクを元の画像サイズのマスクに配置する
  （枠の外は背景）
- 検出できない場合は前景の無いマスク（空マスク）
"""

from __future__ import annotations

import numpy as np

from tias_judge_model.segmentation.detection import Rect, full_rect
from tias_judge_model.segmentation.pipeline import (
    SegmentationParams,
    empty_masks,
    place_masks,
    segment_in_rect,
)
from tias_judge_model.segmentation.resunet import CropMasks, ResUNetRunner

PARAMS = SegmentationParams(foreground_channel=0, min_area=500, mask_threshold=127)


def test_place_masks_puts_crop_at_rect() -> None:
    crop = np.ones((3, 4), bool)
    masks = CropMasks(post=crop, no_post=crop, other_channel=~crop, postprocess_applied=True)
    placed = place_masks(masks, Rect(2, 1, 6, 4), height=6, width=8)
    assert placed.post.shape == (6, 8)
    assert placed.post.sum() == 12
    assert placed.post[1:4, 2:6].all()
    assert not placed.other_channel.any()
    assert placed.postprocess_applied is True


def test_empty_masks_have_no_foreground() -> None:
    m = empty_masks(5, 7)
    for arr in (m.post, m.no_post, m.other_channel):
        assert arr.shape == (5, 7)
        assert not arr.any()
    assert m.postprocess_applied is False


def _scene() -> tuple[np.ndarray, np.ndarray]:
    bgr = np.full((120, 160, 3), (150, 120, 110), np.uint8)  # 青みがかった背景
    yy, xx = np.mgrid[0:120, 0:160]
    mask = ((yy - 60) / 30) ** 2 + ((xx - 90) / 40) ** 2 <= 1
    bgr[mask] = (105, 95, 200)  # 赤（BGR）
    return bgr, mask


def _iou(a: np.ndarray, b: np.ndarray) -> float:
    return (a & b).sum() / (a | b).sum()


def test_segment_whole_image_with_dummy_model(dummy_seg_models: dict) -> None:
    runner = ResUNetRunner(dummy_seg_models["resunet"], 256)
    bgr, mask = _scene()
    m = segment_in_rect(bgr, full_rect(120, 160), runner, PARAMS)
    assert m.post.shape == mask.shape
    assert _iou(m.post, mask) > 0.9


def test_segment_in_sub_rect_keeps_outside_background(dummy_seg_models: dict) -> None:
    runner = ResUNetRunner(dummy_seg_models["resunet"], 256)
    bgr, mask = _scene()
    rect = Rect(40, 20, 140, 100)
    m = segment_in_rect(bgr, rect, runner, PARAMS)
    outside = np.ones_like(mask)
    outside[20:100, 40:140] = False
    assert not m.post[outside].any()
    assert not m.no_post[outside].any()
    assert _iou(m.post, mask) > 0.9