"""tias-train cv のクラス不均衡の対策の結合テスト（SPEC_02「クラス不均衡」→「対策の設定」）.

- `metrics.json` に使った対策（トップレベル）と fold ごとのクラス重み（または抜き出し確率）が入る
- 既存のキー名・構造は変わらない（`scripts/summarize_cv.py` 等が依存する）
- 既定のままなら従来と同じ学習: 対策のキーを書かない設定と，既定値を明示した設定で指標が一致する
- 回収対象の出力（ログを含む）に患者 ID・ファイル名が出ない

合成データ（scripts/make_synthetic_dataset.py）・resnet18・pretrained なし・CPU・1 エポックで回す．
"""

from __future__ import annotations

import contextlib
import importlib.util
import io
import json
import logging
import sys
from pathlib import Path
from types import ModuleType

import numpy as np
import pytest
import yaml

import tias_judge_model.cli.train as cli
from tias_judge_model.data.class_balance import compute_class_weights
from tias_judge_model.data.classes import CLASS_SLUGS, NUM_CLASSES
from tias_judge_model.data.manifest import build_manifest_records, write_manifest

pytestmark = pytest.mark.slow

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

# 変更前（対策の導入前）の metrics.json のキー
LEGACY_TOP_KEYS = {
    "run_id",
    "git_commit",
    "classes",
    "severities_trained",
    "subsets",
    "n_folds",
    "best_epoch",
    "mean_std",
    "pooled",
    "confusion_matrix_sum",
    "folds",
    "timing",
    "environment",
}
LEGACY_FOLD_KEYS = {
    "fold",
    "seed",
    "best_epoch",
    "best_val_macro_f1",
    "counts",
    "timing",
    "subsets",
}


