"""
STRAND 1 — BAYESIAN MODEL COMPARISON: DFD vs LCDM on the joint real-data sky.

Goes beyond the existing BIC proxy (joint_fit.py) to a real nested-sampling Bayesian EVIDENCE
ln Z and Bayes factor, over the joint likelihood:
    CMB plik-lite (full 613x613 Fortran covariance, cobaya-validated)
  + Planck 2018 CMB LENSING (full 9-bin bandpower, cobaya native CMBlikes)   <-- NEW (strand 2)
  + BAO (12 consensus points)
  + Pantheon+ SNe (1580, full stat+sys covariance)

Two models, SAME data, SAME likelihood code:
  LCDM : free {ombh2, omch2, H0, ln10As, ns, tau}                       (6 cosmological)
  DFD  : H0, ns, As-norm, omch2(via Omega_chi misalignment) DERIVED;
         free {ombh2, tau, theta_i}                                     (3 cosmological)
         + the frozen optical screen A=2sqrt(alpha)R/(1+R) (no free knob).

Tiers:
  A) BIC/AIC from the joint best-fit chi^2 + honest parameter counts (instant).  Now folds in
     the strand-2 lensing chi^2 -- which CHANGES the verdict relative to the proxy-lensing joint_fit.
  B) dynesty nested-sampling ln Z for each model -> ln B = lnZ_DFD - lnZ_LCDM (the publication number).

HONESTY-GUARD: report the HONEST joint result and each sector. DFD's structural advantage is FEWER free
parameters (Occam/evidence rewards this); its disadvantage is a worse CMB-lensing chi^2 (low Omega_m).
Do NOT cherry-pick the data combination. The joint number is the headline.

Run with py3.12 (cobaya+camb+dynesty):
  python3.12 bayesian_model_comparison.py --tier A
  ... --tier B --nlive 200      (heavy; launch in background)
"""
import os
import sys
import time
import argparse
import numpy as np
import camb
from scipy.io import FortranFile
from scipy.integrate import cumulative_trapezoid, quad

import dfd_params as P
import background as B
from cobaya.likelihoods.planck_2018_lensing import native as LensLik

C = P.C_KMS
HERE = os.path.dirname(os.path.abspath(__file__))
PKG = "~/cobaya_packages"
DATL = os.path.join(PKG, "data/planck_2018_pliklite_native")
PD = os.path.join(HERE, "pantheon_data")
LMAX = 2500

# DFD frozen screen (kernel-split optical distance screen; derived f_b gate, no refit).
A_SCR = 0.06667

# =============================================================================== plik-lite CMB
NTT, NTE, NEE = 215, 199, 199
NB = NTT + NTE + NEE
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, NB))
ICMB = np.linalg.inv(np.tril(_cov) + np.tril(_cov, -1).T)


def cmb_chi2(tt, te, ee, profile_amp=True):
    # binning mirrors the cobaya-validated joint_fit.py (CMB chi2 593.5/585.5 reproduced):
    # per-spectrum local bin index into the global blmin/blmax/W (concatenated weights).
    u = np.empty(NB)
    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
    if profile_amp:
        uCu = u @ ICMB @ u
        uCX = u @ ICMB @ Xcmb
        XCX = Xcmb @ ICMB @ Xcmb
        return XCX - uCX ** 2 / uCu          # amplitude-profiled (A_s free normalization)
    r = u - Xcmb
    return r @ ICMB @ r


# =============================================================================== Pantheon+ SNe
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)
zc = zsn[ssel]; mbc = mb[ssel]
ISN = np.linalg.inv(covsn[np.ix_(ssel, ssel)])


