import numpy as np, json, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
lo=np.vstack([np.load('zp_pairs.npy'),np.load('zp_pairs_lowextra.npy')])
FILES={'1e12':('zeros3',267653395647),'1e21':('zeros4',144176897509546973000),'1e22':('zeros5',1370919909931995300000)}
def pairs_all(tag):
    f,off=FILES[tag]; g=np.array([float(p[0]) for p in (L.split() for L in open(f)) if len(p)==1 and p[0].replace('.','').isdigit()])
    a=np.load(f'zp_all_{tag}.npy'); assert len(a)==len(g); z=np.abs(a[:,1]); T=float(off)+g.mean(); L=np.log(T/(2*np.pi))
    s=np.diff(g)*L/(2*np.pi); return np.c_[s,z[:-1],z[1:],np.full(len(s),T)]
G={'main (t<=1.13e6)':lo}
for k in ['1e12','1e21','1e22']: G['#'+k]=pairs_all(k)
edges=np.array([0,0.05,0.1,0.15,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0])
def stats(s,x,y):
    out=[]
    for a,b in zip(edges[:-1],edges[1:]):
        m=(s>=a)&(s<b); n=m.sum()
        if n<20: out.append([np.nan]*8+[n]); continue
        X,Y=x[m],y[m]; C=np.cov(X,Y); r=C[0,1]/np.sqrt(C[0,0]*C[1,1])
        # bootstrap err on cov
        rng=np.random.default_rng(0); bs=[np.cov(X[k],Y[k])[0,1] for k in (rng.integers(0,n,n) for _ in range(300))]
        out.append([np.median(s[m]),C[0,1],np.std(bs),C[0,0],C[1,1],r,np.var(X-Y),np.var(X+Y)/4,n])
    return np.array(out)
R={}
for k,a in G.items():
    s,z1,z2,t=a.T; L=np.log(t/(2*np.pi)); x=np.log(z1/L); y=np.log(z2/L); R[k]=stats(s,x,y)
for N in [12,25,46]:
    a=np.load(f'cue_N{N}.npy'); s,l1,l2=a.T; R[f'CUE N={N}']=stats(s,l1-np.log(N),l2-np.log(N))
cols=['s_med','cov','cov_err','var1','var2','corr','var_diff','var_mean','n']
json.dump({k:[dict(zip(cols,map(float,r))) for r in v] for k,v in R.items()},open('cov_vs_s.json','w'),indent=1)
# fit cov = a + b log s on s<0.5
fits={}
for k,v in R.items():
    m=(v[:,0]<0.5)&np.isfinite(v[:,1])&(v[:,8]>=20); w=1/v[m,2]**2
    A=np.vstack([np.ones(m.sum()),np.log(v[m,0])]).T; cv=np.linalg.inv(A.T@(A*w[:,None])); p=cv@A.T@(w*v[m,1])
    chi=float((((v[m,1]-A@p))**2*w).sum())
    # var_diff ~ c s^g
    md=m&(v[:,6]>0)&(v[:,8]>=20); g=np.polyfit(np.log(v[md,0]),np.log(v[md,6]),1)
    fits[k]=dict(cov_a=p[0],cov_b=p[1],cov_b_err=np.sqrt(cv[1,1]),chi2=chi,dof=int(m.sum()-2),vardiff_exp=g[0],corr_smallest=float(v[0,5]),corr_s1=float(v[-1,5]))
    print(k,{q:round(float(x),3) for q,x in fits[k].items()})
json.dump(fits,open('cov_fits.json','w'),indent=1,default=float)
fig,ax=plt.subplots(1,3,figsize=(17,5))
for i,(k,v) in enumerate(R.items()):
    ls='-' if 'CUE' not in k else '--'; c=f'C{i}'
    ax[0].errorbar(v[:,0],v[:,1],v[:,2],fmt='o'+ls,ms=3,color=c,label=k)
    ax[1].plot(v[:,0],v[:,5],'o'+ls,ms=3,color=c,label=k)
    ax[2].loglog(v[:,0],v[:,6],'o'+ls,ms=3,color=c,label=k)
xx=np.linspace(.03,1); ax[0].plot(xx,-0.5*np.log(xx),'k:',label='-½ log s (Argon)')
ax[0].set_xscale('log'); ax[0].set_xlabel('s'); ax[0].set_title("Cov(log|Z'(γ1)|, log|Z'(γ2)| | s)"); ax[0].legend(fontsize=6)
ax[1].set_xlabel('s'); ax[1].set_title('correlation | s'); ax[2].set_xlabel('s'); ax[2].set_title("Var(log|Z'1| − log|Z'2| | s)")
ax[2].loglog(xx,2*xx**2,'k:',label='∝ s²'); ax[2].legend(fontsize=6)
plt.tight_layout(); plt.savefig('cov_vs_s.png',dpi=130)