#!/usr/bin/env python3
"""Summarize repeated LAB-Bench runs by averaging metrics."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from statistics import mean, stdev
from typing import Any


def load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def summarize_group(paths: list[Path]) -> dict[str, Any]:
    payloads = [load_json(path) for path in paths]
    metric_keys = sorted(payloads[0]["metrics_all"])
    summary = {
        "run_count": len(payloads),
        "files": [str(path) for path in paths],
        "metrics_all_mean": {},
        "metrics_all_std": {},
    }
    for key in metric_keys:
        values = [payload["metrics_all"][key] for payload in payloads if isinstance(payload["metrics_all"].get(key), (int, float))]
        if not values:
            continue
        summary["metrics_all_mean"][key] = mean(values)
        summary["metrics_all_std"][key] = stdev(values) if len(values) > 1 else 0.0
    return summary


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--results-root", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()

    groups: dict[str, list[Path]] = {}
    for path in sorted(args.results_root.glob("*.json")):
        stem = path.stem
        if "_run" not in stem:
            continue
        group = stem.rsplit("_run", 1)[0]
        groups.setdefault(group, []).append(path)

    payload = {"results_root": str(args.results_root), "groups": {}}
    for group, paths in sorted(groups.items()):
        payload["groups"][group] = summarize_group(paths)

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
    print(json.dumps(payload["groups"], indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())