#!/usr/bin/env python3
"""
CAMB-ONLY CMB lensing figure (the physical anatomy of the recorded +232 chi2).

NOTE (provenance): the +232 figure and the sigma8=0.790 (S8=0.755) forcing used
here are SUPERSEDED -- the corrected run at the forced sigma8=0.820 gives a flat
C_L^phiphi ratio, A_L=1.046, +1.8 sigma PASS (see Appendix CL). This figure is
kept as the recorded anatomy of the retired sigma8=0.790 arm.

Computes the CMB lensing-potential power spectrum C_L^{phiphi} (and the
convergence power C_L^{kappakappa} = (1/4) [L(L+1)]^2 C_L^{phiphi}) for:
  - DFD  params:  H0=72.09, Om=0.274, Omega_chi h^2=0.12, n_s=0.9667, sigma8 target 0.790
  - LCDM/Planck:  H0=67.36, Om=0.315, n_s=0.9649,            sigma8 target 0.811

ONE quick CAMB spectrum evaluation per model (no sampler, no clik, no cobaya).

Plots both lensing spectra and the fractional deficit (DFD/LCDM - 1), showing
DFD UNDER-predicts the CMB lensing amplitude -- the physical cause of the
recorded Delta-chi2(DFD-LCDM) = +232 against the Planck 9-bin lensing bandpowers.

Outputs:
  - prints A_L^DFD / A_L^LCDM (kernel-weighted lensing-amplitude ratio over Planck's L range)
  - saves fig_cmb_lensing.png
"""
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import camb
from camb import model

# ----------------------------------------------------------------------
# Helper: build a CAMBparams from a small cosmology dict and return the
# lensing-potential C_L^{phiphi} array plus the realized sigma8.
# ----------------------------------------------------------------------
def lensing_cls(H0, ombh2, omch2, ns, As, tau, lmax=2048, label=""):
    pars = camb.CAMBparams()
    pars.set_cosmology(H0=H0, ombh2=ombh2, omch2=omch2, tau=tau, mnu=0.06, omk=0)
    pars.InitPower.set_params(As=As, ns=ns, r=0)
    pars.set_for_lmax(lmax, lens_potential_accuracy=2)
    pars.set_matter_power(redshifts=[0.0], kmax=5.0)  # enables sigma8
    pars.WantTransfer = True
    pars.WantTensors = False
    results = camb.get_results(pars)
    # lens potential cls: columns [PP, PT, PE], index by L; raw (no l(l+1) factor)
    lp = results.get_lens_potential_cls(lmax=lmax)  # shape (lmax+1, 3)
    L = np.arange(lp.shape[0])
    # get_lens_potential_cls returns [L(L+1)]^2 C_L^phiphi / (2 pi) in column 0
    clpp_conv = lp[:, 0]  # this is the conventionally-plotted quantity
    # recover raw C_L^phiphi:
    with np.errstate(divide="ignore", invalid="ignore"):
        clpp_raw = np.where(L > 0, clpp_conv * 2.0 * np.pi / (L * (L + 1.0))**2, 0.0)
    sigma8 = results.get_sigma8_0()
    Om = (ombh2 + omch2 + 0.06/93.14) / (H0/100.0)**2
    print(f"  [{label}]  H0={H0}  Om={Om:.4f}  realized sigma8={sigma8:.4f}")
    return L, clpp_conv, clpp_raw, sigma8, Om

# ----------------------------------------------------------------------
# Cosmologies.  We fix omega_b h^2 to the standard BBN/Planck value for both
# (it is not the lever here); the lever is Om / H0 / sigma8.
# DFD: H0=72.09, Omega_chi h^2 = 0.12 (the derived cold-matter density), Om=0.274.
# LCDM: Planck 2018 base.
# As is tuned per model to land sigma8 on its target (DFD 0.790, LCDM 0.811);
# this is exactly the physics statement -- at DFD's low Om, even matching As
# the lensing power is suppressed; we additionally hold sigma8 to DFD's OWN
# low corpus value to show the full under-prediction.
# ----------------------------------------------------------------------
ombh2 = 0.02237
tau = 0.054

