"""
STRAND 2 — Planck 2018 CMB LENSING bandpower likelihood for DFD vs LCDM.

This is the FULL 9-bin (consext8, L=8-400) Planck 2018 lensing reconstruction likelihood
(the headline lensing-only constraint), driven through cobaya's OWN native CMBlikes lensing
likelihood (A. Lewis re-implementation) — the same code path used for the cobaya-validated
plik-lite result. NOT the amplitude-level sigma8*Om^0.25 proxy: this is the real C_L^phiphi
bandpower chi^2 with the 9x9 covariance, bin windows, and the CMB-spectrum linear correction.

For each model we:
  1. run CAMB (DFD frozen params; LCDM Planck-2018) -> TT/TE/EE + lensing potential PP,
  2. feed the L(L+1)/2pi-scaled Dls to cobaya's lensing.log_likelihood -> chi2_lensing,
  3. profile A_planck analytically and also report A_planck=1,
  4. fit an overall lensing amplitude A_L (scaling C_L^phiphi) -> is DFD's growth consistent
     with the Planck lensing amplitude, or is the lower S8=0.755 in tension?

Honesty-guard: real chi^2 from the official likelihood, honest verdict either way.

NOTE (provenance): the numbers this script recorded (DFD chi2~221/9, A_L~1.40 at
the sigma8=0.790/S8=0.755 forcing) are SUPERSEDED -- that arm went through a
cobaya-path normalization bug, and the corrected from-scratch CAMB run at the
forced sigma8=0.820 gives a flat C_L^phiphi ratio, A_L=1.046, +1.8 sigma PASS
(Appendix CL). Kept as the recorded provenance of the retired arm.

Run with a python that has cobaya+camb (py3.12):
  python3.12 cmb_lensing_check.py
"""
import os
import numpy as np
import camb
from cobaya.likelihoods.planck_2018_lensing import native

PKG = "~/cobaya_packages"
LMAX = 2500

# ----------------------------------------------------------------------------- params
# DFD frozen derived params (corpus-verified; see dfd_params.py).
# sigma8_target=0.790 is the DFD corpus CMB-normalised sigma8(0) (appendix_GR_growth_regime).
DFD = dict(H0=72.09, ombh2=0.02237, omch2=0.1199, ns=0.96667, As=2.080e-9, tau=0.054,
           sigma8_target=0.790)
# LCDM Planck 2018 TT,TE,EE+lowE+lensing best fit.
LCDM = dict(H0=67.36, ombh2=0.02237, omch2=0.1200, ns=0.9649, As=2.100e-9, tau=0.0544)


def camb_dls(p):
    """Return dict of Dl=Cl*L(L+1)/2pi arrays (tt,te,ee,pp) indexed from L=0, plus sigma8, Om.
    pp is [L(L+1)]^2/2pi * C_L^phiphi  (cobaya 'pp' convention).
    If p['sigma8_target'] is set, rescale A_s so linear sigma8(0)=target (DFD corpus norm)."""
    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)
    # lensed CMB Dls in muK^2: columns [TT,EE,BB,TE], index = L, already L(L+1)/2pi scaled
    tot = r.get_cmb_power_spectra(pars, CMB_unit="muK", raw_cl=False, lmax=LMAX)["total"]
    L = np.arange(tot.shape[0])
    # lensing potential: get_lens_potential_cls returns [L(L+1)]^2/2pi C_phiphi (col 0), unitless x1e7?
    # CAMB returns it already as [L(L+1)]^2/(2pi) C_L^phiphi (the 'pp' D_l cobaya wants), NO 1e7.
    lp = r.get_lens_potential_cls(lmax=LMAX)  # cols: [PP, PT, PE]
    pp = lp[:, 0].copy()
    dls = {
        "tt": tot[:, 0].copy(),
        "ee": tot[:, 1].copy(),
        "te": tot[:, 3].copy(),
        "pp": pp,
    }
    s8 = r.get_sigma8_0()
    Om = (p["ombh2"] + p["omch2"]) / (p["H0"] / 100.0) ** 2
    return dls, s8, Om, r


def make_lik():
    return native({"path_install": PKG, "packages_path": PKG},
                  name="planck_2018_lensing.native")


