"""
DFD-CLASS — Perturbation / growth module (module 4).

Tests the combined-tension hypothesis: does DFD's (high derived H0 + low Omega_m
+ mu-law modified growth) SIMULTANEOUSLY relieve H0 AND S8/growth/lensing -- the joint
tension that is ΛCDM's real weak point?

Pieces:
  - sigma8, S8 = sigma8 sqrt(Om/0.3)  vs Planck-CMB (0.834) AND weak lensing DES+KiDS (~0.76).
  - fsigma8(z) vs RSD compilation (BOSS/eBOSS).
  - C_ell^phiphi (CMB lensing) and P(k) shape.
DFD growth law (corpus_extract.md §2): D'' + [2 + dlnH/dlna] D' - 1.5 (Geff/G) Om(a) D = 0,
  Geff/G = 1/mu(xbar), mu(x)=x/(1+x). On FRW the smooth field sits at xbar -> 0 (runaway)
  UNLESS the (asserted, program-grade) Hubble-EFE a_ext=cH gives xbar=1/(2 sqrt(alpha))=5.85
  -> Geff = 1/mu(5.85) = 1.171 G (a constant ~17% ENHANCEMENT, gated OFF at high z/Newtonian).
HONEST: that enhancement is the WRONG sign for S8 (it raises clustering). The low Omega_m is the
  RIGHT sign. We compute the NET, bracketing Geff in [1 (GR/Newtonian-screened), 1.171 (full EFE)].
Honesty-guard both ways: report real chi^2; the verdict falls where the numbers land.
"""

import numpy as np
import camb
from scipy.integrate import solve_ivp
import dfd_params as P

ZS = [0.0, 0.067, 0.15, 0.38, 0.51, 0.61, 0.70, 0.85, 1.48]

# RSD fsigma8 compilation (BOSS/eBOSS DR16 + low-z); approximate consensus values.
FS8 = [
    ("6dFGS",     0.067, 0.423, 0.055),
    ("SDSS MGS",  0.15,  0.530, 0.160),
    ("BOSS z1",   0.38,  0.497, 0.045),
    ("BOSS z2",   0.51,  0.459, 0.038),
    ("BOSS z3",   0.61,  0.436, 0.034),
    ("eBOSS LRG", 0.70,  0.473, 0.041),
    ("eBOSS ELG", 0.85,  0.315, 0.095),
    ("eBOSS QSO", 1.48,  0.462, 0.045),
]
# S8 measurements
S8_PLANCK = (0.834, 0.016)      # Planck 2018 TT,TE,EE+lowE+lensing
S8_WL     = (0.766, 0.020)      # DES-Y3 + KiDS-1000 joint (weak lensing), the "low" side


def run_camb(H0, ombh2, omch2, ns, As, tau=0.054):
    pars = camb.CAMBparams()
    pars.set_cosmology(H0=H0, ombh2=ombh2, omch2=omch2, mnu=0.06, omk=0.0, tau=tau)
    pars.InitPower.set_params(As=As, ns=ns, r=0.0)
    pars.set_matter_power(redshifts=ZS[::-1], kmax=5.0)
    pars.set_for_lmax(2500, lens_potential_accuracy=1)
    pars.NonLinear = camb.model.NonLinear_none
    res = camb.get_results(pars)
    s8_0 = res.get_sigma8_0()
    s8z = res.get_sigma8()[::-1]               # sigma8(z) for ZS (CAMB returns ascending a)
    fs8z = res.get_fsigma8()[::-1]             # f*sigma8(z), GR growth
    Om = (ombh2 + omch2 + 0.06/93.14) / (H0/100.)**2
    cl = res.get_lens_potential_cls(lmax=2000)  # [L,4]: phiphi, phiT, phiE in [L(L+1)]^2/2pi C
    return dict(res=res, s8_0=s8_0, s8z=dict(zip(ZS, s8z)), fs8z=dict(zip(ZS, fs8z)),
                Om=Om, clpp=cl[:, 0])


def E_of_a(a, Om, Or, OL):
    return np.sqrt(Om * a**-3 + Or * a**-4 + OL)


def dlnH_dlna(a, Om, Or, OL):
    E2 = Om * a**-3 + Or * a**-4 + OL
    dE2 = -3 * Om * a**-3 - 4 * Or * a**-4
    return 0.5 * a * dE2 / E2 / a   # = 0.5 dlnE2/dlna


