"""
=====================================================================================
 canonical_joint_chi2.py  --  SUB-BUILD A frozen canonical joint chi^2 (DFD vs LCDM)
=====================================================================================
ONE entry point that RUNS the install-free canonical joint and prints the frozen
numbers, plus the +12.5 CMB-gap isolation proof. Pure numpy/scipy/camb (no cobaya,
no dynesty). If python3.12 (with cobaya) is used the REAL 9-bin lensing row also runs.

RECONCILIATION RESULT (SUB-BUILD A):
  * dfd_likelihoods.py            DFD CMB chi^2 = 593.5   <- CANONICAL (cobaya-validated)
  * bayesian_model_comparison.py  DFD CMB chi^2 = 606.0   <- ARTIFACT of one line:
        sigma8_target=0.790 forcing in the OLD Tier A breakdown.
    The +12.5 is 100% that line (proven by the isolation toggle: [3]-[1] = +12.5, and
    [2]==[3], [1]==[4] to the milli-chi^2). The plik-lite CMB chi^2 profiles A_s as a free
    overall norm, so pinning sigma8=0.790 (=> A_s down ~7%) double-counts the amplitude AND
    changes the lensing-smoothing SHAPE the profile cannot remove -> spurious +12.5.
    The live bayesian_model_comparison.py is ALREADY fixed (uses P.A_S, no sigma8 forcing);
    its 0.790 path survives only as a labeled, NOT-used sensitivity row.

CANONICAL DECISION:  DFD CMB chi^2 = 593.5 ; the 606.0 is retired as an amplitude
double-count artifact. dfd_likelihoods.py is the SINGLE SOURCE OF TRUTH.

FROZEN CANONICAL JOINT (one number for BOTH models, consistent amplitude treatment):
  proxy lensing  (install-free):  DFD 1993.2  LCDM 1987.5   Delta +5.7   N=2206
  real 9-bin lens (cobaya):        DFD 2043.2  LCDM 1995.3   Delta +47.9  N=2214

Run (install-free, proxy lensing headline):
  python3 canonical_joint_chi2.py
Run (with a Python that has cobaya installed -> adds the real 9-bin lensing row):
  python3.12 canonical_joint_chi2.py
"""
import os
import sys
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import dfd_params as P
import dfd_likelihoods as L

THSCR = L.THSCR


