"""
TARGET D — UNIFICATION verdict: does ONE DFD normalization fit
   [CMB-lensing phiphi] + [cosmic-shear S8] + [fsigma8 RSD]  SIMULTANEOUSLY,
or is there an IRREDUCIBLE SPLIT?

Real CAMB throughout. DFD frozen DERIVED norm (no re-fit):
   H0=72.09, omch2=0.1199, ns=0.96667, As=2.080e-9 (-> sigma8=0.820, Om=0.274, S8=0.783).
LCDM-Planck best fit: H0=67.36, omch2=0.1200, ns=0.9649, As=2.10e-9.

Three growth/lensing probes, same model normalization each:
  (1) phiphi : native Planck-2018 9-bandpower lensing chi^2 (full 9x9 cov + windows
               + CMB linear corr). Reuses lensing_chi2_native.
  (2) S8     : cosmic-shear DES-Y3+KiDS-1000 joint, S8=0.766+/-0.020 (1 dof).
  (3) fsigma8: BOSS/eBOSS RSD, 8 points (GR growth from CAMB get_fsigma8).

Sum the chi^2 over the joint growth+lensing sector; quantify NET Delta chi^2.
Then fold into the FULL #1 joint evidence (ln B), with + without growth+lensing,
to see if growth/lensing NET-HELPS or NET-HURTS DFD's case.
"""
import os
import numpy as np
import camb
import lensing_chi2_native as L

# ---------------- model definitions (frozen) ----------------
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)

# ---------------- data ----------------
S8_WL = (0.766, 0.020)        # DES-Y3 + KiDS-1000 joint cosmic shear (the "low S8")
FS8 = [   # BOSS/eBOSS DR16 + low-z RSD compilation (z, obs, err)
    (0.067, 0.423, 0.055), (0.15, 0.530, 0.160), (0.38, 0.497, 0.045),
    (0.51, 0.459, 0.038), (0.61, 0.436, 0.034), (0.70, 0.473, 0.041),
    (0.85, 0.315, 0.095), (1.48, 0.462, 0.045),
]
ZS = sorted(set([z for z, _, _ in FS8]))


