"""
=====================================================================================
 dfd_likelihoods.py  --  SHARED per-sector cosmology chi^2 for DFD vs LambdaCDM
=====================================================================================
SINGLE SOURCE OF TRUTH for the per-sector chi^2 used by:
    joint_fit.py
    cmb_model_comparison.py
    paramcount_sensitivity.py

Before this module existed each of those scripts carried its own copy of the
plik-lite binning / BAO / Pantheon+ / lensing chi^2, and paramcount_sensitivity.py
HARDCODED the numbers -- which can silently drift from the live values. Everything
is now defined ONCE here and imported, so the BIC/AIC tables can never disagree
with the live likelihoods.

----------------------------------------------------------------------------------
WHAT RUNS LIVE IN THIS ENVIRONMENT vs WHAT IS COBAYA-GATED
----------------------------------------------------------------------------------
LIVE (pure CAMB + numpy/scipy, no cobaya):
    * chi2_cmb()            plik-lite TT/TE/EE, full 613x613 Fortran covariance,
                            overall amplitude (A_s/A_planck norm) analytically profiled.
    * chi2_bao()            12 consensus BAO points (DFD frozen kernel-split screen
                            for DFD; standard FRW for LCDM).
    * chi2_sne()            Pantheon+ 1580 SNe, full STAT+SYS cov, abs-mag M marginalised.
    * chi2_lensing_proxy()  amplitude-level lensing proxy  ((sigma8 Om^0.25 - 0.589)/0.020)^2.
                            This is the historical 'joint_fit.py' lensing line that
                            produces the DFD 1993.2 / LCDM 1987.5 headline.

COBAYA-GATED (NOT computable in this environment -- cobaya/clik not installed):
    * chi2_lensing_real()   the REAL Planck 2018 9-bin lensing bandpower likelihood
                            (cobaya native CMBlikes). If cobaya is importable it runs
                            live; otherwise it raises LensingUnavailable so callers
                            fall back EXPLICITLY (never a silent fabricated number).

So every number a script prints is either (a) recomputed live here, or (b) honestly
flagged as cobaya-gated -- but in BOTH cases sourced from THIS one module, so the
three scripts can never drift apart.

HONESTY-GUARD: this is a pure DRY/no-drift refactor. No chi^2 value was changed to make
DFD look better; the live numbers reproduce the pre-refactor headline exactly.
"""
import os
import numpy as np
from scipy.io import FortranFile
from scipy.integrate import cumulative_trapezoid, quad

import dfd_params as P
import background as B

C_KMS = P.C_KMS

# DFD frozen optical screen (derived; NO refit).  THSCR = 100*theta* target the DFD
# CMB Cl are rescaled to; A_SCR = screen amplitude in the BAO/SNe distance gate.
A_SCR = 0.06667
THSCR = 1.04138

_HERE = os.path.dirname(os.path.abspath(__file__))
_PKG = os.environ.get("COBAYA_PACKAGES", os.path.expanduser("~/cobaya_packages"))  # set COBAYA_PACKAGES to your cobaya packages dir
_DATL = os.path.join(_PKG, "data", "planck_2018_pliklite_native")
_PD = os.path.join(_HERE, "pantheon_data")

# Planck-2018 lensing amplitude proxy target (sigma8 * Om^0.25), 1-sigma.
LENS_PROXY = (0.589, 0.020)


# =====================================================================================
# 1.  CMB  plik-lite likelihood  (full 613x613 Fortran covariance, cobaya-validated)
#     LIVE: pure numpy on the Planck data products; no cobaya needed.
# =====================================================================================
_NTT, _NTE, _NEE = 215, 199, 199
NB_CMB = _NTT + _NTE + _NEE        # 613 bandpowers
_OFF = 30

_d = np.loadtxt(os.path.join(_DATL, "cl_cmb_plik_v22.dat"))
_Xcmb = _d[:, 1]
_blmin = np.loadtxt(os.path.join(_DATL, "blmin.dat")).astype(int) + _OFF
_blmax = np.loadtxt(os.path.join(_DATL, "blmax.dat")).astype(int) + _OFF
_w = np.loadtxt(os.path.join(_DATL, "bweight.dat"))
_ls = np.arange(len(_w)) + _OFF
_W = np.hstack((np.zeros(_OFF), _w * 2 * np.pi / _ls / (_ls + 1)))
_cov = FortranFile(os.path.join(_DATL, "c_matrix_plik_v22.dat"), "r").read_reals(dtype=float).reshape((NB_CMB, NB_CMB))
_ICMB = np.linalg.inv(np.tril(_cov) + np.tril(_cov, -1).T)


