"""患者単位分割のテスト（SPEC_02「データ分割」・SPEC_01「patient_id について」）.

最重要: 同一 patient_id が複数の分割に跨らないこと．1 患者に複数枚ある manifest でも同様．
"""

import logging
import random
from collections import Counter

import pytest

from tias_judge_model.data.classes import (
    ALL_DIR_NAMES,
    CLASS_SLUGS,
    VALID_SEVERITIES,
    parse_dir_name,
)
from tias_judge_model.data.manifest import ManifestRecord
from tias_judge_model.data.split import (
    Fold,
    SplitResult,
    count_by_class,
    count_by_class_severity,
    count_by_group,
    filter_by_severity,
    make_kfold_splits,
    patient_groups,
    patient_labels,
    select_records,
    split_records,
    split_train_val_test,
)

SEED = 42


def make_records(
    patients_per_class: int | dict[str, int], images_per_patient: int = 1, prefix: str = "pt"
) -> list[ManifestRecord]:
    """クラスごとに指定人数の患者を持つ合成 manifest を作る（1 患者 images_per_patient 枚）."""
    if isinstance(patients_per_class, int):
        patients_per_class = dict.fromkeys(CLASS_SLUGS, patients_per_class)
    records = []
    for label, n_patients in patients_per_class.items():
        severities = VALID_SEVERITIES[label]
        for p in range(n_patients):
            patient_id = f"{prefix}_{label}_{p:04d}"
            for k in range(images_per_patient):
                severity = severities[(p + k) % len(severities)]
                records.append(
                    ManifestRecord(
                        image_path=f"{label}_{severity}/{patient_id}_{k}.png",
                        patient_id=patient_id,
                        label=label,
                        severity=severity,
                    )
                )
    return records


def make_group_records(
    patients_per_group: dict[str, int], images_per_patient: int = 1, prefix: str = "g"
) -> list[ManifestRecord]:
    """層化グループ（`<label>_<severity>`）ごとに指定人数の患者を持つ合成 manifest を作る."""
    records = []
    for group, n_patients in patients_per_group.items():
        label, severity = parse_dir_name(group)
        for p in range(n_patients):
            patient_id = f"{prefix}_{group}_{p:04d}"
            for k in range(images_per_patient):
                records.append(
                    ManifestRecord(
                        image_path=f"{group}/{patient_id}_{k}.png",
                        patient_id=patient_id,
                        label=label,
                        severity=severity,
                    )
                )
    return records


def shuffled(records: list[ManifestRecord], seed: int) -> list[ManifestRecord]:
    out = list(records)
    random.Random(seed).shuffle(out)
    return out


def all_patients(records: list[ManifestRecord]) -> set[str]:
    return {r.patient_id for r in records}


def labels_in(patient_ids: tuple[str, ...], records: list[ManifestRecord]) -> Counter[str]:
    mapping = patient_labels(records)
    return Counter(mapping[p] for p in patient_ids)


# --- 単一分割 (train / val / test) ---


@pytest.mark.parametrize("images_per_patient", [1, 3])
def test_single_split_keeps_patients_disjoint(images_per_patient: int) -> None:
    records = make_records(20, images_per_patient=images_per_patient)
    result = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)

    train, val, test = set(result.train), set(result.val), set(result.test)
    assert not train & val
    assert not train & test
    assert not val & test
    assert train | val | test == all_patients(records)
    # 各分割は患者 ID の重複を含まない
    for part in (result.train, result.val, result.test):
        assert len(part) == len(set(part))


def test_single_split_puts_all_images_of_a_patient_in_same_split() -> None:
    records = make_records(12, images_per_patient=4)
    result = split_train_val_test(records, val_ratio=0.2, test_ratio=0.2, seed=SEED)

    selected = {
        name: select_records(records, ids)
        for name, ids in (("train", result.train), ("val", result.val), ("test", result.test))
    }
    assert sum(len(v) for v in selected.values()) == len(records)
    for name, recs in selected.items():
        for r in recs:
            images_of_patient = [x for x in records if x.patient_id == r.patient_id]
            assert all(x in recs for x in images_of_patient), name


