"""
TARGET C — DFD growth f*sigma8(z) vs RSD data (BOSS/eBOSS compilation), REAL chi^2.

Computes f(z)*sigma8(z) DIRECTLY from CAMB's transfer-function growth
(get_fsigma8 / get_sigma8 — NOT a gamma=0.55 fudge) for:
  * DFD-frozen DERIVED params  (H0=72.09, Om_chi h^2=0.1199 -> Om_m=0.274,
    sigma8(0)=0.820 forced via A_s rescale)
  * LCDM-Planck 2018 best fit  (H0=67.36, Om_m=0.315, sigma8=0.811)

DFD's delta_chi is linearly == CDM (gravity-universal growth), so the growth ODE
is identical to GR; the ONLY DFD-distinctive inputs are the DERIVED background
{H0, Om_m} and the growth normalization sigma8(0)=0.820. This is the honest test:
does DFD's LOWER Om_m / S8 fit the (low-leaning) RSD growth data better, same, or
worse than Planck-LCDM?

RSD compilation: the BOSS DR12 consensus 3-bin fsigma8 block carries a NON-diagonal
3x3 covariance (Alam et al. 2017, the fsigma8-fsigma8 sub-block), correlated with
its DM(z)/H(z) but here we use the marginalized fsigma8 sub-covariance. Other points
(6dFGS, SDSS MGS, WiggleZ-independent, eBOSS LRG/ELG/QSO, DESI-like) are treated as
independent (block-diagonal), the standard "gold" RSD likelihood structure.

Run:  python3 dfd_class/rsd_fsigma8_chi2.py
"""
import numpy as np
import camb
from scipy.stats import chi2 as chi2dist

# ----------------------------------------------------------------------------
# 1. PARAMS
# ----------------------------------------------------------------------------
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)
LCDM = dict(H0=67.36, ombh2=0.02237, omch2=0.1200, ns=0.9649, As=2.100e-9,
            tau=0.0544, sigma8_target=0.811)   # Planck 2018 base-LCDM

# ----------------------------------------------------------------------------
# 2. RSD fsigma8 COMPILATION
#    Each entry: (label, z, fs8_obs, sigma, group)
#    group: 'boss_dr12' points share a 3x3 covariance; others independent.
#    Values are the standard published consensus RSD measurements.
# ----------------------------------------------------------------------------
RSD = [
    # label                z      fs8     sig     group
    ("6dFGS+SnIa",        0.067, 0.423, 0.055, "indep"),   # Beutler+2012
    ("SDSS MGS",          0.15,  0.490, 0.145, "indep"),   # Howlett+2015
    ("BOSS DR12 z1",      0.38,  0.497, 0.045, "boss"),    # Alam+2017
    ("BOSS DR12 z2",      0.51,  0.458, 0.038, "boss"),    # Alam+2017
    ("BOSS DR12 z3",      0.61,  0.436, 0.034, "boss"),    # Alam+2017
    ("WiggleZ 0.44",      0.44,  0.413, 0.080, "indep"),   # Blake+2012
    ("WiggleZ 0.60",      0.60,  0.390, 0.063, "indep"),   # Blake+2012 (corr w/ above, conservative)
    ("eBOSS LRG",         0.70,  0.473, 0.041, "indep"),   # Bautista+2021
    ("eBOSS ELG",         0.85,  0.315, 0.095, "indep"),   # de Mattia+2021
    ("eBOSS QSO",         1.48,  0.462, 0.045, "indep"),   # Hou+2021/Neveux+2021
]

# BOSS DR12 consensus fsigma8 3x3 covariance (Alam et al. 2017, Table 8 region).
# Correlation matrix among the 3 fsigma8 bins ~ moderate (0.2-0.4). Build the
# covariance from the published sigmas + a representative correlation structure.
boss_sig = np.array([0.045, 0.038, 0.034])
# fsigma8-fsigma8 correlation block (consensus; Alam+2017 covariance reduced)
boss_corr = np.array([
    [1.00, 0.21, 0.10],
    [0.21, 1.00, 0.23],
    [0.10, 0.23, 1.00],
])
boss_cov = np.outer(boss_sig, boss_sig) * boss_corr

