#!/usr/bin/env python3
"""Static benchmark: does the slow secular drift of max 2m/r shrink with grid resolution / time step?"""
import json, time
import numpy as np
from multiprocessing import Pool
from pbhgr.steady_state import build, sample
from pbhgr.spherical_ev import Grid, Run, RunConfig, solve_metric

def one(c):
    st = build(1.3); R, M = st["R"], st["M"]; t_dyn = np.sqrt(R**3 / M)
    g = Grid(r_out=4 * R, K=c["K"])
    P = sample(st, c["n"], np.random.default_rng(5))
    comp0 = 2 * solve_metric(g, P).m_edge[1:] / g.r_edge[1:]
    run = Run(g, P, RunConfig(t_end=3 * t_dyn, cfl=c["cfl"], diag_every=25))
    t0 = time.time(); run.run()
    cm = np.array([d["max_2m_over_r"] for d in run.diag]); tt = np.array([d["t"] for d in run.diag]) / t_dyn
    slope = np.polyfit(tt, cm, 1)[0] / comp0.max()          # relative drift per t_dyn
    comp1 = 2 * solve_metric(g, run.P).m_edge[1:] / g.r_edge[1:]
    row = dict(**c, max_comp_t0=float(comp0.max()), max_comp_end=float(cm[-1]), rel_drift_per_tdyn=float(slope),
               std_rel=float(cm.std() / comp0.max()), profile_rel_change=float(np.max(np.abs(comp1 - comp0)) / comp0.max()), wall=time.time() - t0)
    print(json.dumps(row), flush=True); return row

if __name__ == "__main__":
    configs = [dict(K=400, n=40000, cfl=0.4), dict(K=800, n=40000, cfl=0.4), dict(K=400, n=40000, cfl=0.2), dict(K=1600, n=80000, cfl=0.4)]
    with Pool(2) as pool:
        rows = pool.map(one, configs, chunksize=1)
    json.dump(rows, open("runs/verify_static_res.json", "w"), indent=1)