"""
STRAND 3 — HONEST 2-HALO LOS SCREEN AMPLITUDE + SENSITIVITY  (no tuning)
========================================================================
Goal: replace the single-halo NFW+splashback toy (excess ~0.034) with a proper halo-model /
realization LOS calculation, and report the screen amplitude with an HONEST error budget — does it
(a) bracket 0.030, (b) reach the data-preferred 0.039, or (c) miss?

WHAT THE 2-HALO UPGRADE REVEALS (the real result, honesty-guarded both ways):
  The amplitude is the product of two pieces:
     excess = kappa_gal  x  Q_screen ,   Q_screen = < 1/mu(x) >_LOS ,  x = g_local / a0(z)
  * kappa_gal = b_g r_gm sigma_kappa is now ANCHORED to the REAL data: the per-LOS convergence RMS
    from the BOSS pipeline (boss_data/kappa_fg_dataset.npz) = 0.0029 (random LOS), b_g-enhanced.
    -> kappa_gal ~ 0.0052. This is FIRM (measured), and ~the toy's 0.006.
  * Q_screen is the DFD mu-law boost, and it is ENTIRELY controlled by ONE undetermined quantity:
    the External-Field-Effect (EFE) floor x_ext = g_ext/a0. THIS is the crux the single-halo toy hid:
      - g_ext = peculiar/tidal LSS accel (uniform Hubble field gauged away)  -> x_ext~0.003, Q~300 (BLOWS UP)
      - g_ext = Hubble accel cH (the corpus 'AC' hypothesis)                 -> x_ext=5.85,  Q=1.17 (TOO SMALL)
      - intermediate void-skin field                                        -> x_ext~0.2-0.5, Q~3-5 (IN RANGE)
    The corpus itself flags x_ext=5.85 as 'an admitted hypothesis the action does not produce'
    (growth_regime_adjudication.py:291). So the EFE floor is NOT derived; the screen amplitude is a
    FUNCTION OF IT. We map excess(x_ext) and state where 0.030 / 0.039 land. NO knob tuned to the target.

Halo-model machinery (Tinker08 MF + Tinker10 bias, 1h+2h) is built and used to (i) get b_eff for kappa_gal,
(ii) cross-check the volume/path-weighted boost against an independent lognormal-PDF realization.
"""
import numpy as np
from scipy.integrate import quad
import camb
import dfd_params as P

# ----------------------------------------------------------------- constants / cosmology at z
G = 6.674e-11; Msun = 1.989e30; Mpc = 3.0857e22; c_si = 2.998e8
H0 = P.H0 * 1e3 / Mpc; h = P.H0 / 100.0
DELTA_C = 1.686; RHO_CRIT0 = 3 * H0**2 / (8 * np.pi * G)
Z_EVAL = 0.5
Ez = np.sqrt(P.OMEGA_M * (1 + Z_EVAL)**3 + P.OMEGA_LAMBDA); Hz = H0 * Ez
a0_z = 2 * np.sqrt(P.ALPHA) * c_si * Hz
g_H = c_si * Hz                                # Hubble accel (the corpus AC EFE candidate)
X_HUBBLE = g_H / a0_z                          # = 1/(2 sqrt(alpha)) = 5.85
rho_mean_z = P.OMEGA_M * (1 + Z_EVAL)**3 * RHO_CRIT0

pars = camb.CAMBparams()
pars.set_cosmology(H0=P.H0, ombh2=P.OMEGA_B_H2, omch2=P.OMEGA_CHI_H2, mnu=0.06, tau=0.054)
pars.InitPower.set_params(As=P.A_S, ns=P.N_S)
pars.set_matter_power(redshifts=[Z_EVAL], kmax=100.0)
res = camb.get_results(pars)
PKlin = camb.get_matter_power_interpolator(pars, nonlinear=False, kmax=100, zmax=1.0)
PKnl  = camb.get_matter_power_interpolator(pars, nonlinear=True,  kmax=100, zmax=1.0)
sigma8_z = res.get_sigma8()[0]
def Plin(k): return PKlin.P(Z_EVAL, k)
def Pnl(k):  return PKnl.P(Z_EVAL, k)

