"""
Iteration 6: robustness battery on the CORRECTED (SDSS redshift-truncated foreground) test.
Is the ~-1 sigma negative slope ROBUST (consistent mild LCDM lean) or NOISE (scatters +/-)?
Vary: z-cuts, host-mass control, lensing-kernel vs uniform-front weighting, smoothing scale, residual model.
Each: multivariate GLS slope b + jackknife significance. Maps built once per smoothing scale.
"""
import numpy as np, healpy as hp, glob, os
from scipy.integrate import quad
import dfd_params as P, background as B
from skymap_discriminator import load_sne, dfd_residuals

NSIDE=256; d2r=np.pi/180; FD="foreground_data"
def chi(z): return (P.C_KMS/P.H0)*quad(lambda x:1/B.E(x),0,z,limit=60)[0]

def build_maps(fg_deg, bg_deg=10.0):
    bins=[]
    for f in sorted(glob.glob(os.path.join(FD,"sdss_z*.csv"))):
        g=np.loadtxt(f,delimiter=',',ndmin=2)
        if g.size==0: continue
        ra,dec,z=g[:,0],g[:,1],g[:,2]; zc=z.mean()
        cnt=np.bincount(hp.ang2pix(NSIDE,(90-dec)*d2r,ra*d2r),minlength=hp.nside2npix(NSIDE)).astype(float)
        nfg=hp.smoothing(cnt,fwhm=fg_deg*d2r); nbg=hp.smoothing(cnt,fwhm=bg_deg*d2r)
        dg=np.where(nbg>0.05,nfg/np.maximum(nbg,1e-6)-1.0,0.0)
        bins.append((zc,chi(zc),dg,nbg>0.05))
    foot=np.all([b[3] for b in bins[:3]],axis=0)
    return bins, foot

def kappa(bins,foot,ra,dec,zsn,kernel='lensing'):
    pix=hp.ang2pix(NSIDE,(90-dec)*d2r,ra*d2r); chis=chi(zsn); k=0.0
    for zc,chil,dg,cov in bins:
        if zc>=zsn or chil>=chis: break
        wkern = chil*(chis-chil)/chis*(1+zc) if kernel=='lensing' else 1.0
        k += wkern*dg[pix]
    return k, foot[pix]

def lcdm_residuals(d):
    z=d['zHD']; from scipy.integrate import quad as Q
    pl=P.PLANCK; Om=pl['Omega_m']; H0=pl['H0']; Or=P.OMEGA_R_H2/(H0/100)**2; OL=1-Om-Or
    mu=np.array([5*np.log10((1+x)*(P.C_KMS/H0)*Q(lambda t:1/np.sqrt(Om*(1+t)**3+Or*(1+t)**4+OL),0,x,limit=60)[0])+25 for x in z])
    e=d['m_b_corr_err_DIAG']; w=1/e**2; res=d['m_b_corr']-mu; return res-np.sum(w*res)/np.sum(w), e

def fit(bins,foot,zlo,zhi,kernel='lensing',use_mass=True,resid='dfd'):
    d=load_sne(zlo,zhi); dmu,err=(dfd_residuals(d) if resid=='dfd' else lcdm_residuals(d))
    ra,dec,hm=d['RA'],d['DEC'],d['HOST_LOGMASS']
    kk=np.array([kappa(bins,foot,ra[i],dec[i],d['zHD'][i],kernel)[0] for i in range(len(d))])
    inf=np.array([kappa(bins,foot,ra[i],dec[i],d['zHD'][i],kernel)[1] for i in range(len(d))])
    sel=inf&np.isfinite(kk)&(np.abs(kk)>0)
    y,k,e=dmu[sel],kk[sel],err[sel]; k=(k-k.mean())/k.std(); w=1/e**2
    cols=[np.ones_like(k),k]+([np.where(hm[sel]>0,hm[sel]-10,0.0)] if use_mass else [])
    X=np.column_stack(cols)
    def gls(X,y,w): A=(X.T*w)@X; return np.linalg.solve(A,(X.T*w)@y)
    b=gls(X,y,w)[1]; ras=ra[sel]; pat=(ras//18).astype(int); sl=[]
    for pp in np.unique(pat):
        m=pat!=pp; sl.append(gls(X[m],y[m],w[m])[1])
    sl=np.array(sl); se=np.sqrt((len(sl)-1)/len(sl)*np.sum((sl-sl.mean())**2))
    return b,se,sel.sum()

print("="*82); print("ROBUSTNESS — corrected SDSS foreground test: is the ~-1sigma NEGATIVE robust or noise?"); print("="*82)
bins,foot=build_maps(0.5)
print(f"\n  {'variant':40s} {'N':>5s} {'slope b':>9s} {'jk_se':>7s} {'signif':>8s}")
def row(name,**kw):
    b,se,n=fit(bins,foot,**kw); print(f"  {name:40s} {n:5d} {b:+9.4f} {se:7.4f} {b/se:+7.1f}σ"); return b/se
sigs=[]
sigs.append(row("baseline z0.08-0.6 lensing +mass",zlo=0.08,zhi=0.6))
sigs.append(row("z>0.12",zlo=0.12,zhi=0.6))
sigs.append(row("z>0.20",zlo=0.20,zhi=0.6))
sigs.append(row("z 0.08-0.6 NO host-mass term",zlo=0.08,zhi=0.6,use_mass=False))
sigs.append(row("uniform front-count (no W kernel)",zlo=0.08,zhi=0.6,kernel='uniform'))
sigs.append(row("LCDM-model residuals",zlo=0.08,zhi=0.6,resid='lcdm'))
b2,foot2=build_maps(1.0);
b,se,n=fit(b2,foot2,zlo=0.08,zhi=0.6); print(f"  {'smoothing 1.0 deg':40s} {n:5d} {b:+9.4f} {se:7.4f} {b/se:+7.1f}σ"); sigs.append(b/se)
sigs=np.array(sigs)
print("\n  "+"-"*78)
print(f"  significances: {np.array2string(sigs,precision=1)}")
print(f"  range [{sigs.min():+.1f}, {sigs.max():+.1f}]σ; all negative? {(sigs<0).all()}; any |sig|>2? {(np.abs(sigs)>2).any()}")
print("  Read: consistently ~-1sigma (all negative, none significant) = ROBUST mild LCDM lean (not noise,")
print("  not a detection). Sign flips = noise. >2sigma anywhere = a real detection of that sign.")