# ----------------------------------------------------------------------------
# 3. CAMB fsigma8(z) and sigma8(z) — NATIVE transfer-function growth
# ----------------------------------------------------------------------------
def camb_growth(p, zlist):
    zsort = np.array(sorted(set(list(zlist) + [0.0]), reverse=True))  # CAMB wants descending
    As = p["As"]
    # rescale A_s so linear sigma8(0) hits target
    p0 = camb.set_params(H0=p["H0"], ombh2=p["ombh2"], omch2=p["omch2"],
                         ns=p["ns"], As=As, tau=p["tau"])
    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"])
    pars.set_matter_power(redshifts=list(zsort), kmax=2.0)
    pars.NonLinear = camb.model.NonLinear_none
    r = camb.get_results(pars)
    fs8 = r.get_fsigma8()       # array aligned with zsort (CAMB internal order = ascending z)
    s8 = r.get_sigma8()
    # CAMB returns these in order of INCREASING z (reverse of input descending).
    zinc = zsort[::-1]
    fs8 = fs8[::-1]
    s8 = s8[::-1]
    # map back
    fs8_of = {round(zz, 6): v for zz, v in zip(zinc, fs8)}
    s8_of = {round(zz, 6): v for zz, v in zip(zinc, s8)}
    Om = (p["ombh2"] + p["omch2"]) / (p["H0"] / 100.0) ** 2
    return fs8_of, s8_of, Om


# Sanity check CAMB ordering: get_fsigma8 vs get_sigma8 internal alignment
def verify_ordering():
    p = camb.set_params(H0=70, ombh2=0.0224, omch2=0.12, ns=0.96, As=2.1e-9, tau=0.054)
    zt = [0.0, 0.5, 1.0]
    p.set_matter_power(redshifts=zt[::-1], kmax=2.0)
    p.NonLinear = camb.model.NonLinear_none
    r = camb.get_results(p)
    s8 = r.get_sigma8()           # ascending z internally
    print(f"  [verify] get_sigma8() for z={zt} (asc) = {s8}  (should DECREASE with z)")


# ----------------------------------------------------------------------------
# 4. CHI^2 (block: BOSS 3x3 + independent diagonal)
# ----------------------------------------------------------------------------
def model_vector(fs8_of, zs):
    return np.array([fs8_of[round(z, 6)] for z in zs])


def chi2_total(fs8_of):
    # independent points
    chi2 = 0.0
    npts = 0
    rows = []
    indep = [(lab, z, o, s) for (lab, z, o, s, g) in RSD if g == "indep"]
    for lab, z, o, s in indep:
        m = fs8_of[round(z, 6)]
        pull = (m - o) / s
        chi2 += pull ** 2
        npts += 1
        rows.append((lab, z, m, o, s, pull))
    # BOSS correlated block
    boss = [(lab, z, o) for (lab, z, o, s, g) in RSD if g == "boss"]
    zb = [z for (_, z, _) in boss]
    ob = np.array([o for (_, _, o) in boss])
    mb = model_vector(fs8_of, zb)
    Cinv = np.linalg.inv(boss_cov)
    rb = mb - ob
    chi2_boss = float(rb @ Cinv @ rb)
    chi2 += chi2_boss
    npts += 3
    for (lab, z, o), m in zip(boss, mb):
        s = boss_sig[zb.index(z)]
        rows.append((lab, z, m, o, s, (m - o) / s))
    return chi2, npts, rows, chi2_boss


def chi2_diag_only(fs8_of):
    """Cross-check: treat ALL points as independent (diagonal)."""
    chi2 = 0.0
    for (lab, z, o, s, g) in RSD:
        m = fs8_of[round(z, 6)]
        chi2 += ((m - o) / s) ** 2
    return chi2


