"""
Iteration 2: DECOMPOSE the +49 residual shape gap (after calibration) in the real Planck-spectrum test.
Is it (i) the theta* +1sigma misalignment (frozen A=0.0667 gives theta*=1.04138 vs the fitted 0.0685 -> 1.04109),
or (ii) genuine shape (n_s tilt 0.9667 vs Planck 0.9649, densities)?  Honesty-guard: report what each closes.
"""
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}
def load(f):
    d=np.loadtxt(os.path.join(DDIR,f)); return d[:,0],d[:,1],0.5*(d[:,2]+d[:,3]),d[:,4]
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']
    return np.arange(cl.shape[0]),cl,r.get_derived_params()['thetastar']

def cal_chi2(ell,cl,theta_target,th_raw):
    s=th_raw/theta_target; num=den=0.0; data={}
    for spec,f in FILES.items():
        l,Dl,err,best=load(f); t=np.interp(l/s,ell,cl[:,COL[spec]])
        num+=np.sum(Dl*t/err**2); den+=np.sum(t*t/err**2); data[spec]=(l,Dl,err,best,t)
    c=num/den; out={}; tot=0.0
    for spec,(l,Dl,err,best,t) in data.items():
        ch=np.sum(((Dl-c*t)/err)**2); out[spec]=ch; tot+=ch
    return c,tot,out

# LCDM reference chi^2 (BestFit column)
lcdm_tot=0; lcdm={}
for spec,f in FILES.items():
    l,Dl,err,best=load(f); lcdm[spec]=np.sum(((Dl-best)/err)**2); lcdm_tot+=lcdm[spec]

print("="*80)
print("ITER 2 — DECOMPOSE the residual shape gap (calibration-marginalized chi^2 vs LCDM)")
print("="*80)
print(f"\n  LCDM reference: TT {lcdm['TT']:.0f}  TE {lcdm['TE']:.0f}  EE {lcdm['EE']:.0f}  TOTAL {lcdm_tot:.0f}\n")

# CORRECTION (2026-07 tension hunt): the 'frozen theta*' cases below use the legacy
# frozen-A=0.0667 gate value 1.04138. The frozen<->fitted equivalence claim is RETRACTED:
# A is one fitted amplitude (A* = 0.0908, theta* +1.0 sigma); the frozen gate value gives
# theta* +12.2 sigma at the consistent convention. Case values kept: numeric logic untouched.
cases = [
  ("derived ns, frozen theta* (1.04138)", P.N_S, 1.04138),
  ("derived ns, PERFECT theta* (1.04109)", P.N_S, 1.04109),
  ("Planck ns=0.9649, frozen theta*",      0.9649, 1.04138),
  ("Planck ns=0.9649, PERFECT theta*",     0.9649, 1.04109),
]
print(f"  {'case':40s} {'cal':>6s} {'TT':>7s} {'TE':>6s} {'EE':>6s} {'TOTAL':>7s} {'-LCDM':>7s}")
ell_cache={}
for name,ns,th_t in cases:
    key=round(ns,4)
    if key not in ell_cache:
        ell_cache[key]=camb_cl(P.H0,P.OMEGA_B_H2,P.OMEGA_CHI_H2,ns,P.A_S)
    ell,cl,th_raw=ell_cache[key]
    c,tot,out=cal_chi2(ell,cl,th_t,th_raw)
    print(f"  {name:40s} {c:6.3f} {out['TT']:7.0f} {out['TE']:6.0f} {out['EE']:6.0f} {tot:7.0f} {tot-lcdm_tot:+7.0f}")

# where in TT does the residual live (baseline derived, frozen theta*)?
ell,cl,th_raw=ell_cache[round(P.N_S,4)]; s=th_raw/1.04138
l,Dl,err,best=load(FILES['TT']); t=np.interp(l/s,ell,cl[:,0])
num=np.sum(Dl*t/err**2); den=np.sum(t*t/err**2); c=num/den; res=((Dl-c*t)/err)**2; resL=((Dl-best)/err)**2
print(f"\n  TT residual by multipole region (calibrated DFD vs LCDM), c={c:.3f}:")
print(f"  {'ell range':14s} {'Nbin':>5s} {'chi2 DFD':>9s} {'chi2 LCDM':>10s}")
for lo,hi in [(0,500),(500,1200),(1200,2000),(2000,3000)]:
    m=(l>=lo)&(l<hi)
    print(f"  {f'{lo}-{hi}':14s} {m.sum():5d} {res[m].sum():9.1f} {resL[m].sum():10.1f}")
print("\n  Read: if PERFECT theta* collapses the gap, the +49 was the frozen-A +1sigma (fixable by the")
print("  2.7% larger fitted A). If it persists, it is genuine shape (n_s/densities/damping).")
