"""ResUNet の前処理・後処理（`segmentation/resunet.py`）のテスト.

仕様（DeepTIAS 2.7 の再現．実装サマリーを仕様として扱う）:

- 前処理: 切り出し画像を `cv2.resize`（線形）で S×S → BGR のまま /255 → `(1, 3, S, S)` float32
- 後処理: 出力 ch（既定 0）にシグモイド → ×255 を uint8（切り捨て）→ 大津の 2 値化 →
  輪郭が 0 個または 2 個超なら小領域除去（面積 min_area 未満の前景の島を消し，次に同じく小さい
  背景の穴を埋める．8 近傍）→ 枠のサイズに線形で戻して `> mask_threshold` で 2 値化
- `ResUNetRunner` は入力 `(1, 3, S, S)` → 出力 `(1, 2, S, S)` の ONNX を CPU で推論する．
  ファイルが無い・入力サイズが設定と違うと ValueError
"""

from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np
import pytest

from tias_judge_model.segmentation.resunet import (
    ResUNetRunner,
    binarize_otsu,
    count_contours,
    needs_small_region_removal,
    output_to_masks,
    preprocess_resunet,
    probability_map_u8,
    remove_small_regions,
    resize_mask,
)

BIG = 30.0


# --- 前処理 ---


def test_preprocess_shape_dtype_and_range() -> None:
    rng = np.random.default_rng(0)
    bgr = rng.integers(0, 256, size=(37, 53, 3), dtype=np.uint8)
    x = preprocess_resunet(bgr, 64)
    assert x.shape == (1, 3, 64, 64)
    assert x.dtype == np.float32
    assert 0.0 <= x.min() and x.max() <= 1.0


def test_preprocess_matches_linear_resize_divided_by_255() -> None:
    rng = np.random.default_rng(1)
    bgr = rng.integers(0, 256, size=(90, 70, 3), dtype=np.uint8)
    x = preprocess_resunet(bgr, 32)
    expected = cv2.resize(bgr, (32, 32), interpolation=cv2.INTER_LINEAR).astype(np.float32) / 255
    np.testing.assert_array_equal(x[0], expected.transpose(2, 0, 1))


def test_preprocess_keeps_bgr_channel_order() -> None:
    bgr = np.zeros((10, 10, 3), dtype=np.uint8)
    bgr[..., 0] = 255  # B
    bgr[..., 2] = 51  # R
    x = preprocess_resunet(bgr, 16)
    assert np.all(x[0, 0] == 1.0)  # ch0 = B
    assert np.all(x[0, 1] == 0.0)
    assert np.allclose(x[0, 2], 51 / 255)  # ch2 = R


@pytest.mark.parametrize(
    "bad",
    [
        np.zeros((10, 10), dtype=np.uint8),
        np.zeros((10, 10, 4), dtype=np.uint8),
        np.zeros((10, 10, 3), dtype=np.float32),
        np.zeros((0, 10, 3), dtype=np.uint8),
        np.zeros((10, 0, 3), dtype=np.uint8),
    ],
    ids=["2d", "4ch", "float", "empty_h", "empty_w"],
)
def test_preprocess_rejects_bad_input(bad: np.ndarray) -> None:
    with pytest.raises(ValueError):
        preprocess_resunet(bad, 16)


# --- シグモイドと丸め ---


def _output(ch0: np.ndarray, ch1: np.ndarray | None = None) -> np.ndarray:
    ch1 = -ch0 if ch1 is None else ch1
    return np.stack([ch0, ch1])[None].astype(np.float32)


def test_probability_map_truncates_sigmoid_times_255() -> None:
    # sigmoid(0) = 0.5 → 127.5 → 127（切り捨て）
    logits = np.zeros((2, 2), dtype=np.float32)
    assert np.all(probability_map_u8(_output(logits), 0) == 127)


@pytest.mark.parametrize("target", [1.9, 64.99, 128.5, 200.2, 254.7])
def test_probability_map_floor_for_various_values(target: float) -> None:
    p = target / 255.0
    logit = np.log(p / (1 - p))
    out = probability_map_u8(_output(np.full((1, 1), logit, np.float64)), 0)
    assert out.dtype == np.uint8
    assert int(out[0, 0]) == int(np.floor(target))


def test_probability_map_extremes_do_not_overflow() -> None:
    logits = np.array([[1000.0, -1000.0, 50.0, -50.0]], dtype=np.float32)
    out = probability_map_u8(_output(logits), 0)
    assert out.tolist() == [[255, 0, 255, 0]]


