"""Fit d ln g/dt = -c s^alpha jointly on the b2 0.98 and 0.999 N=82 runs (steps 800..end), then re-run the model."""
import numpy as np, csv
from scipy.optimize import minimize
def load(fn):
    r=list(csv.DictReader(open(fn))); return {k:np.array([float(x[k]) for x in r]) for k in r[0]}
A=load('update_size_b2-0.98_N82_eps1e-8.csv'); B=load('update_size_b2-0.999_N82_eps1e-8.csv')
def simulate(c,al,b2,eps,g0,v0,T):
    g,v=g0,v0; out=[]
    for t in range(T+1):
        s=g/(np.sqrt(v)+eps); out.append((g,np.sqrt(v),s)); g=g*np.exp(-c*s**al); v=v+(1-b2)*(g*g-v)
    return np.array(out)
def loss(p):
    c,al=p; tot=0
    for D,b2 in ((A,0.98),(B,0.999)):
        i0=int(np.where(D['step']==800)[0][0]); sim=simulate(c,al,b2,1e-8,D['rms_grad'][i0],D['rms_sqrt_vhat'][i0]**2,int(D['step'].max())-800)
        idx=(D['step'][i0:]-800).astype(int); tot+=np.mean((np.log(sim[idx,0])-np.log(D['rms_grad'][i0:]))**2)
    return tot
res=minimize(loss,[0.01,1.0],method='Nelder-Mead'); c,al=res.x
print(f"joint fit: c={c:.4f}, alpha={al:.3f}, rms log error {np.sqrt(res.fun/2):.3f}")
for name,D,b2 in (('b2=0.98',A,0.98),('b2=0.999',B,0.999)):
    i0=int(np.where(D['step']==800)[0][0]); sim=simulate(c,al,b2,1e-8,D['rms_grad'][i0],D['rms_sqrt_vhat'][i0]**2,int(D['step'].max())-800)
    print(f"\n{name}: step | g meas | g model | sqrt(v)/g meas | model | step/lr meas | model")
    for t in range(800,int(D['step'].max())+1,400):
        i=int(np.where(D['step']==t)[0][0]); j=t-800
        print(f"{t} | {D['rms_grad'][i]:.2e} | {sim[j,0]:.2e} | {D['rms_sqrt_vhat'][i]/D['rms_grad'][i]:.1f} | {sim[j,1]/sim[j,0]:.1f} | {D['rms_step_over_lr'][i]:.3f} | {sim[j,2]:.3f}")
print("\nPrediction from g0=1e-5, v0=g0^2: steps for ln(g0/g) to reach 7 / 11 / 15; mean s to T7")
print("beta2 | eps | T7 | T11 | T15 | mean s to T7")
for b2 in (0.9,0.95,0.98,0.99,0.995,0.999):
    for eps in (1e-8,1e-16):
        sim=simulate(c,al,b2,eps,1e-5,1e-10,60000); lng=np.log(1e-5/sim[:,0])
        T=[int(np.argmax(lng>=k)) if (lng>=k).any() else -1 for k in (7,11,15)]
        print(f"{b2} | {eps:g} | "+" | ".join(str(x) if x>0 else '>60000' for x in T)+f" | {sim[:T[0],2].mean():.3f}")