"""
REAL Planck-spectrum test (first pass, binned bandpowers, diagonal errors).
Data: Planck 2018 (PR3) binned TT/TE/EE from the PLA (planck_data/*.txt); columns l, Dl, -dDl, +dDl, BestFit.
Theory: DFD CAMB on derived background params, with the screen applied ONLY as the theta* angular rescale.
CORRECTION (2026-07 tension hunt): the 'FROZEN, NO refit of A' framing is retracted -- A is ONE FITTED
amplitude (A* = 0.0908, theta* +1.0 sigma); the frozen 2sqrt(alpha)-gate value 0.0667 (-> 100theta*=1.04138,
used below) gives theta* +12.2 sigma and plik-lite dchi2 = +202/+172 at the consistent convention.
Numeric logic untouched. NO foreground/calibration fudge.
Compares chi^2(DFD) vs chi^2(LCDM Planck best-fit) against the SAME real data. Honesty-guard: real numbers.
Honest scope: diagonal errors only (full plik-lite covariance is the next step); theory at bin centers.
"""
import numpy as np, camb, os
import dfd_params as P

DDIR = os.path.join(os.path.dirname(__file__), 'planck_data')
FILES = {'TT':'COM_PowerSpect_CMB-TT-binned_R3.01.txt',
         'TE':'COM_PowerSpect_CMB-TE-binned_R3.02.txt',
         'EE':'COM_PowerSpect_CMB-EE-binned_R3.02.txt'}
COL = {'TT':0, 'EE':1, 'TE':3}            # CAMB total-cl column index
THETA_SCREENED = 1.04138                  # frozen-A DFD theta* (gate theorem); Planck 1.04109

def load(f):
    d = np.loadtxt(os.path.join(DDIR, f))
    l, Dl, dlo, dhi, best = d[:,0], d[:,1], d[:,2], d[:,3], d[:,4]
    return l, Dl, 0.5*(dlo+dhi), best

def camb_cl(H0, ombh2, omch2, ns, As, tau=0.054):
    pars = camb.set_params(H0=H0, ombh2=ombh2, omch2=omch2, ns=ns, As=As, tau=tau, lmax=2700)
    r = camb.get_results(pars)
    cl = r.get_cmb_power_spectra(pars, CMB_unit='muK', raw_cl=False)['total']  # Dl, [:,0]TT 1EE 2BB 3TE
    return np.arange(cl.shape[0]), cl, r.get_derived_params()['thetastar']

print("="*78)
print("REAL PLANCK-SPECTRUM TEST — DFD (frozen derived screen) vs LCDM, real PR3 binned data")
print("="*78)
ell, cl_dfd, th_dfd = camb_cl(P.H0, P.OMEGA_B_H2, P.OMEGA_CHI_H2, P.N_S, P.A_S)
s = th_dfd/THETA_SCREENED                  # ell-axis rescale from the optical screen
print(f"\n  DFD raw 100*theta* = {th_dfd:.4f}; frozen screen -> {THETA_SCREENED} (Planck 1.04109).")
print(f"  screen ell-rescale factor s = {s:.5f}  (peaks shift by {100*(s-1):+.2f}%)")
print(f"  DFD derived params: H0={P.H0}, ombh2={P.OMEGA_B_H2}, omch2={P.OMEGA_CHI_H2:.4f}, "
      f"ns={P.N_S:.4f}, As={P.A_S:.3e}, tau=0.054 (no foreground/calibration nuisance fit)\n")

print(f"  {'spec':5s} {'Nbins':>6s} {'chi2 DFD':>11s} {'chi2 LCDM':>11s} {'chi2/N DFD':>11s} {'chi2/N LCDM':>12s}")
tot_dfd = tot_lcdm = ntot = 0
for spec, f in FILES.items():
    l, Dl, err, best = load(f)
    dfd_obs = np.interp(l/s, ell, cl_dfd[:, COL[spec]])      # observed = raw at ell/s
    c_dfd = np.sum(((Dl-dfd_obs)/err)**2)
    c_lcdm = np.sum(((Dl-best)/err)**2)
    n = len(l); tot_dfd += c_dfd; tot_lcdm += c_lcdm; ntot += n
    print(f"  {spec:5s} {n:6d} {c_dfd:11.1f} {c_lcdm:11.1f} {c_dfd/n:11.2f} {c_lcdm/n:12.2f}")
print(f"  {'TOTAL':5s} {ntot:6d} {tot_dfd:11.1f} {tot_lcdm:11.1f} {tot_dfd/ntot:11.2f} {tot_lcdm/ntot:12.2f}")

print("\n" + "-"*78)
print(f"  DFD total chi^2 = {tot_dfd:.0f} / {ntot} bins (chi^2/N={tot_dfd/ntot:.2f})")
print(f"  LCDM total chi^2 = {tot_lcdm:.0f} / {ntot} bins (chi^2/N={tot_lcdm/ntot:.2f})  [Planck best-fit ref]")
d = tot_dfd-tot_lcdm
print(f"  Delta chi^2 (DFD - LCDM) = {d:+.0f}.  "
      + ("DFD competitive." if d < 2*ntot**0.5 else "DFD worse than LCDM by >2sqrt(N) -- diagnose."))
print("  HONEST: diagonal errors (no plik-lite covariance yet), theory at bin centres, derived A_s/tau")
print("  (no nuisance fit). Next: full covariance + nuisance marginalization (plik-lite).")

# DIAGNOSTIC: how much of the gap is an overall amplitude/calibration offset (which plik MARGINALIZES)?
print("\n  DIAGNOSTIC — marginalize a single overall calibration c (Dl -> c*Dl), as plik does:")
def cal_chi2(specs):
    num=den=0.0
    for spec,f in FILES.items():
        if spec not in specs: continue
        l,Dl,err,best = load(f); t = np.interp(l/s, ell, cl_dfd[:,COL[spec]])
        num += np.sum(Dl*t/err**2); den += np.sum(t*t/err**2)
    c = num/den; chi2=0.0
    for spec,f in FILES.items():
        if spec not in specs: continue
        l,Dl,err,best = load(f); t = c*np.interp(l/s, ell, cl_dfd[:,COL[spec]])
        chi2 += np.sum(((Dl-t)/err)**2)
    return c, chi2
cc, ck = cal_chi2(['TT','TE','EE'])
ctt, ktt = cal_chi2(['TT'])
print(f"    best overall calibration c = {cc:.4f} ({100*(cc-1):+.1f}%); TT-only c = {ctt:.4f} ({100*(ctt-1):+.1f}%)")
print(f"    DFD chi^2 after calibration: total {ck:.0f} (chi^2/N={ck/ntot:.2f}) vs LCDM {tot_lcdm:.0f} ({tot_lcdm/ntot:.2f})")
print(f"    => calibration absorbs {tot_dfd-ck:.0f} of the {tot_dfd-tot_lcdm:.0f} gap; residual SHAPE gap = {ck-tot_lcdm:+.0f}.")
print(f"    (A_s offset: DFD {P.A_S:.3e} vs Planck 2.1e-9 = {100*(P.A_S/2.1e-9-1):+.1f}%, partly the calibration.)")