def test_single_split_with_uneven_images_per_patient_stays_disjoint() -> None:
    # 1 患者の枚数がばらつく manifest（将来の複数枚追加を想定）
    base = make_records(10, images_per_patient=1, prefix="one")
    many = make_records(10, images_per_patient=5, prefix="many")
    records = base + many
    result = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)

    groups = [set(result.train), set(result.val), set(result.test)]
    for patient in all_patients(records):
        assert sum(patient in g for g in groups) == 1


def test_single_split_is_reproducible_with_same_seed() -> None:
    records = make_records(20)
    a = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    b = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    assert a == b


def test_single_split_changes_with_different_seed() -> None:
    records = make_records(30)
    a = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=0)
    b = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=1)
    assert a != b


@pytest.mark.parametrize("images_per_patient", [1, 3])
def test_single_split_does_not_depend_on_input_order(images_per_patient: int) -> None:
    records = make_records(20, images_per_patient=images_per_patient)
    expected = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    for shuffle_seed in range(3):
        actual = split_train_val_test(
            shuffled(records, shuffle_seed), val_ratio=0.15, test_ratio=0.15, seed=SEED
        )
        assert actual == expected
    assert (
        split_train_val_test(list(reversed(records)), val_ratio=0.15, test_ratio=0.15, seed=SEED)
        == expected
    )


def test_single_split_is_stratified_by_class() -> None:
    # 実データに近い不均衡な分布（TEST_01 の 1/10 程度）
    per_class = {"pale": 10, "light_red": 34, "red": 12, "dark_red": 38, "purple": 15}
    records = make_records(per_class)
    result = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)

    for part in (result.train, result.val, result.test):
        counts = labels_in(part, records)
        assert set(counts) == set(CLASS_SLUGS), "各分割に全クラスが含まれること"

    # 各クラスで val / test の人数が比率どおり（端数は 1 人以内）
    for part in (result.val, result.test):
        counts = labels_in(part, records)
        for label, n in per_class.items():
            assert abs(counts[label] - n * 0.15) <= 1, label


def test_single_split_sizes_follow_ratios() -> None:
    # SPEC_02: 層化グループ（クラス × 程度の 9 グループ）ごとに「患者数 × 比率」を四捨五入する．
    # make_records(20) は light_red_normal が 20 人，他の 8 グループが 10 人ずつ．
    records = make_records(20)
    result = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    # light_red_normal: 20 × 0.15 = 3 人．他の 8 グループ: 10 × 0.15 = 1.5 → 四捨五入で 2 人
    expected = 3 + 8 * 2
    assert len(result.val) == expected
    assert len(result.test) == expected
    assert len(result.train) == 100 - 2 * expected


def test_single_split_rounds_half_up_per_group() -> None:
    # 2.5 → 3（偶数丸めの 2 ではない）．各グループ 10 人 × 0.25 = 2.5
    records = make_group_records(dict.fromkeys(ALL_DIR_NAMES, 10))
    result = split_train_val_test(records, val_ratio=0.25, test_ratio=0.25, seed=SEED)
    groups = patient_groups(records)
    for part in (result.val, result.test):
        counts = Counter(groups[p] for p in part)
        assert counts == dict.fromkeys(ALL_DIR_NAMES, 3)


def test_single_split_to_dict_contains_all_parts() -> None:
    result = split_train_val_test(make_records(10), val_ratio=0.2, test_ratio=0.2, seed=SEED)
    d = result.to_dict()
    assert set(d) == {"train", "val", "test"}
    assert d["train"] == list(result.train)
    assert isinstance(result, SplitResult)


def test_single_split_rejects_empty_records() -> None:
    with pytest.raises(ValueError):
        split_train_val_test([], val_ratio=0.15, test_ratio=0.15, seed=SEED)


