"""学習の設定（`model`・`train` 節）のテスト（SPEC_02「学習の設定」・SPEC_01「設定ファイル」）.

- 既定値が SPEC_02 の表どおり
- `model`・`train` を省略した既存の設定（`tias-dataset` 用）も読める
- 未知のキー・不正な値は拒否する
- リポジトリの `configs/cv_*.yaml` がすべて読め，モデル比較用の設定はモデル名以外が元の設定と同一
"""

from __future__ import annotations

import re
from pathlib import Path

import pytest
import yaml
from pydantic import ValidationError

from tias_judge_model.schema import ExperimentConfig, load_config

CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs"
CV_CONFIGS = sorted(CONFIGS_DIR.glob("cv_*.yaml"))
NO_ENV: dict[str, str] = {}


def _minimal() -> dict:
    return {"data": {"root": "somewhere"}}


def test_model_defaults_follow_spec() -> None:
    c = ExperimentConfig.model_validate(_minimal())
    assert c.model.name == "convnext_tiny"
    assert c.model.pretrained is True


def test_train_defaults_follow_spec() -> None:
    t = ExperimentConfig.model_validate(_minimal()).train
    assert t.image_size == 224
    assert t.batch_size == 32
    assert t.max_epochs == 40
    assert t.patience == 8
    assert t.lr == pytest.approx(1e-4)
    assert t.weight_decay == pytest.approx(0.05)
    assert t.warmup_epochs == 2
    assert t.inner_val_ratio == pytest.approx(0.15)
    assert t.num_workers == 4
    assert t.amp is True


@pytest.mark.parametrize("name", ["synthetic.yaml", "newton.yaml"])
def test_existing_configs_without_model_train_still_load(name: str) -> None:
    c = load_config(CONFIGS_DIR / name, environ=NO_ENV)
    assert c.model.name == "convnext_tiny"
    assert c.train.image_size == 224


@pytest.mark.parametrize(
    "raw",
    [
        {"model": {"nme": "resnet18"}},
        {"train": {"epochs": 3}},
        {"train": {"batch_size": 0}},
        {"train": {"inner_val_ratio": 0.0}},
        {"train": {"inner_val_ratio": 1.0}},
        {"train": {"lr": 0}},
        {"model": {"name": ""}},
    ],
)
def test_invalid_training_settings_are_rejected(raw: dict) -> None:
    with pytest.raises(ValidationError):
        ExperimentConfig.model_validate({**_minimal(), **raw})


def test_cv_configs_exist() -> None:
    names = {p.name for p in CV_CONFIGS}
    assert {"cv_with_mild.yaml", "cv_without_mild.yaml", "cv_synthetic.yaml"} <= names


@pytest.mark.parametrize("path", CV_CONFIGS, ids=lambda p: p.name)
def test_every_cv_config_loads(path: Path) -> None:
    load_config(path, environ=NO_ENV)


def test_with_and_without_mild_differ_only_in_severities() -> None:
    w = load_config(CONFIGS_DIR / "cv_with_mild.yaml", environ=NO_ENV)
    wo = load_config(CONFIGS_DIR / "cv_without_mild.yaml", environ=NO_ENV)
    assert w.data.severities == ("strong", "mild", "normal")
    assert wo.data.severities == ("strong", "normal")
    dw = w.model_dump()
    dwo = wo.model_dump()
    dw["data"].pop("severities")
    dwo["data"].pop("severities")
    assert dw == dwo


def test_production_cv_configs_use_spec_defaults() -> None:
    c = load_config(CONFIGS_DIR / "cv_with_mild.yaml", environ=NO_ENV)
    assert c.model == ExperimentConfig.model_validate(_minimal()).model
    assert c.train == ExperimentConfig.model_validate(_minimal()).train
    assert c.split.n_folds == 5


def test_synthetic_cv_config_does_not_download_weights() -> None:
    c = load_config(CONFIGS_DIR / "cv_synthetic.yaml", environ=NO_ENV)
    assert c.model.pretrained is False


def _model_comparison_pairs() -> list[tuple[Path, Path]]:
    pairs = []
    for p in CV_CONFIGS:
        if p.name in {"cv_with_mild.yaml", "cv_without_mild.yaml", "cv_synthetic.yaml"}:
            continue
        suffix = "with_mild" if p.stem.endswith("_with_mild") else "without_mild"
        pairs.append((p, CONFIGS_DIR / f"cv_{suffix}.yaml"))
    return pairs


@pytest.mark.parametrize(("variant", "base"), _model_comparison_pairs(), ids=lambda p: Path(p).name)
def test_model_comparison_config_differs_only_in_model_name(variant: Path, base: Path) -> None:
    v = load_config(variant, environ=NO_ENV).model_dump()
    b = load_config(base, environ=NO_ENV).model_dump()
    assert v["model"]["name"] != b["model"]["name"]
    v["model"].pop("name")
    b["model"].pop("name")
    assert v == b


