"""
HALO-MODEL LOS (the clean first-principles attempt the referee asked for: compute the galaxy-weighted <x>
from the REAL halo population, NOT inferred). The screen source is the trace ~rho (MASS-weighted), so this
also SETTLES the weighting concern: does the AQUAL boost <1/mu(x)>, x=g(r)/a0, survive when weighted by rho
(which lives in dense Newtonian cores), or die? We integrate the NFW profile of each halo (real g(r), real
rho(r)), weight by the Sheth-Tormen mass function, and report the mass-weighted (=trace-source) <1/mu> and the
galaxy-weighted version (HOD ~ favors group/cluster halos). honesty-guard: x is COMPUTED from g(r)/a0, never chosen.
"""
import numpy as np
from scipy.integrate import quad
import camb
import dfd_params as P

G=6.674e-11; Msun=1.989e30; Mpc=3.0857e22
c_si=2.998e8; H0=P.H0*1e3/Mpc
a0=2*np.sqrt(P.ALPHA)*c_si*H0                      # 1.197e-10 m/s^2
z=0.5; Ez=np.sqrt(P.OMEGA_M*(1+z)**3+P.OMEGA_LAMBDA)
rho_crit_z=3*(H0*Ez)**2/(8*np.pi*G)               # kg/m^3 at z
rho_m_z=P.OMEGA_M*(1+z)**3*3*H0**2/(8*np.pi*G)

# ---- CAMB sigma(M) ----
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],kmax=80.0)
res=camb.get_results(pars); PK=camb.get_matter_power_interpolator(pars,nonlinear=False,kmax=80,zmax=1)
h=P.H0/100
def sigma_M(Msun_mass):
    M=Msun_mass*Msun
    R=(3*M/(4*np.pi*rho_m_z))**(1/3)              # Lagrangian radius [m]
    R_Mpc=R/Mpc
    def integ(kM):                                # kM in 1/Mpc
        kR=kM*R_Mpc; W=3*(np.sin(kR)-kR*np.cos(kR))/kR**3 if kR>1e-3 else 1.0
        return PK.P(z,kM)*W**2*kM**2/(2*np.pi**2)
    return np.sqrt(quad(integ,1e-4,80,limit=200)[0])

dc=1.686
def st_f(nu):                                     # Sheth-Tormen
    A,a,p=0.3222,0.707,0.3
    return A*np.sqrt(2*a/np.pi)*(1+(1/(a*nu**2))**p)*nu*np.exp(-a*nu**2/2)

# ---- NFW per-halo mass-weighted <1/mu> ----
def Rvir(M): return (3*M/(4*np.pi*200*rho_crit_z))**(1/3)   # R200c [m]
def conc(Msun_mass): return 5.0*(Msun_mass/1e14)**(-0.1)    # rough c(M), c~5 at 1e14
def halo_boost(Msun_mass):
    M=Msun_mass*Msun; Rv=Rvir(M); cc=conc(Msun_mass); rs=Rv/cc
    mc=np.log(1+cc)-cc/(1+cc)
    rr=np.logspace(np.log10(rs*1e-2),np.log10(Rv),300)
    x_nfw=rr/rs
    Mr=M*(np.log(1+x_nfw)-x_nfw/(1+x_nfw))/mc      # M(<r)
    g=G*Mr/rr**2; xacc=g/a0
    rho=1.0/(x_nfw*(1+x_nfw)**2)                    # NFW shape (un-normalized)
    invmu=(1+xacc)/xacc
    w=rho*rr**3                                     # rho * r^2 dr with dlnr (rr^3 for log grid)
    return np.trapezoid(invmu*w,np.log(rr))/np.trapezoid(w,np.log(rr)), np.average(xacc,weights=w)

print("="*86)
print("HALO-MODEL LOS — trace(rho)-weighted <1/mu(x)>, x=g(r)/a0 COMPUTED from NFW (z=0.5)")
print("="*86)
print(f"  a0={a0:.3e} m/s^2.  BAO-need Q~5 (x~0.2); referee/corpus inferred x~0.1-0.24.\n")
print(f"  {'M[Msun]':>9s} {'sigma(M)':>9s} {'<1/mu>_halo':>11s} {'<x>_halo':>9s}")
masses=np.logspace(11,15.3,18)
boosts=[]; xbars=[]; sigs=[]
for Mh in masses:
    s=sigma_M(Mh); b,xb=halo_boost(Mh); boosts.append(b); xbars.append(xb); sigs.append(s)
    if abs(np.log10(Mh)-round(np.log10(Mh)))<0.13:
        print(f"  {Mh:9.1e} {s:9.3f} {b:11.2f} {xb:9.3f}")
boosts=np.array(boosts); xbars=np.array(xbars); sigs=np.array(sigs)
nu=dc/sigs
# mass function weight dn/dlnM ~ (rho_m/M) f(nu) |dln sigma/dlnM|
dlnsig=np.gradient(np.log(sigs),np.log(masses))
dndlnM=(rho_m_z/(masses*Msun))*st_f(nu)*np.abs(dlnsig)
massw=dndlnM*masses          # mass-weighted (trace source ~ rho ~ M per halo)
# galaxy-weighted: HOD N_gal ~ (M/M_min) for M>M_min~1e12, favors groups/clusters
Ngal=np.where(masses>1e12,(masses/1e12)**0.9,0.0)
galw=dndlnM*Ngal

Q_mass=np.trapezoid(boosts*massw,np.log(masses))/np.trapezoid(massw,np.log(masses))
x_mass=np.trapezoid(xbars*massw,np.log(masses))/np.trapezoid(massw,np.log(masses))
Q_gal=np.trapezoid(boosts*galw,np.log(masses))/np.trapezoid(galw,np.log(masses))
x_gal=np.trapezoid(xbars*galw,np.log(masses))/np.trapezoid(galw,np.log(masses))
print("\n  POPULATION AVERAGES (computed, not chosen):")
print(f"    MASS-weighted (trace source ~rho):  <x>={x_mass:.3f}  Q=<1/mu>={Q_mass:.2f}  -> A=kappa_gal*Q={0.006*Q_mass:.3f}")
print(f"    GALAXY-weighted (HOD, SNe hosts):   <x>={x_gal:.3f}  Q=<1/mu>={Q_gal:.2f}  -> A={0.006*Q_gal:.3f}")
print(f"\n  BAO need A~0.03 (Q~5).  kappa_gal baseline ~0.006.")
print("  HONEST READ: this is COMPUTED from the NFW halo population (x from g(r)/a0, never chosen). If Q lands")
print("  ~5 (A~0.03) under the TRACE(mass) weighting, the boost SURVIVES the source-weighting concern and the")
print("  amplitude is derived-in-range first-principles. Residuals: NFW c(M), HOD, 1-halo only (no filaments/")
print("  2-halo), lognormal vs N-body. NOT the final word, but x is no longer inferred-to-fit.")
