"""回収した交差検証の結果（runs/*_cv_*/metrics.json）を集計して Markdown の表にする.

`tias-remote pull` で手元に回収した実験ディレクトリを読む（metrics.json と config.yaml は
集計値のみで個人情報を含まない）．モデル × 学習率 × 条件（ややあり込み／除外）ごとに，
全 seed × 5 fold の fold ごとの値の平均 ± 標準偏差と，同じ seed・同じ fold どうしの
対応比較（除外 − 込み）を出す．結果は標準出力（Markdown）と runs/_cv_summary.json に書く．

使い方:
    uv run python scripts/summarize_cv.py [runs]
"""

import json
import math
import sys
from collections import defaultdict
from pathlib import Path

import yaml

CLASSES = ["pale", "light_red", "red", "dark_red", "purple"]
JP = {"pale": "淡白", "light_red": "淡紅", "red": "紅", "dark_red": "暗赤紅", "purple": "紫"}
MODEL_LABEL = {
    "convnext_tiny": "ConvNeXt-Tiny",
    "efficientnet_b0": "EfficientNet-B0",
    "resnet50": "ResNet-50",
    "vit_small_patch16_224": "ViT-Small",
}
MODEL_ORDER = list(MODEL_LABEL)
CONDITIONS = ("with_mild", "without_mild")
CONDITION_LABEL = {"with_mild": "込み", "without_mild": "除外"}

MeanStd = tuple[float, float, int]


def mean_std(xs: list[float]) -> MeanStd:
    """平均・不偏標準偏差・件数を返す."""
    xs = [x for x in xs if x is not None and not math.isnan(x)]
    if not xs:
        return (float("nan"), float("nan"), 0)
    m = sum(xs) / len(xs)
    s = math.sqrt(sum((x - m) ** 2 for x in xs) / (len(xs) - 1)) if len(xs) > 1 else 0.0
    return (m, s, len(xs))


def fmt(ms: MeanStd, digits: int = 3) -> str:
    """平均 ± 標準偏差を文字列にする."""
    m, s, n = ms
    return "—" if n == 0 else f"{m:.{digits}f} ± {s:.{digits}f}"


def label(key: tuple[str, float]) -> str:
    """（モデル名, 学習率）の表示名."""
    name, lr = key
    return f"{MODEL_LABEL[name]}（lr {lr:g}）"


def load_runs(root: Path) -> list[dict]:
    """実験ディレクトリを読む（対象外のモデルは飛ばす）."""
    runs = []
    for mpath in sorted(root.glob("*_cv_*/metrics.json")):
        metrics = json.loads(mpath.read_text(encoding="utf-8"))
        doc = yaml.safe_load((mpath.parent / "config.yaml").read_text(encoding="utf-8"))
        conf = doc.get("config", doc)
        model = metrics["environment"]["model_name"]
        if model not in MODEL_LABEL:
            continue
        cond = "with_mild" if "mild" in metrics["severities_trained"] else "without_mild"
        runs.append(
            {
                "run_id": metrics["run_id"],
                "model": (model, float(conf["train"]["lr"])),
                "cond": cond,
                "seed": conf["split"]["seed"],
                "m": metrics,
            }
        )
    return runs


def fold_values(run: dict, subset: str, key: str) -> list[float]:
    """fold ごとの指標の値."""
    return [f["subsets"][subset][key] for f in run["m"]["folds"]]


def fold_class_recall(run: dict, subset: str, cls: str) -> list[float]:
    """fold ごとのクラス別再現率（正解に現れた fold のみ）."""
    out = []
    for f in run["m"]["folds"]:
        per_class = f["subsets"][subset].get("per_class", {})
        if cls in per_class:
            out.append(per_class[cls]["recall"])
    return out


def mild_to_normal(run: dict) -> float:
    """mild の画像を淡紅と予測した割合（全 fold をまとめた値）."""
    return run["m"]["pooled"]["mild"]["prediction_breakdown"]["overall"]["light_red"]


def pooled_folds(runs: list[dict], subset: str, key: str) -> MeanStd:
    """複数 run の fold ごとの値をまとめて平均 ± 標準偏差にする."""
    return mean_std([v for r in runs for v in fold_values(r, subset, key)])


