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.5; 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); 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)
print(f"LTB mu=0.5: t_C(0)={L.tC(1e-6)/t_H:.3f} t_H; t_AH(r) for r/rm=0.1,0.2,0.3,0.5: " + ", ".join(f"{L.t_AH(f*rm)/t_H:.2f}" for f in (0.1,0.2,0.3,0.5)))
times = np.array([4, 6, 7, 7.5, 8, 8.5, 9, 9.5])*t_H; k = [0]
def cb(run, d):
    if k[0] < len(times) and run.t >= times[k[0]]:
        k[0] += 1
        met = solve_metric(grid, bg, part, run.t); R = met.R_edge[1:]; comp = 2*met.m_edge[1:]/R; RK = R*met.Kth_edge[1:]
        j = np.argmax(comp[:400]); jk = np.argmax(RK)
        Rp = met.a*part.x; o = np.argsort(Rp)[:3]
        W = np.sqrt(1+part.P**2); dRdtau = part.P/met.p_A - Rp*met.p_Kth*W
        print(f"t/t_H={run.t/t_H:5.2f}: max 2m/R={comp[j]:.4f} at x/rm={R[j]/met.a/rm:.4f} (RKth there={RK[j]:+.4f}, alpha={met.alpha_edge[j+1]:.4f}, A={met.A_edge[j+1]:.3f}); max RKth={RK[jk]:+.4f} at x/rm={R[jk]/met.a/rm:.3f}; "
              f"inner shells x/rm={Rp[o]/met.a/rm} P={part.P[o]} dR/dtau={dRdtau[o]} 2m/R@shell={2*met.p_m[o]/Rp[o]}; alpha_c[0]={met.alpha_c[0]:.4f}")
run = CosmoRun(grid, bg, part, CosmoRunConfig(cfl=0.4, t_end=9.5*t_H, diag_every=5), bg.t_i)
run.run(cb)