# ----------------------------------------------------------------- mass function + bias (for b_eff)
def Rofm(Mm): return (3 * Mm * Msun / (4 * np.pi * rho_mean_z))**(1/3) / Mpc
def sigma_M(Mm):
    R = Rofm(Mm)
    def ig(lnk):
        k = np.exp(lnk); kR = k * R
        W = 3 * (np.sin(kR) - kR * np.cos(kR)) / kR**3 if kR > 1e-3 else 1.0
        return Plin(k) * W**2 * k**3 / (2 * np.pi**2)
    return np.sqrt(quad(ig, np.log(1e-4), np.log(100), limit=200)[0])
def tinker08_f(sig, Delta=200.0):
    A0, a0_, b0, c0 = 0.186, 1.47, 2.57, 1.19
    alz = 10 ** (-(0.75 / np.log10(Delta / 75.0))**1.2)
    return (A0*(1+Z_EVAL)**-0.14) * (((sig/(b0*(1+Z_EVAL)**-alz))**(-a0_*(1+Z_EVAL)**-0.06))+1) * np.exp(-c0/sig**2)
def tinker10_bias(nu, Delta=200.0):
    y = np.log10(Delta)
    A = 1.0+0.24*y*np.exp(-(4/y)**4); a = 0.44*y-0.88; B=0.183; b=1.5
    C = 0.019+0.107*y+0.19*np.exp(-(4/y)**4); c=2.4
    return 1 - A*nu**a/(nu**a+DELTA_C**a) + B*nu**b + C*nu**c

def b_eff_tracer():
    masses = np.logspace(11, 15.3, 26)
    sig = np.array([sigma_M(m) for m in masses]); nu = DELTA_C/sig
    dndlnM = (rho_mean_z/(masses*Msun))*tinker08_f(sig)*np.abs(np.gradient(np.log(sig), np.log(masses)))
    occ = dndlnM*masses  # mass-weighted occupation (CMASS-like)
    return np.sum(tinker10_bias(nu)*occ)/np.sum(occ)

# ----------------------------------------------------------------- the boost as a function of x_ext
def mu_inv(x): return (1 + x) / x

def sigmaR_nl(R):
    def ig(lnk):
        k = np.exp(lnk); kR = k*R
        W = 3*(np.sin(kR)-kR*np.cos(kR))/kR**3 if kR>1e-3 else 1.0
        return Pnl(k)*W**2*k**3/(2*np.pi**2)
    return np.sqrt(quad(ig, np.log(1e-4), np.log(100), limit=200)[0])

def Q_screen(x_ext, R_los=5.0, gal_weight=True, bg=1.8):
    """Volume/path-weighted <1/mu> along the LOS, sampling the lognormal density field at smoothing R_los.
    The local x = sqrt(x_self^2 + x_ext^2): x_self from the cell's coherent self-accel, x_ext = EFE floor.
    Weight by the convergence each cell adds (galaxy-density-weighted for the BAO tracer LOS)."""
    sig = sigmaR_nl(R_los); sL = np.sqrt(np.log(1 + sig**2))
    gg = np.linspace(-5*sL, 5*sL, 4000); opd = np.exp(gg - sL**2/2); delta = opd - 1
    Pg = np.exp(-gg**2/(2*sL**2))/(np.sqrt(2*np.pi)*sL)
    x_self = (4/3)*np.pi*G*rho_mean_z*np.abs(delta)*(R_los*Mpc)/a0_z   # coherent self accel / a0
    x = np.sqrt(x_self**2 + x_ext**2); x = np.maximum(x, 1e-3)
    w = Pg*(np.maximum(opd, 1e-6)*bg if gal_weight else np.maximum(delta, 0))
    w = w/np.trapezoid(w, gg)
    return np.trapezoid(mu_inv(x)*w, gg), np.trapezoid(x*w, gg)