def growth_ratio(Om, Q_late):
    """Solve growth ODE for Geff/G = Q(a); return D(a) and f(a)=dlnD/dlna, plus GR baseline.
    Q(a): 1 at high z (Newtonian screened), ramps to Q_late by today (a0~cH regime)."""
    Or = P.OMEGA_R_H2 / P.h**2
    OL = 1.0 - Om - Or
    a_grid = np.linspace(1e-3, 1.0, 4000)

    def Q_of_a(a):
        # smooth ramp: Newtonian (1) early -> Q_late late, transition near a~0.5 (z~1, when a0~cH bites)
        return 1.0 + (Q_late - 1.0) * 0.5 * (1 + np.tanh((a - 0.55) / 0.12))

    def rhs(a, y, Qfun):
        D, Dp = y
        E2 = Om * a**-3 + Or * a**-4 + OL
        dlnE2 = (-3 * Om * a**-3 - 4 * Or * a**-4) / E2
        Oma = Om * a**-3 / E2
        ddD = -(3.0 / a + 0.5 * dlnE2 / a) * Dp + 1.5 * Qfun(a) * Oma * D / a**2
        return [Dp, ddD]

    def solve(Qfun):
        sol = solve_ivp(rhs, (a_grid[0], 1.0), [a_grid[0], 1.0], args=(Qfun,),
                        t_eval=a_grid, rtol=1e-8, atol=1e-10)
        D = sol.y[0]
        f = a_grid * np.gradient(D, a_grid) / D
        return D, f
    D_gr, f_gr = solve(lambda a: 1.0)
    D_q, f_q = solve(Q_of_a)
    return a_grid, D_gr, f_gr, D_q, f_q


