"""Aggregate results/*.json into summary tables (+ figures if matplotlib is
present). Run on any box with access to the shared NFS home.
Usage: python experiments/analyze.py [results_dir]
"""
import glob
import json
import os
import sys
from collections import defaultdict

RES = sys.argv[1] if len(sys.argv) > 1 else "results"

try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    HAVE_MPL = True
except Exception:
    HAVE_MPL = False


def load_all(res):
    runs = []
    for p in sorted(glob.glob(os.path.join(res, "*.json"))):
        try:
            with open(p) as f:
                d = json.load(f)
            d["_path"] = p
            d["_tag"] = os.path.basename(p)[:-5]
            runs.append(d)
        except Exception as e:
            print("skip", p, e)
    return runs


def mean_final_cos(run):
    f = run.get("final", {})
    c = f.get("cos_r_negg")
    return sum(c) / len(c) if c else None


def group(runs, keyfn):
    g = defaultdict(list)
    for r in runs:
        k = keyfn(r)
        if k is not None:
            g[k].append(r)
    return g


def agg(vals):
    vals = [v for v in vals if v is not None]
    if not vals:
        return (None, None, 0)
    m = sum(vals) / len(vals)
    sd = (sum((v - m) ** 2 for v in vals) / len(vals)) ** 0.5
    return (m, sd, len(vals))


def fmt(m, sd, n):
    if m is None:
        return "   n/a  "
    return f"{m:.3f}±{sd:.3f}(n{n})"


def table_main(runs):
    print("\n================ MAIN COMPARISON (mnist depth3 width256) ================")
    sel = [r for r in runs if r["args"].get("depth") == 3 and r["args"].get("width") == 256
           and r["args"].get("dataset") == "mnist" and r["args"].get("nuis_rho", 0) == 0
           and r["_tag"].startswith(("main_", "A_main_"))]
    by_mode = group(sel, lambda r: r["args"]["mode"])
    print(f"{'mode':6s} {'test_acc':22s} {'mean_cos(r,-g)':22s}")
    for mode in ["bp", "dfa", "sdil"]:
        rs = by_mode.get(mode, [])
        acc = agg([r["final"].get("test_acc") for r in rs])
        cos = agg([mean_final_cos(r) for r in rs])
        print(f"{mode:6s} {fmt(*acc):22s} {fmt(*cos):22s}")


def table_nuisance(runs):
    print("\n================ RESIDUALIZATION UNDER NUISANCE ================")
    print("(sdil; does subtracting the soma-predictable baseline preserve alignment as rho grows?)")
    sel = [r for r in runs if r["_tag"].startswith(("nuis_", "A_nuis_"))]
    by = group(sel, lambda r: (r["args"].get("nuis_rho"), r["args"].get("use_residual")))
    rhos = sorted({k[0] for k in by})
    print(f"{'rho':>5s} | {'residual=1 acc':22s} {'cos':16s} | {'residual=0 acc':22s} {'cos':16s}")
    for rho in rhos:
        r1 = by.get((rho, 1), [])
        r0 = by.get((rho, 0), [])
        a1 = agg([r["final"].get("test_acc") for r in r1]); c1 = agg([mean_final_cos(r) for r in r1])
        a0 = agg([r["final"].get("test_acc") for r in r0]); c0 = agg([mean_final_cos(r) for r in r0])
        print(f"{rho:5.1f} | {fmt(*a1):22s} {fmt(*c1):16s} | {fmt(*a0):22s} {fmt(*c0):16s}")


def table_trainedA(runs):
    print("\n================ TRAINED vs FIXED APICAL PATHWAY ================")
    fixed = [r for r in runs if "fixedA" in r["_tag"]]
    af = agg([r["final"].get("test_acc") for r in fixed]); cf = agg([mean_final_cos(r) for r in fixed])
    print(f"{'fixed-A (DFA-like)':22s} acc {fmt(*af):22s} cos {fmt(*cf):16s}")
    tr = [r for r in runs if "trainedA_nd" in r["_tag"]]
    by = group(tr, lambda r: r["args"].get("pert_ndirs"))
    for nd in sorted(by):
        rs = by[nd]
        a = agg([r["final"].get("test_acc") for r in rs]); c = agg([mean_final_cos(r) for r in rs])
        print(f"trained-A ndirs={nd:<3d}     acc {fmt(*a):22s} cos {fmt(*c):16s}")


