import numpy as np, json, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
lo=np.load('zp_pairs.npy'); H={k:np.load(f'zp_pairs_{k}.npy')[:,:4] for k in ['1e12','1e21','1e22']}
G={'t<1e5':lo[lo[:,3]<1e5],'1e5-5e5':lo[(lo[:,3]>=1e5)&(lo[:,3]<5e5)],'5e5-1.13e6':lo[lo[:,3]>=5e5],'#1e12':H['1e12'],'#1e21':H['1e21'],'#1e22':H['1e22']}
def zn(a): s,z1,z2,t=a.T; return s,np.sqrt(z1*z2)/np.log(t/(2*np.pi)),t
rng=np.random.default_rng(0)
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)
    B=np.array([np.polyfit(x[k],y[k],1) for k in (rng.integers(0,len(x),len(x)) for _ in range(1000))]); e=B.std(0)
    r=y-x; C1=np.exp(r.mean()); 
    return dict(alpha=p[0],alpha_err=e[0],lnC=p[1],lnC_err=e[1],C=np.exp(p[1]),C1=C1,C1_err=C1*r.std()/np.sqrt(len(r)),sd=r.std(),n=int(m.sum()),L=float(np.median(np.log(t[m]/(2*np.pi)))))
R={k:fit(v) for k,v in G.items()}; R['pooled']=fit(np.vstack(list(G.values())))
ks=list(G)
for q in ['alpha','C1']:
    v=np.array([R[k][q] for k in ks]); e=np.array([R[k][q+'_err'] for k in ks]); 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,mean=float(m))
L=np.array([R[k]['L'] for k in ks]); C=np.array([R[k]['C1'] for k in ks]); e=np.array([R[k]['C1_err'] for k in ks])
p,cv=np.polyfit(1/L,C,1,w=1/e,cov='unscaled'); chi_a=float((((C-np.polyval(p,1/L))/e)**2).sum())
q,cq=np.polyfit(np.log(L),np.log(C),1,w=C/e,cov='unscaled'); chi_b=float((((np.log(C)-np.polyval(q,np.log(L)))*C/e)**2).sum())
R['model_inv']=dict(C_inf=p[1],C_inf_err=np.sqrt(cv[1,1]),B_times_Cinf=p[0],err=np.sqrt(cv[0,0]),chi2=chi_a,dof=len(L)-2)
R['model_pow']=dict(A=np.exp(q[1]),kappa=q[0],kappa_err=np.sqrt(cq[0,0]),chi2=chi_b,dof=len(L)-2)
json.dump(R,open('fit_per_height.json','w'),indent=1,default=float)
for k in ks+['pooled']: r=R[k]; print(f"{k:12s} L={r['L']:.1f} n={r['n']} alpha={r['alpha']:.3f}±{r['alpha_err']:.3f} C={r['C']:.3f} C1={r['C1']:.3f}±{r['C1_err']:.3f} sd={r['sd']:.2f}")
print({k:R[k] for k in ['chi2_alpha','chi2_C1','model_inv','model_pow']})
xx=np.logspace(-2.4,np.log10(.3)); P=R['pooled']
plt.figure(figsize=(7.5,6))
for (k,a),c in zip(G.items(),['C0','C1','C2','C3','C4','C5']):
    s,z,_=zn(a); plt.loglog(s,z,'.',ms=2.5,alpha=.4,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=7,markerscale=4); plt.title("|Z'| at close pairs vs gap, heights 1e4 to 1.4e21"); plt.tight_layout(); plt.savefig('zprime_vs_gap.png',dpi=130)
fig,ax=plt.subplots(1,2,figsize=(13,5))
for i,k in enumerate(ks):
    s,z,_=zn(G[k]); m=s<.3; ax[0].loglog(s[m],z[m],'.',ms=2,alpha=.3,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}, C1={R[k]['C1']:.3f}")
ax[0].legend(fontsize=7); ax[0].set_xlabel('s'); ax[0].set_title('per-height fits (s<0.3)')
ax[1].errorbar(L,[R[k]['alpha'] for k in ks],[R[k]['alpha_err'] for k in ks],fmt='o',label='α (free fit)')
ax[1].errorbar(L,C,e,fmt='s',label='C (α fixed at 1)')
LL=np.linspace(8,50,200); ax[1].plot(LL,p[1]+p[0]/LL,'g-',label=f"C∞+B'/L: C∞={p[1]:.3f}±{np.sqrt(cv[1,1]):.3f} (χ²={chi_a:.1f}/{len(L)-2})")
ax[1].plot(LL,np.exp(q[1])*LL**q[0],'m:',label=f"A L^κ: κ={q[0]:.3f} (χ²={chi_b:.1f}/{len(L)-2})")
ax[1].axhline(1,ls='--',c='gray'); ax[1].set_xlabel('median log(t/2π)'); ax[1].legend(fontsize=7); ax[1].set_title('α and C vs height')
plt.tight_layout(); plt.savefig('zprime_per_height.png',dpi=130)