"""患者単位分割のテスト（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 CLASS_SLUGS, VALID_SEVERITIES
from tias_judge_model.data.manifest import ManifestRecord
from tias_judge_model.data.split import (
    Fold,
    SplitResult,
    count_by_class,
    count_by_class_severity,
    make_kfold_splits,
    patient_labels,
    select_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 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:
    records = make_records(20)
    result = split_train_val_test(records, val_ratio=0.15, test_ratio=0.15, seed=SEED)
    # 各クラス 20 人 × 0.15 = 3 人
    assert len(result.val) == 15
    assert len(result.test) == 15
    assert len(result.train) == 70


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:
    records = make_records(6, prefix="one") + make_records(6, 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"