def main(root: Path) -> None:
    """集計して表を出力する."""
    runs = load_runs(root)
    groups: dict = defaultdict(list)
    for r in runs:
        groups[(r["model"], r["cond"])].append(r)
    keys = sorted({r["model"] for r in runs}, key=lambda k: (MODEL_ORDER.index(k[0]), k[1]))
    seeds = ", ".join(sorted({str(r["seed"]) for r in runs}))

    lines = [f"対象: {len(runs)} 本（seed {seeds}）\n"]
    lines.append(
        "| モデル | 条件 | seed 数 | common macro F1 | common 正解率 | all macro F1"
        " | mild→淡紅 | 1 本の所要（秒） |"
    )
    lines.append("| --- | --- | --- | --- | --- | --- | --- | --- |")
    summary = {}
    for model in keys:
        for cond in CONDITIONS:
            rs = groups.get((model, cond), [])
            if not rs:
                continue
            row = {
                "common_f1": pooled_folds(rs, "common", "macro_f1"),
                "common_acc": pooled_folds(rs, "common", "accuracy"),
                "all_f1": pooled_folds(rs, "all", "macro_f1"),
                "mild_to_normal": mean_std([mild_to_normal(r) for r in rs]),
                "seconds": mean_std([r["m"]["timing"]["total_seconds"] for r in rs]),
            }
            cells = [
                label(model),
                CONDITION_LABEL[cond],
                str(len(rs)),
                fmt(row["common_f1"]),
                fmt(row["common_acc"]),
                fmt(row["all_f1"]),
                fmt(row["mild_to_normal"], 2),
                fmt(row["seconds"], 0),
            ]
            lines.append("| " + " | ".join(cells) + " |")
            summary[f"{model[0]}/lr{model[1]:g}/{cond}"] = row
    lines.append(
        "\n※ macro F1・正解率は全 seed × 5 fold の fold ごとの値の平均 ± 標準偏差．"
        "mild→淡紅 は run ごとの（全 fold をまとめた）割合の平均 ± 標準偏差．\n"
    )

    lines.append("| モデル | 条件 | " + " | ".join(JP[c] for c in CLASSES) + " |")
    lines.append("| --- | --- | " + " | ".join("---" for _ in CLASSES) + " |")
    for model in keys:
        for cond in CONDITIONS:
            rs = groups.get((model, cond), [])
            if not rs:
                continue
            cells = [
                fmt(mean_std([v for r in rs for v in fold_class_recall(r, "common", c)]), 2)
                for c in CLASSES
            ]
            lines.append(
                f"| {label(model)} | {CONDITION_LABEL[cond]} | " + " | ".join(cells) + " |"
            )
    lines.append(
        "\n※ common（「あり」と淡紅）での各クラスの再現率（fold ごとの値の平均 ± 標準偏差）．\n"
    )

    lines.append(
        "| モデル | 対応のある fold 数 | common F1 の差（除外 − 込み） | 除外が上回った fold |"
    )
    lines.append("| --- | --- | --- | --- |")
    paired = {}
    for model in keys:
        with_mild = {r["seed"]: r for r in groups.get((model, "with_mild"), [])}
        without = {r["seed"]: r for r in groups.get((model, "without_mild"), [])}
        diffs = []
        for seed in sorted(set(with_mild) & set(without)):
            a = fold_values(with_mild[seed], "common", "macro_f1")
            b = fold_values(without[seed], "common", "macro_f1")
            diffs += [y - x for x, y in zip(a, b, strict=True)]
        if diffs:
            wins = sum(1 for d in diffs if d > 0)
            n = len(diffs)
            lines.append(f"| {label(model)} | {n} | {fmt(mean_std(diffs))} | {wins} / {n} |")
            paired[f"{model[0]}/lr{model[1]:g}"] = {
                "diff": mean_std(diffs),
                "wins": wins,
                "n": n,
            }
    lines.append(
        "\n※ 同じ seed なら「込み」と「除外」の分割は完全に同じ（SPEC_02）なので，"
        "fold ごとの差を取っている．\n"
    )

    print("\n".join(lines))
    out = {"summary": summary, "paired": paired, "runs": [r["run_id"] for r in runs]}
    (root / "_cv_summary.json").write_text(
        json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8"
    )


if __name__ == "__main__":
    main(Path(sys.argv[1] if len(sys.argv) > 1 else "runs"))