# ----------------------------------------------------------------- single-halo cross-check (the toy, EFE-free)
def Rdef_200m(Mm): return (3*Mm*Msun/(4*np.pi*200*rho_mean_z))**(1/3)/Mpc
def cMz(Mm, cm=1.0): return cm*5.71*(Mm/(2e12/h))**(-0.084)*(1+Z_EVAL)**(-0.47)
def toy_Q(x_ext, rsp=1.5, slope=1.5, cm=1.0):
    """Halo-profile-weighted <1/mu> WITH the EFE floor (the toy was x_ext=0)."""
    masses = np.logspace(11, 15.3, 22)
    sig = np.array([sigma_M(m) for m in masses]); nu = DELTA_C/sig
    dndlnM = (rho_mean_z/(masses*Msun))*tinker08_f(sig)*np.abs(np.gradient(np.log(sig), np.log(masses)))
    X=[]; W=[]
    for i, Mm in enumerate(masses):
        M=Mm*Msun; Rb=Rdef_200m(Mm); cc=cMz(Mm,cm); rs=Rb/cc; mc=np.log(1+cc)-cc/(1+cc)
        rr=np.logspace(np.log10(rs*1e-2), np.log10(Rb*rsp), 400); xn=rr/rs
        rho_s=M/(4*np.pi*(rs*Mpc)**3*mc)
        rho=np.where(rr<=Rb, rho_s/(xn*(1+xn)**2), rho_s/((Rb/rs)*(1+Rb/rs)**2)*(rr/Rb)**(-slope))
        Menc=M*(np.log(1+xn)-xn/(1+xn))/mc; out=rr>Rb
        if out.any():
            ro=rr[out]*Mpc; rb=rho_s/((Rb/rs)*(1+Rb/rs)**2)*(rr[out]/Rb)**(-slope)
            dM=4*np.pi*.5*(rb[1:]*ro[1:]**2+rb[:-1]*ro[:-1]**2)*np.diff(ro)
            Menc=Menc.copy(); Menc[out]=M+np.concatenate([[0], np.cumsum(dM)])
        x_self=G*Menc/(rr*Mpc)**2/a0_z; x=np.sqrt(x_self**2+x_ext**2)
        w=rho*(rr*Mpc)**2*np.gradient(rr*Mpc)*dndlnM[i]*masses[i]
        X.append(np.maximum(x,1e-3)); W.append(w)
    X=np.concatenate(X); W=np.concatenate(W); W/=W.sum()
    return np.sum(mu_inv(X)*W), np.sum(X*W)


