"""データセット集計のテスト（SPEC_01「データセット集計の出力」・SPEC_02「程度（ややあり）の扱い」）.

- 出力には集計値だけを入れる．**画像単位の値・患者 ID・ファイル名を出さない**（ログも同様）
- 読めない画像・欠けたファイルがあっても止めず，件数として数える
- 舌領域（RGB がすべて 0 でない画素）の Lab を画像ごとに平均し，グループごとに集計する
- 各 mild グループについて「淡紅（light_red_normal）との ΔE」と「同じクラスの strong との ΔE」を
  並べる
- 条件（ややあり込み／除外）ごとの件数と，既定の設定で分割したときの各分割・グループの件数

画像はすべて tmp_path に作る合成画像（臨床データは使わない）．
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

import numpy as np
import pytest
from PIL import Image

from tias_judge_model.data.classes import ALL_DIR_NAMES, CLASS_SLUGS, parse_dir_name
from tias_judge_model.data.color import delta_e_cie76, rgb8_to_lab
from tias_judge_model.data.manifest import ManifestRecord
from tias_judge_model.data.survey import (
    ImageStats,
    collect_image_stats,
    measure_pixels,
    summarize_dataset,
)

SECRET = "SECRETPT"
SPLIT_KW: dict[str, Any] = {"seed": 42, "val_ratio": 0.15, "test_ratio": 0.15, "n_folds": 5}

# 合成データ（scripts/make_synthetic_dataset.py）に近い基準色
BASE = {
    "pale": (235, 205, 200),
    "light_red": (225, 140, 140),
    "red": (215, 40, 45),
    "dark_red": (120, 20, 30),
    "purple": (110, 50, 140),
}


def _uniform(rgb: tuple[int, int, int], *, width: int = 20, height: int = 10) -> np.ndarray:
    """全面が 1 色の舌領域の画像（`(H, W, 3)`, uint8）."""
    pixels = np.zeros((height, width, 3), dtype=np.uint8)
    pixels[:, :] = rgb
    return pixels


def _half_background(rgb: tuple[int, int, int]) -> np.ndarray:
    """左半分が背景（黒），右半分が 1 色の画像."""
    pixels = np.zeros((10, 20, 3), dtype=np.uint8)
    pixels[:, 10:] = rgb
    return pixels


def _stats(pixels: np.ndarray) -> ImageStats:
    return measure_pixels(pixels, mode="RGB", has_alpha=False)


def _record(group: str, index: int, image_index: int = 0) -> ManifestRecord:
    label, severity = parse_dir_name(group)
    patient_id = f"{SECRET}{group}{index:04d}"
    return ManifestRecord(
        image_path=f"{group}/{patient_id}_{image_index}.png",
        patient_id=patient_id,
        label=label,
        severity=severity,
    )


def _blend(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]:
    """a と b を (1 - t) : t で混ぜた色."""
    return tuple(int(round((1 - t) * x + t * y)) for x, y in zip(a, b, strict=True))  # type: ignore[return-value]


def _dataset(per_group: int = 6) -> tuple[list[ManifestRecord], list[ImageStats], dict]:
    """全 9 グループに per_group 人（1 人 1 枚）の合成データ．グループごとの色を返す."""
    normal = BASE["light_red"]
    colors: dict[str, tuple[int, int, int]] = {}
    for group in ALL_DIR_NAMES:
        label, severity = parse_dir_name(group)
        if severity == "mild":
            # red_mild・purple_mild は淡紅寄り，pale_mild・dark_red_mild は strong 寄り
            t = 0.8 if label in ("red", "purple") else 0.1
            colors[group] = _blend(BASE[label], normal, t)
        else:
            colors[group] = BASE[label]
    records: list[ManifestRecord] = []
    stats: list[ImageStats] = []
    for group in ALL_DIR_NAMES:
        for i in range(per_group):
            records.append(_record(group, i))
            stats.append(_stats(_uniform(colors[group], width=20 + i, height=10 + i)))
    return records, stats, colors


def _summarize(records: list[ManifestRecord], stats: list[ImageStats], **kw: Any) -> dict:
    return summarize_dataset(
        records, stats, **{**SPLIT_KW, **kw}, configured_severities=("strong", "mild", "normal")
    )


def _walk_lists(value: Any):
    if isinstance(value, dict):
        for v in value.values():
            yield from _walk_lists(v)
    elif isinstance(value, list):
        yield value
        for v in value:
            yield from _walk_lists(v)


# ---------------------------------------------------------------------------
# 個人情報・画像単位の値を出さない
# ---------------------------------------------------------------------------


class TestNoPersonalData:
    def test_summary_has_no_patient_id_or_filename(self) -> None:
        records, stats, _ = _dataset()
        summary = _summarize(records, stats)
        text = json.dumps(summary, ensure_ascii=False)
        assert SECRET not in text
        assert ".png" not in text
        for r in records:
            assert r.patient_id not in text
            assert r.image_path not in text
        # 分割の確認も実際に行われていること（空振りで個人情報が出ないのではない）
        assert "error" not in summary["split"]["single"]
        assert all("error" not in fold for fold in summary["split"]["kfold"])

    def test_summary_has_no_per_image_values(self) -> None:
        # 画像単位の値の列（画像数と同じ長さのリスト等）を含まない
        records, stats, _ = _dataset()
        summary = _summarize(records, stats)
        for lst in _walk_lists(summary):
            assert len(lst) < len(records)
        # JSON にできる（集計値だけの dict）
        json.dumps(summary)

    def test_logs_have_no_patient_id_or_filename(
        self, tmp_path: Path, caplog: pytest.LogCaptureFixture
    ) -> None:
        caplog.set_level(logging.DEBUG)
        records, _, colors = _dataset(per_group=2)
        images = tmp_path / "images"
        for r in records[:6]:
            path = images / r.image_path
            path.parent.mkdir(parents=True, exist_ok=True)
            Image.fromarray(_uniform(colors[f"{r.label}_{r.severity}"])).save(path)
        broken = images / records[6].image_path
        broken.parent.mkdir(parents=True, exist_ok=True)
        broken.write_bytes(b"not a png")
        # records[7:] は欠けたファイル
        stats = collect_image_stats(records, images)
        _summarize(records, stats)
        assert SECRET not in caplog.text
        assert ".png" not in caplog.text


# ---------------------------------------------------------------------------
# 読み込み（読めない画像・欠けたファイルで止まらない）
# ---------------------------------------------------------------------------


class TestLoading:
    def _write_images(self, images: Path) -> list[ManifestRecord]:
        recs = [_record("red_strong", i) for i in range(6)]
        paths = [images / r.image_path for r in recs]
        paths[0].parent.mkdir(parents=True)
        Image.fromarray(_uniform((200, 50, 50))).save(paths[0])  # RGB
        rgba = np.zeros((10, 20, 4), dtype=np.uint8)
        rgba[..., 0] = 200
        rgba[..., 3] = 255
        Image.fromarray(rgba, mode="RGBA").save(paths[1])  # アルファ付き
        Image.fromarray(np.full((10, 20), 128, dtype=np.uint8), mode="L").save(paths[2])  # 濃淡
        paths[3].write_bytes(b"\x89PNG\r\n\x1a\n broken")  # 壊れた PNG
        paths[4].write_bytes(b"")  # 空ファイル
        # paths[5] は作らない（欠けたファイル）
        return recs

    def test_unreadable_and_missing_are_counted_not_raised(self, tmp_path: Path) -> None:
        images = tmp_path / "images"
        records = self._write_images(images)
        stats = collect_image_stats(records, images)
        assert len(stats) == len(records)
        summary = _summarize(records, stats)
        loading = summary["loading"]
        assert loading["readable"] == 3
        assert loading["unreadable"] == 3
        assert loading["missing"] == 1
        assert loading["with_alpha"] == 1
        assert loading["non_rgb_mode"] == 2  # RGBA と L
        # 読めない画像も件数（manifest 上の画像数）には入る
        assert summary["counts"]["groups"]["red_strong"]["images"] == 6
        # 色の集計は読めた画像だけ
        assert summary["lab"]["groups"]["red_strong"]["n"] == 3

    def test_size_summary(self) -> None:
        records = [_record("red_strong", i) for i in range(3)]
        stats = [
            _stats(_uniform((200, 50, 50), width=w, height=h))
            for w, h in [(10, 40), (30, 20), (50, 30)]
        ]
        size = _summarize(records, stats)["size"]
        assert size["width"] == pytest.approx({"min": 10, "max": 50, "mean": 30, "median": 30})
        assert size["height"] == pytest.approx({"min": 20, "max": 40, "mean": 30, "median": 30})

    def test_tongue_ratio_excludes_background(self) -> None:
        records = [_record("red_strong", 0), _record("red_strong", 1)]
        stats = [_stats(_half_background((200, 50, 50))), _stats(_uniform((200, 50, 50)))]
        ratio = _summarize(records, stats)["tongue_ratio"]
        assert ratio["total"] == pytest.approx({"min": 0.5, "mean": 0.75, "max": 1.0})


# ---------------------------------------------------------------------------
# 色（Lab）と色差
# ---------------------------------------------------------------------------


class TestLabAndDeltaE:
    def test_background_does_not_affect_lab(self) -> None:
        records = [_record("red_strong", 0)]
        stats = [_stats(_half_background((255, 255, 255)))]
        lab = _summarize(records, stats)["lab"]["groups"]["red_strong"]
        assert lab["n"] == 1
        assert lab["L"]["mean"] == pytest.approx(100.0, abs=1e-3)
        assert lab["a"]["mean"] == pytest.approx(0.0, abs=1e-3)
        assert lab["b"]["mean"] == pytest.approx(0.0, abs=1e-3)

    def test_group_statistics_over_image_means(self) -> None:
        # 画像ごとの平均 Lab をグループで集計する．灰色 3 枚（L* が既知）
        grays = [50, 119, 200]
        records = [_record("pale_strong", i) for i in range(3)]
        stats = [_stats(_uniform((v, v, v))) for v in grays]
        expected = [float(rgb8_to_lab(np.array([v, v, v], dtype=np.uint8))[0]) for v in grays]
        lab = _summarize(records, stats)["lab"]["groups"]["pale_strong"]
        assert lab["n"] == 3
        assert lab["L"]["mean"] == pytest.approx(np.mean(expected), abs=1e-3)
        assert lab["L"]["median"] == pytest.approx(expected[1], abs=1e-3)
        assert lab["L"]["p10"] == pytest.approx(np.percentile(expected, 10), abs=1e-3)
        assert lab["L"]["p90"] == pytest.approx(np.percentile(expected, 90), abs=1e-3)
        assert lab["L"]["std"] is not None and lab["L"]["std"] > 0
        assert lab["a"]["mean"] == pytest.approx(0.0, abs=1e-3)
        # 画像の無いグループは件数 0
        assert _summarize(records, stats)["lab"]["groups"]["red_mild"]["n"] == 0

    def test_mild_comparison_with_known_colors(self) -> None:
        records, stats, colors = _dataset()
        summary = _summarize(records, stats)
        lab_of = {g: rgb8_to_lab(np.array(c, dtype=np.uint8)) for g, c in colors.items()}
        items = {item["group"]: item for item in summary["delta_e"]["mild_comparison"]}
        assert set(items) == {"pale_mild", "red_mild", "dark_red_mild", "purple_mild"}
        for group, item in items.items():
            label, _ = parse_dir_name(group)
            to_normal = delta_e_cie76(lab_of[group], lab_of["light_red_normal"])
            to_strong = delta_e_cie76(lab_of[group], lab_of[f"{label}_strong"])
            assert item["delta_e_to_normal"] == pytest.approx(to_normal, abs=1e-3), group
            assert item["delta_e_to_strong"] == pytest.approx(to_strong, abs=1e-3), group
        # 淡紅寄りに作った mild は淡紅に，strong 寄りに作った mild は strong に近い
        assert items["red_mild"]["closer_to"] == "light_red_normal"
        assert items["purple_mild"]["closer_to"] == "light_red_normal"
        assert items["pale_mild"]["closer_to"] == "pale_strong"
        assert items["dark_red_mild"]["closer_to"] == "dark_red_strong"

    def test_group_mean_is_average_of_images(self) -> None:
        # red_mild に色の異なる 2 枚．グループ平均は画像ごとの Lab の平均で，ΔE はそれとの距離
        normal = (225, 140, 140)
        c1, c2 = (220, 100, 100), (200, 60, 60)
        records = [
            _record("light_red_normal", 0),
            _record("red_mild", 0),
            _record("red_mild", 1),
            _record("red_strong", 0),
        ]
        stats = [_stats(_uniform(c)) for c in (normal, c1, c2, BASE["red"])]
        item = next(
            i
            for i in _summarize(records, stats)["delta_e"]["mild_comparison"]
            if i["group"] == "red_mild"
        )
        lab = {c: rgb8_to_lab(np.array(c, dtype=np.uint8)) for c in (normal, c1, c2, BASE["red"])}
        mild_mean = (lab[c1] + lab[c2]) / 2
        assert item["delta_e_to_normal"] == pytest.approx(
            delta_e_cie76(mild_mean, lab[normal]), abs=1e-3
        )
        assert item["delta_e_to_strong"] == pytest.approx(
            delta_e_cie76(mild_mean, lab[BASE["red"]]), abs=1e-3
        )

    def test_matrix_is_symmetric_with_zero_diagonal(self) -> None:
        records, stats, _ = _dataset(per_group=2)
        matrix = _summarize(records, stats)["delta_e"]["matrix"]
        for g1 in ALL_DIR_NAMES:
            assert matrix[g1][g1] == pytest.approx(0.0)
            for g2 in ALL_DIR_NAMES:
                assert matrix[g1][g2] == pytest.approx(matrix[g2][g1])


# ---------------------------------------------------------------------------
# 件数・条件・分割の確認
# ---------------------------------------------------------------------------


class TestCounts:
    def test_counts_by_group_class_severity_and_condition(self) -> None:
        records, stats, _ = _dataset(per_group=6)
        # 1 患者 2 枚（strong と mild）の患者を足す
        extra = [
            ManifestRecord("red_strong/x_0.png", f"{SECRET}mix", "red", "strong"),
            ManifestRecord("red_mild/x_1.png", f"{SECRET}mix", "red", "mild"),
        ]
        records = [*records, *extra]
        stats = [*stats, _stats(_uniform((200, 50, 50))), _stats(_uniform((210, 90, 90)))]
        counts = _summarize(records, stats)["counts"]

        assert counts["total"] == {"images": 56, "patients": 55}
        assert counts["groups"]["red_strong"]["images"] == 7
        assert counts["groups"]["red_mild"]["images"] == 7
        assert counts["groups"]["pale_mild"] == {
            "label": "pale",
            "severity": "mild",
            "images": 6,
            "patients": 6,
        }
        assert set(counts["groups"]) == set(ALL_DIR_NAMES)
        assert counts["classes"]["red"] == {"images": 14, "patients": 13}
        assert counts["classes"]["light_red"] == {"images": 6, "patients": 6}
        assert counts["severities"]["mild"]["images"] == 25
        assert counts["severities"]["normal"]["images"] == 6

        with_mild = counts["conditions"]["with_mild"]
        without_mild = counts["conditions"]["without_mild"]
        assert with_mild["images"] == 56
        assert without_mild["images"] == 56 - 25
        # mix 患者は strong の画像が残るため患者数に入る
        assert without_mild["patients"] == 4 * 6 + 6 + 1
        # 除外後も 5 クラスすべてが残る
        assert set(without_mild["classes"]) == set(CLASS_SLUGS)
        assert all(n > 0 for n in without_mild["classes"].values())
        assert without_mild["classes"]["red"] == 7

    def test_split_counts_follow_default_settings(self) -> None:
        records, stats, _ = _dataset(per_group=10)
        split = _summarize(records, stats)["split"]
        assert (split["seed"], split["n_folds"]) == (42, 5)
        single = split["single"]
        # 各グループの件数を分割で足すと全体の件数に一致し，各分割に 9 グループが入る
        for group in ALL_DIR_NAMES:
            assert sum(single["groups"][part][group] for part in ("train", "val", "test")) == 10
            for part in ("train", "val", "test"):
                assert single["groups"][part][group] > 0, (part, group)
        # val / test はグループごとに 10 × 0.15 = 1.5 → 2 人
        for part in ("val", "test"):
            assert all(n == 2 for n in single["groups"][part].values())
        assert sum(single["patients"].values()) == len(records)
        assert len(split["kfold"]) == 5

    def test_split_conditions_are_subsets(self) -> None:
        # 「ややあり除外」の各分割の件数 = 「ややあり込み」の件数 − その分割の mild の件数
        records, stats, _ = _dataset(per_group=10)
        split = _summarize(records, stats)["split"]
        cases = [(split["single"], ("train", "val", "test"))]
        cases += [(fold, ("train", "held_out")) for fold in split["kfold"]]
        for summary, parts in cases:
            for part in parts:
                groups = summary["groups"][part]
                mild = sum(n for g, n in groups.items() if parse_dir_name(g)[1] == "mild")
                with_mild = summary["conditions"]["with_mild"][part]["images"]
                without_mild = summary["conditions"]["without_mild"][part]["images"]
                assert with_mild == sum(groups.values())
                assert without_mild == with_mild - mild

    def test_split_error_does_not_stop_summary(self) -> None:
        # 分割できない構成（グループの人数が少ない）でも集計は止めない
        records, stats, _ = _dataset(per_group=2)
        summary = _summarize(records, stats)
        assert "error" in summary["split"]["kfold"]
        assert SECRET not in json.dumps(summary, ensure_ascii=False)

    def test_configured_severities_are_recorded(self) -> None:
        records, stats, _ = _dataset(per_group=6)
        summary = summarize_dataset(
            records, stats, **SPLIT_KW, configured_severities=("strong", "normal")
        )
        assert summary["configured_severities"] == ["strong", "normal"]

    def test_length_mismatch_is_error(self) -> None:
        records, stats, _ = _dataset(per_group=2)
        with pytest.raises(ValueError):
            _summarize(records, stats[:-1])