@pytest.mark.parametrize(
    ("val_ratio", "test_ratio"), [(0.5, 0.5), (0.6, 0.5), (0.0, 0.2), (0.2, -0.1), (1.0, 0.1)]
)
def test_single_split_rejects_invalid_ratios(val_ratio: float, test_ratio: float) -> None:
    with pytest.raises(ValueError):
        split_train_val_test(make_records(20), val_ratio, test_ratio, seed=SEED)


def test_single_split_rejects_patient_with_conflicting_labels() -> None:
    records = make_records(10)
    conflict = ManifestRecord(
        image_path="purple_mild/extra.png",
        patient_id=records[0].patient_id,
        label="purple",
        severity="mild",
    )
    assert records[0].label != "purple"
    with pytest.raises(ValueError):
        split_train_val_test([*records, conflict], val_ratio=0.15, test_ratio=0.15, seed=SEED)


def test_split_logs_have_no_patient_id(caplog: pytest.LogCaptureFixture) -> None:
    records = make_records(10, prefix="SECRETPT")
    with caplog.at_level(logging.DEBUG):
        split_train_val_test(records, val_ratio=0.2, test_ratio=0.2, seed=SEED)
        make_kfold_splits(records, n_folds=5, seed=SEED)
    assert "SECRETPT" not in caplog.text


# --- K-fold 交差検証 ---


@pytest.mark.parametrize("images_per_patient", [1, 3])
def test_kfold_held_out_covers_all_patients_exactly_once(images_per_patient: int) -> None:
    records = make_records(12, images_per_patient=images_per_patient)
    folds = make_kfold_splits(records, n_folds=5, seed=SEED)

    assert len(folds) == 5
    assert [f.index for f in folds] == list(range(5))
    held_out_all = [p for f in folds for p in f.held_out]
    assert len(held_out_all) == len(set(held_out_all)), "held_out が fold 間で重複しない"
    assert set(held_out_all) == all_patients(records)


@pytest.mark.parametrize("images_per_patient", [1, 3])
def test_kfold_train_and_held_out_are_disjoint(images_per_patient: int) -> None:
    records = make_records(12, images_per_patient=images_per_patient)
    for fold in make_kfold_splits(records, n_folds=5, seed=SEED):
        assert isinstance(fold, Fold)
        assert not set(fold.train) & set(fold.held_out)
        assert set(fold.train) | set(fold.held_out) == all_patients(records)


def test_kfold_with_uneven_images_per_patient_stays_disjoint() -> None:
    # 各層化グループの患者数が fold 数（5）以上になる構成（SPEC_02: 満たないグループがあれば停止）．
    # "one" は 1 患者 1 枚（strong / mild が 5 人ずつ），"many" は 1 患者 4 枚（strong と mild が
    # 混在し，最も強い strong で層化される）
    records = make_records(10, prefix="one") + make_records(10, images_per_patient=4, prefix="many")
    for fold in make_kfold_splits(records, n_folds=5, seed=SEED):
        train_recs = select_records(records, fold.train)
        held_recs = select_records(records, fold.held_out)
        assert not {r.patient_id for r in train_recs} & {r.patient_id for r in held_recs}
        assert len(train_recs) + len(held_recs) == len(records)


def test_kfold_is_reproducible_with_same_seed() -> None:
    records = make_records(12)
    assert make_kfold_splits(records, 5, seed=SEED) == make_kfold_splits(records, 5, seed=SEED)


def test_kfold_changes_with_different_seed() -> None:
    records = make_records(12)
    assert make_kfold_splits(records, 5, seed=0) != make_kfold_splits(records, 5, seed=1)


@pytest.mark.parametrize("images_per_patient", [1, 3])
def test_kfold_does_not_depend_on_input_order(images_per_patient: int) -> None:
    records = make_records(12, images_per_patient=images_per_patient)
    expected = make_kfold_splits(records, 5, seed=SEED)
    for shuffle_seed in range(3):
        assert make_kfold_splits(shuffled(records, shuffle_seed), 5, seed=SEED) == expected