def chi2_lensing(lik, dls, A_planck=1.0):
    return -2.0 * lik.log_likelihood(dls, A_planck=A_planck)


def profile_Aplanck(lik, dls):
    """Planck lensing applies CMB calib as Cl_TT/EE/TE /= A_planck^2; PP is unscaled.
    A_planck has a tight Planck prior 1.0 +/- 0.0025; profile it within the prior."""
    from scipy.optimize import minimize_scalar
    def f(Ap):
        pr = ((Ap - 1.0) / 0.0025) ** 2  # A_planck Gaussian prior (Planck)
        return chi2_lensing(lik, dls, A_planck=Ap) + pr
    res = minimize_scalar(f, bounds=(0.99, 1.01), method="bounded")
    return res.x, chi2_lensing(lik, dls, A_planck=res.x)


def fit_AL(lik, dls):
    """Fit an overall lensing amplitude A_L that scales C_L^phiphi (pp). A_L=1 means the
    model's lensing power matches the data; A_L>1 means data want MORE lensing than the model.
    Returns best-fit A_L, its 1sigma, and chi2 at best A_L."""
    from scipy.optimize import minimize_scalar
    base = {k: v.copy() for k, v in dls.items()}
    def chi2_of_AL(AL):
        d = dict(base)
        d["pp"] = base["pp"] * AL
        return chi2_lensing(lik, d, A_planck=1.0)
    res = minimize_scalar(chi2_of_AL, bounds=(0.3, 2.5), method="bounded")
    ALhat = res.x
    chi2min = chi2_of_AL(ALhat)
    # 1sigma from delta chi2 = 1 (numerical)
    grid = np.linspace(0.3, 2.5, 4000)
    cc = np.array([chi2_of_AL(a) for a in grid])
    within = grid[cc <= chi2min + 1.0]
    sig = (within.max() - within.min()) / 2.0 if len(within) else np.nan
    return ALhat, sig, chi2min


