"""Summarise grid.py output. Usage: grid_summary.py <grid.jsonl>"""
import json
import sys

import pandas as pd

rows = []
for line in open(sys.argv[1]):
    r = json.loads(line)
    if "error" in r:
        print("ERROR", r["key"], r["error"])
        continue
    rows.append({**r["key"], "acc": r["accuracy"], "f1_macro": r["f1_macro"], "f1_w": r["f1_weighted"],
                 "epochs": r["epochs_completed"], "best": r["best_epoch"], "classes": r["classes_predicted"],
                 "collapsed": r["classes_predicted"] == 1, "pred": "/".join(map(str, r["pred_counts"])),
                 "min_recall": min(r["recall"])})
df = pd.DataFrame(rows)
df["acts"] = df.act1 + "/" + df.act2 + "/" + df.act3
pd.set_option("display.width", 220)
pd.set_option("display.max_rows", 300)

print(f"configs: {len(df)}  collapsed (1 class): {df.collapsed.sum()}  "
      f"all 5 classes predicted: {(df.classes == 5).sum()}  macro F1 >= 0.70: {(df.f1_macro >= 0.70).sum()}")
print("\nTOP 20 by macro F1")
print(df.sort_values("f1_macro", ascending=False).head(20)[
    ["optimizer", "batch_size", "projection_dim", "acts", "acc", "f1_macro", "f1_w", "min_recall", "best", "epochs", "pred"]
].round(3).to_string(index=False))
for factor in ["optimizer", "batch_size", "projection_dim", "act1", "act2", "act3"]:
    g = df.groupby(factor).agg(n=("f1_macro", "size"), f1_mean=("f1_macro", "mean"), f1_max=("f1_macro", "max"),
                               collapsed=("collapsed", "sum"), working=("f1_macro", lambda s: int((s >= 0.70).sum())))
    print(f"\nBY {factor}\n" + g.round(3).to_string())