def test_probability_map_uses_requested_channel() -> None:
    ch0 = np.full((2, 2), 100.0, np.float32)
    out = _output(ch0)
    assert np.all(probability_map_u8(out, 0) == 255)
    assert np.all(probability_map_u8(out, 1) == 0)


def test_binarize_otsu_separates_bimodal_map() -> None:
    prob = np.full((10, 10), 20, np.uint8)
    prob[3:7, 3:7] = 230
    binary = binarize_otsu(prob)
    assert set(np.unique(binary)) == {0, 255}
    assert np.all((binary == 255) == (prob == 230))


# --- 輪郭の数と分岐 ---


def _canvas(size: int = 64) -> np.ndarray:
    return np.zeros((size, size), np.uint8)


def test_count_contours_blank_is_zero() -> None:
    assert count_contours(_canvas()) == 0


def test_count_contours_counts_separate_blobs() -> None:
    img = _canvas()
    cv2.circle(img, (20, 20), 8, 255, -1)
    assert count_contours(img) == 1
    cv2.circle(img, (45, 45), 6, 255, -1)
    assert count_contours(img) == 2
    cv2.rectangle(img, (5, 50), (12, 58), 255, -1)
    assert count_contours(img) == 3


def test_sigmoid_of_moderate_logit_is_truncated_below_255() -> None:
    # sigmoid(30) × 255 = 254.99999… → 254（四捨五入ではなく切り捨て）
    out = probability_map_u8(_output(np.full((1, 1), 30.0, np.float32)), 0)
    assert int(out[0, 0]) == 254


def test_count_contours_full_foreground_is_zero() -> None:
    # 全面が前景ならエッジが無く輪郭 0 個（DeepTIAS の条件では小領域除去の対象）
    assert count_contours(np.full((32, 32), 255, np.uint8)) == 0


@pytest.mark.parametrize(
    ("n", "expected"), [(0, True), (1, False), (2, False), (3, True), (10, True)]
)
def test_needs_small_region_removal(n: int, expected: bool) -> None:
    assert needs_small_region_removal(n) is expected


# --- 小領域除去 ---


def test_removes_small_island_and_keeps_boundary_area() -> None:
    img = _canvas(40)
    img[2:12, 2:12] = 255  # 面積 100（大）
    img[20:22, 20:24] = 255  # 面積 8（min_area 未満）
    img[30:33, 30:33] = 255  # 面積 9（ちょうど min_area）
    out = remove_small_regions(img, 9)
    assert np.all(out[2:12, 2:12] == 255)
    assert np.all(out[20:22, 20:24] == 0)
    assert np.all(out[30:33, 30:33] == 255)  # 未満だけを除く


def test_fills_small_hole_and_keeps_boundary_hole() -> None:
    img = np.full((40, 40), 255, np.uint8)
    img[5:7, 5:9] = 0  # 穴 8（未満）
    img[20:23, 20:23] = 0  # 穴 9（ちょうど）
    out = remove_small_regions(img, 9)
    assert np.all(out[5:7, 5:9] == 255)
    assert np.all(out[20:23, 20:23] == 0)


def test_islands_use_8_connectivity() -> None:
    img = _canvas(20)
    img[5, 5] = 255
    img[6, 6] = 255  # 斜めに隣接 → 8 近傍で 1 つの成分（面積 2）
    out = remove_small_regions(img, 2)
    assert out[5, 5] == 255 and out[6, 6] == 255
    out = remove_small_regions(img, 3)
    assert out[5, 5] == 0 and out[6, 6] == 0


def test_holes_use_8_connectivity() -> None:
    img = np.full((20, 20), 255, np.uint8)
    img[5, 5] = 0
    img[6, 6] = 0  # 斜めに隣接する穴 → 面積 2 の 1 成分
    assert np.all(remove_small_regions(img, 2) == img)
    assert np.all(remove_small_regions(img, 3) == 255)


def test_islands_are_removed_before_holes_are_filled() -> None:
    # 細い輪（前景 16 画素）が穴（背景 9 画素）を囲む．min_area 20:
    # 島の除去が先なら輪（16 < 20）が消えて穴は外の背景とつながり，全面が背景になる．
    # 穴埋めが先だと輪＋穴 = 25 ≥ 20 で残ってしまう．
    img = _canvas(30)
    img[10:15, 10:15] = 255
    img[11:14, 11:14] = 0
    out = remove_small_regions(img, 20)
    assert np.all(out == 0)