def sne_chi2_marg(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                   # M (abs mag) analytically marginalized


# =============================================================================== lensing likelihood
_lens = LensLik({"path_install": PKG, "packages_path": PKG}, name="planck_2018_lensing.native")
NLENS = _lens.nbins_used


def lensing_chi2(dls, A_planck=1.0):
    return -2.0 * _lens.log_likelihood(dls, A_planck=A_planck)


# =============================================================================== CAMB driver
LMAX_CMB = 2600       # CMB Dls need length >= 2509 for the TT binning (matches joint_fit.py)


def camb_run(H0, ombh2, omch2, ns, As, tau, sigma8_target=None):
    if sigma8_target:
        p0 = camb.set_params(H0=H0, ombh2=ombh2, omch2=omch2, ns=ns, As=As, tau=tau, lmax=LMAX_CMB)
        p0.set_matter_power(redshifts=[0.0], kmax=2.0)
        s0 = camb.get_results(p0).get_sigma8_0()
        As = As * (sigma8_target / s0) ** 2
    pars = camb.set_params(H0=H0, ombh2=ombh2, omch2=omch2, ns=ns, As=As, tau=tau, lmax=LMAX_CMB)
    pars.set_matter_power(redshifts=[0.0], kmax=2.0)
    r = camb.get_results(pars)
    tot = r.get_cmb_power_spectra(pars, CMB_unit="muK", raw_cl=False)["total"]   # length LMAX_CMB+1
    lp = r.get_lens_potential_cls(lmax=LMAX)                                       # PP up to LMAX
    dp = r.get_derived_params()
    th = dp["thetastar"]
    rd = dp["rdrag"]
    s8 = r.get_sigma8_0()
    dls = {"tt": tot[:, 0].copy(), "ee": tot[:, 1].copy(), "te": tot[:, 3].copy(), "pp": lp[:, 0].copy()}
    return dls, th, s8, rd


def screen_cmb(dls, thetastar, thscr):
    """Apply DFD frozen optical screen as the theta* ell-rescale on the CMB Dls (TT/TE/EE)."""
    s = thetastar / thscr
    nC = len(dls["tt"])
    Lr = np.arange(nC)
    Lq = Lr / s
    return {
        "tt": np.interp(Lq, Lr, dls["tt"]),
        "te": np.interp(Lq, Lr, dls["te"]),
        "ee": np.interp(Lq, Lr, dls["ee"]),
        "pp": dls["pp"],          # lensing PP not theta*-rescaled (screen is an optical distance bias)
    }


# =============================================================================== BAO + SNe distances
_zg = np.linspace(0, P.Z_STAR, 40000)


def dfd_bao_sne_chi2():
    cHd = C / (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)

    DMd_opt = cumulative_trapezoid(np.exp(A_SCR * shp(_zg)) * cHd, _zg, initial=0.0)

    def cHd_at(z):
        return C / (P.H0 * B.E(z))

    bao = 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
        bao += ((pred - b["val"]) / b["err"]) ** 2
    mu = 5 * np.log10((1 + zc) * np.interp(zc, _zg, DMd_opt)) + 25
    return bao, sne_chi2_marg(mu)


def lcdm_bao_sne_chi2(Om, H0, rd):
    Or = P.OMEGA_R_H2 / (H0 / 100) ** 2
    OL = 1 - Om - Or

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

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

    bao = 0.0
    for b in B.BAO:
        z = b["z"]; dm, dh = DM(z), DH(z)
        pred = dm / rd if b["kind"] == "DM" else dh / rd if b["kind"] == "DH" else (z * dm ** 2 * dh) ** (1 / 3) / rd
        bao += ((pred - b["val"]) / b["err"]) ** 2
    mu = np.array([5 * np.log10((1 + z) * DM(z)) + 25 for z in zc])
    return bao, sne_chi2_marg(mu)


# =============================================================================== model chi2 at fixed best-fit
def chi2_breakdown_dfd():
    # CANONICAL (reconciled with dfd_likelihoods.py): do NOT force sigma8=0.790.
    # The CMB chi^2 is amplitude-profiled (A_s/A_planck is a free overall norm), so pinning A_s
    # via a sigma8 target double-counts the amplitude. Worse, forcing sigma8=0.790 lowers A_s ~7%,
    # which WEAKENS the lensing peak-smoothing -- a SHAPE change the profiling cannot remove --
    # and INFLATES the DFD CMB chi^2 by +12.5 (593.5 -> 606.0, spurious). We use the CAMB-native
    # sigma8 (= DFD's profiled prediction, ~0.820) consistently for BOTH the CMB and the lensing.
    # See reconcile_cmb_chi2.py for the one-toggle isolation proof.
    dls, th, s8, rd = camb_run(P.H0, P.OMEGA_B_H2, P.OMEGA_CHI_H2, P.N_S, P.A_S, P.TAU_REION)
    THSCR = 1.04138
    sdls = screen_cmb(dls, th, THSCR)
    cmb = cmb_chi2(sdls["tt"], sdls["te"], sdls["ee"])          # = 593.5 canonical
    # LENSING: feed the UNSCREENED physical CMB spectra to the N1 linear correction (the optical
    # screen is an apparent-distance bias, not a change to the true CMB, so it must NOT absorb the
    # lensing reconstruction correction). PP is never screened. At native sigma8 this is ~50/9.
    lens = lensing_chi2(dls)
    # sensitivity row: what the (retired) corpus sigma8=0.790 forcing WOULD have given, for the record
    dls_790, th_790, s8_790, _ = camb_run(P.H0, P.OMEGA_B_H2, P.OMEGA_CHI_H2, P.N_S, P.A_S,
                                          P.TAU_REION, sigma8_target=0.790)
    lens_790 = lensing_chi2(dls_790)
    bao, sne = dfd_bao_sne_chi2()
    return dict(CMB=cmb, lensing=lens, lensing_790=lens_790, s8_790=0.790,
                BAO=bao, SNe=sne, s8=s8,
                Om=(P.OMEGA_B_H2 + P.OMEGA_CHI_H2) / (P.H0 / 100) ** 2)


def chi2_breakdown_lcdm():
    pl = P.PLANCK
    dls, th, s8, rd = camb_run(pl["H0"], pl["omega_b_h2"], pl["omega_c_h2"], pl["n_s"], pl["A_s"], 0.0544)
    cmb = cmb_chi2(dls["tt"], dls["te"], dls["ee"])
    lens = lensing_chi2(dls)
    bao, sne = lcdm_bao_sne_chi2(pl["Omega_m"], pl["H0"], pl["r_drag"])
    return dict(CMB=cmb, lensing=lens, BAO=bao, SNe=sne, s8=s8, Om=pl["Omega_m"])


# =============================================================================== TIER A: BIC / AIC
def tier_A():
    print("=" * 84)
    print("TIER A — JOINT chi^2 + BIC/AIC (now INCLUDING the full Planck CMB-LENSING bandpowers)")
    print("=" * 84)
    print("  running 2 CAMB models (DFD frozen, LCDM Planck)...")
    d = chi2_breakdown_dfd()
    l = chi2_breakdown_lcdm()
    N = {"CMB": 613, "lensing": NLENS, "BAO": 12, "SNe": len(zc)}
    print(f"\n  {'sector':12s} {'N':>5s} {'chi2 DFD':>10s} {'chi2 LCDM':>10s} {'Delta':>8s}")
    TD = TL = NT = 0
    for k in ["CMB", "lensing", "BAO", "SNe"]:
        print(f"  {k:12s} {N[k]:5d} {d[k]:10.1f} {l[k]:10.1f} {d[k]-l[k]:+8.1f}")
        TD += d[k]; TL += l[k]; NT += N[k]
    print(f"  {'TOTAL':12s} {NT:5d} {TD:10.1f} {TL:10.1f} {TD-TL:+8.1f}")
    print(f"\n  [lensing sensitivity] DFD with CAMB-native sigma8={d['s8']:.3f} -> chi2_lens={d['lensing']:.0f} "
          f"(used above, self-consistent);")
    print(f"                        DFD with retired corpus sigma8=0.790 -> "
          f"chi2_lens={d['lensing_790']:.0f} (amplitude double-count, NOT used).")
    print(f"                        DFD shows lensing tension either way (LCDM {l['lensing']:.0f}/9).")

    print("\n" + "-" * 84)
    print(f"  Joint chi^2:  DFD {TD:.1f}   LCDM {TL:.1f}   Delta(DFD-LCDM) = {TD-TL:+.1f}  over N={NT}")
    print(f"  (chi^2 scatter sqrt(2N) = {np.sqrt(2*NT):.0f})")
    # parameter-count scenarios (from FALSIFY.md / paramcount_sensitivity.py)
    scenarios = {
        "fair (DFD derives H0,ns,As-norm,omch2-mech)": (4, 6),
        "skeptic (DFD derivations all rejected)":      (8, 6),
    }
    print("\n  free-parameter scenarios and information criteria:")
    for tag, (kD, kL) in scenarios.items():
        aicD, aicL = TD + 2 * kD, TL + 2 * kL
        bicD, bicL = TD + kD * np.log(NT), TL + kL * np.log(NT)
        fav_aic = "DFD" if aicD < aicL else "LCDM"
        fav_bic = "DFD" if bicD < bicL else "LCDM"
        print(f"   [{tag}]  k_DFD={kD}, k_LCDM={kL}")
        print(f"      dAIC(DFD-LCDM) = {aicD-aicL:+.1f} -> favors {fav_aic};   "
              f"dBIC = {bicD-bicL:+.1f} -> favors {fav_bic}")
    print("\n  HONEST: with the FULL lensing bandpowers folded in, the CMB sector now carries DFD's")
    print("  lensing tension (low Omega_m). This is the honesty-guard-correct joint number; the earlier")
    print("  joint_fit.py used the amplitude-level lensing proxy that MISSED this. See VERDICT below.")
    return TD, TL, NT, d, l


# =============================================================================== TIER B: nested evidence
def joint_loglike_lcdm(theta):
    ombh2, omch2, H0, ln10As, ns, tau = theta
    As = 1e-10 * np.exp(ln10As)
    if not (0.005 < tau < 0.12):
        return -np.inf
    try:
        dls, th, s8, rd = camb_run(H0, ombh2, omch2, ns, As, tau)
    except Exception:
        return -np.inf
    if not np.all(np.isfinite(dls["tt"])):
        return -np.inf
    # CMB amplitude profiled analytically (fair plik-lite treatment: A_s/A_planck is a free
    # normalization; the verification record established the free-amplitude profile is the fair number).
    cmb = cmb_chi2(dls["tt"], dls["te"], dls["ee"], profile_amp=True)
    lens = lensing_chi2(dls)
    Om = (ombh2 + omch2) / (H0 / 100) ** 2
    bao, sne = lcdm_bao_sne_chi2(Om, H0, rd)
    tau_prior = ((tau - 0.054) / 0.007) ** 2          # Planck lowE tau prior
    return -0.5 * (cmb + lens + bao + sne + tau_prior)


def joint_loglike_dfd(theta):
    """DFD: only {ombh2, tau, theta_i} free; everything else derived from alpha.
    Internally consistent treatment: A_s is DFD's derived 32pi*alpha^5; the overall CMB amplitude
    is profiled (same fair treatment as LCDM), so sigma8 floats at its CAMB-native value and feeds
    the lensing self-consistently (NO sigma8=0.790 forcing here -- that double-counts the amplitude).
    The corpus sigma8=0.790 sensitivity is shown separately in cmb_lensing_check.py."""
    ombh2, tau, theta_i = theta
    if not (0.005 < tau < 0.12) or not (0.05 < theta_i < 0.6):
        return -np.inf
    omch2 = 1.62 * theta_i ** 2
    H0 = 72.09; ns = P.N_S; As = P.A_S
    try:
        dls, th, s8, rd_camb = camb_run(H0, ombh2, omch2, ns, As, tau)
    except Exception:
        return -np.inf
    if not np.all(np.isfinite(dls["tt"])):
        return -np.inf
    sdls = screen_cmb(dls, th, 1.04138)
    cmb = cmb_chi2(sdls["tt"], sdls["te"], sdls["ee"], profile_amp=True)
    lens = lensing_chi2(dls)            # UNSCREENED CMB for the N1 correction (conservative)
    # BAO+SNe with DFD screen at THIS ombh2 (recompute r_drag at the sampled ombh2)
    par = _DFDpar(ombh2, omch2)
    rd = B.r_drag(par)
    bao, sne = _dfd_bao_sne_at(par, rd)
    tau_prior = ((tau - 0.054) / 0.007) ** 2
    return -0.5 * (cmb + lens + bao + sne + tau_prior)


class _DFDpar:
    """Lightweight param object overriding ombh2/omch2 for r_drag + distances."""
    def __init__(self, ombh2, omch2):
        for k in dir(P):
            if k.isupper():
                setattr(self, k, getattr(P, k))
        self.OMEGA_B_H2 = ombh2
        self.OMEGA_CHI_H2 = omch2
        h = self.H0 / 100
        self.OMEGA_B = ombh2 / h ** 2
        self.OMEGA_CHI = omch2 / h ** 2
        self.OMEGA_M = self.OMEGA_B + self.OMEGA_CHI
        self.OMEGA_LAMBDA = 1 - self.OMEGA_M - self.OMEGA_R_H2 / h ** 2


def _dfd_bao_sne_at(par, rd):
    zg = np.linspace(0, par.Z_STAR, 20000)
    cHd = C / (par.H0 * B.E(zg, par))
    OL = par.OMEGA_LAMBDA

    def shp(z):
        return OL / (OL + par.OMEGA_M * (1 + z) ** 3)

    DMd_opt = cumulative_trapezoid(np.exp(A_SCR * shp(zg)) * cHd, zg, initial=0.0)
    bao = 0.0
    for b in B.BAO:
        z = b["z"]; DMo = np.interp(z, zg, DMd_opt); DHo = np.exp(A_SCR * shp(z)) * C / (par.H0 * B.E(z, par))
        pred = DMo / rd if b["kind"] == "DM" else DHo / rd if b["kind"] == "DH" else (z * DMo ** 2 * DHo) ** (1 / 3) / rd
        bao += ((pred - b["val"]) / b["err"]) ** 2
    mu = 5 * np.log10((1 + zc) * np.interp(zc, zg, DMd_opt)) + 25
    return bao, sne_chi2_marg(mu)


def tier_B(nlive=200, model="both", outdir=None):
    import dynesty
    from dynesty import utils as dyfunc
    outdir = outdir or HERE
    results = {}

    # priors (unit cube -> params)
    def ptform_lcdm(u):
        x = np.empty(6)
        x[0] = 0.019 + u[0] * (0.025 - 0.019)     # ombh2
        x[1] = 0.10 + u[1] * (0.14 - 0.10)        # omch2
        x[2] = 60 + u[2] * (80 - 60)              # H0
        x[3] = 2.9 + u[3] * (3.2 - 2.9)           # ln10As
        x[4] = 0.92 + u[4] * (1.00 - 0.92)        # ns
        x[5] = 0.03 + u[5] * (0.10 - 0.03)        # tau
        return x

    def ptform_dfd(u):
        x = np.empty(3)
        x[0] = 0.019 + u[0] * (0.025 - 0.019)     # ombh2
        x[1] = 0.03 + u[1] * (0.10 - 0.03)        # tau
        x[2] = 0.15 + u[2] * (0.45 - 0.15)        # theta_i
        return x

    jobs = []
    if model in ("both", "lcdm"):
        jobs.append(("LCDM", joint_loglike_lcdm, ptform_lcdm, 6))
    if model in ("both", "dfd"):
        jobs.append(("DFD", joint_loglike_dfd, ptform_dfd, 3))

    for name, loglike, ptform, ndim in jobs:
        print(f"\n[{name}] nested sampling (ndim={ndim}, nlive={nlive})...", flush=True)
        t0 = time.time()
        sampler = dynesty.NestedSampler(loglike, ptform, ndim, nlive=nlive, bound="multi", sample="rwalk")
        sampler.run_nested(dlogz=0.5, print_progress=True)
        res = sampler.results
        lnZ = res.logz[-1]; lnZerr = res.logzerr[-1]
        results[name] = (lnZ, lnZerr)
        np.savez(os.path.join(outdir, f"nested_{name}.npz"),
                 logz=res.logz, logzerr=res.logzerr, samples=res.samples,
                 logwt=res.logwt, logl=res.logl)
        print(f"[{name}] ln Z = {lnZ:.2f} +/- {lnZerr:.2f}   ({(time.time()-t0)/60:.1f} min)", flush=True)

    if "DFD" in results and "LCDM" in results:
        lnB = results["DFD"][0] - results["LCDM"][0]
        err = np.hypot(results["DFD"][1], results["LCDM"][1])
        print("\n" + "=" * 70)
        print(f"  ln B (DFD - LCDM) = {lnB:+.2f} +/- {err:.2f}")
        scale = ("DECISIVE for " + ("DFD" if lnB > 0 else "LCDM")) if abs(lnB) > 5 else \
                ("STRONG for " + ("DFD" if lnB > 0 else "LCDM")) if abs(lnB) > 2.5 else \
                ("substantial for " + ("DFD" if lnB > 0 else "LCDM")) if abs(lnB) > 1 else "inconclusive"
        print(f"  Jeffreys scale: {scale}")
        print("=" * 70)
    return results


# =============================================================================== main
def verdict(TD, TL, NT):
    print("\n" + "=" * 84)
    print("  STRAND-1 VERDICT (joint Bayesian/IC model comparison)")
    print("=" * 84)
    print(f"  Joint chi^2 over {NT} real points: DFD {TD:.0f} vs LCDM {TL:.0f}, Delta {TD-TL:+.0f}.")
    print("  The full CMB-lensing bandpowers (strand 2) move the joint chi^2 toward LCDM because")
    print("  DFD's high derived H0 -> low Omega_m under-predicts the lensing amplitude (A_L~1.4).")
    print("  Occam/BIC still credits DFD for deriving H0,n_s,A_s-norm,Omega_chi; whether the net")
    print("  Bayes factor favors DFD depends on (a) the lensing penalty vs (b) the parsimony reward.")
    print("  Tier B (dynesty ln Z) resolves (a) vs (b) self-consistently -- run it for the headline.")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--tier", default="A", choices=["A", "B", "AB"])
    ap.add_argument("--nlive", type=int, default=200)
    ap.add_argument("--model", default="both", choices=["both", "dfd", "lcdm"])
    args = ap.parse_args()

    if args.tier in ("A", "AB"):
        TD, TL, NT, d, l = tier_A()
        verdict(TD, TL, NT)
    if args.tier in ("B", "AB"):
        tier_B(nlive=args.nlive, model=args.model)
