"""汇总消融实验结果：对比 baseline / +moe / +llm / +moe+llm 的 nested CV 指标。

读取每个 run 目录下的 summary.json（nested_cv 产出），输出对比表 CSV 并打印。

用法：
    python scripts/summarize_ablation.py \
        --run baseline=models/nested_cv/baseline/<timestamp> \
        --run +moe=models/nested_cv/moe/<timestamp> \
        --run +llm=models/nested_cv/llm/<timestamp> \
        --run +moe+llm=models/nested_cv/both/<timestamp> \
        --out reports/ablation_summary.csv
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import pandas as pd


def load_summary_stats(run_dir: Path) -> dict:
    """读取 run 目录的 summary.json，返回 summary_stats。"""
    summary_path = run_dir / "summary.json"
    if not summary_path.exists():
        raise FileNotFoundError(f"summary.json not found in {run_dir}")
    with open(summary_path) as f:
        return json.load(f).get("summary_stats", {})


def main() -> None:
    parser = argparse.ArgumentParser(description="Summarize ablation nested-CV runs")
    parser.add_argument(
        "--run", action="append", default=[],
        help="LABEL=run_dir，可重复传入多个变体",
    )
    parser.add_argument("--out", type=str, default="reports/ablation_summary.csv")
    args = parser.parse_args()

    if not args.run:
        parser.error("至少需要一个 --run LABEL=run_dir")

    rows = []
    for spec in args.run:
        label, sep, run_dir = spec.partition("=")
        if not sep:
            parser.error(f"--run 需为 LABEL=run_dir 形式，收到: {spec}")
        stats = load_summary_stats(Path(run_dir))
        row = {"variant": label, "run_dir": run_dir}
        for task, metrics in stats.items():
            for metric_name, value in metrics.items():
                row[f"{task}.{metric_name}"] = value
        rows.append(row)

    df = pd.DataFrame(rows)
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(out, index=False)

    print(df.to_string(index=False))
    print(f"\nSaved ablation summary to {out}")


if __name__ == "__main__":
    main()