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

`tias-remote pull` で手元に回収した実験ディレクトリを読む（metrics.json と config.yaml は
集計値のみで個人情報を含まない）．モデル × 学習率 × 条件（ややあり込み／除外）ごとに，
全 seed × 5 fold の fold ごとの値の平均 ± 標準偏差と，同じ seed・同じ fold どうしの
対応比較（除外 − 込み）を出す．結果は標準出力（Markdown）と runs/_cv_summary.json に書く．
色特徴分類器（tias-train color-cv．モデル名 color_*）の run も同じ表に並べる（学習率を持たない
ため，学習率は表示せず集計のキーでは 0 とする）．

使い方:
    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",
    "color_logistic": "色特徴＋ロジスティック回帰",
    "color_hist_gb": "色特徴＋勾配ブースティング",
}
# 学習率を持たないモデル（色特徴分類器．config.yaml の train.lr は使われない既定値）
MODELS_WITHOUT_LR = frozenset({"color_logistic", "color_hist_gb"})
MODEL_ORDER = list(MODEL_LABEL)
CONDITIONS = ("with_mild", "without_mild")
MEASURES = ("none", "inverse", "sqrt_inverse", "balanced")
MEASURE_LABEL = {
    "none": "",
    "inverse": "・重み inverse",
    "sqrt_inverse": "・重み sqrt",
    "balanced": "・balanced",
}
# 不均衡対策の対応比較で recall の差を示すクラス（少数クラス → 多数クラスの順）
COMPARED_RECALL_CLASSES = ("pale", "purple", "light_red", "dark_red")
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]) -> str:
    """（モデル名, 学習率, 不均衡対策）の表示名."""
    name, lr, measure = key
    if name in MODELS_WITHOUT_LR:
        extra = MEASURE_LABEL[measure].removeprefix("・")
        return MODEL_LABEL[name] + (f"（{extra}）" if extra else "")
    return f"{MODEL_LABEL[name]}（lr {lr:g}{MEASURE_LABEL[measure]}）"


def key_name(key: tuple[str, float, str]) -> str:
    """_cv_summary.json のキー（モデル名/lr/対策．学習率を持たないモデルは lr を省く）."""
    name, lr, measure = key
    return f"{name}/{measure}" if name in MODELS_WITHOUT_LR else f"{name}/lr{lr:g}/{measure}"


def imbalance_measure(metrics: dict) -> str:
    """不均衡対策の種類（対策の記録が無い古い run は対策なし）."""
    imb = metrics.get("imbalance") or {}
    if imb.get("sampler") == "balanced":
        return "balanced"
    weighting = imb.get("class_weighting", "none")
    return weighting if weighting in MEASURES else "none"


def sort_key(key: tuple[str, float, str]) -> tuple[int, float, int]:
    """表の並び順（モデル → 学習率 → 対策）."""
    return (MODEL_ORDER.index(key[0]), key[1], MEASURES.index(key[2]))


def run_lr(model: str, conf: dict) -> float:
    """集計のキーに使う学習率（学習率を持たないモデルは 0）."""
    return 0.0 if model in MODELS_WITHOUT_LR else float(conf["train"]["lr"])


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, run_lr(model, conf), imbalance_measure(metrics)),
                "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 measure_comparison(keys: list, groups: dict) -> list[str]:
    """不均衡対策あり − 対策なし の対応比較（同じモデル・学習率・条件・seed・fold）."""
    recall_headers = [f"{JP[c]} recall の差" for c in COMPARED_RECALL_CLASSES]
    headers = ["モデル", "条件", "fold 数", "common F1 の差", "上回った fold", *recall_headers]
    lines = [
        "| " + " | ".join(headers) + " |",
        "|" + " --- |" * len(headers),
    ]
    found = False
    for model in keys:
        if model[2] == "none":
            continue
        base_key = (model[0], model[1], "none")
        for cond in CONDITIONS:
            treated = {r["seed"]: r for r in groups.get((model, cond), [])}
            base = {r["seed"]: r for r in groups.get((base_key, cond), [])}
            seeds = sorted(set(treated) & set(base))
            if not seeds:
                continue
            found = True
            f1_diffs: list[float] = []
            recall_diffs: dict[str, list[float]] = {c: [] for c in COMPARED_RECALL_CLASSES}
            for seed in seeds:
                a = fold_values(base[seed], "common", "macro_f1")
                b = fold_values(treated[seed], "common", "macro_f1")
                f1_diffs += [y - x for x, y in zip(a, b, strict=True)]
                for c in COMPARED_RECALL_CLASSES:
                    ra = fold_class_recall(base[seed], "common", c)
                    rb = fold_class_recall(treated[seed], "common", c)
                    if len(ra) == len(rb):
                        recall_diffs[c] += [y - x for x, y in zip(ra, rb, strict=True)]
            wins = sum(1 for d in f1_diffs if d > 0)
            n = len(f1_diffs)
            cells = [label(model), CONDITION_LABEL[cond], str(n), fmt(mean_std(f1_diffs))]
            cells.append(f"{wins} / {n}")
            cells += [fmt(mean_std(recall_diffs[c]), 2) for c in COMPARED_RECALL_CLASSES]
            lines.append("| " + " | ".join(cells) + " |")
    if not found:
        return []
    lines.append(
        "\n※ 不均衡対策あり − 対策なし（同じモデル・学習率・条件・seed・fold どうし）．"
        "recall は common でのクラス別再現率．\n"
    )
    return lines


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=sort_key)
    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"{key_name(model)}/{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[key_name(model)] = {
                "diff": mean_std(diffs),
                "wins": wins,
                "n": n,
            }
    lines.append(
        "\n※ 同じ seed なら「込み」と「除外」の分割は完全に同じ（SPEC_02）なので，"
        "fold ごとの差を取っている．\n"
    )

    lines += measure_comparison(keys, groups)

    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"))