"""Newtonian spherical shell code started from the exact LTB state at 0.9 t_C(0) (LTB shell dynamics are Newtonian in
form, so this is exact until the first crossing).  After crossing: multi-stream Newtonian gravity a = -m_in r/(r^2+eps^2)^{3/2}
with m_in = mass of shells inside (+ half self).  Diagnostic: max 2 m_in/r over the core shells (would be a trapped surface in GR)."""
import sys, numpy as np
sys.path.insert(0, '.')
from pbhgr.ltb import LTBGaussian
mu = float(sys.argv[1]) if len(sys.argv) > 1 else 0.1
eps = float(sys.argv[2]) if len(sys.argv) > 2 else 0.05
L = LTBGaussian(mu); tH = L.t_H; rm = L.r_m; tC0 = L.tC(1e-6)
K, s, npc = 2000, 4.0, 4
xe = 2*rm*np.sinh(s*np.arange(K+1)/K)/np.sinh(s); xe[0] = 0
off = (np.arange(npc)+0.5)/npc
xl = (xe[:-1,None] + (off[None,:]-0.5/npc)*np.diff(xe)[:,None]).ravel(); xr = (xe[:-1,None] + (off[None,:]+0.5/npc)*np.diff(xe)[:,None]).ravel()
lab = (0.5*(xl**3+xr**3))**(1/3); keep = lab < 0.4*rm
lab, xl, xr = lab[keep], xl[keep], xr[keep]
N = L.m(xr) - L.m(xl); m0 = L.m(lab); E0 = L.E(lab); n = len(lab)
t = 0.9*tC0
R = np.array([L.R(t, r) for r in lab]); Rd = np.array([L.Rdot(t, r) for r in lab]); v = Rd.copy()
print(f"mu={mu}, eps={eps}: {n} shells (labels < 0.4 r_m), start t={t/tH:.2f} t_H = 0.9 t_C(0); t_C(0)={tC0/tH:.2f}; LTB t_AH at labels 0.05,0.1,0.15,0.2: " + ", ".join(f"{L.t_AH(f*rm)/tH:.2f}" for f in (0.05,0.1,0.15,0.2)))
def accel(R):
    order = np.argsort(R); Ns = N[order]
    m_in = np.cumsum(Ns) - 0.5*Ns
    Rs = R[order]
    a = np.empty(n); a[order] = -m_in*Rs/(Rs**2 + eps**2)**1.5
    comp = np.empty(n); comp[order] = 2*m_in/np.sqrt(Rs**2 + eps**2)
    return a, comp, m_in[np.argsort(order)]
a, comp, m_in = accel(R)
t_end = 1.25*tC0; step = 0; next_print = 0.96*tC0; peak = 0.0; peak_t = 0
first_trap = None
while t < t_end:
    dt = min(0.05*np.min(np.sqrt((R**2+eps**2)**1.5/np.maximum(m_in, 1e-9))), 0.05*np.min((np.abs(R)+eps)/np.maximum(np.abs(v),1e-9)), 1e-3*t, t_end-t)
    v += 0.5*dt*a; R += dt*v
    neg = R < 0; R[neg] = -R[neg]; v[neg] = -v[neg]
    a, comp, m_in = accel(R); v += 0.5*dt*a
    t += dt; step += 1
    c = comp.max()
    if c > peak: peak, peak_t, peak_lab = c, t, lab[np.argmax(comp)]
    if first_trap is None and c >= 1.0: first_trap = (t, lab[np.argmax(comp)], R[np.argmax(comp)], m_in[np.argmax(comp)])
    if t >= next_print:
        j = np.argmax(comp)
        print(f"t/t_H={t/tH:7.3f} (t/tC0={t/tC0:.3f}, {step} steps) max 2m/r={c:.3f} at label {lab[j]/rm:.3f} r_m (r={R[j]:.3f}, m_in={m_in[j]:.3f}); mass inside r<1: {N[R<1].sum():.3f}, r<2.5: {N[R<2.5].sum():.3f}, r<5: {N[R<5].sum():.3f}, r<10: {N[R<10].sum():.3f}; crossed(v>0 & label<0.1): {np.mean(v[lab<0.1*rm]>0):.2f}", flush=True)
        next_print += 0.02*tC0
print(f"peak 2m/r = {peak:.3f} at t/t_H={peak_t/tH:.3f} (label {peak_lab/rm:.3f}); first 2m/r>=1: {None if first_trap is None else (round(first_trap[0]/tH,3), round(first_trap[1]/rm,3), round(first_trap[2],3), round(first_trap[3],3))}")