"""
THE DECISIVE SN-lensing dataset: calibrated foreground convergence kappa_fg per SN.
kappa_fg_i = sum_{z_l < z_SN,i} (3 H0^2 Om/2c^2) chi_l (chi_s,i - chi_l)/chi_s,i (1+z_l) delta_m(n_i,z_l) dchi_l
with delta_m = delta_g / b(z),  delta_g = (N_d - alpha N_r)/(alpha N_r)  [PROPER randoms, not a count proxy],
N_d weighted by the BOSS completeness weights, randoms = official BOSS randoms (subsampled), bias-corrected,
redshift-TRUNCATED per SN. Tracer: BOSS DR12 LSS (LOWZ z~0.15-0.43 b~1.85; CMASS z~0.43-0.7 b~2.0).
Output: the calibrated kappa_fg dataset + the multivariate regression delta_mu = a + b*kappa_fg + d*hostmass.
"""
import numpy as np, healpy as hp, glob, os
from astropy.io import fits
from scipy.integrate import quad
import dfd_params as P, background as B
from skymap_discriminator import load_sne, dfd_residuals

BD="boss_data"; NSIDE=128; SMOOTH=1.0; d2r=np.pi/180
def chi(z): return (P.C_KMS/P.H0)*quad(lambda x:1/B.E(x),0,max(z,1e-4),limit=60)[0]
PRE=3*(P.H0/P.C_KMS)**2*P.OMEGA_M/2.0      # 3 H0^2 Om / 2c^2  [1/Mpc^2]

def load_gal(f):
    d=fits.open(f, ignore_missing_end=True)[1].data
    ra=np.asarray(d['RA'],float); dec=np.asarray(d['DEC'],float); z=np.asarray(d['Z'],float)
    try: w=np.asarray(d['WEIGHT_SYSTOT'],float)*(np.asarray(d['WEIGHT_CP'],float)+np.asarray(d['WEIGHT_NOZ'],float)-1.0)
    except Exception: w=np.ones(len(ra))
    return ra,dec,z,w
def load_rand(f):
    # truncated FITS: header claims full NAXIS2 but only a prefix is present. Patch NAXIS2 -> recovered rows.
    with open(f,'rb') as fh: raw=bytearray(fh.read())
    h=fits.open(f, ignore_missing_end=True); off=int(h[1]._data_offset); nb1=int(h[1].header['NAXIS1']); h.close()
    nrow=(len(raw)-off)//nb1
    idx=raw.find(b'NAXIS2'); raw[idx:idx+80]=("NAXIS2  = %20d" % nrow).ljust(80).encode()
    tmp=f+'.patched'; open(tmp,'wb').write(bytes(raw[:off+nrow*nb1]))
    d=fits.open(tmp, ignore_missing_end=True)[1].data
    ra=np.asarray(d['RA'],float); dec=np.asarray(d['DEC'],float); z=np.asarray(d['Z'],float)
    m=np.isfinite(ra)&np.isfinite(dec)&np.isfinite(z); return ra[m],dec[m],z[m]

# gather available BOSS samples + LOAD ONCE (galaxies + randoms cached)
SAMP=[]
for tag,bias in [("CMASS_N",2.0),("LOWZ_N",1.85),("CMASS_S",2.0),("LOWZ_S",1.85)]:
    gf=f"{BD}/galaxy_{tag}.fits.gz"; rf=f"{BD}/rand_{tag}.fits"
    if os.path.exists(gf) and os.path.exists(rf):
        g=load_gal(gf); r=load_rand(rf); SAMP.append((tag,bias,g,r))
        print(f"  {tag}: {len(g[0])} gals, {len(r[0])} randoms (ratio {len(r[0])/len(g[0]):.1f})")
print("BOSS samples:", [s[0] for s in SAMP])