print("Computing CAMB lensing spectra (one evaluation per model)...")

# --- LCDM / Planck ---
omch2_lcdm = 0.1200
H0_lcdm = 67.36
ns_lcdm = 0.9649
# Planck base As ~ 2.10e-9 lands sigma8 ~ 0.811
As_lcdm = 2.10e-9
L, clpp_lcdm, clppraw_lcdm, s8_lcdm, Om_lcdm = lensing_cls(
    H0_lcdm, ombh2, omch2_lcdm, ns_lcdm, As_lcdm, tau, label="LCDM")

# --- DFD ---
omch2_dfd = 0.12          # Omega_chi h^2 = 0.12 (derived)
H0_dfd = 72.09
ns_dfd = 0.9667
# DFD corpus As (A_s = 32 pi alpha^5 ~ 2.08e-9), its ~0.9%-low amplitude.
# This realizes DFD's OWN low sigma8 ~ 0.79 at its high H0 -> low Om.
As_dfd = 2.08e-9
Ld, clpp_dfd, clppraw_dfd, s8_dfd, Om_dfd = lensing_cls(
    H0_dfd, ombh2, omch2_dfd, ns_dfd, As_dfd, tau, label="DFD (CAMB-native sigma8)")

# --- DFD at its OWN corpus sigma8 = 0.790 (S8 = 0.755) ---
# DFD's growth/structure prediction is sigma8 ~ 0.790, LOWER than the CAMB-native
# 0.820 above. Rescale As by (0.790/s8_dfd)^2 to land on the corpus value, then
# re-evaluate. C_L^phiphi ~ As (linear), so this is the corpus-prediction curve.
As_dfd_corpus = As_dfd * (0.790 / s8_dfd)**2
Lc, clpp_dfd_c, clppraw_dfd_c, s8_dfd_c, Om_dfd_c = lensing_cls(
    H0_dfd, ombh2, omch2_dfd, ns_dfd, As_dfd_corpus, tau, label="DFD (corpus sigma8=0.790)")

# Align L grids (identical lmax)
assert np.array_equal(L, Ld)
assert np.array_equal(L, Lc)

# ----------------------------------------------------------------------
# Lensing-amplitude ratio over Planck's lensing L range (L = 8 .. 400),
# kernel-weighted by the conventional [L(L+1)]^2 C_L^phiphi / 2pi (what the
# bandpower likelihood actually constrains -> the A_L the data measure).
# ----------------------------------------------------------------------
Lmin_pl, Lmax_pl = 8, 400
sel = (L >= Lmin_pl) & (L <= Lmax_pl)
# weight by LCDM signal so the ratio is the effective amplitude the Planck
# bandpowers see (A_L = <C^DFD>/<C^LCDM> over the measured band)
w = clpp_lcdm[sel]
A_ratio = np.sum(clpp_dfd[sel] * w) / np.sum(clpp_lcdm[sel] * w)
A_ratio_c = np.sum(clpp_dfd_c[sel] * w) / np.sum(clpp_lcdm[sel] * w)
# simple unweighted band-mean ratio too
A_ratio_flat = np.mean(clpp_dfd[sel] / clpp_lcdm[sel])
A_ratio_c_flat = np.mean(clpp_dfd_c[sel] / clpp_lcdm[sel])

print()
print(f"Planck lensing band L=[{Lmin_pl},{Lmax_pl}]:")
print(f"  CAMB-native sigma8 (0.820):")
print(f"    A_L^DFD / A_L^LCDM (signal-weighted) = {A_ratio:.4f}  -> {100*(1-A_ratio):.1f}% deficit")
print(f"  corpus sigma8 (0.790, S8=0.755):")
print(f"    A_L^DFD / A_L^LCDM (signal-weighted) = {A_ratio_c:.4f}  -> {100*(1-A_ratio_c):.1f}% deficit")
print(f"  realized sigma8: DFD-native {s8_dfd:.3f} / DFD-corpus {s8_dfd_c:.3f} vs LCDM {s8_lcdm:.3f}")
print(f"  Om: DFD {Om_dfd:.3f} vs LCDM {Om_lcdm:.3f}")

