"""
Iteration 3: REAL Pantheon+ supernova test, full STAT+SYS covariance, frozen DFD screen (no refit).
1701 SNe; cosmology sample zHD>0.01, non-calibrators. mu = m_b_corr - M, M marginalized analytically.
DFD: D_L^opt=(1+z)*D_M^opt with the integrated-kernel screen delta-psi=A*Omega_L(z), A=0.0667. vs LCDM.
Honesty-guard: real chi^2, full covariance, report DFD vs LCDM honestly.
"""
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

d = np.genfromtxt(os.path.join(PD,'PantheonPlusSH0ES.dat'), names=True, dtype=None, encoding=None)
z = d['zHD']; mb = d['m_b_corr']; iscal = d['IS_CALIBRATOR']
cov = np.loadtxt(os.path.join(PD,'PantheonPlusSH0ES_STATSYS.cov'), skiprows=1).reshape(len(z), len(z))
sel = (z > 0.01) & (iscal == 0)
print(f"loaded Pantheon+: {len(z)} SNe, cosmology sample (z>0.01, non-cal) = {sel.sum()}")
zc, mbc = z[sel], mb[sel]
Cs = cov[np.ix_(sel, sel)]
Ci = np.linalg.inv(Cs)

# distance-modulus models (H0/M absorbed by marginalized offset; only the SHAPE matters)
zg = np.linspace(0, zc.max()*1.05, 40000)
OL = 1 - P.OMEGA_M - P.OMEGA_R_H2/(P.H0/100)**2
def shp(z): return OL/(OL + P.OMEGA_M*(1+z)**3)
cH_dfd = C/(P.H0*np.array([B.E(zz) for zz in zg]))
DM_dfd = cumulative_trapezoid(np.exp(A_SCR*shp(zg))*cH_dfd, zg, initial=0.0)   # screened transverse
def mu_dfd(zz):
    DMo = np.interp(zz, zg, DM_dfd); DL = (1+zz)*DMo; return 5*np.log10(DL) + 25
pl = P.PLANCK; OmL = pl['Omega_m']; OrL = P.OMEGA_R_H2/(pl['H0']/100)**2; OLL = 1-OmL-OrL
from scipy.integrate import quad
def mu_lcdm(zz):
    DM = (C/pl['H0'])*quad(lambda x: 1/np.sqrt(OmL*(1+x)**3+OrL*(1+x)**4+OLL), 0, zz, limit=80)[0]
    return 5*np.log10((1+zz)*DM) + 25

def chi2_marg(mu_model):
    dvec = mbc - mu_model
    one = np.ones_like(dvec)
    A = dvec@Ci@dvec; Bb = one@Ci@dvec; Cc = one@Ci@one
    return A - Bb*Bb/Cc                          # marginalized over the M offset

print("="*72)
print("REAL PANTHEON+ SNe — DFD (frozen screen) vs LCDM, full STAT+SYS covariance")
print("="*72)
mu_d = np.array([mu_dfd(zz) for zz in zc])
mu_l = np.array([mu_lcdm(zz) for zz in zc])
c_d = chi2_marg(mu_d); c_l = chi2_marg(mu_l); n = sel.sum()
print(f"\n  DFD  chi^2 = {c_d:.1f} / {n}  (chi^2/N = {c_d/n:.3f})")
print(f"  LCDM chi^2 = {c_l:.1f} / {n}  (chi^2/N = {c_l/n:.3f})")
print(f"  Delta chi^2 (DFD - LCDM) = {c_d-c_l:+.1f}")
print(f"  screen at SNe: delta-psi(0)={A_SCR*shp(0):.4f}, (0.5)={A_SCR*shp(0.5):.4f}, (1)={A_SCR*shp(1):.4f}")
print(f"\n  {'verdict':8s}: " + ("DFD ties/beats LCDM on real SNe" if c_d-c_l < 2*np.sqrt(2*n) else
       "DFD worse -- diagnose"))
print("  (H0 & abs-mag M marginalized; SNe constrain the distance-redshift SHAPE, which the frozen")
print("   screen modifies. Same A=0.0667 as the CMB/BAO test, no refit.)")
