"""
DFD vs LCDM SKY-MAP DISCRIMINATOR (real data).
Cross-correlate Pantheon+ SN Hubble residuals delta_mu with the Planck 2018 lensing convergence kappa
sampled at each SN sightline. DFD predicts slope d(delta_mu)/d(kappa) > 0 (overdense LOS -> objects look
FARTHER); LCDM lensing predicts < 0 (magnification -> CLOSER). OPPOSITE SIGN = the discriminant.

Honest significance: SN residuals and kappa both have spatial structure, so the naive t-test over-counts.
We use a ROTATION NULL (rotate the kappa map by random angles, re-sample, recompute slope) for the p-value.
kappa-source caveat: Planck kappa is CMB lensing (source z~1100), a DILUTED but SIGN-PRESERVING foreground
tracer (kappa_CMB shares the foreground modes with kappa_fg). z-cuts probe where the foreground fraction grows.

Parametrized so the robustness battery can vary lmax, smoothing, z-cut, mask, residual def.
"""
import numpy as np, os
from scipy.integrate import cumulative_trapezoid
import dfd_params as P, background as B

PD = os.path.join(os.path.dirname(__file__),'pantheon_data')
C, A_SCR = P.C_KMS, 0.06667

def load_sne(zmin=0.01, zmax=2.5, drop_cal=True):
    d=np.genfromtxt(os.path.join(PD,'PantheonPlusSH0ES.dat'),names=True,dtype=None,encoding=None)
    sel=(d['zHD']>zmin)&(d['zHD']<zmax)
    if drop_cal: sel &= (d['IS_CALIBRATOR']==0)
    return d[sel]

def dfd_residuals(d):
    """delta_mu = m_b_corr - mu_DFD(z) - M_marg (frozen-screen model; M marginalized to zero-mean)."""
    z=d['zHD']
    zg=np.linspace(0,z.max()*1.05,40000); cH=C/(P.H0*np.array([B.E(x) for x in zg]))
    OL=1-P.OMEGA_M-P.OMEGA_R_H2/(P.H0/100)**2
    DM=cumulative_trapezoid(np.exp(A_SCR*(OL/(OL+P.OMEGA_M*(1+zg)**3)))*cH,zg,initial=0.0)
    mu=np.array([5*np.log10((1+x)*np.interp(x,zg,DM))+25 for x in z])
    err=d['m_b_corr_err_DIAG']; w=1.0/err**2
    res=d['m_b_corr']-mu; M=np.sum(w*res)/np.sum(w)            # marginalized offset
    return res-M, err

def kappa_at(klm_file, ra, dec, lmax=400, fwhm_deg=1.0, nside=1024, mask_file=None):
    import healpy as hp
    klm=hp.read_alm(klm_file)
    lmax_alm=hp.Alm.getlmax(len(klm))
    klm=hp.almxfl(klm, (np.arange(lmax_alm+1)<=lmax).astype(float))   # low-pass
    kmap=hp.alm2map(klm, nside)
    kmap=hp.smoothing(kmap, fwhm=np.radians(fwhm_deg))
    kmap=hp.Rotator(coord=['G','C']).rotate_map_pixel(kmap)   # Planck is GALACTIC -> rotate to EQUATORIAL (SN RA/Dec)
    theta=np.radians(90.0-dec); phi=np.radians(ra)
    pix=hp.ang2pix(nside, theta, phi)
    kap=kmap[pix]; good=np.ones(len(ra),bool)
    if mask_file and os.path.exists(mask_file):
        m=hp.ud_grade(hp.read_map(mask_file), nside); good=m[pix]>0.5
    return kap, good, kmap, nside

def wls_slope(x, y, w):
    xb=np.sum(w*x)/np.sum(w); yb=np.sum(w*y)/np.sum(w)
    b=np.sum(w*(x-xb)*(y-yb))/np.sum(w*(x-xb)**2)
    a=yb-b*xb; resid=y-(a+b*x); se=np.sqrt(np.sum(w*resid**2)/np.sum(w)/np.sum(w*(x-xb)**2))
    return b, se

def rotation_null(kmap, nside, ra, dec, dmu, w, nrot=500, seed_list=None):
    import healpy as hp
    theta=np.radians(90.0-dec)
    slopes=[]
    # rotate by deterministic angle grid (no RNG; Math.random unavailable-style determinism)
    for i in range(nrot):
        dphi=2*np.pi*(i+0.5)/nrot
        pix=hp.ang2pix(nside, theta, np.radians(ra)+dphi)
        b,_=wls_slope(kmap[pix], dmu, w); slopes.append(b)
    return np.array(slopes)

def run(klm_file, lmax=400, fwhm_deg=1.0, zmin=0.01, zmax=2.5, mask_file=None, label="baseline"):
    d=load_sne(zmin,zmax); dmu,err=dfd_residuals(d); w=1.0/err**2
    kap,good,kmap,nside=kappa_at(klm_file, d['RA'], d['DEC'], lmax, fwhm_deg, mask_file=mask_file)
    x=kap[good]; y=dmu[good]; ww=w[good]
    b,se=wls_slope(x,y,ww)
    null=rotation_null(kmap, nside, d['RA'][good], d['DEC'][good], y, ww)
    nm,ns=null.mean(), null.std()
    z_signif=(b-nm)/ns
    p=2*min((null>=b).mean(),(null<=b).mean())
    r=np.corrcoef(x,y)[0,1]
    print(f"[{label}] N={good.sum()}  slope d(dmu)/d(kappa) = {b:+.3f} +/- {se:.3f}  "
          f"(rotation-null {nm:+.3f}+/-{ns:.3f}); signif {z_signif:+.1f} sigma, p={p:.3f}; Pearson r={r:+.3f}")
    print(f"   sign = {'POSITIVE -> DFD' if b>0 else 'NEGATIVE -> LCDM'} "
          f"({'significant' if abs(z_signif)>2 else 'not significant'})")
    return dict(label=label, N=int(good.sum()), slope=b, se=se, signif=z_signif, p=p, r=r, sign=np.sign(b))

if __name__=="__main__":
    import sys
    klm=sys.argv[1] if len(sys.argv)>1 else "lensing_map/MV/dat_klm.fits"
    print("="*78); print("DFD vs LCDM SKY-MAP DISCRIMINATOR — Pantheon+ residuals x Planck lensing kappa")
    print("="*78)
    run(klm, label="baseline (lmax400, 1deg, all z)")
