"""Spike/period analysis of probe.csv + log.csv.  Usage: python period_analysis.py DIR BETA2 [DIR BETA2 ...]
Spike = probe row where rms_sqrt_vhat jumps >3x above the previous row OR rms_grad jumps >100x (rows are log_every apart).
peak  = max rms_sqrt_vhat in the spike row and the next 2 rows; floor = rms_sqrt_vhat in the row just before the spike.
slope = least-squares slope of ln rms_sqrt_vhat vs step from (peak row + 2) to (row before next spike), per step.
Only spikes after the first 500 steps (initial transient) are used for cycle statistics."""
import sys, csv, math, numpy as np
def load(path):
    r = list(csv.DictReader(open(path)))
    return {k: np.array([float(x[k]) for x in r]) for k in r[0]}
def first_step(lg, thr):
    i = np.where(lg['test_acc'] >= thr)[0]
    return int(lg['step'][i[0]]) if len(i) else None
def analyse(d, b2, verbose=True, tmin=500):
    p = load(f"{d}/probe.csv"); lg = load(f"{d}/log.csv")
    s, g, v, st = p['step'], p['rms_grad'], p['rms_sqrt_vhat'], p['rms_step_over_lr']
    idx = [i for i in range(1, len(s)) if s[i] > tmin and (v[i] > 3*v[i-1] or g[i] > 100*g[i-1])]
    # merge consecutive rows into one spike
    sp = []
    for i in idx:
        if sp and i - sp[-1][-1] <= 2: sp[-1].append(i)
        else: sp.append([i])
    rows = []
    for k, grp in enumerate(sp):
        i0 = grp[0]
        pk_i = i0 + int(np.argmax(v[i0:i0+3]))
        nxt = sp[k+1][0] if k+1 < len(sp) else None
        floor_next = v[nxt-1] if nxt else None
        slope = None
        if nxt and nxt-1 - (pk_i+2) >= 2:
            a, b = pk_i+2, nxt
            slope = np.polyfit(s[a:b], np.log(v[a:b]), 1)[0]
        rows.append(dict(step=int(s[i0]), gpeak=g[i0:i0+3].max(), vpeak=v[pk_i], gpre=g[i0-1], floor_next=floor_next,
                         period=(int(s[nxt]-s[i0]) if nxt else None), slope=slope,
                         step_pre=st[nxt-1] if nxt else None))
    per = np.array([r['period'] for r in rows if r['period']])
    lnr = np.array([math.log(r['vpeak']/r['floor_next']) for r in rows if r['floor_next']])
    slp = np.array([r['slope'] for r in rows if r['slope'] is not None])
    gfl = np.array([r['gpre'] for r in rows])
    out = dict(dir=d, beta2=b2, nspikes=len(rows), steps=int(s[-1]), log_every=int(s[2]-s[1]),
               per_mean=per.mean() if len(per) else float('nan'), per_sd=per.std() if len(per) else float('nan'),
               per_med=np.median(per) if len(per) else float('nan'),
               lnr=lnr.mean() if len(lnr) else float('nan'), lnr_sd=lnr.std() if len(lnr) else float('nan'),
               vpeak=np.median([r['vpeak'] for r in rows]) if rows else float('nan'),
               vfloor=np.median([r['floor_next'] for r in rows if r['floor_next']]) if len(lnr) else float('nan'),
               gpeak=np.median([r['gpeak'] for r in rows]) if rows else float('nan'),
               gfloor=np.median(gfl) if len(gfl) else float('nan'),
               slope=slp.mean() if len(slp) else float('nan'), slope_sd=slp.std() if len(slp) else float('nan'),
               slope_th=math.log(b2)/2, step_mean=st[s > tmin].mean(), step_med=np.median(st[s > tmin]),
               step_pre=np.median([r['step_pre'] for r in rows if r['step_pre'] is not None]) if len(per) else float('nan'),
               acc_end=lg['test_acc'][-1], acc_max=lg['test_acc'].max(),
               t50=first_step(lg, .5), t90=first_step(lg, .9), t99=first_step(lg, .99))
    out['pred'] = 2*out['lnr']/(1-b2)
    out['pred_slope'] = out['lnr']/(-out['slope']) if len(slp) else float('nan')
    if verbose:
        print(f"\n### {d} (beta2={b2}, last step {out['steps']}, log_every {out['log_every']})")
        print("| spike step | period to next | peak rms_grad | peak rms_sqrt_vhat | rms_grad just before | floor rms_sqrt_vhat before next | ln(peak/floor) | slope ln sqrt(vhat)/step | step/lr before next |")
        print("|---|---|---|---|---|---|---|---|---|")
        for r in rows:
            f = lambda x, fm: (fm % x) if x is not None else "-"
            print(f"| {r['step']} | {f(r['period'],'%d')} | {r['gpeak']:.2e} | {r['vpeak']:.2e} | {r['gpre']:.1e} | {f(r['floor_next'],'%.2e')} | "
                  f"{f(math.log(r['vpeak']/r['floor_next']) if r['floor_next'] else None,'%.2f')} | {f(r['slope'],'%.5f')} | {f(r['step_pre'],'%.3f')} |")
    return out
if __name__ == "__main__":
    a = sys.argv[1:]; res = [analyse(a[i], float(a[i+1])) for i in range(0, len(a), 2)]
    print("\n| beta2 | run | steps | spikes | mean period (sd) [median] | predicted 2ln(peak/floor)/(1-b2) | ln(peak/floor) (sd) | median peak / floor sqrt(vhat) | median rms_grad spike / floor | decay slope meas. (sd) vs ln(b2)/2 | mean (median) step/lr | step/lr before spike | test acc end (max) | first step test acc 0.5/0.9/0.99 |")
    print("|---|---|---|---|---|---|---|---|---|---|---|---|---|---|")
    for o in res:
        print(f"| {o['beta2']} | {o['dir']} | {o['steps']} | {o['nspikes']} | {o['per_mean']:.0f} ({o['per_sd']:.0f}) [{o['per_med']:.0f}] | {o['pred']:.0f} | {o['lnr']:.2f} ({o['lnr_sd']:.2f}) | "
              f"{o['vpeak']:.1e} / {o['vfloor']:.1e} | {o['gpeak']:.1e} / {o['gfloor']:.1e} | {o['slope']:.4f} ({o['slope_sd']:.4f}) vs {o['slope_th']:.4f} | {o['step_mean']:.3f} ({o['step_med']:.4f}) | {o['step_pre']:.3f} | "
              f"{o['acc_end']:.3f} ({o['acc_max']:.3f}) | {o['t50']}/{o['t90']}/{o['t99']} |")