"""
SUB-BUILD B — NATIVE Planck-2018 CMB-lensing C_L^phiphi bandpower chi^2 (NO cobaya).

Pure numpy/scipy + CAMB. Loads the on-disk Planck 2018 lensing reconstruction
bandpowers, the full 9x9 bin covariance, the bin-window files, AND the CMB-spectrum
linear (N1/normalization) correction, then forms the real Gaussian
    chi^2 = (d - t)^T C^-1 (d - t)
for DFD and LCDM. This kills both the cobaya dependency and the 1-dof A_L proxy.

Dataset: smicadx12_Dec5_ftl_mv2_ndclpp_p_teb_consext8 — the published Planck 2018
"conservative" lensing-only likelihood (9 bandpowers, L=8-400). Identical data +
binning the cobaya native CMBlikes lensing likelihood uses, so the LCDM number here
reproduces the cobaya-validated chi2 ~ 8.7/9.

CMBlikes native theory vector (exact reproduction):
    t_b = sum_L W_b^PP(L) * Cpp(L)                      [bin window contraction]
        + sum_{X in TT,EE,TE,PP} sum_L dW_b^X(L)*Cx(L)  [CMB linear correction]
        - fid_corr_b                                    [fiducial value of that corr]
  where Cpp = [L(L+1)]^2/(2pi) C_L^phiphi (DIMENSIONLESS, CAMB pp col 0), and the
  TT/EE/TE delta-window spectra Cx = [L(L+1)]/(2pi) C_L^X are in muK^2 (CAMB
  CMB_unit="muK", raw_cl=False). VERIFIED: at the Planck fiducial model,
  sum_X dW_b^X*Cx reproduces the on-disk fiducial-correction file to <0.5% per bin
  (bin1 2.853e-7 vs 2.837e-7, bin9 7.565e-8 vs 7.590e-8), so the NET correction
  (corr - fid_corr) is a small physical N1/normalization term, not a blow-up.
  Window-only LCDM chi2=8.74 already reproduces the cobaya-validated 8.7/9.

Run:  python3 lensing_chi2_native.py
"""
import os
import numpy as np
import camb

LENS_DIR = (os.environ.get("COBAYA_PACKAGES", os.path.expanduser("~/cobaya_packages")) + "/data/planck_supp_data_and_covmats/"
            "lensing/2018")
ROOT = "smicadx12_Dec5_ftl_mv2_ndclpp_p_teb_consext8"   # 9-bin, L=8-400
NBINS = 9
LMAX = 2500

# ----- params -----------------------------------------------------------------
# DFD frozen DERIVED params (corpus). A_s, n_s are the explicit derived values
# (A_s = 32 pi alpha^5 = 2.08e-9; n_s = 1-2/Omega = 0.96667). H0=72.09, Om_chi h^2
# = 0.1199 (chi dark sector == CDM under {Om_cdm->Om_chi, H0->72.09}+optical theta*).
DFD = dict(H0=72.09, ombh2=0.02237, omch2=0.1199, ns=0.96667, As=2.080e-9, tau=0.054,
           sigma8_target=0.820)  # forced sigma8(0)=0.820 (App GR/PK growth norm)
# LCDM Planck 2018 TT,TE,EE+lowE+lensing best fit (the fits ARE its free params).
LCDM = dict(H0=67.36, ombh2=0.02237, omch2=0.1200, ns=0.9649, As=2.100e-9, tau=0.0544)


def load_data():
    bp = np.loadtxt(os.path.join(LENS_DIR, f"{ROOT}_bandpowers.dat"))  # bin Lmin Lmax Lav PP Err Ahat
    d = bp[:NBINS, 4].copy()
    cov = np.loadtxt(os.path.join(LENS_DIR, f"{ROOT}_cov.dat"))[:NBINS, :NBINS]
    Cinv = np.linalg.inv(cov)
    fid = np.loadtxt(os.path.join(LENS_DIR, f"{ROOT}_lensing_fiducial_correction.dat"))[:NBINS, 1]
    win, dwin = [], []
    for b in range(1, NBINS + 1):
        w = np.loadtxt(os.path.join(LENS_DIR, f"{ROOT}_window", f"window{b}.dat"))
        win.append((w[:, 0].astype(int), w[:, 1]))
        dw = np.loadtxt(os.path.join(LENS_DIR, f"{ROOT}_lens_delta_window",
                                     f"window{b}.dat"))  # cols: L wTT wEE wTE wPP
        dwin.append(dw)
    return d, cov, Cinv, fid, win, dwin


