"""Lagrangian-shell test of the PM code against the exact LTB solution (spherical Yoo profile, mu = 0.1).
Usage: python3 pm3d_shelltest.py N Np smooth a_end"""
import sys, numpy as np; sys.path.insert(0, '.')
from pbhgr.pm3d import PM3D, H_I, T_I
from pbhgr.profiles import profile_table
from pbhgr.ltb import LTBGaussian
N, Np, smooth, a_end = int(sys.argv[1]), int(sys.argv[2]), float(sys.argv[3]), float(sys.argv[4])
mu = 0.1; L = LTBGaussian(mu); rm = L.r_m; tC0 = L.tC(1e-6)
pm = PM3D(N, 6 * rm, Np, smooth=smooth).make_ic(mu, profile_table("yoo")["spline"])
shells = [(0.05, 0.1), (0.1, 0.15), (0.15, 0.2), (0.25, 0.3), (0.4, 0.5)]
sel = [(pm.r_ell_q >= lo * rm) & (pm.r_ell_q < hi * rm) for lo, hi in shells]
print(f"N={N} Np={Np} smooth={smooth}: shells (r_q/r_m) " + ", ".join(f"{lo}-{hi} ({s.sum()} p)" for (lo, hi), s in zip(shells, sel)))
print("t/tC0  | PM mean comoving radius / initial, per shell   ||  LTB R/(a e^zeta r) per shell")
a = 1.0
for a_t in np.array([200, 400, 600, 800, 1000, 1150, 1300]):
    if a_t > a_end: break
    while a < a_t:
        da = min(0.02 * a, a_t - a); a, _ = pm.step(a, da)
    t = T_I * a**1.5
    c = pm.L / 2
    pmv, ltbv = [], []
    for (lo, hi), s in zip(shells, sel):
        rel = (pm.x[s] - c + pm.L / 2) % pm.L - pm.L / 2
        r = np.sqrt((rel**2).sum(1)); pmv.append(r.mean() / pm.r_ell_q[s].mean())
        rq = 0.5 * (lo + hi) * rm
        ltbv.append(L.R(t, rq) / (a * rq * np.exp(L.zeta(rq))) if t < L.tC(rq) else 0.0)   # LTB areal radius R = a e^zeta r (1 - ...): compare per unit Lagrangian q
    print(f"{t/tC0:5.3f}  | " + " ".join(f"{v:6.3f}" for v in pmv) + "   ||  " + " ".join(f"{v:6.3f}" for v in ltbv), flush=True)