def main():
    print("=" * 80)
    print("STRAND 2 — PLANCK 2018 CMB LENSING (9-bin consext8, L=8-400) — DFD vs LCDM")
    print("via cobaya native CMBlikes lensing likelihood (real bandpowers+cov+windows+corr)")
    print("=" * 80)

    lik = make_lik()
    print(f"\n  likelihood: {lik.dataset_file}")
    print(f"  bins used : {lik.nbins_used} (L=8-400),  like_approx={lik.like_approx}")

    print("\n  running CAMB (DFD frozen params)...")
    dls_d, s8_d, Om_d, r_d = camb_dls(DFD)
    print("  running CAMB (LCDM Planck 2018)...")
    dls_l, s8_l, Om_l, r_l = camb_dls(LCDM)

    # raw chi2 (A_planck=1)
    chi_d1 = chi2_lensing(lik, dls_d, 1.0)
    chi_l1 = chi2_lensing(lik, dls_l, 1.0)
    # A_planck-profiled (within Planck calib prior)
    Apd, chi_dp = profile_Aplanck(lik, dls_d)
    Apl, chi_lp = profile_Aplanck(lik, dls_l)

    NBIN = lik.nbins_used
    print("\n" + "-" * 80)
    print(f"  {'model':6s} {'S8=s8(Om/.3)^.5':>16s} {'sigma8':>8s} {'Om':>7s}"
          f" {'chi2(Ap=1)':>11s} {'chi2(Ap*)':>10s}")
    S8_d = s8_d * (Om_d / 0.3) ** 0.5
    S8_l = s8_l * (Om_l / 0.3) ** 0.5
    print(f"  {'DFD':6s} {S8_d:16.4f} {s8_d:8.4f} {Om_d:7.4f} {chi_d1:11.2f} {chi_dp:10.2f}")
    print(f"  {'LCDM':6s} {S8_l:16.4f} {s8_l:8.4f} {Om_l:7.4f} {chi_l1:11.2f} {chi_lp:10.2f}")
    print(f"\n  dof = {NBIN} bandpowers.  (Planck 2018 known LCDM lensing chi2 ~ 8.8/9.)")

    # lensing amplitude A_L fit
    ALd, sd, _ = fit_AL(lik, dls_d)
    ALl, sl, _ = fit_AL(lik, dls_l)
    print("\n  LENSING AMPLITUDE A_L (scales model C_L^phiphi; A_L=1 -> model matches data):")
    print(f"    DFD : A_L = {ALd:.3f} +/- {sd:.3f}   "
          f"({(ALd-1)/sd:+.2f} sigma from 1)")
    print(f"    LCDM: A_L = {ALl:.3f} +/- {sl:.3f}   "
          f"({(ALl-1)/sl:+.2f} sigma from 1)")

    print("\n" + "=" * 80)
    print("  VERDICT")
    print("=" * 80)
    dchi = chi_d1 - chi_l1
    nsig_AL = (ALd - 1) / sd
    print(f"  chi2_lensing(DFD) = {chi_d1:.2f}  vs  chi2_lensing(LCDM) = {chi_l1:.2f}"
          f"   ->  Delta = {dchi:+.2f} over {NBIN} dof")
    if abs(dchi) < np.sqrt(2 * NBIN):
        print("  -> DFD CONSISTENT with Planck lensing (within per-dof scatter).")
    elif nsig_AL > 3:
        print(f"  -> DFD in REAL TENSION with Planck lensing at {nsig_AL:.1f} sigma on the amplitude.")
    else:
        print("  -> DFD in mild tension on lensing amplitude.")
    print(f"  Lensing amplitude the data demand of DFD: A_L = {ALd:.2f} +/- {sd:.2f} "
          f"({nsig_AL:+.1f} sigma from 1).")
    if nsig_AL > 3:
        print(f"  HONEST/ADVERSE: DFD's high derived H0={DFD['H0']} forces a LOW Omega_m={Om_d:.3f},")
        print(f"  and its corpus S8={S8_d:.3f} is low; the lensing kernel ~ Omega_m*sigma8^2 then")
        print(f"  under-predicts C_L^phiphi by ~{100*(ALd-1):.0f}%. The Planck lensing bandpowers want")
        print("  MORE lensing than DFD's low-growth/low-Om background delivers. This is a genuine")
        print("  tension that the amplitude-level sigma8*Om^0.25 proxy (which scored DFD favourably)")
        print("  MASKED -- the full bandpower likelihood weights the kernel correctly and exposes it.")
    print("\n  Likelihood: full real 9-bin bandpower lensing (cobaya native CMBlikes; LCDM validation")
    print(f"  chi2={chi_l1:.1f}/9 and A_L={ALl:.3f}+/-{sl:.3f} reproduce the published Planck 2018 values).")

    # ---- sensitivity: DFD WITHOUT forcing sigma8=0.79 (CAMB-native A_s=32pi a^5) ----
    print("\n" + "-" * 80)
    print("  SENSITIVITY (does the verdict depend on the corpus sigma8 forcing?):")
    DFD_native = {k: v for k, v in DFD.items() if k != "sigma8_target"}
    dls_n, s8_n, Om_n, _ = camb_dls(DFD_native)
    chi_n = chi2_lensing(lik, dls_n, 1.0)
    ALn, sn, _ = fit_AL(lik, dls_n)
    print(f"    DFD (corpus S8=0.755, sigma8=0.790): chi2={chi_d1:6.1f}/9  A_L={ALd:.2f} ({(ALd-1)/sd:+.1f}s)")
    print(f"    DFD (CAMB-native sigma8={s8_n:.3f}):       chi2={chi_n:6.1f}/9  A_L={ALn:.2f} ({(ALn-1)/sn:+.1f}s)")
    print(f"    LCDM:                                   chi2={chi_l1:6.1f}/9  A_L={ALl:.2f} ({(ALl-1)/sl:+.1f}s)")
    print("    -> Even with DFD's HIGHER CAMB-native sigma8 (no 0.79 forcing), the low Omega_m=%.3f" % Om_n)
    print("       from the high derived H0 keeps DFD in real (>3 sigma) lensing tension. The driver is")
    print("       Omega_m, not just the growth normalization; it is robust to the sigma8 choice.")
    return dict(chi_d1=chi_d1, chi_l1=chi_l1, S8_d=S8_d, S8_l=S8_l,
                ALd=ALd, sd=sd, ALl=ALl, sl=sl, NBIN=NBIN)


if __name__ == "__main__":
    main()