def _bin(tt, te, ee):
    """Bin model D_l (TT,TE,EE) onto the 613 plik-lite bandpowers."""
    u = np.empty(NB_CMB)
    ix = 0
    for Dl, nb in [(tt, _NTT), (te, _NTE), (ee, _NEE)]:
        for i in range(nb):
            u[ix] = np.dot(Dl[_blmin[i]:_blmax[i] + 1], _W[_blmin[i]:_blmax[i] + 1])
            ix += 1
    return u


def chi2_cmb(tt, te, ee):
    """plik-lite chi^2 with the overall amplitude (A_s / A_planck normalisation)
    analytically profiled out -> the fair plik-lite number (A_s is a free norm).
    LIVE in this environment (CAMB + numpy)."""
    u = _bin(tt, te, ee)
    uCu = u @ _ICMB @ u
    uCX = u @ _ICMB @ _Xcmb
    XCX = _Xcmb @ _ICMB @ _Xcmb
    return XCX - uCX ** 2 / uCu


def chi2_cmb_fixed_amp(tt, te, ee):
    """plik-lite chi^2 at A_planck = 1 (no amplitude freedom). LIVE."""
    u = _bin(tt, te, ee)
    return float((_Xcmb - u) @ _ICMB @ (_Xcmb - u))


# =====================================================================================
# 2.  BAO  (12 consensus points).  LIVE.
# =====================================================================================
_zg = np.linspace(0, P.Z_STAR, 80000)
_cHd = C_KMS / (P.H0 * np.array([B.E(z) for z in _zg]))
_rd = B.r_drag()
_OL = 1 - P.OMEGA_M - P.OMEGA_R_H2 / (P.H0 / 100) ** 2


def _shp(z):
    return _OL / (_OL + P.OMEGA_M * (1 + z) ** 3)


# DFD optical comoving distance integrand with the frozen screen folded in.
_DMd_opt = cumulative_trapezoid(np.exp(A_SCR * _shp(_zg)) * _cHd, _zg, initial=0.0)


def _cHd_at(z):
    return C_KMS / (P.H0 * B.E(z))


def chi2_bao_dfd():
    """BAO chi^2 for DFD (frozen kernel-split optical screen, NO refit). LIVE."""
    c = 0.0
    for b in B.BAO:
        z = b["z"]
        DMo = np.interp(z, _zg, _DMd_opt)
        DHo = np.exp(A_SCR * _shp(z)) * _cHd_at(z)
        pred = DMo / _rd if b["kind"] == "DM" else DHo / _rd if b["kind"] == "DH" else (z * DMo ** 2 * DHo) ** (1 / 3) / _rd
        c += ((pred - b["val"]) / b["err"]) ** 2
    return c


def chi2_bao_lcdm(Om=None, H0=None, rdL=None):
    """BAO chi^2 for LCDM (standard FRW). Defaults to the Planck-2018 reference. LIVE."""
    if Om is None:
        Om = P.PLANCK["Omega_m"]
    if H0 is None:
        H0 = P.PLANCK["H0"]
    if rdL is None:
        rdL = P.PLANCK["r_drag"]
    Or = P.OMEGA_R_H2 / (H0 / 100) ** 2
    OLL = 1 - Om - Or

    def DM(z):
        return (C_KMS / H0) * quad(lambda x: 1 / np.sqrt(Om * (1 + x) ** 3 + Or * (1 + x) ** 4 + OLL), 0, z, limit=80)[0]

    def DH(z):
        return C_KMS / (H0 * np.sqrt(Om * (1 + z) ** 3 + Or * (1 + z) ** 4 + OLL))

    c = 0.0
    for b in B.BAO:
        z = b["z"]
        dm, dh = DM(z), DH(z)
        pred = dm / rdL if b["kind"] == "DM" else dh / rdL if b["kind"] == "DH" else (z * dm ** 2 * dh) ** (1 / 3) / rdL
        c += ((pred - b["val"]) / b["err"]) ** 2
    return c