def test_min_area_zero_changes_nothing() -> None:
    rng = np.random.default_rng(0)
    img = (rng.random((30, 30)) > 0.5).astype(np.uint8) * 255
    np.testing.assert_array_equal(remove_small_regions(img, 0), img)


def test_nonzero_values_are_treated_as_foreground() -> None:
    img = _canvas(20)
    img[2:18, 2:18] = 1
    out = remove_small_regions(img, 5)
    assert out.dtype == np.uint8
    assert set(np.unique(out)) <= {0, 255}
    assert np.all(out[2:18, 2:18] == 255)


# --- 枠のサイズへの戻し ---


def test_resize_mask_shape_and_threshold() -> None:
    mask = np.zeros((8, 8), np.uint8)
    mask[:, 4:] = 255
    out = resize_mask(mask, 16, 4, 127)
    assert out.shape == (4, 16)
    assert out.dtype == bool
    assert not out[:, :7].any()
    assert out[:, 9:].all()


def test_resize_mask_threshold_is_strictly_greater() -> None:
    mask = np.full((4, 4), 127, np.uint8)
    assert not resize_mask(mask, 4, 4, 127).any()
    assert resize_mask(mask, 4, 4, 126).all()


# --- output_to_masks ---


def _disk_logits(size: int, spots: int) -> np.ndarray:
    img = np.zeros((size, size), np.uint8)
    cv2.circle(img, (size // 2, size // 2), size // 4, 255, -1)
    for k in range(spots):
        img[4 + 8 * k : 8 + 8 * k, 4:8] = 255  # 面積 16 の点
    return np.where(img > 0, BIG, -BIG).astype(np.float32)


def test_output_to_masks_without_spots_skips_postprocess() -> None:
    logits = _disk_logits(64, spots=0)
    m = output_to_masks(
        _output(logits), 64, 64, foreground_channel=0, min_area=50, mask_threshold=127
    )
    assert not m.postprocess_applied
    np.testing.assert_array_equal(m.post, m.no_post)
    np.testing.assert_array_equal(m.post, logits > 0)
    np.testing.assert_array_equal(m.other_channel, ~(logits > 0))


def test_output_to_masks_one_spot_keeps_it_since_two_contours() -> None:
    # 輪郭 2 個（舌＋点 1 つ）では小領域除去をしない（DeepTIAS と同じ条件）
    logits = _disk_logits(64, spots=1)
    m = output_to_masks(
        _output(logits), 64, 64, foreground_channel=0, min_area=50, mask_threshold=127
    )
    assert not m.postprocess_applied
    assert m.post[4:8, 4:8].all()


def test_output_to_masks_many_spots_are_removed() -> None:
    logits = _disk_logits(64, spots=2)
    m = output_to_masks(
        _output(logits), 64, 64, foreground_channel=0, min_area=50, mask_threshold=127
    )
    assert m.postprocess_applied
    assert not m.post[4:16, 4:8].any()  # 点は消える
    assert m.no_post[4:8, 4:8].all()  # 後処理なしには残る
    disk = _disk_logits(64, spots=0) > 0
    np.testing.assert_array_equal(m.post, disk)


def test_output_to_masks_resizes_to_box_and_respects_foreground_channel() -> None:
    logits = _disk_logits(64, spots=0)
    m0 = output_to_masks(
        _output(logits), 30, 20, foreground_channel=0, min_area=10, mask_threshold=127
    )
    m1 = output_to_masks(
        _output(logits), 30, 20, foreground_channel=1, min_area=10, mask_threshold=127
    )
    assert m0.post.shape == (20, 30)
    np.testing.assert_array_equal(m1.post, m0.other_channel)
    np.testing.assert_array_equal(m1.other_channel, m0.post)


# --- ResUNetRunner（ダミーの ONNX） ---


def test_runner_runs_dummy_model(dummy_seg_models: dict) -> None:
    runner = ResUNetRunner(dummy_seg_models["resunet"], 256)
    bgr = np.zeros((50, 80, 3), np.uint8)
    bgr[..., 2] = 200  # 赤
    out = runner.run(bgr)
    assert out.shape == (1, 2, 256, 256)
    assert np.all(out[0, 0] > 0)  # ダミーは赤みを前景（ch0 > 0）とする
    assert np.all(out[0, 1] < 0)


def test_runner_rejects_missing_file(tmp_path: Path) -> None:
    with pytest.raises(ValueError):
        ResUNetRunner(tmp_path / "none.onnx", 256)


def test_runner_rejects_input_size_mismatch(dummy_seg_models: dict) -> None:
    with pytest.raises(ValueError):
        ResUNetRunner(dummy_seg_models["resunet"], 128)