"""交差検証の fold 内の検証分割のテスト（SPEC_02「交差検証の実行」手順 2-1・「データ分割」）.

- 学習側から `inner_val_ratio` の患者を検証用に取る（患者単位・9 グループ層化）
- 学習用と検証用の患者が重ならず，合わせると学習側の全患者になる
- 同じ seed で同じ結果．manifest の行順に依存しない
- held_out と組み合わせても，held_out・学習・検証の患者が互いに重ならない
"""

from __future__ import annotations

import random
from collections import Counter

import pytest

from tias_judge_model.data.classes import ALL_DIR_NAMES, parse_dir_name
from tias_judge_model.data.manifest import ManifestRecord
from tias_judge_model.data.split import make_kfold_splits, select_records, split_inner_val

SEED = 42


def make_records(per_group: int = 20, images_per_patient: int = 1) -> list[ManifestRecord]:
    records = []
    for group in ALL_DIR_NAMES:
        label, severity = parse_dir_name(group)
        for p in range(per_group):
            pid = f"pt_{group}_{p:03d}"
            for k in range(images_per_patient):
                records.append(ManifestRecord(f"{group}/{pid}_{k}.png", pid, label, severity))
    return records


def _group_of(records: list[ManifestRecord]) -> dict[str, str]:
    return {r.patient_id: f"{r.label}_{r.severity}" for r in records}


def test_train_and_val_are_disjoint_and_cover_all_patients() -> None:
    records = make_records()
    train, val = split_inner_val(records, 0.15, SEED)
    assert set(train).isdisjoint(val)
    assert set(train) | set(val) == {r.patient_id for r in records}


def test_val_ratio_is_stratified_by_group() -> None:
    records = make_records(per_group=20)
    _, val = split_inner_val(records, 0.15, SEED)
    group = _group_of(records)
    counts = Counter(group[p] for p in val)
    # 20 × 0.15 = 3 人ずつ，9 グループすべてから
    assert counts == dict.fromkeys(ALL_DIR_NAMES, 3)


def test_multiple_images_per_patient_stay_together() -> None:
    records = make_records(per_group=10, images_per_patient=3)
    train, val = split_inner_val(records, 0.2, SEED)
    assert set(train).isdisjoint(val)
    val_images = select_records(records, val)
    assert len(val_images) == 3 * len(val)


def test_same_seed_same_result_and_order_independent() -> None:
    records = make_records()
    shuffled = records[:]
    random.Random(0).shuffle(shuffled)
    assert split_inner_val(records, 0.15, SEED) == split_inner_val(shuffled, 0.15, SEED)


def test_different_seed_changes_result() -> None:
    records = make_records()
    assert split_inner_val(records, 0.15, 1) != split_inner_val(records, 0.15, 2)


@pytest.mark.parametrize("ratio", [0.0, 1.0, -0.1, 1.5])
def test_invalid_ratio_raises(ratio: float) -> None:
    with pytest.raises(ValueError):
        split_inner_val(make_records(), ratio, SEED)


def test_held_out_train_val_are_mutually_disjoint_in_every_fold() -> None:
    records = make_records(per_group=15)
    all_patients = {r.patient_id for r in records}
    folds = make_kfold_splits(records, 5, SEED)
    for fold in folds:
        train_side = select_records(records, fold.train)
        train, val = split_inner_val(train_side, 0.15, SEED + fold.index)
        held = set(fold.held_out)
        assert held.isdisjoint(train)
        assert held.isdisjoint(val)
        assert set(train).isdisjoint(val)
        assert held | set(train) | set(val) == all_patients