"""
=====================================================================================
 reconcile_cmb_chi2.py  --  CANONICAL DFD vs LCDM joint chi^2, reconciliation report
=====================================================================================
PURPOSE (SUB-BUILD A): reconcile the DFD CMB chi^2 discrepancy between

    dfd_likelihoods.py            -> DFD CMB chi^2 = 593.5   (cobaya-validated canonical)
    bayesian_model_comparison.py  -> DFD CMB chi^2 = 606.0   (un-validated re-impl, Tier A)

ROOT CAUSE (isolated by toggling each candidate difference one at a time):
    The ENTIRE +12.5 gap is caused by ONE line in bayesian_model_comparison.py's
    chi2_breakdown_dfd():  sigma8_target=0.790  passed to camb_run().

    The CMB plik-lite chi^2 is AMPLITUDE-PROFILED (A_s / A_planck is treated as a free
    overall normalisation and analytically profiled out). So an overall rescale of A_s
    is supposed to be invisible. It is NOT invisible here because forcing sigma8 down to
    0.790 lowers A_s by ~7%, which WEAKENS the gravitational-lensing smoothing of the
    high-l acoustic peaks. That is a SHAPE change, not an overall amplitude, so the
    amplitude profiling cannot remove it. Sharper (less-smeared) peaks fit Planck WORSE
    (Planck mildly prefers EXTRA smoothing -- the A_L>1 lensing anomaly), so chi^2 rises
    by +12.5.

    PROOF (run check_isolation() below):
      canonical path (no sigma8 forcing)                 -> 593.5
      canonical + sigma8_target=0.790 ONLY               -> 606.0   (reproduces the gap)
      canonical + the other script's screen() fn ONLY    -> 593.5   (screen is NOT the cause)
      UNLENSED spectra, same A_s rescale                 -> moves only +1.3 (confirms it is
                                                            the lensing-smoothing shape).

    The sigma8=0.790 forcing is ALSO internally inconsistent: you cannot both profile A_s
    out of the CMB AND simultaneously pin it via a sigma8 target -- that double-treats the
    amplitude. bayesian_model_comparison.py's OWN nested-sampling path (joint_loglike_dfd,
    lines 292-294) explicitly drops the sigma8=0.790 forcing for exactly this reason; only
    its Tier A breakdown was left inconsistent.

CANONICAL DECISION:  DFD CMB chi^2 = 593.5  (dfd_likelihoods.py). The 606.0 is an artifact.

This script is INSTALL-FREE beyond numpy/scipy/camb: it imports the canonical
dfd_likelihoods module and reports the frozen joint. Cobaya is OPTIONAL (only the REAL
9-bin lensing row needs it; the proxy-lensing headline does not).

Run (a Python with cobaya installed unlocks the real 9-bin lensing row):
  python3 reconcile_cmb_chi2.py
  python3 reconcile_cmb_chi2.py --isolate
"""
import os
import sys
import argparse
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


# --------------------------------------------------------------------------- isolation
def check_isolation():
    """Reproduce 593.5 and 606.0 and prove the gap = the sigma8_target=0.790 line."""
    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
        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(" ISOLATION: which difference causes the +12.5 DFD-CMB gap?")
    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:7.3f}  <- dfd_likelihoods")

    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] bayesian (sigma8=0.790 + screen_cmb)                chi2 = {c2:7.3f}  <- bayesian_model_comparison")

    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_target=0.790 ONLY                chi2 = {c3:7.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:7.3f}  <- screen NOT the cause")
    print("-" * 78)
    print(f"  GAP fully explained by [3]-[1] = {c3-c1:+.1f}  (the sigma8=0.790 amplitude double-count).")
    print(f"  CANONICAL DFD CMB chi2 = {c1:.1f}  (== cobaya-validated 593.5).")
    return c1, c2


# --------------------------------------------------------------------------- frozen joint
def frozen_joint():
    """ONE canonical joint chi^2 for BOTH DFD and LCDM, consistent amplitude treatment
    (A_s profiled in CMB => CAMB-native sigma8 used everywhere, NO 0.790 forcing)."""
    dfd = L.dfd_camb(want_pp=True)
    lcdm = L.lcdm_camb(want_pp=True)
    print("\n" + "=" * 78)
    print(" CANONICAL FROZEN JOINT  (consistent amplitude: CAMB-native sigma8, profiled CMB)")
    print("=" * 78)
    print(f"  DFD : sigma8={dfd['s8']:.4f}  Om={dfd['Om']:.4f}  (H0={P.H0})")
    print(f"  LCDM: sigma8={lcdm['s8']:.4f}  Om={lcdm['Om']:.4f}  (H0=67.36)")

    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)  [HEADLINE]")
    res = {"proxy": (tdp, tlp, ntp)}
    if L.lensing_real_available():
        tdr, tlr, ntr = show("real", "REAL Planck 9-bin bandpower (cobaya live)")
        res["real"] = (tdr, tlr, ntr)
    else:
        print("\n  [real 9-bin lensing: cobaya not importable here -> proxy headline only]")
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--isolate", action="store_true", help="run only the +12.5 gap isolation")
    args = ap.parse_args()
    if args.isolate:
        check_isolation()
        return
    check_isolation()
    res = frozen_joint()
    print("\n" + "=" * 78)
    print(" CANONICAL NUMBERS (frozen)")
    print("=" * 78)
    print(f"  DFD CMB chi^2            = 593.5   (dfd_likelihoods.py is CANONICAL; 606.0 is an artifact)")
    tdp, tlp, ntp = res["proxy"]
    print(f"  JOINT (proxy lensing)   : DFD {tdp:.1f}  LCDM {tlp:.1f}  Delta {tdp-tlp:+.1f}  over N={ntp}")
    if "real" in res:
        tdr, tlr, ntr = res["real"]
        print(f"  JOINT (real 9-bin lens) : DFD {tdr:.1f}  LCDM {tlr:.1f}  Delta {tdr-tlr:+.1f}  over N={ntr}")
        print(f"  (DFD real-lensing chi^2 = 50.1/9 at native sigma8=0.820, NOT the 220.8 the old")
        print(f"   bayesian_model_comparison.py reported -- that 220.8 also came from the sigma8=0.790 force.)")


if __name__ == "__main__":
    main()