# =====================================================================================
# 3.  Pantheon+ SNe  (1580 cosmology sample, full STAT+SYS cov, abs-mag M marginalised). LIVE.
# =====================================================================================
_sn = np.genfromtxt(os.path.join(_PD, "PantheonPlusSH0ES.dat"), names=True, dtype=None, encoding=None)
_zsn = _sn["zHD"]
_mb = _sn["m_b_corr"]
_cal = _sn["IS_CALIBRATOR"]
_covsn = np.loadtxt(os.path.join(_PD, "PantheonPlusSH0ES_STATSYS.cov"), skiprows=1).reshape(len(_zsn), len(_zsn))
_ssel = (_zsn > 0.01) & (_cal == 0)
SN_Z = _zsn[_ssel]                  # the 1580 cosmology-sample redshifts
_mbc = _mb[_ssel]
_ISN = np.linalg.inv(_covsn[np.ix_(_ssel, _ssel)])
N_SNE = int(_ssel.sum())


def chi2_sne(mu):
    """Pantheon+ chi^2 for a model distance modulus array mu(SN_Z), with the absolute
    magnitude M analytically marginalised. LIVE."""
    mu = np.asarray(mu)
    dv = _mbc - mu
    one = np.ones_like(dv)
    A = dv @ _ISN @ dv
    Bb = one @ _ISN @ dv
    Cc = one @ _ISN @ one
    return A - Bb * Bb / Cc


def mu_dfd(zz):
    """DFD distance modulus (optical-screened comoving distance)."""
    DMo = np.interp(zz, _zg, _DMd_opt)
    return 5 * np.log10((1 + zz) * DMo) + 25


def mu_lcdm(zz, Om=None, H0=None):
    """LCDM distance modulus. Scalar or array zz. Defaults to Planck-2018 reference."""
    if Om is None:
        Om = P.PLANCK["Omega_m"]
    if H0 is None:
        H0 = P.PLANCK["H0"]
    Or = P.OMEGA_R_H2 / (H0 / 100) ** 2
    OLL = 1 - Om - Or

    def _mu(z):
        DM = (C_KMS / H0) * quad(lambda x: 1 / np.sqrt(Om * (1 + x) ** 3 + Or * (1 + x) ** 4 + OLL), 0, z, limit=80)[0]
        return 5 * np.log10((1 + z) * DM) + 25

    if np.ndim(zz) == 0:
        return _mu(float(zz))
    return np.array([_mu(z) for z in np.asarray(zz)])


def chi2_sne_dfd():
    """Pantheon+ chi^2 for DFD (frozen optical screen). LIVE."""
    return chi2_sne(np.array([mu_dfd(z) for z in SN_Z]))


def chi2_sne_lcdm(Om=None, H0=None):
    """Pantheon+ chi^2 for LCDM. Defaults to Planck-2018 reference. LIVE."""
    return chi2_sne(mu_lcdm(SN_Z, Om=Om, H0=H0))


# =====================================================================================
# 4a.  Planck lensing -- amplitude PROXY  ((sigma8 Om^0.25 - 0.589)/0.020)^2.  LIVE.
#      This is the historical joint_fit.py lensing line; it produces the
#      DFD 1993.2 / LCDM 1987.5 four-sector headline.  Amplitude-level only.
# =====================================================================================
def chi2_lensing_proxy(s8, Om):
    """Amplitude-level Planck-lensing proxy. LIVE (needs only sigma8 and Om from CAMB).
    NOTE: this is a PROXY. The real 9-bin bandpower likelihood (chi2_lensing_real)
    weights the lensing kernel ~ Omega_m sigma8^2 correctly and shows DFD in tension;
    that one is cobaya-gated. See FALSIFY.md."""
    return ((s8 * Om ** 0.25 - LENS_PROXY[0]) / LENS_PROXY[1]) ** 2


# =====================================================================================
# 4b.  Planck lensing -- REAL 9-bin bandpower likelihood (cobaya native CMBlikes).
#      COBAYA-GATED: only runs if cobaya/clik is importable in this environment.
# =====================================================================================
class LensingUnavailable(RuntimeError):
    """Raised when the real Planck lensing bandpower likelihood cannot be computed
    because cobaya/clik is not installed in this environment."""


_LENSLIK = None
_LENSLIK_ERR = None


def _get_lenslik():
    """Lazily construct the cobaya lensing likelihood. Raises LensingUnavailable if
    cobaya is not importable (the case in this environment)."""
    global _LENSLIK, _LENSLIK_ERR
    if _LENSLIK is not None:
        return _LENSLIK
    if _LENSLIK_ERR is not None:
        raise LensingUnavailable(_LENSLIK_ERR)
    try:
        from cobaya.likelihoods.planck_2018_lensing import native as LensLik
    except Exception as e:                                      # noqa: BLE001
        _LENSLIK_ERR = f"cobaya/clik not available: {e!r} -- requires cobaya plik-lite to recompute live"
        raise LensingUnavailable(_LENSLIK_ERR)
    _LENSLIK = LensLik({"path_install": _PKG, "packages_path": _PKG}, name="planck_2018_lensing.native")
    return _LENSLIK