def isolation():
    """Prove the entire +12.5 CMB gap = the sigma8_target=0.790 forcing (not the screen,
    not the binning, not the cov handling). Pure camb + the canonical chi2_cmb()."""
    import camb

    def camb_raw(As, tau, sigma8_target=None, lmax=2600):
        if sigma8_target:
            p0 = camb.set_params(H0=P.H0, ombh2=P.OMEGA_B_H2, omch2=P.OMEGA_CHI_H2,
                                 ns=P.N_S, 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=P.H0, ombh2=P.OMEGA_B_H2, omch2=P.OMEGA_CHI_H2,
                               ns=P.N_S, 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"]
        return cl, r.get_derived_params()["thetastar"]

    def screen_inline(cl, th, lmax=2600):     # dfd_likelihoods.camb_model branch
        s = th / THSCR
        Lr = np.arange(cl.shape[0]); Lq = np.arange(lmax) / s
        return (np.interp(Lq, Lr, cl[:, 0]), np.interp(Lq, Lr, cl[:, 3]),
                np.interp(Lq, Lr, cl[:, 1]))

    def screen_bmc(cl, th):                    # bayesian_model_comparison.screen_cmb branch
        s = th / THSCR
        Lr = np.arange(cl.shape[0]); Lq = Lr / s
        return (np.interp(Lq, Lr, cl[:, 0]), np.interp(Lq, Lr, cl[:, 3]),
                np.interp(Lq, Lr, cl[:, 1]))

    print("=" * 78)
    print(" SUB-BUILD A  ISOLATION: what causes the +12.5 DFD-CMB chi^2 gap (593.5 vs 606)?")
    print("=" * 78)
    cl, th = camb_raw(P.A_S, 0.054, sigma8_target=None)
    c1 = L.chi2_cmb(*screen_inline(cl, th))
    print(f"  [1] canonical  (no sigma8 forcing, inline screen)   chi2 = {c1:8.3f}  dfd_likelihoods.py")
    cl, th = camb_raw(P.A_S, P.TAU_REION, sigma8_target=0.790)
    c2 = L.chi2_cmb(*screen_bmc(cl, th))
    print(f"  [2] old BMC    (sigma8=0.790 + screen_cmb)          chi2 = {c2:8.3f}  old bayesian Tier A")
    cl, th = camb_raw(P.A_S, 0.054, sigma8_target=0.790)
    c3 = L.chi2_cmb(*screen_inline(cl, th))
    print(f"  [3] canonical + sigma8=0.790 ONLY                   chi2 = {c3:8.3f}  the +{c3-c1:.1f} culprit")
    cl, th = camb_raw(P.A_S, 0.054, sigma8_target=None)
    c4 = L.chi2_cmb(*screen_bmc(cl, th))
    print(f"  [4] canonical + screen_cmb() ONLY                   chi2 = {c4:8.3f}  screen is NOT the cause")
    print("-" * 78)
    print(f"  gap [3]-[1] = {c3-c1:+.2f}  (sigma8=0.790 amplitude double-count)")
    print(f"  [2]=[3]={c2:.3f}=={c3:.3f}? {abs(c2-c3)<1e-3} ;  [1]=[4]={c1:.3f}=={c4:.3f}? {abs(c1-c4)<1e-3}")
    print(f"  => the screen impl is provably identical; the whole gap is the sigma8 forcing.")
    print(f"  CANONICAL DFD CMB chi^2 = {c1:.1f}")
    return c1


def joint():
    """ONE canonical frozen joint for BOTH models (consistent CAMB-native sigma8)."""
    dfd = L.dfd_camb(want_pp=(L.lensing_real_available()))
    lcdm = L.lcdm_camb(want_pp=(L.lensing_real_available()))
    print("\n" + "=" * 78)
    print(" CANONICAL FROZEN JOINT chi^2  (A_s profiled => CAMB-native sigma8 everywhere)")
    print("=" * 78)
    print(f"  DFD : sigma8={dfd['s8']:.4f}  Om={dfd['Om']:.4f}  (H0={P.H0}, derived)")
    print(f"  LCDM: sigma8={lcdm['s8']:.4f}  Om={lcdm['Om']:.4f}  (H0=67.36, Planck)")

    def show(mode, label):
        secs = L.per_sector_chi2(lensing=mode, dfd_model=dfd, lcdm_model=lcdm)
        print(f"\n  --- lensing = {label} ---")
        print(f"  {'sector':34s}{'N':>5s}{'DFD':>10s}{'LCDM':>10s}{'D(DFD-LCDM)':>13s}")
        TD = TL = NT = 0.0
        for k, s in secs.items():
            print(f"  {k:34s}{s['N']:5d}{s['dfd']:10.1f}{s['lcdm']:10.1f}{s['dfd']-s['lcdm']:+13.1f}")
            TD += s['dfd']; TL += s['lcdm']; NT += s['N']
        print(f"  {'TOTAL':34s}{int(NT):5d}{TD:10.1f}{TL:10.1f}{TD-TL:+13.1f}")
        return TD, TL, int(NT)

    tdp, tlp, ntp = show("proxy", "amplitude proxy (sigma8 Om^0.25)  [INSTALL-FREE HEADLINE]")
    if L.lensing_real_available():
        tdr, tlr, ntr = show("real", "REAL Planck 9-bin bandpower (cobaya)")
    print("\n" + "=" * 78)
    print(" CANONICAL NUMBERS (frozen, SUB-BUILD A)")
    print("=" * 78)
    print(f"  DFD CMB chi^2          = 593.5  (canonical; 606.0 retired as artifact)")
    print(f"  JOINT (proxy lensing)  : DFD {tdp:.1f}  LCDM {tlp:.1f}  Delta {tdp-tlp:+.1f}  N={ntp}")
    if L.lensing_real_available():
        print(f"  JOINT (real 9-bin lens): DFD {tdr:.1f}  LCDM {tlr:.1f}  Delta {tdr-tlr:+.1f}  N={ntr}")


if __name__ == "__main__":
    isolation()
    joint()