if __name__ == "__main__":
    print("=" * 96)
    print("STRAND 3 — 2-HALO LOS SCREEN AMPLITUDE vs the EFE floor x_ext (honest; nothing tuned to 0.030)")
    print("=" * 96)
    KAPPA_RAND = 0.0029            # MEASURED per-LOS convergence RMS (boss_data/kappa_fg_dataset.npz)
    beff = b_eff_tracer()
    bg = 1.8; r_gm = 1.0
    kappa_gal = bg*r_gm*KAPPA_RAND
    print(f"  z={Z_EVAL}  sigma8(z)={sigma8_z:.4f}  a0(z)={a0_z:.3e} m/s^2")
    print(f"  b_eff(Tinker10, CMASS-like) = {beff:.2f}    kappa_gal = b_g r_gm * 0.0029 = {kappa_gal:.4f}")
    print(f"  EFE candidates: x_ext(peculiar LSS)~0.003 [Q~300] ; x_ext(Hubble cH)={X_HUBBLE:.2f} [Q={mu_inv(X_HUBBLE):.2f}]")
    print(f"  Target FIXED: 0.030 (BAO diff); 0.039 (data-pref). Required Q = target/kappa_gal = "
          f"{0.030/kappa_gal:.1f} -> needs x_ext = 1/(Q-1) = {1/(0.030/kappa_gal-1):.3f}\n")

    print("  [1] BOOST vs EFE FLOOR — two INDEPENDENT LOS estimators (halo-profile & lognormal-PDF path):")
    print(f"  {'x_ext':>7s} {'Q_toy(halo)':>12s} {'exc_halo':>9s} {'Q_pdf(LOS)':>11s} {'exc_pdf':>9s}")
    xext_grid = [0.0, 0.1, 0.2, 0.24, 0.3, 0.5, 1.0, 2.0, X_HUBBLE]
    for xe in xext_grid:
        Qt,_ = toy_Q(xe); Qp,_ = Q_screen(xe)
        tag = "  <- corpus Hubble-EFE" if abs(xe-X_HUBBLE)<1e-6 else ("  <- screen module xbar~0.24" if abs(xe-0.24)<1e-6 else "")
        print(f"  {xe:7.3f} {Qt:12.2f} {kappa_gal*Qt:9.4f} {Qp:11.2f} {kappa_gal*Qp:9.4f}{tag}")

    print("\n  [2] WHERE DO THE TARGETS LAND? (solve x_ext that yields each amplitude, halo estimator)")
    from scipy.optimize import brentq
    for tgt in (0.030, 0.039):
        try:
            xsol = brentq(lambda xe: kappa_gal*toy_Q(xe)[0] - tgt, 0.01, 50.0)
            print(f"     excess={tgt}: requires x_ext = {xsol:.3f}  (Q={tgt/kappa_gal:.2f})")
        except Exception as e:
            print(f"     excess={tgt}: no solution in range ({e})")

    print("\n  [3] PRE-REGISTERED INGREDIENT SENSITIVITY at the screen-module floor x_ext=0.24 (un-tuned):")
    print(f"  {'axis':24s} {'value':>8s} {'Q':>7s} {'excess':>9s} {'x/0.030':>8s}")
    base_xe = 0.24
    scans = [
        ("splashback R_sp/R200m", [1.0,1.5,2.0], lambda v: toy_Q(base_xe, rsp=v)[0]),
        ("infall slope",          [1.0,1.5,2.0], lambda v: toy_Q(base_xe, slope=v)[0]),
        ("concentration xc",      [0.7,1.0,1.4], lambda v: toy_Q(base_xe, cm=v)[0]),
        ("LOS smooth R [Mpc]",    [2.0,5.0,8.0], lambda v: Q_screen(base_xe, R_los=v)[0]),
        ("galaxy bias b_g",       [1.5,1.8,2.1], None),
        ("kappa_rand RMS",        [0.0023,0.0029,0.0045], None),
    ]
    lo=hi=kappa_gal*toy_Q(base_xe)[0]
    for name, vals, fn in scans:
        for v in vals:
            if name=="galaxy bias b_g":
                Q=toy_Q(base_xe)[0]; e=(v*r_gm*KAPPA_RAND)*Q
            elif name=="kappa_rand RMS":
                Q=toy_Q(base_xe)[0]; e=(bg*r_gm*v)*Q
            else:
                Q=fn(v); e=kappa_gal*Q
            lo=min(lo,e); hi=max(hi,e)
            print(f"  {name:24s} {str(v):>8s} {Q:7.2f} {e:9.4f} {e/0.030:8.2f}")
    print("\n  " + "-"*92)
    print(f"  AMPLITUDE BAND at the screen-module floor x_ext=0.24, over un-tuned ingredients: "
          f"[{lo:.4f}, {hi:.4f}]")
    print(f"  target 0.030 in band? {'YES' if lo<=0.030<=hi else 'NO'} ; "
          f"data-pref 0.039 in band? {'YES' if lo<=0.039<=hi else 'NO'}")
    print("\n  VERDICT: the 2-halo upgrade does NOT pin the amplitude; it shows the amplitude is a steep")
    print("  function of the UNDERIVED EFE floor x_ext. At the corpus Hubble-EFE floor (x=5.85) the screen")
    print("  is 5x TOO SMALL (~0.006). The in-range result requires x_ext~0.2-0.3 (weak void-skin field),")
    print("  which the corpus admits is a hypothesis, not action-derived. So the 2-halo calc WEAKENS the")
    print("  'amplitude brackets 0.030' claim: it survives ONLY in the weak-EFE void-skin regime.")
