"""
REAL SN-residual x foreground-convergence test (the DFD decisive test, with REAL lensing maps).
DFD predicts slope b>0 (SN looks FARTHER through foreground mass); LCDM cosmological constant predicts b~0.
Usage: python3 sn_lensing_real.py [planck|des_ks|des_tomo1|des_tomo2]
Pipeline: Pantheon+ Hubble residuals delta-mu  x  convergence kappa sampled at each SN sightline.
Coordinate care: Planck maps are GALACTIC -> rotate to Equatorial; DES maps are already Equatorial.
Rotation-null: randomly rotate the map and re-fit -> the noise floor (should be ~0).

Data (public, NOT bundled with this deposit -- place under the working directory):
  pantheon_data/PantheonPlusSH0ES.dat                  (Pantheon+SH0ES public release)
  lensing_map/COM_Lensing_4096_R3.00/MV/dat_klm.fits   (Planck 2018 lensing, PLA)
  des_y3/KS_full.fits, wiener_tomo*.fits, glimpse_mask.fits  (DES Y3 mass maps)
"""
import sys, numpy as np
import healpy as hp
from scipy.integrate import quad

MAP = sys.argv[1] if len(sys.argv)>1 else 'planck'

# ---------- Pantheon+ residuals ----------
rows=[l.split() for l in open('pantheon_data/PantheonPlusSH0ES.dat')]
H=rows[0]; iz=H.index('zHD'); iMU=H.index('MU_SH0ES'); iE=H.index('MU_SH0ES_ERR_DIAG'); iRA=H.index('RA'); iD=H.index('DEC')
z=np.array([float(r[iz]) for r in rows[1:]]); mu=np.array([float(r[iMU]) for r in rows[1:]])
emu=np.array([float(r[iE]) for r in rows[1:]]); ra=np.array([float(r[iRA]) for r in rows[1:]]); dec=np.array([float(r[iD]) for r in rows[1:]])
# fiducial flat LCDM distance modulus
c=299792.458; Om=0.31
def dL(zz,H0=70.0):
    return (1+zz)*(c/H0)*quad(lambda x:1/np.sqrt(Om*(1+x)**3+1-Om),0,zz,limit=80)[0]
sel=(z>0.02)&(z<0.8)&(emu>0)&(emu<1.0)
z,mu,emu,ra,dec=z[sel],mu[sel],emu[sel],ra[sel],dec[sel]
mumod=np.array([5*np.log10(dL(zz))+25 for zz in z])
res=mu-mumod; w=1/emu**2; res=res-np.average(res,weights=w)   # marginalize absolute scale
print(f"Pantheon+: {len(z)} SNe (0.02<z<0.8), <z>={z.mean():.3f}, residual rms={res.std():.3f} mag")

# ---------- load convergence map, sample at SN sightlines ----------
def sample(mapfile, galactic, nside_out=1024, fwhm_deg=0.5, des=False, field=0):
    if mapfile.endswith('_klm.fits') or 'klm' in mapfile:
        klm=hp.read_alm(mapfile); k=hp.alm2map(klm,nside=nside_out)
    else:
        k=hp.read_map(mapfile, field=field)
        if hp.get_nside(k)!=nside_out: k=hp.ud_grade(k,nside_out)
    k=hp.smoothing(k, fwhm=np.radians(fwhm_deg))
    if galactic:
        k=hp.Rotator(coord=['G','C']).rotate_map_pixel(k)   # Galactic -> Equatorial
    bad=(k==hp.UNSEEN)|~np.isfinite(k)
    if bad.any(): k=np.where(bad,np.nan,k)
    th=np.radians(90-dec); ph=np.radians(ra)
    pix=hp.ang2pix(nside_out, th, ph)
    return k[pix], k

desmask=None
if MAP=='planck':
    kap, kmap = sample('lensing_map/COM_Lensing_4096_R3.00/MV/dat_klm.fits', galactic=True)
    label="Planck 2018 CMB lensing (full-sky; WRONG kernel = diluted foreground)"
else:
    mfile={'des_ks':'des_y3/KS_full.fits','des_tomo1':'des_y3/wiener_tomo1.fits','des_tomo2':'des_y3/wiener_tomo2.fits'}[MAP]
    label={'des_ks':'DES Y3 KS convergence (all sources)','des_tomo1':'DES Y3 Wiener tomo1 (low-z src = FOREGROUND lens)','des_tomo2':'DES Y3 Wiener tomo2'}[MAP]
    kap, kmap = sample(mfile, galactic=False)
    m=hp.read_map('des_y3/glimpse_mask.fits'); m=hp.ud_grade(m, 1024) if hp.get_nside(m)!=1024 else m
    th=np.radians(90-dec); ph=np.radians(ra); desmask=m[hp.ang2pix(1024,th,ph)]>0.5  # inside DES footprint

good=np.isfinite(kap)
if desmask is not None: good=good & desmask
zS,resS,wS,kS,raS,decS=z[good],res[good],w[good],kap[good],ra[good],dec[good]
print(f"MAP: {label}\n  {good.sum()} SNe land on valid map pixels (footprint overlap)")
if good.sum()<25:
    print("  TOO FEW SNe on the map footprint for a meaningful fit."); sys.exit()

# standardize kappa (zero-mean unit handled by slope units); weighted regression res ~ b*kappa
kbar=np.average(kS,weights=wS); rbar=np.average(resS,weights=wS)
Sxx=np.sum(wS*(kS-kbar)**2); b=np.sum(wS*(kS-kbar)*(resS-rbar))/Sxx
sb=np.sqrt(1/Sxx)*np.sqrt(np.sum(wS*(resS-rbar-b*(kS-kbar))**2)/(good.sum()-2))
print(f"\n  SLOPE  delta-mu = b * kappa :  b = {b:+.3f} +/- {sb:.3f}   ({b/sb:+.2f} sigma from zero)")
print(f"  SIGN:  DFD predicts b>0 (farther through mass);  LCDM predicts b~0 / slightly <0.")
print(f"  -> {'DFD-sign (positive)' if b>0 else 'LCDM-sign (<=0)'}, significance {abs(b/sb):.2f} sigma.")

# rotation-null: random map rotations give the noise floor
rng=np.random.default_rng(7); nulls=[]
for s in range(200):
    dphi=rng.uniform(0,360)
    pix=hp.ang2pix(hp.get_nside(kmap), np.radians(90-decS), np.radians((raS+dphi)%360))
    kk=kmap[pix]; m=np.isfinite(kk)
    if m.sum()<25: continue
    kb=np.average(kk[m],weights=wS[m]); bn=np.sum(wS[m]*(kk[m]-kb)*(resS[m]-np.average(resS[m],weights=wS[m])))/np.sum(wS[m]*(kk[m]-kb)**2)
    nulls.append(bn)
nulls=np.array(nulls); print(f"\n  ROTATION-NULL: b_null = {nulls.mean():+.3f} +/- {nulls.std():.3f} (200 random rotations) -> empirical sigma")
print(f"  empirical significance = {b/nulls.std():+.2f} sigma (vs the rotation noise floor)")
