"""
KERNEL-SPLIT TEST (the referee's proposed wall-breaker). Honesty-guard BOTH ways.

The theta*-vs-BAO 'shape wall' came from my treating CMB theta* and BAO as ONE ruler under ONE flat
multiplicative screen e^{Dpsi(z)} applied to D_M at every z. The referee's point: the optical screen is a
LINE-OF-SIGHT effect, so transverse and radial distances are screened DIFFERENTLY:
   D_M^opt(z) = INTEGRAL_0^z e^{dpsi(z')} (c/H) dz'   (transverse = accumulated)
   D_H^opt(z) = e^{dpsi(z)} (c/H)(z)                  (radial = LOCAL)
Consequence: theta* depends on the screen INTEGRATED to z*=1090 (so a LOW-z screen barely moves it),
while low-z BAO feels the LOCAL screen strongly. That could fix BAO without wrecking theta*.

HONEST TEST: does a LOW-PARAMETER, PHYSICALLY-SHAPED screen (concentrated at low z, like late-time
structure/DE) close theta* AND the 12 BAO points TOGETHER under the kernel split, where the FLAT
screen could not? If it only closes with a flexible/overfit screen, that is NOT a win (just hidden
params) -- I will report the param count and shape honestly.
"""
import numpy as np
from scipy.integrate import cumulative_trapezoid
from scipy.optimize import minimize
import dfd_params as P
import background as B

C, H0 = P.C_KMS, P.H0
THETA_PLANCK, THETA_ERR = 1.04109, 0.00030      # 100*theta* (Planck)

# raw DFD background
zg = np.linspace(0.0, P.Z_STAR, 120000)
cH = C/(H0*np.array([B.E(z) for z in zg]))       # c/H(z) [Mpc]
DM_raw = cumulative_trapezoid(cH, zg, initial=0.0)
rs_star = B.sound_horizon(P.Z_STAR)              # photon-baryon sound horizon at recombination
rd      = B.r_drag()                             # baryon drag ruler
DMstar_raw = DM_raw[-1]
theta_raw = 100*rs_star/DMstar_raw

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

# physical low-z screen shapes (1 amplitude param A each). f(z) concentrated at low z.
SHAPES = {
    'OmegaLambda(z)': lambda z: (1-P.OMEGA_M)/((1-P.OMEGA_M)+P.OMEGA_M*(1+z)**3),  # DE energy fraction
    '1/(1+z)':        lambda z: 1.0/(1+z),
    'exp(-z)':        lambda z: np.exp(-z),
}

def predict_chi2(A, shapef, mode):
    dpsi_g = A*shapef(zg)
    if mode == 'kernel':      # integrated transverse, local radial (external review)
        DM_opt = cumulative_trapezoid(np.exp(dpsi_g)*cH, zg, initial=0.0)
    else:                     # flat: multiply D_M(z) by e^{dpsi(z)} everywhere (my old way)
        DM_opt = np.exp(dpsi_g)*DM_raw
    DMstar_opt = DM_opt[-1]
    theta = 100*rs_star/DMstar_opt
    chi2 = ((theta-THETA_PLANCK)/THETA_ERR)**2
    for b in B.BAO:
        z = b["z"]
        DMo = np.interp(z, zg, DM_opt)
        DHo = np.exp(A*shapef(z))*cH_at(z) if mode=='kernel' else np.exp(A*shapef(z))*cH_at(z)
        pred = DMo/rd if b["kind"]=="DM" else DHo/rd if b["kind"]=="DH" else (z*DMo**2*DHo)**(1/3)/rd
        chi2 += ((pred-b["val"])/b["err"])**2
    return chi2, theta

def best(shapef, mode):
    r = minimize(lambda a: predict_chi2(a[0], shapef, mode)[0], x0=[0.05],
                 method="Nelder-Mead", options=dict(xatol=1e-5, fatol=1e-4))
    A = r.x[0]; c, th = predict_chi2(A, shapef, mode)
    return A, c, th

