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.3; 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, 1000); 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)
order = np.argsort(x); Mi = np.empty(len(x)); cs = np.cumsum(N[order]); Mi[order] = cs - 0.5*N[order]
rL = np.linspace(1e-6, x_out*1.05, 200001); mL = L.m(rL); EL = L.E(rL)
MrestL = np.concatenate([[0], np.cumsum(np.diff(mL)/np.sqrt(1+2*0.5*(EL[1:]+EL[:-1])))])
r_lab = np.interp(Mi, MrestL, rL); m0 = L.m(r_lab); E0 = L.E(r_lab)
met0 = solve_metric(grid, bg, part, bg.t_i); R0 = met0.a*part.x
tau_off = LTBGaussian.shell_tau(m0, E0, R0, np.ones_like(R0))
sel = [int(np.argmin(np.abs(x - f*rm))) for f in (0.02, 0.05, 0.1, 0.2, 0.5, 1.0)]
print("shell x0/rm:", [f"{x[i]/rm:.3f}" for i in sel], " t_C(0)/t_H=%.2f" % (L.tC(1e-6)/t_H), " t_C(shell)/t_H:", [f"{L.tC(r_lab[i])/t_H:.2f}" for i in sel])
def report(run, d):
met = solve_metric(grid, bg, part, run.t); R = met.a*part.x
tau = tau_off + part.tau
Rp = LTBGaussian.shell_R_of_tau(m0[sel], E0[sel], tau[sel])
Rd = LTBGaussian.shell_R_of_tau(m0[sel], E0[sel], tau[sel]*1.001)
al = np.interp(R[sel], met.R_edge, met.alpha_edge)
print(f"t/t_H={run.t/t_H:6.3f} " + " | ".join(f"R/RL={R[i]/Rp[j]-1:+.1e} P={part.P[i]:+.1e} a={al[j]:.3f} RKth/(HR)={met.p_Kth[i]*R[i]/(bg.H(run.t)*R[i]):+.3f} dRdtau_L={(Rd[j]-Rp[j])/(0.001*tau[i]):+.2e} dRdtau_sim={part.P[i]/met.p_A[i]-R[i]*met.p_Kth[i]*np.sqrt(1+part.P[i]**2):+.2e}" for j, i in enumerate(sel)))
run = CosmoRun(grid, bg, part, CosmoRunConfig(cfl=0.4, t_end=6.0*t_H, diag_every=300), bg.t_i)
run.run(report)