"""Aggregate the timing fields already recorded by the final experiments."""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

import pandas as pd


def _identity(path: Path, report: dict) -> tuple[str, str, int | None]:
    value = str(path).lower()
    protocol = "without_vb_in_range" if "no-vb-in-range" in value else "with_vb_in_range"
    match = re.search(r"seed[-_](\d+)", value)
    seed = int(match.group(1)) if match else report.get("config", {}).get("seed")
    if "b2_1_random_forest" in value:
        model = "B2.1 Random Forest"
    elif "b2_2_xgboost" in value:
        model = "B2.2 XGBoost"
    elif "baseline3" in value:
        model = "B3 EdgeMLP"
    elif "graphsage" in value:
        model = "GraphSAGE"
    elif report.get("baseline") == "B1":
        model, protocol, seed = "B1 PEP 440", "rule", None
    else:
        raise ValueError(f"Cannot identify timing report: {path}")
    return model, protocol, seed


def _timing(report: dict) -> tuple[float, float, float | None]:
    if report.get("baseline") == "B1":
        inference = report["inference"]
    elif "inference" in report.get("test", {}):
        inference = report["test"]["inference"]
    else:
        inference = report["inference"]["test"]
    fit = report.get("fit_seconds")
    if fit is None:
        fit = report.get("training", {}).get("fit_seconds")
    return float(inference["seconds_per_pair"]), float(inference["total_seconds"]), fit


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--metrics-root", action="append", type=Path, default=[])
    parser.add_argument("--metrics-file", action="append", type=Path, default=[])
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()

    paths = list(args.metrics_file)
    for root in args.metrics_root:
        paths.extend(root.glob("**/metrics.json"))
    rows = []
    for path in sorted(set(paths)):
        report = json.loads(path.read_text(encoding="utf-8"))
        try:
            seconds_per_pair, total_seconds, fit_seconds = _timing(report)
            model, protocol, seed = _identity(path, report)
        except (KeyError, TypeError, ValueError):
            continue
        environment = report.get("environment", {})
        rows.append(
            {
                "model": model,
                "protocol": protocol,
                "seed": seed,
                "seconds_per_pair": seconds_per_pair,
                "test_total_seconds": total_seconds,
                "fit_seconds": fit_seconds,
                "device": environment.get("device", "cpu"),
                "python": environment.get("python"),
                "source": str(path),
            }
        )
    frame = pd.DataFrame(rows)
    if frame.empty:
        raise ValueError("No timing records found")
    grouped = frame.groupby(["model", "protocol"], dropna=False)
    summary = grouped.agg(
        runs=("seconds_per_pair", "size"),
        seconds_per_pair_mean=("seconds_per_pair", "mean"),
        seconds_per_pair_std=("seconds_per_pair", "std"),
        seconds_per_pair_min=("seconds_per_pair", "min"),
        seconds_per_pair_max=("seconds_per_pair", "max"),
        fit_seconds_mean=("fit_seconds", "mean"),
        fit_seconds_std=("fit_seconds", "std"),
    ).reset_index()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    frame.to_csv(args.output_dir / "per_run.csv", index=False)
    summary.to_csv(args.output_dir / "summary.csv", index=False)
    groups = [
        {key: (None if pd.isna(value) else value) for key, value in record.items()}
        for record in summary.to_dict("records")
    ]
    payload = {
        "scope": "descriptive aggregation of timings recorded during original runs",
        "comparability_warning": (
            "B2 ran on CPU and B3/GraphSAGE on CUDA; Python and library versions also differ. "
            "Use these values descriptively, not as a controlled hardware benchmark."
        ),
        "groups": groups,
    }
    (args.output_dir / "summary.json").write_text(
        json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8"
    )
    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()