def run_camb(p):
    """One CAMB run at the model's frozen norm. Returns sigma8(0), Om, S8,
    fsigma8 at the RSD redshifts, and CMB+lensing spectra for the bandpower code."""
    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=L.LMAX + 100)
        p0.set_matter_power(redshifts=ZS + [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=L.LMAX + 100)
    pars.set_matter_power(redshifts=ZS + [0.0], kmax=2.0)
    pars.set_for_lmax(2500, lens_potential_accuracy=1)
    r = camb.get_results(pars)
    s8_0 = r.get_sigma8_0()
    Om = (p["ombh2"] + p["omch2"]) / (p["H0"] / 100.0) ** 2
    S8 = s8_0 * (Om / 0.3) ** 0.5
    # fsigma8(z): CAMB returns ascending a, i.e. descending z; map back to ZS
    z_camb = ZS + [0.0]
    fs8_all = r.get_fsigma8()       # ordered by ascending scale factor (z high->low)
    # get_fsigma8 returns in the order of redshifts passed reversed; build a dict safely
    s8z = r.get_sigma8()
    # CAMB orders results by increasing a (decreasing z); our input was ZS+[0.0]
    z_in = np.array(ZS + [0.0])
    order = np.argsort(-z_in)        # increasing a
    fs8_map = {}
    for idx, zi in zip(order, np.sort(z_in)[::-1]):
        pass
    # robust: re-evaluate via interpolation on the returned grid
    z_sorted_desc = np.sort(z_in)[::-1]    # matches ascending-a order
    fs8_by_z = {float(z): float(f) for z, f in zip(z_sorted_desc, fs8_all)}
    sp = {"tt": None}
    tot = r.get_cmb_power_spectra(pars, CMB_unit="muK", raw_cl=False, lmax=L.LMAX)["total"]
    pp = r.get_lens_potential_cls(lmax=L.LMAX)[:, 0].copy()
    spec = {"tt": tot[:, 0].copy(), "ee": tot[:, 1].copy(), "te": tot[:, 3].copy(), "pp": pp}
    return dict(s8_0=s8_0, Om=Om, S8=S8, fs8_by_z=fs8_by_z, spec=spec, As=As)


def phiphi_chi2(spec):
    d, cov, Cinv, fid, win, dwin = L.load_data()
    t_win = L.theory_bandpowers(spec, fid, win, dwin, use_linear=False)
    t_lin = L.theory_bandpowers(spec, fid, win, dwin, use_linear=True)
    return L.chi2(d, Cinv, t_win), L.chi2(d, Cinv, t_lin)


def s8_chi2(S8):
    return ((S8 - S8_WL[0]) / S8_WL[1]) ** 2


def fs8_chi2(fs8_by_z):
    c = 0.0
    rows = []
    for z, obs, err in FS8:
        # nearest grid z (we ran exactly the RSD z's)
        zk = min(fs8_by_z.keys(), key=lambda k: abs(k - z))
        pred = fs8_by_z[zk]
        c += ((pred - obs) / err) ** 2
        rows.append((z, obs, err, pred))
    return c, rows


def main():
    print("=" * 84)
    print("TARGET D — JOINT growth+lensing UNIFICATION: phiphi + shear-S8 + fsigma8")
    print("  one frozen normalization per model; real CAMB; honesty-guard both ways")
    print("=" * 84)

    out = {}
    for name, p in [("LCDM", LCDM), ("DFD", DFD)]:
        m = run_camb(p)
        c_win, c_lin = phiphi_chi2(m["spec"])
        c_s8 = s8_chi2(m["S8"])
        c_fs8, rows = fs8_chi2(m["fs8_by_z"])
        out[name] = dict(s8_0=m["s8_0"], Om=m["Om"], S8=m["S8"],
                         phi_win=c_win, phi_lin=c_lin, c_s8=c_s8, c_fs8=c_fs8, rows=rows)
        print(f"\n--- {name}  sigma8={m['s8_0']:.4f}  Om={m['Om']:.4f}  S8={m['S8']:.4f} ---")
        print(f"  phiphi-lensing  chi^2 = {c_win:7.2f} /9 (window)   {c_lin:7.2f} /9 (+linear corr)")
        print(f"  shear-S8        chi^2 = {c_s8:7.2f} /1   (S8={m['S8']:.3f} vs WL {S8_WL[0]}+/-{S8_WL[1]})")
        print(f"  fsigma8 (RSD)   chi^2 = {c_fs8:7.2f} /8")

    print("\n" + "=" * 84)
    print("JOINT growth+lensing chi^2 (sum over phiphi + shear-S8 + fsigma8)")
    print("=" * 84)
    for tag, key in [("phiphi WINDOW-only", "phi_win"), ("phiphi +LINEAR-corr", "phi_lin")]:
        for sub in ["full"]:
            jd = out["DFD"][key] + out["DFD"]["c_s8"] + out["DFD"]["c_fs8"]
            jl = out["LCDM"][key] + out["LCDM"]["c_s8"] + out["LCDM"]["c_fs8"]
            print(f"\n[{tag}]   N_dof = 9 + 1 + 8 = 18")
            print(f"  {'probe':22s} {'DFD':>9s} {'LCDM':>9s}   {'DFD-LCDM':>9s}")
            print(f"  {'phiphi-lensing':22s} {out['DFD'][key]:9.2f} {out['LCDM'][key]:9.2f}   "
                  f"{out['DFD'][key]-out['LCDM'][key]:+9.2f}")
            print(f"  {'shear-S8':22s} {out['DFD']['c_s8']:9.2f} {out['LCDM']['c_s8']:9.2f}   "
                  f"{out['DFD']['c_s8']-out['LCDM']['c_s8']:+9.2f}")
            print(f"  {'fsigma8 (RSD)':22s} {out['DFD']['c_fs8']:9.2f} {out['LCDM']['c_fs8']:9.2f}   "
                  f"{out['DFD']['c_fs8']-out['LCDM']['c_fs8']:+9.2f}")
            print(f"  {'-'*52}")
            print(f"  {'JOINT TOTAL':22s} {jd:9.2f} {jl:9.2f}   {jd-jl:+9.2f}")
            net = jd - jl
            who = "DFD WINS" if net < -0.5 else ("LCDM WINS" if net > 0.5 else "TIE")
            print(f"  -> NET Delta chi^2 (DFD-LCDM) = {net:+.2f}  =>  {who} the joint growth+lensing sector")
            out[f"joint_{key}"] = dict(dfd=jd, lcdm=jl, net=net)

    # ---- decompose the SPLIT ----
    print("\n" + "=" * 84)
    print("SPLIT DIAGNOSIS (is the tension irreducible?)")
    print("=" * 84)
    print(f"  shear-S8 + fsigma8 (LOW-side probes): DFD {out['DFD']['c_s8']+out['DFD']['c_fs8']:.2f} "
          f"vs LCDM {out['LCDM']['c_s8']+out['LCDM']['c_fs8']:.2f}  "
          f"-> DFD better by {(out['LCDM']['c_s8']+out['LCDM']['c_fs8'])-(out['DFD']['c_s8']+out['DFD']['c_fs8']):.2f}")
    print(f"  phiphi-lensing (HIGH-side probe, window): DFD {out['DFD']['phi_win']:.2f} "
          f"vs LCDM {out['LCDM']['phi_win']:.2f}  "
          f"-> LCDM better by {out['DFD']['phi_win']-out['LCDM']['phi_win']:.2f}")
    print(f"  phiphi-lensing (+linear corr): DFD {out['DFD']['phi_lin']:.2f} "
          f"vs LCDM {out['LCDM']['phi_lin']:.2f}  "
          f"-> LCDM better by {out['DFD']['phi_lin']-out['LCDM']['phi_lin']:.2f}")

    return out


if __name__ == "__main__":
    main()
