"""Aggregate the sweeps into slopes, ratios and the numbers quoted in pages/."""
import glob
import json
import os
import numpy as np
import rdpg

R = 5


def load(tag):
    with open(f"outputs/main_{tag}.json") as f:
        return json.load(f)


def agg(tag):
    d = load(tag)
    runs = d["runs"]
    ns = sorted({r["n"] for r in runs})
    by = {n: [r for r in runs if r["n"] == n] for n in ns}
    out = {"tag": tag, "ngrid": ns, "reps": [len(by[n]) for n in ns], "dims": d["dims"]}

    def col(f):
        return [float(np.mean([f(r) for r in by[n]])) for n in ns]

    def colmed(f):
        return [float(np.median([f(r) for r in by[n]])) for n in ns]

    out["err_mean"] = {str(dd): col(lambda r, dd=dd: r["err"][str(dd)]) for dd in d["dims"]}
    out["err_se"] = {str(dd): [float(np.std([r["err"][str(dd)] for r in by[n]], ddof=1)
                                     / np.sqrt(len(by[n]))) for n in ns] for dd in d["dims"]}
    out["base_mean"] = col(lambda r: r["base_2inf"])
    out["trail_mean"] = {str(dd): col(lambda r, dd=dd: r["trail_2inf"][str(dd)])
                         for dd in d["dims"] if dd > R}
    out["deloc_all_mean"] = col(lambda r: r["deloc_max_all"])
    out["deloc_all_med"] = colmed(lambda r: r["deloc_max_all"])
    out["deloc_rp1_mean"] = col(lambda r: r["deloc_max_rp1"])
    out["deloc_rp1_med"] = colmed(lambda r: r["deloc_max_rp1"])
    out["signal_2inf_mean"] = col(lambda r: r["signal_2inf"])
    out["min_decomp_slack"] = float(min(min(r["decomp_slack"].values()) for r in runs))
    out["s_hat_mean"] = [float(np.mean([r["s_hat"][i] for r in by[ns[-1]]])) for i in range(8)]

    lg = np.log(np.array(ns, float))
    out["slopes"] = {k: rdpg.loglog_slope(ns, v) for k, v in out["err_mean"].items()}
    out["slopes_trail"] = {k: rdpg.loglog_slope(ns, v) for k, v in out["trail_mean"].items()}
    out["slope_base"] = rdpg.loglog_slope(ns, out["base_mean"])
    out["slope_deloc_all"] = rdpg.loglog_slope(ns, out["deloc_all_mean"])
    out["slope_deloc_rp1"] = rdpg.loglog_slope(ns, out["deloc_rp1_mean"])
    # remove the sqrt(log n) factor that delocalised max-entry statistics carry
    out["slope_deloc_rp1_delogged"] = rdpg.loglog_slope(
        ns, np.array(out["deloc_rp1_mean"]) / np.sqrt(lg))
    out["deloc_rp1_over_sqrt2logn"] = [
        float(v / np.sqrt(2 * np.log(n) / n)) for v, n in zip(out["deloc_rp1_mean"], ns)]
    out["deloc_all_over_sqrt4logn"] = [
        float(v / np.sqrt(4 * np.log(n) / n)) for v, n in zip(out["deloc_all_mean"], ns)]
    out["theorem31_bound_gamma0"] = [float(R ** 2 * np.log(n) ** 4 / np.sqrt(n)) for n in ns]
    out["theorem31_bound_holds"] = bool(all(
        v <= b for v, b in zip(out["deloc_all_mean"], out["theorem31_bound_gamma0"])))
    # over-specification: strip the predicted sqrt(log n)
    out["slopes_delogged"] = {k: rdpg.loglog_slope(ns, np.array(v) / np.sqrt(lg))
                              for k, v in out["err_mean"].items()}
    out["slopes_trail_delogged"] = {k: rdpg.loglog_slope(ns, np.array(v) / np.sqrt(lg))
                                    for k, v in out["trail_mean"].items()}
    # predicted trailing-term constant sqrt(4 sigma k log n) n^{-1/4} check
    return out


if __name__ == "__main__":
    tags = sorted(os.path.basename(p)[5:-5] for p in glob.glob("outputs/main_*.json"))
    A = {t: agg(t) for t in tags}
    with open("outputs/analysis.json", "w") as f:
        json.dump(A, f, indent=1)
    for t in tags:
        a = A[t]
        print(f"\n===== {t}  n={a['ngrid']} reps={a['reps']}")
        print("  d :  slope  R2   | delogged slope |  err(n_max)")
        for dd in a["dims"]:
            s = a["slopes"][str(dd)]
            sd = a["slopes_delogged"][str(dd)]
            print(f"  {dd:2d}: {s[0]:+.3f} {s[2]:.3f} | {sd[0]:+.3f} {sd[2]:.3f} |"
                  f" {a['err_mean'][str(dd)][-1]:.4f}")
        print(f"  base slope {a['slope_base'][0]:+.3f} (R2 {a['slope_base'][2]:.3f})"
              f"  min decomposition slack {a['min_decomp_slack']:.3e}")
        print(f"  deloc(all>r) slope {a['slope_deloc_all'][0]:+.3f}"
              f"  ratio to sqrt(4 log n/n): {[round(x,3) for x in a['deloc_all_over_sqrt4logn']]}")
        print(f"  deloc(r+1)  slope {a['slope_deloc_rp1'][0]:+.3f}"
              f" delogged {a['slope_deloc_rp1_delogged'][0]:+.3f}"
              f"  ratio to sqrt(2 log n/n): {[round(x,3) for x in a['deloc_rp1_over_sqrt2logn']]}")
        print(f"  s_hat(n_max)[:8] {[round(x,1) for x in a['s_hat_mean']]}")