"""
Iteration 2/3: the REAL plik-lite likelihood with the FULL 613x613 covariance (not diagonal).
Exactly replicates cobaya's PlanckPlikLite.get_chi_squared (binning + Fortran covariance + A_planck calib).
Theory = DFD CAMB with the FROZEN screen as the theta* l-rescale (A_DFD=0.0667, no refit). Honesty-guard:
the broad TT residual that cost +49 in diagonal chi^2 may be heavily down-weighted by the real covariance
-- or not. Report the real number both ways (A_planck=1 fixed; and overall-amplitude profiled = pure shape).
"""
import numpy as np, camb, os
from scipy.io import FortranFile
import dfd_params as P

DAT = "~/cobaya_packages/data/planck_2018_pliklite_native"
NTT, NTE, NEE = 215, 199, 199; NB = NTT+NTE+NEE; OFF = 30
THETA_SCREENED = 1.04138

# --- load plik-lite data exactly as cobaya does ---
data = np.loadtxt(os.path.join(DAT,"cl_cmb_plik_v22.dat"))
Xdata = data[:,1]                                   # bandpowers (Cl units), all 613
blmin = np.loadtxt(os.path.join(DAT,"blmin.dat")).astype(int) + OFF
blmax = np.loadtxt(os.path.join(DAT,"blmax.dat")).astype(int) + OFF
w = np.loadtxt(os.path.join(DAT,"bweight.dat"))
ls = np.arange(len(w)) + OFF
w = w * 2*np.pi/ls/(ls+1)                            # work with Dl, not Cl
weights = np.hstack((np.zeros(OFF), w))
cov = FortranFile(os.path.join(DAT,"c_matrix_plik_v22.dat"),"r").read_reals(dtype=float).reshape((NB,NB))
cov = np.tril(cov) + np.tril(cov,-1).T
invcov = np.linalg.inv(cov)
print(f"loaded plik-lite: {NB} bandpowers, cov {cov.shape}, blmin/blmax len {len(blmin)}")

def bin_theory(Dl_tt, Dl_te, Dl_ee):
    cl = np.empty(NB); ix = 0
    for Dl, nbin in [(Dl_tt,NTT),(Dl_te,NTE),(Dl_ee,NEE)]:
        for i in range(nbin):
            cl[ix] = np.dot(Dl[blmin[i]:blmax[i]+1], weights[blmin[i]:blmax[i]+1]); ix += 1
    return cl

def camb_dl(H0,ombh2,omch2,ns,As,tau=0.054,screen=None):
    pars = camb.set_params(H0=H0,ombh2=ombh2,omch2=omch2,ns=ns,As=As,tau=tau,lmax=2600)
    r = camb.get_results(pars); cl = r.get_cmb_power_spectra(pars,CMB_unit='muK',raw_cl=False)['total']
    L = np.arange(cl.shape[0]); th = r.get_derived_params()['thetastar']
    tt,ee,te = cl[:,0].copy(), cl[:,1].copy(), cl[:,3].copy()
    if screen:                                       # apply theta* rescale: Dl_obs(l)=Dl_raw(l/s)
        s = th/screen; Lq = np.arange(2600)/s
        tt = np.interp(Lq,L,cl[:,0]); ee = np.interp(Lq,L,cl[:,1]); te = np.interp(Lq,L,cl[:,3])
    return tt,te,ee,th

def chi2_report(name, tt,te,ee):
    u = bin_theory(tt,te,ee)
    chi2_fix = float((Xdata-u) @ invcov @ (Xdata-u))         # A_planck=1
    # profile overall amplitude beta=1/A^2 (free calibration = pure shape):
    uCu = u@invcov@u; uCX = u@invcov@Xdata; XCX = Xdata@invcov@Xdata
    beta = uCX/uCu; chi2_shape = XCX - uCX**2/uCu
    A = beta**-0.5 if beta>0 else float('nan')
    print(f"  {name:26s} chi2(A_planck=1)={chi2_fix:7.1f}   chi2(amp-profiled)={chi2_shape:7.1f}  (A_eff={A:.4f})")
    return chi2_fix, chi2_shape

print("="*82)
print("REAL plik-lite chi^2 (full covariance), 613 bandpowers, dof~613")
print("="*82)
tt,te,ee,th = camb_dl(P.H0,P.OMEGA_B_H2,P.OMEGA_CHI_H2,P.N_S,P.A_S, screen=THETA_SCREENED)
print(f"  DFD raw 100*theta*={th:.4f}, screened->{THETA_SCREENED}")
dfd = chi2_report("DFD (frozen screen)", tt,te,ee)
lt,le,lee,lth = camb_dl(67.36,0.02237,0.12,0.9649,2.1e-9)     # Planck 2018 LCDM
lcdm = chi2_report("LCDM (Planck 2018)", lt,le,lee)
print("-"*82)
print(f"  Delta chi^2 (DFD-LCDM): A_planck=1 -> {dfd[0]-lcdm[0]:+.1f};  amplitude-profiled(shape) -> {dfd[1]-lcdm[1]:+.1f}")
print(f"  dof = 613.  amplitude-profiled isolates SHAPE (free overall norm). The full covariance is the")
print(f"  honest test; compare the +49 DIAGONAL shape gap from iter-1 to the covariant shape gap here.")
