#!/usr/bin/env python3
"""Convergence of the horizon-approach case (nu=0.3, flat=3, L_rms=0.5) in time step, grid and particles.
Reports the times at which max 2m/r first exceeds 0.5, 8/9 and 0.95, and the ADM drift."""
import json, sys, time
import numpy as np
from pbhgr.spherical_ev import Grid, Run, RunConfig, solve_metric
from pbhgr.initial_data import shell_family

nu, flatness, L_rms = 0.30, 3.0, 0.5
M = 1.0; R = 2 * M / nu; t_dyn = np.sqrt(R**3 / M)
configs = [dict(K=300, n=20000, cfl=0.4), dict(K=300, n=20000, cfl=0.2), dict(K=300, n=20000, cfl=0.1),
           dict(K=300, n=80000, cfl=0.4), dict(K=600, n=40000, cfl=0.4), dict(K=1200, n=80000, cfl=0.4)]
def one(c):
    rng = np.random.default_rng(11)
    g = Grid(r_out=6 * R, K=c["K"])
    P = shell_family(M, R, c["n"], rng, flatness=flatness, L_rms=L_rms)
    for _ in range(4): P.N *= M / solve_metric(g, P).M
    run = Run(g, P, RunConfig(t_end=4.0 * t_dyn, cfl=c["cfl"], diag_every=25, horizon_threshold=0.995))
    t0 = time.time(); out = run.run()
    dg = run.diag; comp = np.array([d["max_2m_over_r"] for d in dg]); tt = np.array([d["t"] for d in dg]) / t_dyn
    Mt = np.array([d["M_total"] for d in dg])
    def tcross(th):
        i = np.argmax(comp > th)
        return float(np.interp(th, comp[i-1:i+1], tt[i-1:i+1])) if comp[i] > th and i > 0 else None
    row = dict(**c, outcome=out, max_comp=float(comp.max()), t_cross_0p5=tcross(0.5), t_cross_8_9=tcross(8/9), t_cross_0p95=tcross(0.95),
               adm_drift_max=float(np.max(np.abs(Mt - Mt[0]))), r_at_max=dg[int(np.argmax(comp))]["r_at_max"] / R, wall=time.time() - t0)
    print(json.dumps(row), flush=True)
    return row

if __name__ == "__main__":
    from multiprocessing import Pool
    with Pool(2) as pool:
        rows = pool.map(one, configs, chunksize=1)
    json.dump(rows, open("runs/verify_bh_convergence.json", "w"), indent=1)