import numpy as np, json, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
lo=np.load('zp_pairs.npy'); hi=np.load('zp_pairs_1e12.npy')[:,:4]
def zn(a): s,z1,z2,t=a.T; return s,np.sqrt(z1*z2)/np.log(t/(2*np.pi)),t
groups={'t<1e5':lo[lo[:,3]<1e5],'1e5-5e5':lo[(lo[:,3]>=1e5)&(lo[:,3]<5e5)],'5e5-1.13e6':lo[lo[:,3]>=5e5],'~2.7e11 (#1e12)':hi}
def fit(a,smax=0.3):
    s,z,t=zn(a); m=s<smax; x,y=np.log(s[m]),np.log(z[m]); p=np.polyfit(x,y,1)
    rng=np.random.default_rng(0); B=[np.polyfit(x[k],y[k],1) for k in (rng.integers(0,len(x),len(x)) for _ in range(500))]
    e=np.std(B,0); return dict(alpha=p[0],alpha_err=e[0],lnC=p[1],lnC_err=e[1],C=np.exp(p[1]),n=int(m.sum()),logt=float(np.median(np.log(t[m]/(2*np.pi)))))
R={k:fit(v) for k,v in groups.items()}; R['pooled']=fit(np.vstack([lo,hi]))
# chi2 consistency
for q in ['alpha','lnC']:
    v=np.array([R[k][q] for k in groups]); e=np.array([R[k][q+'_err'] for k in groups]); w=1/e**2; m=(v*w).sum()/w.sum()
    R['chi2_'+q]=dict(chi2=float(((v-m)**2*w).sum()),dof=len(v)-1,weighted_mean=float(m))
json.dump(R,open('fit_per_height.json','w'),indent=1,default=float); print(json.dumps(R,indent=1,default=float))
P=R['pooled']; xx=np.logspace(-2.4,np.log10(.3))
plt.figure(figsize=(7.5,6))
for (k,a),c in zip([('t ≤ 1.13e6 (first 2M zeros)',lo),('t ≈ 2.7e11 (zeros #1e12+)',hi)],['C0','C3']):
    s,z,_=zn(a); plt.loglog(s,z,'.',ms=2.5,alpha=.45,color=c,label=k)
plt.loglog(xx,P['C']*xx**P['alpha'],'k',lw=2,label=f"pooled fit s<0.3: C={P['C']:.3f}, α={P['alpha']:.3f}±{P['alpha_err']:.3f}")
plt.loglog(xx,xx,'r--',label='C=1, α=1')
plt.xlabel('normalized gap s'); plt.ylabel(r"$\sqrt{|Z'(\gamma_1)Z'(\gamma_2)|}\,/\,\log(t/2\pi)$"); plt.legend(fontsize=8); plt.title("|Z'| at close pairs vs gap (corrected normalization)")
plt.tight_layout(); plt.savefig('zprime_vs_gap.png',dpi=130)
fig,ax=plt.subplots(1,2,figsize=(12,4.8)); ks=list(groups)
for i,k in enumerate(ks):
    s,z,_=zn(groups[k]); m=s<.3; ax[0].loglog(s[m],z[m],'.',ms=2,alpha=.35,color=f'C{i}'); ax[0].loglog(xx,R[k]['C']*xx**R[k]['alpha'],color=f'C{i}',label=f"{k}: α={R[k]['alpha']:.3f}, C={R[k]['C']:.3f}")
ax[0].legend(fontsize=7); ax[0].set_xlabel('s'); ax[0].set_title('per-height fits')
L=[R[k]['logt'] for k in ks]
ax[1].errorbar(L,[R[k]['alpha'] for k in ks],[R[k]['alpha_err'] for k in ks],fmt='o',label='α'); ax[1].errorbar(L,[R[k]['C'] for k in ks],[R[k]['C']*R[k]['lnC_err'] for k in ks],fmt='s',label='C')
ax[1].axhline(1,ls='--',c='gray'); ax[1].set_xlabel('median log(t/2π)'); ax[1].legend(); ax[1].set_title('α and C vs height')
plt.tight_layout(); plt.savefig('zprime_per_height.png',dpi=130)