"""交差検証の手順のテスト（SPEC_02「交差検証の実行」「程度（ややあり）の扱い」）.

最重要:
- 分割は manifest 全体で作り，`data.severities` に関わらず同一（同じ seed なら同じ分割）
- 各 fold で held_out・学習・検証の患者が互いに重ならず，held_out は fold 間で manifest を分割する
- `data.severities` は学習用と検証用にだけ適用し，held_out は全程度のまま評価する

合成データ（scripts/make_synthetic_dataset.py）・極小設定（resnet18・pretrained なし・1 エポック）で
実際に学習を回す（`@pytest.mark.slow`）．
"""

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path
from types import ModuleType

import numpy as np
import pytest

from tias_judge_model.data.manifest import read_manifest
from tias_judge_model.schema import ExperimentConfig
from tias_judge_model.training.cross_validation import CrossValidationResult, run_cross_validation

SCRIPT_PATH = Path(__file__).resolve().parents[2] / "scripts" / "make_synthetic_dataset.py"
N_FOLDS = 5


def _load_script() -> ModuleType:
    spec = importlib.util.spec_from_file_location("make_synthetic_dataset_cv", SCRIPT_PATH)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _config(root: Path, severities: list[str]) -> ExperimentConfig:
    return ExperimentConfig.model_validate(
        {
            "data": {"root": str(root), "severities": severities},
            "split": {"seed": 7, "n_folds": N_FOLDS},
            "model": {"name": "resnet18", "pretrained": False},
            "train": {
                "image_size": 32,
                "batch_size": 8,
                "max_epochs": 1,
                "warmup_epochs": 0,
                "inner_val_ratio": 0.2,
                "num_workers": 0,
                "amp": False,
            },
        }
    )


@pytest.fixture(scope="module")
def cv_runs(
    tmp_path_factory: pytest.TempPathFactory,
) -> tuple[list, CrossValidationResult, CrossValidationResult]:
    root = tmp_path_factory.mktemp("cv") / "dataset"
    _load_script().generate(root, per_dir=6, seed=0)
    records = read_manifest(root / "manifest.csv")
    with_mild = run_cross_validation(
        _config(root, ["strong", "mild", "normal"]), records, root.parent / "ck_with"
    )
    without_mild = run_cross_validation(
        _config(root, ["strong", "normal"]), records, root.parent / "ck_without"
    )
    return records, with_mild, without_mild


pytestmark = pytest.mark.slow


def test_runs_all_folds(cv_runs) -> None:
    _, with_mild, without_mild = cv_runs
    assert [f.index for f in with_mild.folds] == list(range(N_FOLDS))
    assert len(without_mild.folds) == N_FOLDS


def test_patients_disjoint_within_each_fold(cv_runs) -> None:
    records, with_mild, without_mild = cv_runs
    all_patients = {r.patient_id for r in records}
    for result in (with_mild, without_mild):
        for f in result.folds:
            train, val, held = set(f.train_patients), set(f.val_patients), set(f.held_out_patients)
            assert train and val and held
            assert train.isdisjoint(val)
            assert train.isdisjoint(held)
            assert val.isdisjoint(held)
            assert train | val | held == all_patients


def test_held_out_partitions_manifest_across_folds(cv_runs) -> None:
    records, with_mild, _ = cv_runs
    held_all = [p for f in with_mild.folds for p in f.held_out_patients]
    assert len(held_all) == len(set(held_all))
    assert set(held_all) == {r.patient_id for r in records}
    indices = np.concatenate([f.held_out_indices for f in with_mild.folds])
    assert sorted(indices.tolist()) == list(range(len(records)))


def test_split_is_same_regardless_of_severities(cv_runs) -> None:
    _, with_mild, without_mild = cv_runs
    for a, b in zip(with_mild.folds, without_mild.folds, strict=True):
        assert a.held_out_patients == b.held_out_patients
        assert a.train_patients == b.train_patients
        assert a.val_patients == b.val_patients
        np.testing.assert_array_equal(a.held_out_indices, b.held_out_indices)


def test_held_out_keeps_mild_even_when_excluded_from_training(cv_runs) -> None:
    records, _, without_mild = cv_runs
    severities = np.array([r.severity for r in records])
    held_sev = np.concatenate([severities[f.held_out_indices] for f in without_mild.folds])
    assert (held_sev == "mild").sum() == (severities == "mild").sum()
    for f in without_mild.folds:
        held_ids = set(f.held_out_patients)
        expected = {i for i, r in enumerate(records) if r.patient_id in held_ids}
        assert set(f.held_out_indices.tolist()) == expected


def test_severities_applied_only_to_train_and_val(cv_runs) -> None:
    _, with_mild, without_mild = cv_runs
    for f in without_mild.folds:
        for split in ("train", "val"):
            mild_counts = {g: n for g, n in f.group_counts[split].items() if g.endswith("_mild")}
            assert sum(mild_counts.values()) == 0, split
        assert sum(n for g, n in f.group_counts["held_out"].items() if g.endswith("_mild")) > 0
    # ややあり込みでは学習用に mild が入る
    assert any(
        sum(n for g, n in f.group_counts["train"].items() if g.endswith("_mild")) > 0
        for f in with_mild.folds
    )


def test_held_out_probabilities_shape(cv_runs) -> None:
    _, with_mild, _ = cv_runs
    for f in with_mild.folds:
        assert f.held_out_probs.shape == (len(f.held_out_indices), 5)
        np.testing.assert_allclose(f.held_out_probs.sum(axis=1), 1.0, atol=1e-4)
        assert 1 <= f.best_epoch <= len(f.history)


def test_checkpoints_per_fold(cv_runs, tmp_path_factory) -> None:
    records, _, _ = cv_runs
    # fixture と同じ tmp の下に fold<k>.pt が書かれている
    base = tmp_path_factory.getbasetemp()
    found = sorted(p.name for p in base.rglob("ck_with/fold*.pt"))
    assert found == [f"fold{k}.pt" for k in range(N_FOLDS)]