"""Collect per-(method, target, seed) JSON results into a unified table.

Walks results/ for *.json files written by the training scripts and produces:
  - results/summary.csv  : per-method × target mean ± std table
  - results/summary.tex  : LaTeX-formatted version (paper table-ready)

Run after all jobs have finished:
    python eval_collect.py --root results/
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from collections import defaultdict

import numpy as np
import pandas as pd


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="results/")
    ap.add_argument("--out_csv", default="results/summary.csv")
    ap.add_argument("--out_tex", default="results/summary.tex")
    args = ap.parse_args()

    rows = []
    for path in Path(args.root).rglob("seed*.json"):
        try:
            with open(path) as fp:
                data = json.load(fp)
            rows.append(data)
        except Exception as e:
            print(f"[skip] {path}: {e}")

    if not rows:
        print("[warn] no results found.")
        return

    df = pd.DataFrame(rows)
    grp = df.groupby(["method", "target"])
    summary = grp.agg(
        r2_mean=("r2", "mean"),
        r2_std=("r2", "std"),
        rmse_mean=("rmse", "mean"),
        rmse_std=("rmse", "std"),
        mae_mean=("mae", "mean"),
        mae_std=("mae", "std"),
        n_params=("n_params", "first"),
        n_seeds=("seed", "count"),
    ).reset_index()
    summary.to_csv(args.out_csv, index=False)
    print(f"[ok] wrote {args.out_csv} with {len(summary)} rows")

    with open(args.out_tex, "w") as fp:
        fp.write("\\begin{tabular}{llrrrr}\n\\toprule\n")
        fp.write("Method & Target & $R^2$ & RMSE & MAE & Params \\\\\n\\midrule\n")
        for _, r in summary.iterrows():
            fp.write(
                f"{r['method']} & {r['target']} & "
                f"{r['r2_mean']:.3f} $\\pm$ {r['r2_std']:.3f} & "
                f"{r['rmse_mean']:.4g} & {r['mae_mean']:.4g} & "
                f"{int(r['n_params'])} \\\\\n"
            )
        fp.write("\\bottomrule\n\\end{tabular}\n")
    print(f"[ok] wrote {args.out_tex}")


if __name__ == "__main__":
    main()