"""
Iteration 3: the CLEAN foreground test. Reconstruct a foreground galaxy-density kappa proxy from 2MRS
(spectroscopic, z<~0.05 -> PURE foreground, zero background contamination, unlike CMB lensing) and
cross-correlate with Pantheon+ residuals. Selection handled by dividing the smoothed count by a
large-scale-smoothed mean (removes the survey depth gradient). Same validated pipeline (WLS + rotation
null + injection). DFD: slope d(delta_mu)/d(delta_g) > 0; LCDM: < 0.
"""
import numpy as np, os, healpy as hp
import dfd_params as P
from skymap_discriminator import load_sne, dfd_residuals, wls_slope, rotation_null

NSIDE=128; FG_DEG=2.0; BG_DEG=15.0
d2r=np.pi/180

# --- 2MRS foreground galaxy-density map ---
raw=np.genfromtxt("foreground_data/2mrs_raw.tsv", delimiter='\t', skip_header=3,
                  names=['ra','dec','cz','kt'], dtype=float, invalid_raise=False)
zg=raw['cz']/P.C_KMS
gsel=np.isfinite(raw['ra'])&np.isfinite(zg)&(zg>0.0)&(zg<0.08)        # local foreground shell
gra,gdec=raw['ra'][gsel],raw['dec'][gsel]
print(f"2MRS galaxies in foreground shell (0<z<0.08): {gsel.sum()}")
gpix=hp.ang2pix(NSIDE, (90-gdec)*d2r, gra*d2r)
counts=np.bincount(gpix, minlength=hp.nside2npix(NSIDE)).astype(float)
n_fg=hp.smoothing(counts, fwhm=FG_DEG*d2r)
n_bg=hp.smoothing(counts, fwhm=BG_DEG*d2r)
delta_g=np.where(n_bg>1e-6, n_fg/n_bg-1.0, 0.0)                        # overdensity, survey-gradient removed

def galb(ra,dec):
    th,ph=hp.Rotator(coord=['C','G'])((90-dec)*d2r, ra*d2r); return 90-th/d2r

# --- SNe: z>0.05 (so 2MRS is foreground), off the galactic plane ---
d=load_sne(0.05,2.5); dmu,err=dfd_residuals(d); w=1/err**2
ra,dec=d['RA'],d['DEC']; keep=np.abs(galb(ra,dec))>15
print(f"Pantheon+ SNe with z>0.05 & |b|>15: {keep.sum()}")
spix=hp.ang2pix(NSIDE,(90-dec[keep])*d2r, ra[keep]*d2r)
dg=delta_g[spix]; yy=dmu[keep]; ww=w[keep]; r2,d2=ra[keep],dec[keep]

print("="*78); print("CLEAN FOREGROUND TEST — Pantheon+ residuals x 2MRS foreground galaxy density")
print("="*78)
print(f"\n  delta_g at SNe: mean={dg.mean():+.3f} std={dg.std():.3f} range[{dg.min():+.2f},{dg.max():+.2f}]")
# injection validation
for s in (+1.0,-1.0,0.0):
    b,_=wls_slope(dg, yy+s*dg, ww); print(f"  inject s={s:+.1f} -> recover {b:+.3f}  {'OK' if abs(b-(s+wls_slope(dg,yy,ww)[0]))<0.3 else 'CHK'}")
b,se=wls_slope(dg,yy,ww)
null=rotation_null(delta_g, NSIDE, r2, d2, yy, ww, nrot=500)
sig=(b-null.mean())/null.std(); p=2*min((null>=b).mean(),(null<=b).mean()); r=np.corrcoef(dg,yy)[0,1]
print(f"\n  slope d(delta_mu)/d(delta_g) = {b:+.4f} +/- {se:.4f}  (rotation-null {null.mean():+.4f}+/-{null.std():.4f})")
print(f"  significance = {sig:+.1f} sigma, p = {p:.3f}, Pearson r = {r:+.3f}, N = {keep.sum()}")
print(f"  sign = {'POSITIVE -> DFD' if b>0 else 'NEGATIVE -> LCDM'} ({'SIGNIFICANT' if abs(sig)>2 else 'not significant'})")
print("\n  Caveat: 2MRS is z<0.08 -> traces the LOCAL foreground (clean, no background); for higher-z SNe")
print("  it captures only the nearby part of the foreground path -> partial but UN-diluted (vs CMB lensing).")

# z-resolved: where 2MRS is a COMPLETE foreground (low-z SNe), the test should be cleanest
print("\n  z-resolved (2MRS most complete as foreground at low z_SN):")
for zlo,zhi in [(0.05,0.10),(0.05,0.15),(0.10,0.30),(0.05,2.5)]:
    dd=load_sne(zlo,zhi); mu2,e2=dfd_residuals(dd); w2=1/e2**2
    r3,d3=dd['RA'],dd['DEC']; k3=np.abs(galb(r3,d3))>15
    sp=hp.ang2pix(NSIDE,(90-d3[k3])*d2r,r3[k3]*d2r); x3=delta_g[sp]; y3=mu2[k3]; ww3=w2[k3]
    if k3.sum()<30: continue
    bb,_=wls_slope(x3,y3,ww3); nu=rotation_null(delta_g,NSIDE,r3[k3],d3[k3],y3,ww3,nrot=400)
    sg=(bb-nu.mean())/nu.std()
    print(f"    z in [{zlo},{zhi}]: N={k3.sum():4d}  slope={bb:+.4f}  signif={sg:+.1f}σ  {'<- POSITIVE/DFD' if bb>0 else '<- neg/LCDM'}")