def test_kfold_is_stratified_by_class() -> None:
    per_class = {"pale": 10, "light_red": 34, "red": 12, "dark_red": 38, "purple": 15}
    records = make_records(per_class)
    for fold in make_kfold_splits(records, n_folds=5, seed=SEED):
        counts = labels_in(fold.held_out, records)
        assert set(counts) == set(CLASS_SLUGS), "各 fold の held_out に全クラスが含まれること"
        for label, n in per_class.items():
            assert abs(counts[label] - n / 5) <= 1, label


@pytest.mark.parametrize("n_folds", [1, 0, -1])
def test_kfold_rejects_n_folds_below_two(n_folds: int) -> None:
    with pytest.raises(ValueError):
        make_kfold_splits(make_records(10), n_folds=n_folds, seed=SEED)


def test_kfold_rejects_empty_records() -> None:
    with pytest.raises(ValueError):
        make_kfold_splits([], n_folds=5, seed=SEED)


def test_kfold_rejects_patient_with_conflicting_labels() -> None:
    records = make_records(10)
    conflict = ManifestRecord(
        image_path="purple_mild/extra.png",
        patient_id=records[0].patient_id,
        label="purple",
        severity="mild",
    )
    with pytest.raises(ValueError):
        make_kfold_splits([*records, conflict], n_folds=5, seed=SEED)


# --- 集計 ---


def test_patient_labels_maps_each_patient_once() -> None:
    records = make_records(3, images_per_patient=2)
    mapping = patient_labels(records)
    assert len(mapping) == 15
    for r in records:
        assert mapping[r.patient_id] == r.label


def test_count_by_class_lists_all_classes_in_id_order() -> None:
    records = make_records({"red": 2, "purple": 1}, images_per_patient=2)
    counts = count_by_class(records)
    assert list(counts) == list(CLASS_SLUGS)
    assert counts == {"pale": 0, "light_red": 0, "red": 4, "dark_red": 0, "purple": 2}


def test_count_by_class_severity_counts_pairs() -> None:
    records = make_records({"red": 3, "light_red": 2})
    counts = count_by_class_severity(records)
    assert counts == {("light_red", "normal"): 2, ("red", "strong"): 2, ("red", "mild"): 1}
    # クラス ID 順に並ぶ
    assert [label for label, _ in counts][0] == "light_red"


# --- 層化（クラス × 程度の 9 グループ．SPEC_02「データ分割」） ---

# 実データの 1/5 程度で，グループごとに人数が異なる分布
UNEVEN_GROUPS = {
    "pale_strong": 12,
    "pale_mild": 11,
    "light_red_normal": 40,
    "red_strong": 14,
    "red_mild": 13,
    "dark_red_strong": 30,
    "dark_red_mild": 20,
    "purple_strong": 16,
    "purple_mild": 15,
}


def _round_half_up(value: float) -> int:
    return int(value + 0.5)


def test_single_split_is_stratified_by_nine_groups() -> None:
    records = make_group_records(UNEVEN_GROUPS)
    result = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    groups = patient_groups(records)

    for part in (result.train, result.val, result.test):
        counts = Counter(groups[p] for p in part)
        assert set(counts) == set(ALL_DIR_NAMES), "各分割に 9 グループすべてが含まれること"
    # val / test の人数はグループごとに「患者数 × 比率」の四捨五入
    for part in (result.val, result.test):
        counts = Counter(groups[p] for p in part)
        for group, n in UNEVEN_GROUPS.items():
            assert counts[group] == _round_half_up(n * 0.15), group


def test_kfold_is_stratified_by_nine_groups() -> None:
    records = make_group_records(UNEVEN_GROUPS)
    groups = patient_groups(records)
    for fold in make_kfold_splits(records, n_folds=5, seed=SEED):
        counts = Counter(groups[p] for p in fold.held_out)
        assert set(counts) == set(ALL_DIR_NAMES), "各 fold の held_out に 9 グループが含まれること"
        for group, n in UNEVEN_GROUPS.items():
            assert abs(counts[group] - n / 5) < 1, group