def main():
    print("=" * 72)
    print("PERTURBATION / GROWTH MODULE — combined H0 + S8/growth/lensing test")
    print("=" * 72)
    dfd = run_camb(P.H0, P.OMEGA_B_H2, P.OMEGA_CHI_H2, P.N_S, P.A_S)
    lcdm = run_camb(P.PLANCK["H0"], P.PLANCK["omega_b_h2"], P.PLANCK["omega_c_h2"],
                    P.PLANCK["n_s"], P.PLANCK["A_s"])

    # --- sigma8 / S8 (GR growth from CAMB) ---
    def S8(s8, Om):
        return s8 * np.sqrt(Om / 0.3)
    s8_dfd_gr, s8_l = dfd["s8_0"], lcdm["s8_0"]
    S8_dfd_gr = S8(s8_dfd_gr, dfd["Om"])
    S8_l = S8(s8_l, lcdm["Om"])
    print(f"\n[sigma8 / S8, GR growth]")
    print(f"  Omega_m:     DFD {dfd['Om']:.4f}   LCDM {lcdm['Om']:.4f}")
    print(f"  sigma8(0):   DFD {s8_dfd_gr:.4f}   LCDM {s8_l:.4f}")
    print(f"  S8:          DFD {S8_dfd_gr:.4f}   LCDM {S8_l:.4f}")
    print(f"  data: Planck-CMB S8={S8_PLANCK[0]:.3f}+/-{S8_PLANCK[1]:.3f} | "
          f"weak-lensing S8={S8_WL[0]:.3f}+/-{S8_WL[1]:.3f}")

    # --- DFD modified growth: bracket Geff in [1, 1.171] ---
    Q_LATE = 1.0 / (5.85 / 6.85)     # = 1.171
    a_g, D_gr, f_gr, D_q, f_q = growth_ratio(dfd["Om"], Q_LATE)
    ratio0 = D_q[-1] / D_gr[-1]       # growth boost at a=1 from Geff
    s8_dfd_mod = s8_dfd_gr * ratio0
    S8_dfd_mod = S8(s8_dfd_mod, dfd["Om"])
    print(f"\n[DFD modified growth, Geff/G ramps 1 -> {Q_LATE:.3f} (Hubble-EFE, late-time)]")
    print(f"  growth boost D_DFD/D_GR(a=1) = {ratio0:.4f}")
    print(f"  sigma8(0): GR {s8_dfd_gr:.4f} -> modified {s8_dfd_mod:.4f}")
    print(f"  S8:        GR {S8_dfd_gr:.4f} -> modified {S8_dfd_mod:.4f}")
    for tag, s8v, S8v in [("GR-growth", s8_dfd_gr, S8_dfd_gr), ("EFE-growth", s8_dfd_mod, S8_dfd_mod)]:
        pull_wl = (S8v - S8_WL[0]) / S8_WL[1]
        pull_pl = (S8v - S8_PLANCK[0]) / S8_PLANCK[1]
        print(f"    DFD {tag:11s} S8={S8v:.3f}: {pull_wl:+.1f}sig vs WL, {pull_pl:+.1f}sig vs Planck-CMB")
    pull_wl_l = (S8_l - S8_WL[0]) / S8_WL[1]
    print(f"    LCDM-Planck S8={S8_l:.3f}: {pull_wl_l:+.1f}sig vs WL (the S8 tension)")

    # --- fsigma8(z) vs RSD ---
    print(f"\n[fsigma8(z) vs RSD data]")
    print(f"{'survey':10s} {'z':>5s} {'obs':>7s} {'err':>6s} {'DFD-GR':>7s} {'DFD-EFE':>8s} {'LCDM':>7s}")
    chi2 = dict(dfd_gr=0., dfd_efe=0., lcdm=0.)
    f_q_interp = lambda z: np.interp(1/(1+z), a_g, f_q)
    f_gr_interp = lambda z: np.interp(1/(1+z), a_g, f_gr)
    D_q_interp = lambda z: np.interp(1/(1+z), a_g, D_q) / D_q[-1]
    for name, z, obs, err in FS8:
        s8z_dfd = dfd["s8z"][z]
        fs8_dfd_gr = dfd["fs8z"][z]
        # modified: f*sigma8 with EFE growth (rescale sigma8 by D_q/D_gr ratio at z, f by f_q)
        rz = (np.interp(1/(1+z), a_g, D_q) / D_q[-1]) / (np.interp(1/(1+z), a_g, D_gr) / D_gr[-1])
        fs8_dfd_efe = f_q_interp(z) * s8z_dfd * ratio0 * rz / f_gr_interp(z) * f_gr_interp(z)
        fs8_dfd_efe = f_q_interp(z) * (s8z_dfd * ratio0 * rz)
        fs8_lcdm = lcdm["fs8z"][z]
        chi2["dfd_gr"]  += ((fs8_dfd_gr  - obs) / err)**2
        chi2["dfd_efe"] += ((fs8_dfd_efe - obs) / err)**2
        chi2["lcdm"]    += ((fs8_lcdm    - obs) / err)**2
        print(f"{name:10s} {z:5.2f} {obs:7.3f} {err:6.3f} {fs8_dfd_gr:7.3f} {fs8_dfd_efe:8.3f} {fs8_lcdm:7.3f}")
    n = len(FS8)
    print(f"  fsigma8 chi^2/N:  DFD-GR {chi2['dfd_gr']/n:.2f}   DFD-EFE {chi2['dfd_efe']/n:.2f}   "
          f"LCDM {chi2['lcdm']/n:.2f}")

    # --- CMB lensing C_ell^phiphi amplitude ---
    L = np.arange(2, 2001)
    clpp_dfd = dfd["clpp"][2:2001]
    clpp_l = lcdm["clpp"][2:2001]
    band = (L >= 40) & (L <= 400)
    A_lens = np.mean(clpp_dfd[band] / clpp_l[band])
    print(f"\n[CMB lensing C_ell^phiphi] DFD/LCDM amplitude (40<L<400) = {A_lens:.3f}")

    # --- verdict ---
    print("\n" + "-" * 72)
    print("HONEST VERDICT (module 4 — the combined-tension test):")
    print(f"  H0:  DFD 72.09 (SH0ES side, -0.9sig) vs LCDM 67.4 (-5.5sig) -- DFD WINS this axis.")
    win_s8_gr = abs((S8_dfd_gr-S8_WL[0])/S8_WL[1]) < abs(pull_wl_l)
    print(f"  S8:  DFD-GR {S8_dfd_gr:.3f} ({(S8_dfd_gr-S8_WL[0])/S8_WL[1]:+.1f}sig WL) vs LCDM {S8_l:.3f} "
          f"({pull_wl_l:+.1f}sig WL) -- DFD {'CLOSER to lensing' if win_s8_gr else 'NOT closer'}.")
    print(f"  S8 with DFD mu-law growth boost ({Q_LATE:.2f}G): {S8_dfd_mod:.3f} -- the EFE enhancement")
    print(f"       pushes S8 {'UP (wrong way, partially undoing the low-Om gain)' if S8_dfd_mod>S8_dfd_gr else 'down'}.")
    print(f"  fsigma8: DFD-GR chi^2/N {chi2['dfd_gr']/n:.2f} vs LCDM {chi2['lcdm']/n:.2f} "
          f"-- DFD {'BETTER' if chi2['dfd_gr']<chi2['lcdm'] else 'WORSE OR TIE'} on growth.")
    print(f"  lensing A_phiphi = {A_lens:.3f} (vs LCDM 1.00).")
    combined = win_s8_gr and (chi2['dfd_gr'] <= chi2['lcdm']*1.1)
    print(f"  COMBINED: DFD {'SIMULTANEOUSLY high-H0 AND lower-S8/good-growth = beats LCDM on the joint tension' if combined else 'helps H0+S8 via low Om, but mu-law growth + the theta*/BAO screen issue keep it PARTIAL'}.")
    return dict(S8_dfd_gr=S8_dfd_gr, S8_dfd_mod=S8_dfd_mod, S8_lcdm=S8_l,
                fs8_chi2=chi2, A_lens=A_lens, s8_dfd=s8_dfd_gr)


if __name__ == "__main__":
    main()