def camb_spectra(p):
    """Return D_l-scaled tt,ee,te (col=L) in muK^2 and pp = [L(L+1)]^2/2pi C_phiphi
    (dimensionless), plus sigma8, Om. CMB_unit="muK" matches the muK^2 normalization
    the Planck delta-windows expect (verified vs the on-disk fiducial-correction file)."""
    As = p["As"]
    if p.get("sigma8_target"):
        p0 = camb.set_params(H0=p["H0"], ombh2=p["ombh2"], omch2=p["omch2"],
                             ns=p["ns"], As=As, tau=p["tau"], lmax=LMAX + 100)
        p0.set_matter_power(redshifts=[0.0], kmax=2.0)
        s0 = camb.get_results(p0).get_sigma8_0()
        As = As * (p["sigma8_target"] / s0) ** 2
    pars = camb.set_params(H0=p["H0"], ombh2=p["ombh2"], omch2=p["omch2"],
                           ns=p["ns"], As=As, tau=p["tau"], lmax=LMAX + 100)
    pars.set_matter_power(redshifts=[0.0], kmax=2.0)
    r = camb.get_results(pars)
    tot = r.get_cmb_power_spectra(pars, CMB_unit="muK", raw_cl=False, lmax=LMAX)["total"]
    pp = r.get_lens_potential_cls(lmax=LMAX)[:, 0].copy()
    sp = {"tt": tot[:, 0].copy(), "ee": tot[:, 1].copy(), "te": tot[:, 3].copy(),
          "pp": pp}
    Om = (p["ombh2"] + p["omch2"]) / (p["H0"] / 100.0) ** 2
    return sp, r.get_sigma8_0(), Om


def theory_bandpowers(sp, fid, win, dwin, use_linear=True):
    pp = sp["pp"]
    t = np.empty(NBINS)
    for b in range(NBINS):
        L, w = win[b]
        tb = np.sum(w * pp[L])                      # bin-window contraction
        if use_linear:
            dw = dwin[b]
            Lc = dw[:, 0].astype(int)
            corr = (np.sum(dw[:, 1] * sp["tt"][Lc]) + np.sum(dw[:, 2] * sp["ee"][Lc])
                    + np.sum(dw[:, 3] * sp["te"][Lc]) + np.sum(dw[:, 4] * sp["pp"][Lc]))
            tb += corr - fid[b]                     # CMB linear correction (N1/norm)
        t[b] = tb
    return t


def chi2(d, Cinv, t):
    r = d - t
    return float(r @ Cinv @ r)


def fit_AL(d, Cinv, t):
    """Analytic best-fit overall lensing amplitude (t -> A_L t) + 1sigma."""
    tCt = t @ Cinv @ t
    AL = (t @ Cinv @ d) / tCt
    sig = 1.0 / np.sqrt(tCt)
    return AL, sig, chi2(d, Cinv, AL * t)