def lensing_real_available():
    """True iff the cobaya real-lensing likelihood can be constructed here."""
    try:
        _get_lenslik()
        return True
    except LensingUnavailable:
        return False


def nbins_lensing_real():
    """Number of lensing bandpowers used by the real likelihood (9). Cobaya-gated."""
    return _get_lenslik().nbins_used


def chi2_lensing_real(lens_dls):
    """Real Planck 2018 9-bin lensing-reconstruction chi^2 (A_planck=1). lens_dls is a
    dict with the UNSCREENED tt/te/ee (for the N1/linear correction) and physical
    pp = C_L^phiphi. COBAYA-GATED: raises LensingUnavailable if cobaya is absent."""
    lik = _get_lenslik()
    return -2.0 * lik.log_likelihood(lens_dls, A_planck=1.0)


# =====================================================================================
# NOTE on RSD / growth (fsigma8):  paramcount_sensitivity.py historically carried a 5th
# hardcoded row (6.6 / 7.3) for the RSD/growth sector. That number comes from
# perturbations.py's modified-growth pipeline (growth ODE + CAMB fsigma8), NOT from the
# four-sector joint_fit stack, and is NOT part of the DFD 1993.2 / LCDM 1987.5 headline.
# A clean GR-only live recompute does NOT reproduce 6.6/7.3 (perturbations.py uses a
# different growth treatment), so we do NOT fabricate a shared RSD chi^2 here. The four
# sectors that ALL THREE scripts share, and that define the headline, are the ones
# centralised in this module; RSD/growth stays the responsibility of perturbations.py.
# =====================================================================================

# =====================================================================================
# 5.  CAMB driver + convenience (DFD, LCDM) pair helpers.
#     One CAMB call per model -> CMB Cl (optionally theta*-screened) + sigma8 + Om + pp.
# =====================================================================================
def camb_model(H0, ombh2, omch2, ns, As, tau=0.054, screen=None, lmax=2600,
               sigma8_target=None, want_pp=False):
    """Run CAMB once; return dict(tt, te, ee, s8, Om, th[, pp, lens_dls]).
    If screen is set, the primary CMB Cl are rescaled by theta*/screen (the frozen
    optical screen as an apparent-distance bias). The lensing pp is NEVER screened.
    Importing CAMB does NOT require cobaya."""
    import camb
    if sigma8_target:                                  # rescale A_s so linear sigma8(0)=target
        p0 = camb.set_params(H0=H0, ombh2=ombh2, omch2=omch2, ns=ns, As=As, tau=tau, lmax=lmax)
        p0.set_matter_power(redshifts=[0.0], kmax=2.0)
        As = As * (sigma8_target / camb.get_results(p0).get_sigma8_0()) ** 2
    pars = camb.set_params(H0=H0, ombh2=ombh2, omch2=omch2, ns=ns, As=As, tau=tau, lmax=lmax)
    pars.set_matter_power(redshifts=[0.0], kmax=2.0)
    r = camb.get_results(pars)
    cl = r.get_cmb_power_spectra(pars, CMB_unit="muK", raw_cl=False)["total"]
    Lr = np.arange(cl.shape[0])
    th = r.get_derived_params()["thetastar"]           # 100*theta*
    s8 = r.get_sigma8_0()
    Om = (ombh2 + omch2) / (H0 / 100) ** 2
    if screen:                                         # apply frozen screen as theta*-rescale
        s = th / screen
        Lq = np.arange(lmax) / s
        tt = np.interp(Lq, Lr, cl[:, 0]); ee = np.interp(Lq, Lr, cl[:, 1]); te = np.interp(Lq, Lr, cl[:, 3])
    else:
        tt, te, ee = cl[:, 0], cl[:, 3], cl[:, 1]
    out = dict(tt=tt, te=te, ee=ee, s8=s8, Om=Om, th=th)
    if want_pp:
        pp = r.get_lens_potential_cls(lmax=min(lmax, 2500))[:, 0]   # [L(L+1)]^2/2pi C_L^phiphi
        out["pp"] = pp
        # lensing dls feed the cobaya lensing likelihood the UNSCREENED physical CMB + pp.
        out["lens_dls"] = {"tt": cl[:, 0].copy(), "te": cl[:, 3].copy(),
                           "ee": cl[:, 1].copy(), "pp": pp.copy()}
    return out


