"""Put the results of run.py side by side.

    python compare.py [results_dir]

Prints a table of all runs found in results_dir (default results/) and, for
every problem and case run with both backends, the speedup and the relative
difference of the surface fields. The table is also written to results_dir/comparison.md.
"""

import glob
import json
import os
import sys

import numpy as np


def fmt(x, spec):
    return "" if x is None else format(x, spec)


def main():
    rdir = sys.argv[1] if len(sys.argv) > 1 else os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")
    runs = {}
    for f in sorted(glob.glob(os.path.join(rdir, "*", "summary.json"))):
        s = json.load(open(f))
        s["dir"] = os.path.dirname(f)
        runs[(s.get("problem", "room"), s["case"], s["backend"])] = s
    if not runs:
        print(f"no results in {rdir}")
        return

    lines = ["| problem | case | k | points | backend | machine | setup (s) | solve (s) | setup + solve (s) | "
             "iterations | self-test error | energy mismatch |",
             "|---|---|---|---|---|---|---|---|---|---|---|---|"]
    order = {"small": 0, "medium": 1, "large": 2}
    sort_key = lambda pc: (pc[0] != "room", pc[0], order.get(pc[1], 9))
    for (problem, case, backend), s in sorted(runs.items(), key=lambda kv: (sort_key(kv[0][:2]), kv[0][2])):
        m = s["machine"]
        where = m.get("gpu", "").split(",")[0] if backend == "gpu" else f"{m.get('cpu', '?')}, {m.get('cpu_count')} threads"
        lines.append(f"| {problem} | {case} | {s['k']:.4g} | {s['npts']:,} | {backend} | {where} | {s['setup_s']:.2f} | "
                     f"{s['solve_s']:.2f} | {s['setup_plus_solve_s']:.2f} | {s['niter']} | "
                     f"{fmt(s.get('selftest_err'), '.1e')} | {s['energy_mismatch']:.1e} |")

    lines += ["", "| problem | case | speedup (setup + solve) | surface field difference |", "|---|---|---|---|"]
    for problem, case in sorted({pc[:2] for pc in runs}, key=sort_key):
        if (problem, case, "gpu") in runs and (problem, case, "fmm3dbie") in runs:
            g, c = runs[(problem, case, "gpu")], runs[(problem, case, "fmm3dbie")]
            ug = np.load(os.path.join(g["dir"], "field.npz"))["u_wall"]
            uc = np.load(os.path.join(c["dir"], "field.npz"))["u_wall"]
            diff = np.linalg.norm(ug - uc) / np.linalg.norm(uc) if ug.shape == uc.shape else float("nan")
            lines.append(f"| {problem} | {case} | {c['setup_plus_solve_s'] / g['setup_plus_solve_s']:.0f}x | {diff:.1e} |")
    text = "\n".join(lines)
    print(text)
    with open(os.path.join(rdir, "comparison.md"), "w") as f:
        f.write(text + "\n")


if __name__ == "__main__":
    main()