"""Caustic benchmark: cold dust sphere at rest (time-symmetric, weak field), evolved with the GR areal/maximal-slicing
code (StaticBackground) and with a Newtonian multi-stream shell code on IDENTICAL shells.  Until the central caustic
both must agree with the LTB/Newtonian free fall; after it, the multi-stream pile-up (mass inside fixed radii, max 2m/R)
is compared in the regime where the compaction is still small.  Usage: python3 v1_caustic_bench.py [C=2M/R0] [K] [s] [npc]"""
import sys, time, numpy as np
sys.path.insert(0, '.')
from pbhgr.cosmo_ev import StaticBackground, CosmoGrid, CParticles, CosmoRun, CosmoRunConfig, solve_metric
from pbhgr.cosmo_id import sample_shells_quiet
C = float(sys.argv[1]) if len(sys.argv) > 1 else 0.02
K = int(sys.argv[2]) if len(sys.argv) > 2 else 2000
s = float(sys.argv[3]) if len(sys.argv) > 3 else 4.0
npc = int(sys.argv[4]) if len(sys.argv) > 4 else 4
t_end_fac = float(sys.argv[5]) if len(sys.argv) > 5 else 1.3
R0 = 1.0; x_out = 4.0 * R0
M = C * R0 / 2                                        # rho = rho0 exp(-R^2/R0^2): M = pi^{3/2} rho0 R0^3
rho0 = M / (np.pi**1.5 * R0**3)
Rg = np.linspace(0, x_out, 40001)
prof = dict(R=Rg, dN_dR=4 * np.pi * Rg**2 * rho0 * np.exp(-Rg**2 / R0**2), v=np.zeros_like(Rg))
grid = CosmoGrid(x_out=x_out, K=K, stretch=s)
x, P, Lsq, N = sample_shells_quiet(prof, grid.x_edge, npc)
from scipy.special import erf
m_of = lambda R: np.pi**1.5 * rho0 * R0**3 * (erf(R / R0) - 2 * R / (np.sqrt(np.pi) * R0) * np.exp(-R**2 / R0**2))
tC = np.pi / 2 * np.sqrt(x**3 / (2 * m_of(x)))         # Newtonian free-fall time of each shell from rest
tC0 = np.pi / 2 / np.sqrt(8 * np.pi * rho0 / 3)
print(f"C=2M/R0={C}: M={M:.4f}, rho0={rho0:.4e}, t_C(0)={tC0:.3f}, t_C(0.3 R0)/t_C(0)={np.interp(0.3, x, tC)/tC0:.3f}, t_C(R0)/t_C(0)={np.interp(1.0, x, tC)/tC0:.3f}; K={K} s={s} dx_min={grid.dx:.2e}, shells={len(x)}", flush=True)
Rlist = (0.01, 0.03, 0.1, 0.3)
tlist = np.array([0.9, 0.98, 1.0, 1.01, 1.02, 1.05, 1.1, 1.2, 1.3]) * tC0
tlist = tlist[tlist <= t_end_fac * tC0 + 1e-9]
# ---------------- Newtonian shells (leapfrog, softening eps = dx_min)
def newton(eps):
    R = x.copy(); v = np.zeros_like(R); n = len(R)
    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
        mi = np.empty(n); mi[order] = m_in
        return a, mi
    a, mi = accel(R); t = 0.0; out = {}; step = 0; peak = 0.0
    for tt in tlist:
        while t < tt:
            dt = min(0.05 * np.min(np.sqrt((R**2 + eps**2)**1.5 / np.maximum(mi, 1e-12))), 0.05 * np.min((np.abs(R) + eps) / np.maximum(np.abs(v), 1e-12)), tt - t)
            v += 0.5 * dt * a; R += dt * v
            neg = R < 0; R[neg] = -R[neg]; v[neg] = -v[neg]
            a, mi = accel(R); v += 0.5 * dt * a; t += dt; step += 1
            comp = 2 * mi / np.sqrt(R**2 + eps**2); peak = max(peak, comp.max())
        out[tt] = ([N[R < Rf].sum() for Rf in Rlist], (2 * mi / np.sqrt(R**2 + eps**2)).max(), peak, step)
    return out
t0 = time.time(); nw = newton(grid.dx); print(f"Newtonian done in {time.time()-t0:.0f}s ({nw[tlist[-1]][3]} steps)", flush=True)
# ---------------- GR
part = CParticles(x.copy(), P.copy(), Lsq.copy(), N.copy())
bg = StaticBackground()
met0 = solve_metric(grid, bg, part, 0.0)
print(f"GR t=0: M_ADM={met0.m_edge[-1]:.5f} (sum N={N.sum():.5f}), alpha(0)={met0.alpha_c[0]:.5f}, max 2m/R={np.max(2*met0.m_edge[1:]/met0.R_edge[1:]):.4f}", flush=True)
gr = {}; peak = [0.0]
def cb(run, key):
    met = solve_metric(grid, bg, part, run.t); comp = 2 * met.m_edge[1:] / met.R_edge[1:]
    peak[0] = max(peak[0], comp[met.Kth_edge[1:] > 0].max() if np.any(met.Kth_edge[1:] > 0) else 0.0)
    gr[key] = ([float(np.interp(Rf, met.R_edge, met.m_edge)) for Rf in Rlist], float(comp[met.Kth_edge[1:] > 0].max()) if np.any(met.Kth_edge[1:] > 0) else 0.0, peak[0], run.step, float(met.alpha_c[0]), float(np.median(part.tau[np.abs(part.x) < 0.05])), int(met.n_bad))
cfg = CosmoRunConfig(cfl=0.3, t_end=tlist[-1], diag_every=10**9, max_dt_frac=1e9, snapshot_times=tuple(tlist))
run = CosmoRun(grid, bg, part, cfg, 0.0)
t0 = time.time()
# drive the run so that the GR sample is taken at equal CENTRAL PROPER TIME tau_core = tt (alpha < 1 in the core)
core = np.abs(part.x) < 0.05
def tau_core(): return float(np.median(part.tau[np.abs(part.x) < 0.05]))
rate = float(met0.alpha_c[0])
for tt in tlist:
    while True:
        rem = (tt - tau_core()) / rate
        if rem < 1e-7 * tC0: break
        t_a, tau_a = run.t, tau_core()
        run.cfg.t_end = run.t + rem
        run.run(lambda r, d: None)
        if run.t > t_a: rate = max((tau_core() - tau_a) / (run.t - t_a), 0.05)
    cb(run, tt)
print(f"GR done in {time.time()-t0:.0f}s ({run.step} steps)", flush=True)
print("tau/tC0 (Newton: t/tC0) | Newton m(<0.01) m(<0.03) m(<0.1) m(<0.3) max2m/r peak | GR m(<0.01) m(<0.03) m(<0.1) m(<0.3) max2m/R peak | alpha(0) tau_core/t  bad")
for tt in tlist:
    a = nw[tt]; b = gr[tt]
    print(f"{tt/tC0:5.2f} | " + " ".join(f"{v/M:7.4f}" for v in a[0]) + f" {a[1]:6.3f} {a[2]:6.3f} | " + " ".join(f"{v/M:7.4f}" for v in b[0]) + f" {b[1]:6.3f} {b[2]:6.3f} | {b[4]:.4f} {b[5]/tt:.4f} {b[6]}")
gr_t = sorted(gr.keys())