def test_model_comparison_configs_are_present() -> None:
    assert len(_model_comparison_pairs()) >= 2


# --- クラス不均衡の対策（SPEC_02「クラス不均衡」→「対策の設定」） ---

IMBALANCE_CONFIGS = sorted((CONFIGS_DIR / "imbalance").glob("*.yaml"))
REPO_ROOT = CONFIGS_DIR.parent
# 設定ファイル名の末尾と，対策のキー・値の対応
IMBALANCE_SUFFIXES = {
    "_cwinv": ("class_weighting", "inverse"),
    "_cwsqrt": ("class_weighting", "sqrt_inverse"),
    "_balanced": ("sampler", "balanced"),
}


def test_imbalance_defaults_are_no_countermeasure() -> None:
    t = ExperimentConfig.model_validate(_minimal()).train
    assert t.class_weighting == "none"
    assert t.sampler == "random"


@pytest.mark.parametrize(
    "train",
    [
        {"class_weighting": "none"},
        {"class_weighting": "inverse"},
        {"class_weighting": "sqrt_inverse"},
        {"sampler": "random"},
        {"sampler": "balanced"},
        {"class_weighting": "inverse", "sampler": "balanced"},
    ],
    ids=str,
)
def test_valid_imbalance_settings_are_accepted(train: dict) -> None:
    t = ExperimentConfig.model_validate({**_minimal(), "train": train}).train
    for key, value in train.items():
        assert getattr(t, key) == value


@pytest.mark.parametrize(
    "train",
    [
        {"class_weighting": "balanced"},
        {"class_weighting": "Inverse"},
        {"class_weighting": "focal"},
        {"class_weighting": ""},
        {"class_weighting": None},
        {"class_weighting": True},
        {"sampler": "weighted"},
        {"sampler": "inverse"},
        {"sampler": "Balanced"},
        {"sampler": None},
        {"class_weights": "inverse"},
    ],
    ids=str,
)
def test_invalid_imbalance_settings_are_rejected(train: dict) -> None:
    with pytest.raises(ValidationError):
        ExperimentConfig.model_validate({**_minimal(), "train": train})


def _imbalance_base(path: Path) -> tuple[str, Path]:
    # 先頭行のコメントの最後の「（<キー> 以外は <元設定> と同一」が，この設定の元
    first = path.read_text(encoding="utf-8").splitlines()[0]
    found = re.findall(r"（(\w+) 以外は (configs/[\w./-]+\.yaml) と同一", first)
    assert found, f"元設定の記載が無い: {path.name}"
    key, base = found[-1]
    return key, REPO_ROOT / base


def test_imbalance_configs_count() -> None:
    assert len(IMBALANCE_CONFIGS) == 36


@pytest.mark.parametrize("path", IMBALANCE_CONFIGS, ids=lambda p: p.name)
def test_imbalance_config_differs_from_base_only_in_countermeasure(path: Path) -> None:
    suffix = next(s for s in IMBALANCE_SUFFIXES if path.stem.endswith(s))
    key, value = IMBALANCE_SUFFIXES[suffix]
    comment_key, base = _imbalance_base(path)
    assert comment_key == key
    assert base.is_file(), base
    v = load_config(path, environ=NO_ENV).model_dump()
    b = load_config(base, environ=NO_ENV).model_dump()
    assert v["train"][key] == value
    # 元設定は対策なし
    assert b["train"]["class_weighting"] == "none"
    assert b["train"]["sampler"] == "random"
    v["train"].pop(key)
    b["train"].pop(key)
    assert v == b
    # 生の YAML でも，元設定との違いは対策のキー 1 つだけ
    raw_v = yaml.safe_load(path.read_text(encoding="utf-8"))
    raw_b = yaml.safe_load(base.read_text(encoding="utf-8"))
    assert raw_v["train"].pop(key) == value
    assert raw_v == raw_b


def test_imbalance_configs_cover_each_base_with_three_countermeasures() -> None:
    by_base: dict[Path, set[str]] = {}
    for path in IMBALANCE_CONFIGS:
        suffix = next(s for s in IMBALANCE_SUFFIXES if path.stem.endswith(s))
        _, base = _imbalance_base(path)
        by_base.setdefault(base.resolve(), set()).add(suffix)
        # 設定名は「元設定名 + 対策の接尾辞」
        assert path.stem == base.stem + suffix
    assert len(by_base) == 12
    assert all(s == set(IMBALANCE_SUFFIXES) for s in by_base.values())