"""
FIGURE matter_pk — two panels:
 (a) linear matter power P(k) at z=0, DFD vs LCDM (CAMB get_matter_power_spectrum)
 (b) growth f*sigma8(z) for DFD (GR + EFE-modified) vs RSD data vs LCDM.
Confirms S8 ~ 0.785 and fsigma8 chi2 claims; annotates S8 on the plot.
"""
import sys
sys.path.insert(0, ".")

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import camb

import dfd_params as P
import perturbations as pert
from perturbations import FS8, S8_PLANCK, S8_WL, ZS

OUT_PNG = "./figures/matter_pk.png"


def get_pk(H0, ombh2, omch2, ns, As, tau=0.054, kmax=5.0):
    """Linear matter P(k) at z=0 plus sigma8(0)."""
    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=[0.0], kmax=kmax)
    pars.NonLinear = camb.model.NonLinear_none
    res = camb.get_results(pars)
    kh, z, pk = res.get_matter_power_spectrum(minkh=1e-4, maxkh=kmax, npoints=400)
    s8 = res.get_sigma8_0()
    return kh, pk[0], s8


def S8(s8, Om):
    return s8 * np.sqrt(Om / 0.3)


def main():
    # ---- panel (b) inputs: full CAMB growth run (reproduces perturbations.main) ----
    dfd = pert.run_camb(P.H0, P.OMEGA_B_H2, P.OMEGA_CHI_H2, P.N_S, P.A_S)
    lcdm = pert.run_camb(P.PLANCK["H0"], P.PLANCK["omega_b_h2"], P.PLANCK["omega_c_h2"],
                         P.PLANCK["n_s"], P.PLANCK["A_s"])

    # ---- panel (a) inputs: linear P(k) z=0 ----
    kh_d, pk_d, s8_d_pk = get_pk(P.H0, P.OMEGA_B_H2, P.OMEGA_CHI_H2, P.N_S, P.A_S)
    kh_l, pk_l, s8_l_pk = get_pk(P.PLANCK["H0"], P.PLANCK["omega_b_h2"],
                                 P.PLANCK["omega_c_h2"], P.PLANCK["n_s"], P.PLANCK["A_s"])

    # ---- S8 (GR growth from CAMB) ----
    s8_dfd_gr, s8_lcdm = dfd["s8_0"], lcdm["s8_0"]
    S8_dfd_gr = S8(s8_dfd_gr, dfd["Om"])
    S8_lcdm = S8(s8_lcdm, lcdm["Om"])

    # ---- DFD modified (EFE) growth boost, Geff/G ramps 1 -> 1.171 ----
    Q_LATE = 1.0 / (5.85 / 6.85)   # = 1.171
    a_g, D_gr, f_gr, D_q, f_q = pert.growth_ratio(dfd["Om"], Q_LATE)
    ratio0 = D_q[-1] / D_gr[-1]
    s8_dfd_mod = s8_dfd_gr * ratio0
    S8_dfd_mod = S8(s8_dfd_mod, dfd["Om"])

    # ---- fsigma8(z): GR, EFE, LCDM at the data redshifts + chi2 ----
    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)
    chi2 = dict(dfd_gr=0., dfd_efe=0., lcdm=0.)
    zz, obs_a, err_a, fs8_gr_a, fs8_efe_a, fs8_lcdm_a = [], [], [], [], [], []
    for name, z, obs, err in FS8:
        s8z_dfd = dfd["s8z"][z]
        fs8_dfd_gr = dfd["fs8z"][z]
        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)
        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
        zz.append(z); obs_a.append(obs); err_a.append(err)
        fs8_gr_a.append(fs8_dfd_gr); fs8_efe_a.append(fs8_dfd_efe); fs8_lcdm_a.append(fs8_lcdm)
    n = len(FS8)

    # smooth fsigma8(z) curves over a fine z grid for plotting (GR + LCDM via CAMB grid)
    zgrid = np.array(ZS)
    fs8_dfd_grid = np.array([dfd["fs8z"][z] for z in ZS])
    fs8_lcdm_grid = np.array([lcdm["fs8z"][z] for z in ZS])
    order = np.argsort(zgrid)
    zgrid_s = zgrid[order]
    fs8_dfd_grid_s = fs8_dfd_grid[order]
    fs8_lcdm_grid_s = fs8_lcdm_grid[order]

    # ---- console confirmation ----
    print("=== matter_pk figure: confirmation numbers ===")
    print(f"Omega_m: DFD {dfd['Om']:.4f}  LCDM {lcdm['Om']:.4f}")
    print(f"sigma8(0): DFD {s8_dfd_gr:.4f} (P(k) run {s8_d_pk:.4f})  "
          f"LCDM {s8_lcdm:.4f} (P(k) run {s8_l_pk:.4f})")
    print(f"S8 (GR): DFD {S8_dfd_gr:.4f}  LCDM {S8_lcdm:.4f}")
    print(f"S8 (DFD EFE-modified): {S8_dfd_mod:.4f}  (boost {ratio0:.4f})")
    print(f"fsigma8 chi2: DFD-GR {chi2['dfd_gr']:.2f} (chi2/N {chi2['dfd_gr']/n:.2f})  "
          f"DFD-EFE {chi2['dfd_efe']:.2f} ({chi2['dfd_efe']/n:.2f})  "
          f"LCDM {chi2['lcdm']:.2f} ({chi2['lcdm']/n:.2f})")
    print(f"data: Planck-CMB S8={S8_PLANCK[0]} WL S8={S8_WL[0]}")

    # ============================ PLOT ============================
    plt.rcParams.update({"font.size": 11, "axes.grid": True,
                         "grid.alpha": 0.3, "figure.dpi": 150})
    fig, (axA, axB) = plt.subplots(1, 2, figsize=(13, 5.2))

    # ---- panel (a): P(k) ----
    axA.loglog(kh_d, pk_d, color="#c0392b", lw=2.2,
               label=f"DFD ($\\sigma_8$={s8_d_pk:.3f})")
    axA.loglog(kh_l, pk_l, color="#2c3e50", lw=2.0, ls="--",
               label=f"$\\Lambda$CDM Planck ($\\sigma_8$={s8_l_pk:.3f})")
    axA.set_xlabel(r"$k\ \ [h\,\mathrm{Mpc}^{-1}]$")
    axA.set_ylabel(r"$P(k)\ \ [(h^{-1}\mathrm{Mpc})^3]$")
    axA.set_title("(a) Linear matter power spectrum, $z=0$")
    axA.legend(loc="lower left", frameon=True)
    axA.set_xlim(1e-4, 5.0)
    axA.annotate(f"$S_8^{{\\rm DFD}}={S8_dfd_gr:.3f}$\n$S_8^{{\\Lambda\\rm CDM}}={S8_lcdm:.3f}$",
                 xy=(0.97, 0.95), xycoords="axes fraction",
                 fontsize=11, va="top", ha="right",
                 bbox=dict(boxstyle="round", fc="#fdf2e9", ec="#c0392b", alpha=0.9))

    # ---- panel (b): fsigma8(z) ----
    axB.errorbar(zz, obs_a, yerr=err_a, fmt="o", color="black", ms=6, capsize=3,
                 label="RSD data (BOSS/eBOSS)", zorder=5)
    axB.plot(zgrid_s, fs8_dfd_grid_s, color="#c0392b", lw=2.2, marker="s", ms=4,
             label=f"DFD GR-growth ($\\chi^2/N$={chi2['dfd_gr']/n:.2f})")
    axB.plot(zz, fs8_efe_a, color="#e67e22", lw=1.8, ls=":", marker="^", ms=5,
             label=f"DFD EFE-growth ($\\chi^2/N$={chi2['dfd_efe']/n:.2f})")
    axB.plot(zgrid_s, fs8_lcdm_grid_s, color="#2c3e50", lw=2.0, ls="--", marker="D", ms=4,
             label=f"$\\Lambda$CDM ($\\chi^2/N$={chi2['lcdm']/n:.2f})")
    axB.set_xlabel(r"redshift $z$")
    axB.set_ylabel(r"$f\sigma_8(z)$")
    axB.set_title("(b) Growth rate $f\\sigma_8(z)$ vs RSD")
    axB.legend(loc="lower left", frameon=True, fontsize=9.5)
    axB.set_xlim(0, 1.6)

    fig.suptitle("DFD vs $\\Lambda$CDM: matter clustering and growth",
                 fontsize=13, fontweight="bold")
    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fig.savefig(OUT_PNG, dpi=150)
    print(f"saved {OUT_PNG}")

    return dict(S8_dfd_gr=S8_dfd_gr, S8_lcdm=S8_lcdm, S8_dfd_mod=S8_dfd_mod,
                s8_dfd=s8_dfd_gr, s8_lcdm=s8_lcdm, Om_dfd=dfd["Om"], Om_lcdm=lcdm["Om"],
                chi2=chi2, n=n, ratio0=ratio0)


if __name__ == "__main__":
    main()