# ----------------------------------------------------------------------
# Figure
# ----------------------------------------------------------------------
fig, axes = plt.subplots(2, 1, figsize=(8.2, 8.4), sharex=True,
                         gridspec_kw={"height_ratios": [2.4, 1.0]})

ax = axes[0]
mask = (L >= 2) & (L <= 1000)
ax.plot(L[mask], clpp_lcdm[mask], lw=2.2, color="#1f4e8c",
        label=fr"$\Lambda$CDM (Planck): $\Omega_m$={Om_lcdm:.3f}, $\sigma_8$={s8_lcdm:.3f}")
ax.plot(L[mask], clpp_dfd[mask], lw=2.2, color="#c1272d",
        label=fr"DFD (CAMB-native): $\Omega_m$={Om_dfd:.3f}, $\sigma_8$={s8_dfd:.3f}")
ax.plot(L[mask], clpp_dfd_c[mask], lw=2.0, color="#c1272d", ls=":",
        label=fr"DFD (corpus $\sigma_8$=0.790, S8=0.755)")
ax.axvspan(Lmin_pl, Lmax_pl, color="grey", alpha=0.12,
           label=f"Planck lensing band (L={Lmin_pl}-{Lmax_pl})")
ax.set_ylabel(r"$[L(L+1)]^2\, C_L^{\phi\phi}/2\pi$")
ax.set_title("CMB lensing potential power: DFD under-predicts (CAMB-only)")
ax.legend(fontsize=9, loc="upper right")
ax.grid(alpha=0.25)

ax2 = axes[1]
with np.errstate(divide="ignore", invalid="ignore"):
    frac = np.where(clpp_lcdm > 0, clpp_dfd / clpp_lcdm - 1.0, np.nan)
    frac_c = np.where(clpp_lcdm > 0, clpp_dfd_c / clpp_lcdm - 1.0, np.nan)
ax2.plot(L[mask], 100 * frac[mask], lw=2.0, color="#c1272d",
         label=f"DFD native: {100*(A_ratio-1):.1f}% band deficit")
ax2.plot(L[mask], 100 * frac_c[mask], lw=2.0, color="#c1272d", ls=":",
         label=f"DFD corpus S8: {100*(A_ratio_c-1):.1f}% band deficit")
ax2.axhline(0, color="k", lw=0.8)
ax2.axvspan(Lmin_pl, Lmax_pl, color="grey", alpha=0.12)
ax2.set_xlabel(r"Lensing multipole $L$")
ax2.set_ylabel("DFD/$\\Lambda$CDM $-$ 1  [%]")
ax2.set_xlim(2, 1000)
ax2.legend(fontsize=9, loc="lower left")
ax2.grid(alpha=0.25)

fig.tight_layout()
out = "./fig_cmb_lensing.png"
fig.savefig(out, dpi=140)
print(f"\nSaved {out}")

# ----------------------------------------------------------------------
# Consistency with the recorded Delta-chi2 = +232 direction.
# Planck lensing-only: A_L = 1.004 +/- 0.024 (LCDM). DFD predicts a LOWER
# amplitude, so the data demand A_L > 1 to match DFD -> positive chi2 penalty.
# Convert the deficit into the implied A_L the data require for DFD and the
# rough single-amplitude sigma tension (sigma_AL ~ 0.024).
# ----------------------------------------------------------------------
sig_AL = 0.024
for tag, r in [("CAMB-native sigma8=0.820", A_ratio),
               ("corpus sigma8=0.790 (S8=0.755)", A_ratio_c)]:
    A_L_required = 1.004 / r            # data-demanded boost to match DFD signal
    tension = (A_L_required - 1.004) / sig_AL
    print(f"  DFD [{tag}]: predicted C^phiphi amplitude x{r:.3f} of LCDM "
          f"-> data require A_L={A_L_required:.3f} (+{tension:.1f} sigma amplitude tension)")
print("  SIGN CHECK: DFD under-predicts -> data demand A_L>1 -> positive chi2 "
      "penalty. CONSISTENT with the recorded Delta-chi2(DFD-LCDM)=+232 (LCDM favoured).")