def dfd_camb(want_pp=False, sigma8_target=None):
    """One CAMB call for DFD with frozen-derived params + frozen optical screen (NO refit)."""
    return camb_model(P.H0, P.OMEGA_B_H2, P.OMEGA_CHI_H2, P.N_S, P.A_S,
                      screen=THSCR, want_pp=want_pp, sigma8_target=sigma8_target)


def lcdm_camb(want_pp=False, H0=67.36, ombh2=0.02237, omch2=0.12, ns=0.9649, As=2.1e-9):
    """One CAMB call for the Planck-2018 best-fit LCDM reference."""
    return camb_model(H0, ombh2, omch2, ns, As, want_pp=want_pp)


def per_sector_chi2(lensing="proxy", dfd_model=None, lcdm_model=None):
    """Compute the (DFD, LCDM) per-sector chi^2 pair for the standard joint stack,
    LIVE via CAMB + numpy/scipy.  Returns a dict keyed by sector name:
        {sector: {"N": N, "dfd": chi2_dfd, "lcdm": chi2_lcdm, "live": bool}}

    lensing="proxy"  -> amplitude proxy (sigma8 Om^0.25), LIVE. This is the stack that
                        reproduces the DFD 1993.2 / LCDM 1987.5 headline.
    lensing="real"   -> cobaya 9-bin bandpower, COBAYA-GATED; raises LensingUnavailable
                        here if cobaya is absent.
    lensing="none"   -> omit the lensing sector.

    dfd_model / lcdm_model: optional pre-computed camb_model() dicts (to reuse a CAMB
    call); if None they are computed here (with pp only when lensing=='real')."""
    want_pp = (lensing == "real")
    if dfd_model is None:
        dfd_model = dfd_camb(want_pp=want_pp)
    if lcdm_model is None:
        lcdm_model = lcdm_camb(want_pp=want_pp)

    out = {}
    out["CMB TT/TE/EE (plik-lite cov)"] = {
        "N": NB_CMB, "live": True,
        "dfd": chi2_cmb(dfd_model["tt"], dfd_model["te"], dfd_model["ee"]),
        "lcdm": chi2_cmb(lcdm_model["tt"], lcdm_model["te"], lcdm_model["ee"]),
    }
    out["BAO (12 consensus pts)"] = {
        "N": len(B.BAO), "live": True,
        "dfd": chi2_bao_dfd(),
        "lcdm": chi2_bao_lcdm(),
    }
    out["Pantheon+ SNe (full cov)"] = {
        "N": N_SNE, "live": True,
        "dfd": chi2_sne_dfd(),
        "lcdm": chi2_sne_lcdm(),
    }
    if lensing == "proxy":
        out["Planck lensing (sigma8 Om^0.25)"] = {
            "N": 1, "live": True,
            "dfd": chi2_lensing_proxy(dfd_model["s8"], dfd_model["Om"]),
            "lcdm": chi2_lensing_proxy(lcdm_model["s8"], lcdm_model["Om"]),
        }
    elif lensing == "real":
        out["Planck lensing (9-bin C_Lpp)"] = {
            "N": nbins_lensing_real(), "live": True,
            "dfd": chi2_lensing_real(dfd_model["lens_dls"]),
            "lcdm": chi2_lensing_real(lcdm_model["lens_dls"]),
        }
    elif lensing != "none":
        raise ValueError(f"unknown lensing mode {lensing!r}")
    return out


if __name__ == "__main__":
    # Smoke test: print the live per-sector chi^2 (proxy lensing) and the 4-sector total.
    secs = per_sector_chi2(lensing="proxy")
    print(f"{'sector':32s} {'N':>5s} {'chi2 DFD':>10s} {'chi2 LCDM':>10s} {'live':>5s}")
    TD = TL = NT = 0.0
    for name, s in secs.items():
        print(f"{name:32s} {s['N']:5d} {s['dfd']:10.1f} {s['lcdm']:10.1f} {str(s['live']):>5s}")
        TD += s["dfd"]; TL += s["lcdm"]; NT += s["N"]
    print(f"{'TOTAL':32s} {int(NT):5d} {TD:10.1f} {TL:10.1f}")
    print(f"real-lensing (cobaya) available here: {lensing_real_available()}")