print("="*78)
print("KERNEL-SPLIT TEST — can a 1-param physical screen close theta* AND BAO together?")
print("="*78)
print(f"\n  r_s(z*)={rs_star:.2f} Mpc,  r_drag={rd:.2f} Mpc,  N_BAO={len(B.BAO)}")
print(f"  RAW DFD (no screen):  100*theta*={theta_raw:.5f}  "
      f"(Planck {THETA_PLANCK}, off {100*(theta_raw-THETA_PLANCK)/THETA_PLANCK:+.2f}% = "
      f"{(theta_raw-THETA_PLANCK)/THETA_ERR:+.0f}sigma)")
c0,_ = predict_chi2(0.0, SHAPES['1/(1+z)'], 'kernel')
bao0 = c0 - ((theta_raw-THETA_PLANCK)/THETA_ERR)**2
print(f"  RAW DFD BAO chi^2 = {bao0:.1f}/{len(B.BAO)} (chi^2/N={bao0/len(B.BAO):.2f});  "
      f"theta*+BAO total chi^2 = {c0:.1f}  <-- THE WALL")

print("\n  1-parameter screen fits (theta* + 12 BAO = 13 data, 1 free param A):")
print(f"  {'shape':16s} {'mode':8s} {'A':>8s} {'100theta*':>10s} {'total chi2':>11s} {'chi2/dof':>9s}")
for name, f in SHAPES.items():
    for mode in ('flat','kernel'):
        A,c,th = best(f, mode)
        print(f"  {name:16s} {mode:8s} {A:8.4f} {th:10.5f} {c:11.1f} {c/12:9.2f}")

print("\n" + "-"*78)
print("Read: 'flat' = my old single-screen treatment; 'kernel' = the referee's split (integrated")
print("transverse, local radial). dof=13-1=12. theta* pinned if 100theta* ~ 1.04109.")

# ---- consistency check: does the best Omega_Lambda(z) kernel screen also satisfy SNe? ----
# DFD's optical screen IS its dark energy; the SAME screen must keep D_L tracking the data.
# Use LCDM-Planck D_M as the 'fits-everything' reference across the SNe/BAO range.
print("\n" + "="*78)
print("CONSISTENCY of the winning screen (Omega_Lambda(z) shape, kernel mode):")
fOL = SHAPES['OmegaLambda(z)']
A,c,th = best(fOL, 'kernel')
dpsi_g = A*fOL(zg)
DM_opt = cumulative_trapezoid(np.exp(dpsi_g)*cH, zg, initial=0.0)
def DM_lcdm(z):
    pl=P.PLANCK; Om=pl['Omega_m']; Or=P.OMEGA_R_H2/(pl['H0']/100)**2; OL=1-Om-Or
    from scipy.integrate import quad
    return (C/pl['H0'])*quad(lambda zz:1/np.sqrt(Om*(1+zz)**3+Or*(1+zz)**4+OL),0,z,limit=80)[0]
print(f"  best A = {A:.4f} (1 fitted amplitude; shape FIXED, not fitted)")
print(f"  screen dpsi(z):  z=0:{A*fOL(0):.4f}  z=0.5:{A*fOL(0.5):.4f}  z=1:{A*fOL(1):.4f}  "
      f"z=2:{A*fOL(2):.4f}  z*:{A*fOL(P.Z_STAR):.5f}")
print(f"  -> screen at z* is ~{A*fOL(P.Z_STAR):.5f} (negligible): theta* moves via the INTEGRAL, not a local kick.")
print(f"\n  Distance vs LCDM-Planck (which fits Pantheon+/BAO) across the SNe range:")
print(f"  {'z':>6s} {'D_M^opt(DFD)':>13s} {'D_M(LCDM)':>11s} {'ratio-1':>9s}")
for z in (0.05,0.1,0.2,0.35,0.5,0.7,1.0,1.3):
    DMo = np.interp(z, zg, DM_opt); DMl = DM_lcdm(z)
    print(f"  {z:6.2f} {DMo:13.1f} {DMl:11.1f} {100*(DMo/DMl-1):+8.2f}%")
print("\n  VERDICT: if |ratio-1| stays within a few % across 0<z<1.3, the SAME screen that pins")
print("  theta* + BAO also keeps SNe distances on the data-preferred track (SNe absolute scale is")
print("  degenerate with H0/M_B, so a smooth few-% offset is absorbed). Honest residual: the")
print("  amplitude A is still 1 FITTED number -- the open task is to DERIVE A from DFD (baryon-")
print("  inertia f_b / deep-MOND a0(z) / structure screen), which would make it parameter-free.")
