"""Depth / no-BN analysis: accuracy vs depth, and the per-layer 'depth utility'
profile (does the teaching signal reach early layers?). DFA is expected to give
near-random credit (~0.07) at every layer; SDIL should stay aligned, especially
in the hard early layers.
Usage: python experiments/analyze_depth.py <prefix>   e.g. depth  or  deepres
"""
import glob
import json
import os
import sys
from collections import defaultdict

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

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


def load():
    runs = {}
    for p in sorted(glob.glob(os.path.join(RES, f"{PFX}_*.json"))):
        with open(p) as f:
            d = json.load(f)
        runs[os.path.basename(p)[:-5]] = d
    return runs


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 main():
    runs = load()
    if not runs:
        print(f"no runs with prefix {PFX}_")
        return
    # collect datasets, depths, seeds
    metas = []
    for tag, d in runs.items():
        a = d["args"]
        metas.append((a["dataset"], a["mode"], a["depth"], a.get("seed", 0), tag))
    datasets = sorted({m[0] for m in metas})
    depths = sorted({m[2] for m in metas})

    for ds in datasets:
        print(f"\n############ {PFX}  dataset={ds} ############")
        print("\n--- TEST ACCURACY vs DEPTH ---")
        print(f"{'depth':>6s} | {'bp':>18s} | {'dfa':>18s} | {'sdil':>18s} | {'dfa->sdil gap':>14s}")
        for depth in depths:
            cells = {}
            for mode in ["bp", "dfa", "sdil"]:
                accs = [runs[t]["final"].get("test_acc") for (D, M, DP, S, t) in metas
                        if D == ds and M == mode and DP == depth]
                cells[mode] = agg(accs)
            def cell(m):
                mm, sd, n = cells[m]
                return f"{mm:.4f}±{sd:.3f}" if mm is not None else "   n/a  "
            gap = ""
            if cells["dfa"][0] is not None and cells["sdil"][0] is not None:
                gap = f"{cells['sdil'][0] - cells['dfa'][0]:+.4f}"
            print(f"{depth:6d} | {cell('bp'):>18s} | {cell('dfa'):>18s} | {cell('sdil'):>18s} | {gap:>14s}")

        # depth utility: per-layer alignment averaged over seeds, DFA vs SDIL
        print("\n--- DEPTH UTILITY: per-layer cos(r,-g) [input-side ... output-side] ---")
        for depth in depths:
            for mode in ["dfa", "sdil"]:
                perlayer_runs = [runs[t]["final"].get("cos_r_negg") for (D, M, DP, S, t) in metas
                                 if D == ds and M == mode and DP == depth and runs[t]["final"].get("cos_r_negg")]
                if not perlayer_runs:
                    continue
                nL = len(perlayer_runs[0])
                mean_pl = [sum(r[i] for r in perlayer_runs) / len(perlayer_runs) for i in range(nL)]
                early = sum(mean_pl[:max(1, nL // 3)]) / max(1, nL // 3)
                print(f"  d{depth:<2d} {mode:4s} early3rd={early:+.3f} | " + " ".join(f"{v:+.2f}" for v in mean_pl))

    if HAVE_MPL:
        figdir = os.path.join(RES, "figs")
        os.makedirs(figdir, exist_ok=True)
        for ds in datasets:
            plt.figure(figsize=(6, 4))
            for mode, style in [("bp", "k-o"), ("dfa", "r-s"), ("sdil", "b-^")]:
                ys, es = [], []
                for depth in depths:
                    accs = [runs[t]["final"].get("test_acc") for (D, M, DP, S, t) in metas
                            if D == ds and M == mode and DP == depth]
                    m, sd, n = agg(accs)
                    ys.append(m); es.append(sd or 0)
                if any(v is not None for v in ys):
                    xs = [d for d, y in zip(depths, ys) if y is not None]
                    yv = [y for y in ys if y is not None]
                    ev = [e for y, e in zip(ys, es) if y is not None]
                    plt.errorbar(xs, yv, yerr=ev, fmt=style, label=mode, capsize=3)
            plt.xlabel("depth (hidden layers)"); plt.ylabel("test acc")
            plt.title(f"{PFX} {ds}: accuracy vs depth (no-BN)")
            plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
            fp = os.path.join(figdir, f"{PFX}_{ds}_acc_vs_depth.png")
            plt.savefig(fp, dpi=120); plt.close()
            print(f"[fig] {fp}")

            # depth-utility figure: early-layer alignment vs depth (the collapse)
            plt.figure(figsize=(6, 4))
            for mode, style in [("dfa", "r-s"), ("sdil", "b-^")]:
                xs, ys = [], []
                for depth in depths:
                    pl = [runs[t]["final"].get("cos_r_negg") for (D, M, DP, S, t) in metas
                          if D == ds and M == mode and DP == depth and runs[t]["final"].get("cos_r_negg")]
                    if not pl:
                        continue
                    nL = len(pl[0]); k = max(1, nL // 3)
                    early = [sum(r[:k]) / k for r in pl]
                    xs.append(depth); ys.append(sum(early) / len(early))
                if xs:
                    plt.plot(xs, ys, style, label=mode, markersize=7)
            plt.xlabel("depth (hidden layers)"); plt.ylabel("early-layer cos(r, -grad)")
            plt.title(f"{PFX} {ds}: credit reaching early layers vs depth")
            plt.ylim(0, 1); plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
            fp2 = os.path.join(figdir, f"{PFX}_{ds}_depth_utility.png")
            plt.savefig(fp2, dpi=120); plt.close()
            print(f"[fig] {fp2}")


if __name__ == "__main__":
    main()