def main():
    print("=" * 80)
    print("NATIVE Planck-2018 CMB-lensing chi^2 (pure numpy/scipy+CAMB, NO cobaya)")
    print(f"dataset={ROOT}")
    print(f"  {NBINS} bandpowers, L=8-400, full {NBINS}x{NBINS} cov + bin windows + linear corr")
    print("=" * 80)
    d, cov, Cinv, fid, win, dwin = load_data()
    print(f"\nloaded {len(d)} bandpowers; cov {cov.shape}; {len(win)} bin windows; "
          f"{len(dwin)} delta-windows")

    out = {}
    for name, p in [("LCDM", LCDM), ("DFD", DFD)]:
        sp, s8, Om = camb_spectra(p)
        t_lin = theory_bandpowers(sp, fid, win, dwin, use_linear=True)
        t_nol = theory_bandpowers(sp, fid, win, dwin, use_linear=False)
        c_lin = chi2(d, Cinv, t_lin)
        c_nol = chi2(d, Cinv, t_nol)
        AL_f, sig_f, cAL_f = fit_AL(d, Cinv, t_lin)   # full (with linear corr)
        AL_w, sig_w, cAL_w = fit_AL(d, Cinv, t_nol)   # window-only (= record's A_L)
        S8 = s8 * (Om / 0.3) ** 0.5
        out[name] = dict(chi2=c_lin, chi2_nolin=c_nol, AL=AL_f, sig=sig_f,
                         AL_w=AL_w, sig_w=sig_w, s8=s8, Om=Om, S8=S8)
        print(f"\n--- {name}  (sigma8={s8:.4f}, Om={Om:.4f}, S8={S8:.4f}) ---")
        print(f"  chi^2 (window only)      = {c_nol:8.2f}  / {NBINS} bins")
        print(f"     A_L (window only)     = {AL_w:.4f} +/- {sig_w:.4f}   "
              f"({(AL_w-1)/sig_w:+.2f} sigma)   <- matches read-from-record")
        print(f"  chi^2 (+ CMB linear corr)= {c_lin:8.2f}  (corr effect = {c_lin-c_nol:+.2f})")
        print(f"     A_L (full)            = {AL_f:.4f} +/- {sig_f:.4f}   "
              f"({(AL_f-1)/sig_f:+.2f} sigma)")

    print("\n" + "=" * 80)
    print("SUMMARY")
    print(f"  WINDOW-ONLY (headline / record-comparable):")
    print(f"    chi^2(LCDM)={out['LCDM']['chi2_nolin']:.2f}/9  "
          f"A_L={out['LCDM']['AL_w']:.3f}+/-{out['LCDM']['sig_w']:.3f} "
          f"({(out['LCDM']['AL_w']-1)/out['LCDM']['sig_w']:+.2f}s)")
    print(f"    chi^2(DFD) ={out['DFD']['chi2_nolin']:.2f}/9  "
          f"A_L={out['DFD']['AL_w']:.3f}+/-{out['DFD']['sig_w']:.3f} "
          f"({(out['DFD']['AL_w']-1)/out['DFD']['sig_w']:+.2f}s)")
    print(f"  FULL (+ CMB linear/N1 correction):")
    print(f"    chi^2(LCDM)={out['LCDM']['chi2']:.2f}/9  A_L={out['LCDM']['AL']:.3f} "
          f"({(out['LCDM']['AL']-1)/out['LCDM']['sig']:+.2f}s)")
    print(f"    chi^2(DFD) ={out['DFD']['chi2']:.2f}/9  A_L={out['DFD']['AL']:.3f} "
          f"({(out['DFD']['AL']-1)/out['DFD']['sig']:+.2f}s)")
    print("\n  VALIDATION: native LCDM chi2=8.72/9, A_L=1.003+/-0.025 reproduces the")
    print("  cobaya-validated 8.7/9 + A_L=1.004+/-0.024 with NO cobaya (full 9x9 cov,")
    print("  bin windows, CMB linear corr; not the 1-dof amplitude proxy).")
    print("  DFD window-only A_L=%.3f (%.2fs) matches the read-from-record +1.8s 'mild PASS'."
          % (out['DFD']['AL_w'], (out['DFD']['AL_w']-1)/out['DFD']['sig_w']))
    print("  Adding the CMB linear correction (which responds to DFD's high-H0/low-Om")
    print("  background) sharpens DFD to A_L=%.2f (%.1fs) -> real lensing tension, not mild."
          % (out['DFD']['AL'], (out['DFD']['AL']-1)/out['DFD']['sig']))
    print("=" * 80)
    return out


if __name__ == "__main__":
    main()
