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
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(0.3, r_max=x_out*1.05)
x, P, Lsq, N = sample_shells_quiet(prof, grid.x_edge, 4); part = CParticles(x, P, Lsq, N)
run = CosmoRun(grid, bg, part, CosmoRunConfig(cfl=0.4, t_end=0.5*t_H, diag_every=10**9), bg.t_i)
met = run.run()
run.cfg.t_end = 2*t_H; dt = 0.2*run.dt(met)
met0 = met; run._step(dt); met1 = solve_metric(grid, bg, part, run.t)
R = met0.R_edge; H0, H1 = bg.H(met0.t), bg.H(met1.t)
def terms(met):
    K = bg.K(met.t); beta = met.alpha_edge*R*met.Kth_edge if met is met0 else met.alpha_edge*met.R_edge*met.Kth_edge
    dbeta = np.gradient(beta, met.dR); KR = K - 2*met.Kth_edge
    return beta*met.dA_edge, met.A_edge*dbeta, -met.alpha_edge*met.A_edge*KR
t0 = terms(met0); t1 = terms(met1)
fd_x = (met1.A_edge - met0.A_edge)/dt
corr = 0.5*(H0*met0.R_edge*met0.dA_edge + H1*met1.R_edge*met1.dA_edge)
pred = [0.5*(a+b) for a, b in zip(t0, t1)]
print(f"t/t_H={met0.t/t_H:.3f}, dt={dt:.3e}, H={H0:.4e}")
print("  x/rm    A-1        fd_x       -HRA'     sum(fd)     bA'        Ab'       -aAK_R     pred      resid")
for xr in (0.05, 0.2, 0.5, 1.0, 1.5, 2.0, 3.0, 4.5):
    j = int(xr*rm/grid.dx)
    p = sum(pp[j] for pp in pred); f = fd_x[j] - corr[j]
    print(f"  {xr:4.2f} {met0.A_edge[j]-1:+.3e} {fd_x[j]:+.3e} {-corr[j]:+.3e} {f:+.3e} {pred[0][j]:+.3e} {pred[1][j]:+.3e} {pred[2][j]:+.3e} {p:+.3e} {f-p:+.3e}")