def _load_script() -> ModuleType:
    spec = importlib.util.spec_from_file_location("make_synthetic_dataset_train_imb", 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 _train(**extra: str) -> dict:
    return {
        "image_size": 32,
        "batch_size": 8,
        "max_epochs": 2,
        "warmup_epochs": 1,
        "inner_val_ratio": 0.2,
        "num_workers": 0,
        "amp": False,
        **extra,
    }


# 設定名 → (severities, train の追加キー)
RUNS = {
    "cv_default": (["strong", "mild", "normal"], {}),
    "cv_explicit_default": (
        ["strong", "mild", "normal"],
        {"class_weighting": "none", "sampler": "random"},
    ),
    "cv_cwsqrt_nomild": (["strong", "normal"], {"class_weighting": "sqrt_inverse"}),
    "cv_cwinv": (["strong", "mild", "normal"], {"class_weighting": "inverse"}),
    "cv_balanced": (["strong", "mild", "normal"], {"sampler": "balanced"}),
}


@pytest.fixture(scope="module")
def runs(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]:
    tmp = tmp_path_factory.mktemp("train_cli_imb")
    data_root = tmp / "dataset"
    _load_script().generate(data_root, per_dir=PER_DIR, seed=0)
    for n, png in enumerate(sorted((data_root / "images").rglob("*.png"))):
        png.rename(png.with_name(f"{SECRET}{n:04d}.png"))
    write_manifest(build_manifest_records(data_root / "images"), data_root / "manifest.csv")

    out: dict[str, dict] = {}
    mp = pytest.MonkeyPatch()
    mp.delenv("TIAS_DATA_ROOT", raising=False)
    mp.setenv("TIAS_GIT_COMMIT", "abc1234")
    root_logger = logging.getLogger()
    old_level = root_logger.level
    root_logger.setLevel(logging.INFO)
    try:
        for name, (severities, extra) in RUNS.items():
            config = tmp / f"{name}.yaml"
            config.write_text(
                yaml.safe_dump(
                    {
                        "data": {"root": data_root.as_posix(), "severities": severities},
                        "split": {"seed": 3, "n_folds": N_FOLDS},
                        "model": {"name": "resnet18", "pretrained": False},
                        "train": _train(**extra),
                    }
                ),
                encoding="utf-8",
            )
            out_root = tmp / f"runs_{name}"
            stdout, stderr = io.StringIO(), io.StringIO()
            handler = logging.StreamHandler(stderr)
            root_logger.addHandler(handler)
            try:
                with contextlib.redirect_stdout(stdout):
                    rc = cli.main(["cv", "--config", str(config), "--out-root", str(out_root)])
            finally:
                root_logger.removeHandler(handler)
            run_dir = next(p for p in out_root.iterdir() if p.is_dir())
            out[name] = {
                "rc": rc,
                "dir": run_dir,
                "stdout": stdout.getvalue(),
                "stderr": stderr.getvalue(),
                "metrics": json.loads((run_dir / "metrics.json").read_text(encoding="utf-8")),
            }
    finally:
        root_logger.setLevel(old_level)
        mp.undo()
    return out


def test_all_runs_succeed(runs: dict) -> None:
    assert all(r["rc"] == 0 for r in runs.values())


@pytest.mark.parametrize("name", list(RUNS))
def test_legacy_keys_are_kept(runs: dict, name: str) -> None:
    m = runs[name]["metrics"]
    assert set(m) == LEGACY_TOP_KEYS | {"imbalance"}
    for fold in m["folds"]:
        assert set(fold) == LEGACY_FOLD_KEYS | {"class_balance"}
    # summarize_cv.py が読むキー
    assert m["environment"]["model_name"] == "resnet18"
    assert isinstance(m["severities_trained"], list)
    assert "light_red" in m["pooled"]["mild"]["prediction_breakdown"]["overall"]
    assert m["timing"]["total_seconds"] > 0
    for fold in m["folds"]:
        for subset in ("all", "common", "mild"):
            assert {"macro_f1", "accuracy", "per_class"} <= set(fold["subsets"][subset])


@pytest.mark.parametrize(
    ("name", "weighting", "sampler"),
    [
        ("cv_default", "none", "random"),
        ("cv_explicit_default", "none", "random"),
        ("cv_cwsqrt_nomild", "sqrt_inverse", "random"),
        ("cv_cwinv", "inverse", "random"),
        ("cv_balanced", "none", "balanced"),
    ],
)
def test_top_level_records_countermeasure(
    runs: dict, name: str, weighting: str, sampler: str
) -> None:
    imb = runs[name]["metrics"]["imbalance"]
    assert imb["class_weighting"] == weighting
    assert imb["sampler"] == sampler


def _counts(entry: dict) -> list[int]:
    return [entry["class_counts"][s] for s in CLASS_SLUGS]


@pytest.mark.parametrize(
    ("name", "method"), [("cv_cwsqrt_nomild", "sqrt_inverse"), ("cv_cwinv", "inverse")]
)
def test_fold_class_weights_recorded(runs: dict, name: str, method: str) -> None:
    for fold in runs[name]["metrics"]["folds"]:
        cb = fold["class_balance"]
        assert cb["class_weighting"] == method
        assert list(cb["class_weights"]) == list(CLASS_SLUGS)
        assert cb["class_draw_probabilities"] is None
        counts = _counts(cb)
        expected = compute_class_weights(np.repeat(np.arange(NUM_CLASSES), counts), method)
        got = np.array([cb["class_weights"][s] for s in CLASS_SLUGS])
        np.testing.assert_allclose(got, expected, atol=1e-5)
        # 学習用データ（data.severities 適用後）のクラス別枚数
        groups = fold["counts"]["images_by_group"]["train"]
        for slug in CLASS_SLUGS:
            assert cb["class_counts"][slug] == sum(
                n for g, n in groups.items() if g.rsplit("_", 1)[0] == slug
            )


def test_without_mild_counts_exclude_mild(runs: dict) -> None:
    for fold in runs["cv_cwsqrt_nomild"]["metrics"]["folds"]:
        groups = fold["counts"]["images_by_group"]["train"]
        assert sum(n for g, n in groups.items() if g.endswith("_mild")) == 0
        assert sum(fold["class_balance"]["class_counts"].values()) == sum(groups.values())


def test_fold_draw_probabilities_recorded(runs: dict) -> None:
    for fold in runs["cv_balanced"]["metrics"]["folds"]:
        cb = fold["class_balance"]
        assert cb["sampler"] == "balanced"
        assert cb["class_weights"] is None
        probs = np.array([cb["class_draw_probabilities"][s] for s in CLASS_SLUGS])
        np.testing.assert_allclose(probs, 1 / NUM_CLASSES, atol=1e-5)


def test_default_run_records_no_weights(runs: dict) -> None:
    for fold in runs["cv_default"]["metrics"]["folds"]:
        cb = fold["class_balance"]
        # 記録が無い（None）か，対策なしであることが分かる形
        if cb is None:
            continue
        assert cb["class_weighting"] == "none"
        assert cb["sampler"] == "random"
        assert cb["class_weights"] is None
        assert cb["class_draw_probabilities"] is None


def _strip_volatile(m: dict) -> dict:
    m = json.loads(json.dumps(m))
    for key in ("run_id", "timing"):
        m.pop(key)
    for fold in m["folds"]:
        fold.pop("timing")
    return m


def test_explicit_defaults_give_same_results_as_omitted_keys(runs: dict) -> None:
    a = _strip_volatile(runs["cv_default"]["metrics"])
    b = _strip_volatile(runs["cv_explicit_default"]["metrics"])
    assert a == b
    ha = (runs["cv_default"]["dir"] / "history.csv").read_text(encoding="utf-8").splitlines()
    hb = (runs["cv_explicit_default"]["dir"] / "history.csv").read_text(encoding="utf-8")
    # 所要秒数の列を除けば学習の経過も一致する
    header = ha[0].split(",")
    keep = [i for i, c in enumerate(header) if "second" not in c]
    rows_a = [[r.split(",")[i] for i in keep] for r in ha]
    rows_b = [[r.split(",")[i] for i in keep] for r in hb.splitlines()]
    assert rows_a == rows_b


def test_countermeasure_changes_training(runs: dict) -> None:
    # 対策を入れると学習損失の経過が変わる（対策が実際に効いていること）
    def train_losses(name: str) -> list[str]:
        text = (runs[name]["dir"] / "history.csv").read_text(encoding="utf-8").splitlines()
        col = text[0].split(",").index("train_loss")
        return [r.split(",")[col] for r in text[1:]]

    assert train_losses("cv_cwinv") != train_losses("cv_default")
    assert train_losses("cv_balanced") != train_losses("cv_default")


@pytest.mark.parametrize("name", list(RUNS))
def test_outputs_have_no_patient_ids(runs: dict, name: str) -> None:
    r = runs[name]
    for fname in ("metrics.json", "log.txt", "history.csv", "config.yaml"):
        assert SECRET not in (r["dir"] / fname).read_text(encoding="utf-8"), fname
    assert SECRET not in r["stdout"]
    assert SECRET not in r["stderr"]