def table_depth(runs):
    print("\n================ DEPTH SCALING ================")
    sel = [r for r in runs if r["_tag"].startswith("depth_")]
    if not sel:
        return
    by = group(sel, lambda r: (r["args"]["mode"], r["args"]["depth"]))
    depths = sorted({k[1] for k in by})
    print(f"{'depth':>6s} | " + " | ".join(f"{m:>26s}" for m in ["bp", "dfa", "sdil"]))
    for d in depths:
        cells = []
        for m in ["bp", "dfa", "sdil"]:
            rs = by.get((m, d), [])
            a = agg([r["final"].get("test_acc") for r in rs])
            c = agg([mean_final_cos(r) for r in rs])
            cells.append(f"a{fmt(*a)[:5]} c{('%.2f'%c[0]) if c[0] is not None else ' n/a'}")
        print(f"{d:6d} | " + " | ".join(f"{c:>26s}" for c in cells))


def plots(runs):
    if not HAVE_MPL:
        print("\n[matplotlib not available -> skipping figures]")
        return
    figdir = os.path.join(RES, "figs")
    os.makedirs(figdir, exist_ok=True)

    # Fig 1: alignment trajectory sdil vs dfa (seed0 main)
    plt.figure(figsize=(6, 4))
    for tag, lab in [("A_main_sdil_s0", "SDIL (trained A)"), ("A_main_dfa_s0", "DFA (fixed A)"),
                     ("main_sdil_mnist_s0", "SDIL"), ("main_dfa_mnist_s0", "DFA")]:
        r = next((x for x in runs if x["_tag"] == tag), None)
        if not r:
            continue
        xs = [s["step"] for s in r["steps"] if "cos_r_negg" in s]
        ys = [sum(s["cos_r_negg"]) / len(s["cos_r_negg"]) for s in r["steps"] if "cos_r_negg" in s]
        if xs:
            plt.plot(xs, ys, label=lab)
    plt.xlabel("step"); plt.ylabel("mean cos(r, -grad)"); plt.title("Gradient alignment over training")
    plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
    plt.savefig(os.path.join(figdir, "alignment_traj.png"), dpi=120); plt.close()

    # Fig 2: acc & alignment vs rho (residual vs not)
    sel = [r for r in runs if r["_tag"].startswith(("nuis_", "A_nuis_"))]
    if sel:
        by = group(sel, lambda r: (r["args"].get("nuis_rho"), r["args"].get("use_residual")))
        rhos = sorted({k[0] for k in by})
        fig, ax = plt.subplots(1, 2, figsize=(10, 4))
        for ur, lab in [(1, "residual (SDIL)"), (0, "raw apical")]:
            accs = [agg([r["final"].get("test_acc") for r in by.get((rho, ur), [])])[0] for rho in rhos]
            coss = [agg([mean_final_cos(r) for r in by.get((rho, ur), [])])[0] for rho in rhos]
            ax[0].plot(rhos, accs, marker="o", label=lab)
            ax[1].plot(rhos, coss, marker="o", label=lab)
        ax[0].set_xlabel("nuisance rho"); ax[0].set_ylabel("test acc"); ax[0].set_title("Accuracy vs nuisance")
        ax[1].set_xlabel("nuisance rho"); ax[1].set_ylabel("mean cos(r,-grad)"); ax[1].set_title("Alignment vs nuisance")
        for a in ax:
            a.legend(); a.grid(alpha=.3)
        plt.tight_layout(); plt.savefig(os.path.join(figdir, "nuisance.png"), dpi=120); plt.close()

    # Fig 3: alignment vs depth
    sel = [r for r in runs if r["_tag"].startswith("depth_")]
    if sel:
        by = group(sel, lambda r: (r["args"]["mode"], r["args"]["depth"]))
        depths = sorted({k[1] for k in by})
        plt.figure(figsize=(6, 4))
        for m in ["dfa", "sdil"]:
            ys = [agg([mean_final_cos(r) for r in by.get((m, d), [])])[0] for d in depths]
            plt.plot(depths, ys, marker="o", label=m)
        plt.xlabel("hidden depth"); plt.ylabel("mean cos(r,-grad)"); plt.title("Alignment vs depth")
        plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
        plt.savefig(os.path.join(figdir, "depth.png"), dpi=120); plt.close()
    print(f"\n[figures written to {figdir}]")


def main():
    runs = load_all(RES)
    print(f"loaded {len(runs)} runs from {RES}")
    table_main(runs)
    table_nuisance(runs)
    table_trainedA(runs)
    table_depth(runs)
    plots(runs)


if __name__ == "__main__":
    main()