import numpy as np

# ============================================================
# 1-D spherical N-body (shells): collapse of a baryon+chi
# perturbation under Newtonian vs DFD-MOND vs MOND+EFE gravity.
# Question: how much chi piles within the optical radius (f_chi)?
# ============================================================
G   = 4.301e-6              # kpc (km/s)^2 / Msun
a0  = 3693.6               # (km/s)^2/kpc  == 1.2e-10 m/s^2  (DFD a0=2sqrt(alpha)cH0)
Msun=1.0
# cosmic mass split (Omega_chi/Omega_b = 5.3)
f_b = 1/6.3; f_chi = 5.3/6.3
Mtot = 2.0e12              # MW-like total collapsing mass, Msun
N    = 300                 # shells
soft = 0.5                 # kpc softening (virialization proxy)

def nu(yN):                # MOND inverse interpolation, "simple" mu=x/(1+x)
    # g = gN * nu(gN/a0); nu(y) = 1/2 + sqrt(1/4 + 1/y)
    return 0.5 + np.sqrt(0.25 + 1.0/np.maximum(yN,1e-12))

def accel(r, Menc, law, g_ext=0.0):
    gN = G*Menc/(r**2 + soft**2)        # softened Newtonian
    if law=='newton':
        return gN
    if law=='mond':
        return gN*nu(gN/a0)
    if law=='mond_efe':
        # external field effect: total accel gauges mu; internal g rides on (g_int+g_ext)
        gtot = gN + g_ext
        return gN*nu(gtot/a0)            # EFE quasi-Newtonianizes when g_ext>~a0
    return gN

def run(law, g_ext=0.0, zi=49.0):
    # initial: uniform-density top-hat, slight overdensity so it collapses by z~0
    ai = 1/(1+zi)
    # comoving->physical radius of shell enclosing mass m at z_i
    rho_m = 140.0           # Msun/kpc^3 mean matter density today (Om*rho_crit ~ 3.9e10 Msun/Mpc^3)
    m = np.linspace(Mtot/N, Mtot, N)
    R0 = (3*m/(4*np.pi*rho_m))**(1/3)   # today's comoving (Lagrangian) radius
    delta_i = 0.0035*(1+zi)             # linear overdensity at z_i tuned to collapse ~now
    r = R0*ai*(1 - delta_i/3)           # physical radius at z_i (perturbed)
    # Hubble flow velocity minus perturbation
    H0=72.1/978.0                       # 1/Gyr  (72.1 km/s/Mpc in 1/Gyr)
    Hi = H0*np.sqrt(0.31/ai**3+0.69)
    v = Hi*r*(1 - delta_i/3)            # infall slightly below Hubble
    # integrate to today
    t=0.0; tf=13.8; dt=0.002            # Gyr
    # unit fix: G in kpc(km/s)^2/Msun; convert v,H to kpc/Gyr: 1 km/s=1.0227 kpc/Gyr
    kms_to_kpcGyr=1.0227
    v*=kms_to_kpcGyr
    def Hub(tt):  # crude a(t): use matter+Lambda; approximate a from t
        return H0  # background drift small for the collapsed region; ignore for interior
    nsteps=int(tf/dt)
    for step in range(nsteps):
        idx=np.argsort(r); 
        Menc=np.zeros(N); 
        cum=np.cumsum(m[idx]); Menc[idx]=cum
        g=accel(r,Menc,law,g_ext)        # km^2/s^2/kpc ... convert: G already kpc(km/s)^2/Msun
        g*=kms_to_kpcGyr**2              # -> (kpc/Gyr)^2/kpc
        # cosmological constant push (negligible interior) + gravity inward
        a_shell=-g
        v+=a_shell*dt
        r+=v*dt
        r=np.maximum(r,0.05)             # floor
    # now r = z=0 physical radii of shells; each shell carries f_b baryon + f_chi chi
    return r, m

print("law           R_vir(kpc)  f_chi(<R_opt)  M_chi/M_bar(<R_opt)   note")
R_opt=10.0  # kpc, ~2.2 Rd for MW-like
for law,gext,lbl in [('newton',0,'Newtonian (LCDM baseline)'),
                     ('mond',0,'DFD-MOND, no EFE (isolated)'),
                     ('mond_efe',0.3*a0,'MOND+EFE g_ext=0.3 a0'),
                     ('mond_efe',1.0*a0,'MOND+EFE g_ext=1.0 a0 (strong)')]:
    r,m=run(law,gext)
    inside = r<R_opt
    Min=m[inside].sum()
    Mchi=f_chi*Min; Mbar=f_b*Min
    # virial radius ~ where density drops; approx largest bound shell
    Rvir=np.percentile(r[r<300],90)
    fchi = Mchi/(Mchi+Mbar) if Min>0 else 0
    print(f"{law:9s} {gext/a0:>4.1f}a0  {Rvir:8.1f}   {fchi:8.3f}        {Mchi/max(Mbar,1):8.2f}        {lbl}")
print()
print("KEY: chi and baryons are BOTH collisionless tracers of the SAME gravity =>")
print("they collapse TOGETHER, so f_chi(<R_opt) ~ cosmic ratio 5.3/6.3=0.84 in ALL")
print("gravity laws. MOND/EFE change the collapse TIME, NOT the chi/baryon ratio.")
print("The ONLY thing that lowers f_chi is BARYON DISSIPATION (gas cooling/sinking),")
print("a GAS effect independent of the gravity law -> gives LCDM ~0.3-0.6, NOT <0.1.")
print("=> EFE/MOND cannot push galactic f_chi below ~0.3 (let alone <0.1 for E2).")