def test_single_split_with_nine_groups_is_reproducible() -> None:
    records = make_group_records(UNEVEN_GROUPS)
    a = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    b = split_train_val_test(shuffled(records, 3), val_ratio=0.15, test_ratio=0.15, seed=SEED)
    assert a == b
    assert make_kfold_splits(records, 5, seed=SEED) == make_kfold_splits(
        shuffled(records, 3), 5, seed=SEED
    )


def test_kfold_stops_when_a_group_has_fewer_patients_than_folds() -> None:
    per_group = dict.fromkeys(ALL_DIR_NAMES, 10)
    per_group["red_mild"] = 4
    with pytest.raises(ValueError):
        make_kfold_splits(make_group_records(per_group), n_folds=5, seed=SEED)
    # fold 数以上ならちょうど通る
    per_group["red_mild"] = 5
    assert len(make_kfold_splits(make_group_records(per_group), n_folds=5, seed=SEED)) == 5


def test_single_split_stops_when_a_group_leaves_no_train_patient() -> None:
    per_group = dict.fromkeys(ALL_DIR_NAMES, 10)
    per_group["purple_mild"] = 2  # val 1・test 1（2 × 0.3 = 0.6 → 1）で train が 0 人
    with pytest.raises(ValueError):
        split_train_val_test(
            make_group_records(per_group), val_ratio=0.3, test_ratio=0.3, seed=SEED
        )


def test_patient_with_multiple_severities_uses_strongest() -> None:
    records = [
        ManifestRecord("red_mild/a_0.png", "a", "red", "mild"),
        ManifestRecord("red_strong/a_1.png", "a", "red", "strong"),
        ManifestRecord("pale_mild/b_0.png", "b", "pale", "mild"),
        ManifestRecord("light_red_normal/c_0.png", "c", "light_red", "normal"),
    ]
    assert patient_groups(records) == {
        "a": "red_strong",
        "b": "pale_mild",
        "c": "light_red_normal",
    }
    # 行の順序に依存しない
    assert patient_groups(list(reversed(records)))["a"] == "red_strong"


def test_kfold_stratifies_mixed_patients_by_strongest_severity() -> None:
    # red_mild だけの患者は 0 人で，mild と strong の両方を持つ患者が 3 人いる．
    # 最も強い strong で層化されれば red_mild グループは存在せず 5-fold が成り立つ
    # （mild で層化されると 3 人 < 5 fold で停止する）
    per_group = dict.fromkeys(ALL_DIR_NAMES, 10)
    del per_group["red_mild"]
    records = make_group_records(per_group)
    for i in range(3):
        pid = f"mixed_{i}"
        records.append(ManifestRecord(f"red_mild/{pid}_0.png", pid, "red", "mild"))
        records.append(ManifestRecord(f"red_strong/{pid}_1.png", pid, "red", "strong"))
    folds = make_kfold_splits(records, n_folds=5, seed=SEED)
    held = [p for f in folds for p in f.held_out]
    assert {f"mixed_{i}" for i in range(3)} <= set(held)
    # 単一分割でも止まらず，1 患者の画像は同じ分割に入る
    result = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    parts = [set(result.train), set(result.val), set(result.test)]
    for i in range(3):
        assert sum(f"mixed_{i}" in part for part in parts) == 1


# --- 分割後の程度による除外（SPEC_02「程度（ややあり）の扱い」） ---

WITHOUT_MILD = ("strong", "normal")


