#!/usr/bin/env python3
"""Measure wall time per RK4 step and peak RSS vs particle number / grid size, then extrapolate
the Stage-A workload to see where this VPS stops being sufficient."""
import json, os, resource, subprocess, sys, time
import numpy as np

def probe(n_part, K, steps=10):
    code = f"""
import numpy as np, time, resource, json
from pbhgr.spherical_ev import Grid, RunConfig, Run, solve_metric
from pbhgr.initial_data import shell_family
rng = np.random.default_rng(0)
g = Grid(r_out=30.0, K={K})
P = shell_family(1.0, 5.0, {n_part}, rng, flatness=2.0, L_rms=0.3)
P.N *= 1.0/solve_metric(g, P).M
run = Run(g, P, RunConfig(Gamma=0.2, t_end=1e9, cfl=0.4, diag_every=10**9, max_steps={steps}))
t0=time.time(); run.run(); dt=(time.time()-t0)/{steps}
print(json.dumps(dict(n_part={n_part}, K={K}, s_per_step=dt, peak_rss_MB=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024)))
"""
    out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True)
    if out.returncode != 0:
        err = out.stderr.strip().splitlines()[-1][:200] if out.stderr.strip() else f"killed (signal/OOM), returncode {out.returncode}"
        return dict(n_part=n_part, K=K, error=err)
    return json.loads(out.stdout.strip().splitlines()[-1])

rows = []
for n_part, K in [(2e4, 400), (1e5, 400), (1e5, 4000), (1e6, 4000), (3e6, 4000), (1e7, 10000), (2e7, 10000)]:
    r = probe(int(n_part), K, steps=5 if n_part >= 1e6 else 10)
    rows.append(r); print(r, flush=True)
json.dump(rows, open("runs/resource_probe.json", "w"), indent=1)