import sys, numpy as np
sys.path.insert(0, '.')
from pbhgr.cosmo_ev import Background, CosmoGrid, CParticles, CosmoRun, CosmoRunConfig, solve_metric
from pbhgr.cosmo_id import yoo_spherical_cmc, sample_shells_quiet
from pbhgr.ltb import LTBGaussian
mu = 0.36; bg = Background(5.0); rm = np.sqrt(6); t_H = bg.t_i*(np.sqrt(6)/0.2)**3; x_out = 5*rm
grid = CosmoGrid(x_out, 2000, stretch=4.0); prof = yoo_spherical_cmc(mu, r_max=x_out*1.05)
x, P, Lsq, N = sample_shells_quiet(prof, grid.x_edge, 4); part = CParticles(x, P, Lsq, N)
L = LTBGaussian(mu)
sel = [int(np.argmin(np.abs(x - f*rm))) for f in (0.01, 0.03, 0.06, 0.09, 0.12, 0.2)]
print("labels x0/rm:", [f"{x[i]/rm:.3f}" for i in sel], "LTB t_C(0)=%.3f t_H" % (L.tC(1e-6)/t_H), "t_C(labels):", [f"{L.tC(x[i])/t_H:.3f}" for i in sel], "2m_LTB:", [f"{2*L.m(x[i]):.3f}" for i in sel])
print("columns per shell: R | P | alpha | 2m/R ; then global: min_alpha@x/rm, max2m/R(coll)@x/rm, bad, n(x<0.02rm)")
def cb(run, d):
if run.t < 8.3*t_H: return
met = solve_metric(grid, bg, part, run.t); R = met.a*np.abs(part.x)
al = np.interp(R, met.R_edge, met.alpha_edge)
comp = 2*met.p_m/np.maximum(R,1e-300)
j = np.argmin(met.alpha_c); Re = met.R_edge[1:]; cc = 2*met.m_edge[1:]/Re; cc = np.where(met.Kth_edge[1:]>0, cc, -1); jm = np.argmax(cc)
print(f"t={run.t/t_H:6.3f} " + " ".join(f"[{R[i]:7.3f} {part.P[i]:+7.3f} {al[i]:.2e} {comp[i]:5.2f} tau={ (part.tau[i]+bg.t_i)/t_H:5.3f}]" for i in sel)
+ f" | a_min={met.alpha_c.min():.1e}@{met.R_c[j]/met.a/rm:.4f} c_max={cc[jm]:.2f}@{Re[jm]/met.a/rm:.4f} bad={met.n_bad} n_in={(np.abs(part.x)<0.02*rm).sum()}")
run = CosmoRun(grid, bg, part, CosmoRunConfig(cfl=0.4, t_end=12.5*t_H, diag_every=40), bg.t_i)
run.run(cb)