def _mixed_records() -> list[ManifestRecord]:
    # グループ別の 1 患者 1 枚の患者に加え，1 患者に strong と mild が混在する患者を含める
    records = make_group_records(UNEVEN_GROUPS)
    for i in range(6):
        pid = f"mixed_{i}"
        label = ("pale", "red", "dark_red", "purple")[i % 4]
        records.append(ManifestRecord(f"{label}_strong/{pid}_0.png", pid, label, "strong"))
        records.append(ManifestRecord(f"{label}_mild/{pid}_1.png", pid, label, "mild"))
    return records


def _assert_exclusion_is_subset(
    records: list[ManifestRecord], split: SplitResult | Fold, part_names: list[str]
) -> None:
    everything = split_records(records, split)
    with_mild = split_records(records, split, ("strong", "mild", "normal"))
    without_mild = split_records(records, split, WITHOUT_MILD)
    assert list(everything) == part_names
    assert list(without_mild) == part_names
    # 全程度を指定したときは除外しないときと同じ
    assert with_mild == everything
    for name in part_names:
        kept, full = without_mild[name], everything[name]
        # 部分集合であり，除かれたのは mild の画像だけ，残ったのは mild 以外のすべて
        assert set(kept) <= set(full), name
        removed = set(full) - set(kept)
        assert all(r.severity == "mild" for r in removed), name
        assert set(kept) == {r for r in full if r.severity != "mild"}, name


def test_single_split_exclusion_is_subset_of_full_split() -> None:
    records = _mixed_records()
    split = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    _assert_exclusion_is_subset(records, split, ["train", "val", "test"])
    # 全画像がいずれか 1 つの分割に入る
    everything = split_records(records, split)
    assert sum(len(v) for v in everything.values()) == len(records)


def test_kfold_exclusion_is_subset_of_full_split() -> None:
    records = _mixed_records()
    for fold in make_kfold_splits(records, n_folds=5, seed=SEED):
        _assert_exclusion_is_subset(records, fold, ["train", "held_out"])


def test_exclusion_keeps_strong_images_of_mixed_patient_in_same_split() -> None:
    records = _mixed_records()
    split = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    without_mild = split_records(records, split, WITHOUT_MILD)
    for i in range(6):
        pid = f"mixed_{i}"
        homes = [name for name, rs in without_mild.items() if any(r.patient_id == pid for r in rs)]
        assert len(homes) == 1
        (home,) = homes
        assert pid in split.parts()[home]
        assert [r.severity for r in without_mild[home] if r.patient_id == pid] == ["strong"]


def test_mild_exclusion_keeps_all_five_classes() -> None:
    # SPEC_02: 「ややあり除外」でも 5 クラスすべてが残る
    records = make_group_records(UNEVEN_GROUPS)
    split = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    for name, rs in split_records(records, split, WITHOUT_MILD).items():
        assert all(n > 0 for n in count_by_class(rs).values()), name


@pytest.mark.parametrize("severities", [(), ("weak",), ("strong", "unknown")])
def test_split_records_rejects_invalid_severities(severities: tuple[str, ...]) -> None:
    records = make_group_records(UNEVEN_GROUPS)
    split = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    with pytest.raises(ValueError):
        split_records(records, split, severities)


def test_filter_by_severity_keeps_only_given_severities() -> None:
    records = _mixed_records()
    assert filter_by_severity(records, WITHOUT_MILD) == [r for r in records if r.severity != "mild"]
    assert filter_by_severity(records, ("mild",)) == [r for r in records if r.severity == "mild"]


@pytest.mark.parametrize("severities", [(), ("weak",)])
def test_filter_by_severity_rejects_invalid(severities: tuple[str, ...]) -> None:
    with pytest.raises(ValueError):
        filter_by_severity(make_group_records(UNEVEN_GROUPS), severities)


def test_count_by_group_lists_all_nine_groups() -> None:
    records = make_group_records({"red_mild": 2, "light_red_normal": 3}, images_per_patient=2)
    counts = count_by_group(records)
    assert list(counts) == list(ALL_DIR_NAMES)
    assert counts["red_mild"] == 4
    assert counts["light_red_normal"] == 6
    assert sum(counts.values()) == 10