# tomographic delta_m maps with PROPER randoms
ZB=np.arange(0.10,0.75,0.05); npix=hp.nside2npix(NSIDE)
dmaps=[]; foot=np.zeros(npix,bool)
for i in range(len(ZB)-1):
    zlo,zhi=ZB[i],ZB[i+1]; zc=0.5*(zlo+zhi); Nd=np.zeros(npix); Nr=np.zeros(npix); bz=2.0
    for tag,bias,(gra,gdec,gz,gw),(rra,rdec,rz) in SAMP:
        gs=(gz>=zlo)&(gz<zhi); rs=(rz>=zlo)&(rz<zhi)
        if gs.sum()==0 or rs.sum()==0: continue
        np.add.at(Nd,hp.ang2pix(NSIDE,(90-gdec[gs])*d2r,gra[gs]*d2r),gw[gs])
        np.add.at(Nr,hp.ang2pix(NSIDE,(90-rdec[rs])*d2r,rra[rs]*d2r),1.0)
        bz=bias
    if Nr.sum()<100: dmaps.append((zc,chi(zc),None)); continue
    alpha=Nd.sum()/Nr.sum(); cov=Nr>0; foot|=cov
    dg=np.full(npix,0.0); dg[cov]=(Nd[cov]-alpha*Nr[cov])/(alpha*Nr[cov])
    dg=hp.smoothing(dg,fwhm=SMOOTH*d2r)            # delta_g
    dmaps.append((zc, chi(zc), dg/bz))             # delta_m = delta_g / b(z)
    print(f"  zbin {zc:.3f}: Nd={Nd.sum():.0f} Nr={int(Nr.sum())} alpha={alpha:.4f} b={bz} footprint {cov.mean()*100:.1f}%")

def kappa_fg(ra,dec,zsn):
    pix=hp.ang2pix(NSIDE,(90-dec)*d2r,ra*d2r); chis=chi(zsn); k=0.0; dchi=chi(0.15)-chi(0.10)
    for zc,chil,dm in dmaps:
        if dm is None or zc>=zsn or chil>=chis: continue
        k += PRE*chil*(chis-chil)/chis*(1+zc)*dm[pix]*(chi(zc+0.025)-chi(zc-0.025))
    return k, foot[pix]

# SNe in footprint, z>0.15 (BOSS foreground exists), velocity-clean
d=load_sne(0.15,1.0); dmu,err=dfd_residuals(d); ra,dec,hm=d['RA'],d['DEC'],d['HOST_LOGMASS']
kap=np.array([kappa_fg(ra[i],dec[i],d['zHD'][i])[0] for i in range(len(d))])
inf=np.array([kappa_fg(ra[i],dec[i],d['zHD'][i])[1] for i in range(len(d))])
sel=inf&np.isfinite(kap)&(kap!=0)
print(f"\nPantheon+ SNe in BOSS footprint (z>0.15): {sel.sum()}")
y=dmu[sel]; k=kap[sel]; e=err[sel]; w=1/e**2; mass=np.where(hm[sel]>0,hm[sel]-10,0.0)
print(f"calibrated kappa_fg: mean={k.mean():+.4f} std={k.std():.4f} range[{k.min():+.3f},{k.max():+.3f}]")
np.savez("boss_data/kappa_fg_dataset.npz", ra=ra[sel],dec=dec[sel],z=d['zHD'][sel],dmu=y,kappa_fg=k,err=e,hostmass=mass)

X=np.column_stack([np.ones_like(k),k,mass])
def gls(X,y,w): A=(X.T*w)@X; return np.linalg.solve(A,(X.T*w)@y), np.linalg.inv(A)
beta,covb=gls(X,y,w); b_slope=beta[1]
pat=(ra[sel]//18).astype(int); sl=[gls(X[pat!=pp],y[pat!=pp],w[pat!=pp])[0][1] for pp in np.unique(pat)]
sl=np.array(sl); jk=np.sqrt((len(sl)-1)/len(sl)*np.sum((sl-sl.mean())**2))
print("="*78); print("DECISIVE TEST — calibrated BOSS kappa_fg (randoms+bias+truncation) x Pantheon+"); print("="*78)
for s in (+2.,-2.,0.): bb,_=gls(X,y+s*k,w); print(f"  inject {s:+.1f} -> {bb[1]:+.2f}")
print(f"\n  slope d(delta_mu)/d(kappa_fg) = {b_slope:+.3f}  GLS-se {np.sqrt(covb[1,1]):.3f}  jackknife-se {jk:.3f}")
print(f"  significance = {b_slope/jk:+.1f} sigma; sign = {'POSITIVE->DFD' if b_slope>0 else 'NEGATIVE->LCDM'} "
      f"({'SIGNIFICANT' if abs(b_slope/jk)>2 else 'not significant'})")
print("  (kappa_fg is now physically calibrated: proper randoms, bias-corrected, redshift-truncated.)")