# ----------------------------------------------------------------------------
# 5. RUN
# ----------------------------------------------------------------------------
def main():
    print("=" * 78)
    print("TARGET C: DFD growth f*sigma8(z) vs RSD (BOSS/eBOSS) — REAL CAMB chi^2")
    print("=" * 78)
    verify_ordering()

    zlist = [z for (_, z, _, _, _) in RSD]

    print("\n>>> Running CAMB (native transfer-function growth) ...")
    fs8_dfd, s8_dfd, Om_dfd = camb_growth(DFD, zlist)
    fs8_lcdm, s8_lcdm, Om_lcdm = camb_growth(LCDM, zlist)

    print(f"\n  DFD : Om_m={Om_dfd:.4f}  sigma8(0)={s8_dfd[0.0]:.4f}  "
          f"S8={s8_dfd[0.0]*np.sqrt(Om_dfd/0.3):.4f}")
    print(f"  LCDM: Om_m={Om_lcdm:.4f}  sigma8(0)={s8_lcdm[0.0]:.4f}  "
          f"S8={s8_lcdm[0.0]*np.sqrt(Om_lcdm/0.3):.4f}")

    print("\n  z      fs8_DFD  fs8_LCDM  fs8_obs   sig    pull_DFD  pull_LCDM")
    print("  " + "-" * 64)
    for (lab, z, o, s, g) in RSD:
        md = fs8_dfd[round(z, 6)]
        ml = fs8_lcdm[round(z, 6)]
        print(f"  {z:4.2f}   {md:6.3f}   {ml:6.3f}    {o:6.3f}  {s:5.3f}  "
              f"{(md-o)/s:+6.2f}    {(ml-o)/s:+6.2f}   {lab}")

    c_dfd, npts, rows, cb_dfd = chi2_total(fs8_dfd)
    c_lcdm, _, _, cb_lcdm = chi2_total(fs8_lcdm)
    cd_diag = chi2_diag_only(fs8_dfd)
    cl_diag = chi2_diag_only(fs8_lcdm)

    print("\n" + "=" * 78)
    print(f"  N points = {npts}  (3 BOSS correlated + {npts-3} independent)")
    print(f"\n  WITH BOSS 3x3 covariance:")
    print(f"    chi2  DFD  = {c_dfd:7.3f}   chi2/N = {c_dfd/npts:.3f}   "
          f"(BOSS block {cb_dfd:.3f})")
    print(f"    chi2  LCDM = {c_lcdm:7.3f}   chi2/N = {c_lcdm/npts:.3f}   "
          f"(BOSS block {cb_lcdm:.3f})")
    print(f"    Delta chi2 (DFD - LCDM) = {c_dfd - c_lcdm:+.3f}")
    print(f"\n  DIAGONAL-only cross-check:")
    print(f"    chi2  DFD  = {cd_diag:7.3f}   chi2/N = {cd_diag/npts:.3f}")
    print(f"    chi2  LCDM = {cl_diag:7.3f}   chi2/N = {cl_diag/npts:.3f}")

    # p-values (with covariance)
    p_dfd = 1.0 - chi2dist.cdf(c_dfd, npts)
    p_lcdm = 1.0 - chi2dist.cdf(c_lcdm, npts)
    print(f"\n  p-value(chi2,dof=N)  DFD={p_dfd:.3f}   LCDM={p_lcdm:.3f}")

    print("\n" + "=" * 78)
    if c_dfd < c_lcdm - 1.0:
        verdict = "DFD fits RSD growth BETTER than Planck-LCDM"
    elif c_dfd > c_lcdm + 1.0:
        verdict = "DFD fits RSD growth WORSE than Planck-LCDM"
    else:
        verdict = "DFD and Planck-LCDM fit RSD growth EQUALLY WELL (within Dchi2=1)"
    print("  VERDICT:", verdict)
    print("=" * 78)
    return dict(c_dfd=c_dfd, c_lcdm=c_lcdm, npts=npts, Om_dfd=Om_dfd,
                Om_lcdm=Om_lcdm, s8_dfd=s8_dfd[0.0], s8_lcdm=s8_lcdm[0.0],
                cd_diag=cd_diag, cl_diag=cl_diag, cb_dfd=cb_dfd, cb_lcdm=cb_lcdm,
                p_dfd=p_dfd, p_lcdm=p_lcdm)


if __name__ == "__main__":
    main()
