#!/usr/bin/env python3
"""Run representative cases with full snapshot storage for visualisation -> runs/viz/*.pkl"""
import pickle, sys, time
from multiprocessing import Pool
from pathlib import Path
import numpy as np
from pbhgr.spherical_ev import Grid, Run, RunConfig, solve_metric
from pbhgr.initial_data import shell_family

CASES = {
    "bounce_warm":      dict(nu=0.15, flatness=1.0, L_rms=0.5, Gamma_tdyn=0.0),
    "bh_approach_flat": dict(nu=0.30, flatness=3.0, L_rms=0.5, Gamma_tdyn=0.0),
    "decay_warm":       dict(nu=0.15, flatness=1.0, L_rms=0.5, Gamma_tdyn=0.3),
    "cold_caustic":     dict(nu=0.15, flatness=1.0, L_rms=0.0, Gamma_tdyn=0.0),
}

def run_case(name, n_part=20000, K=300, seed=11):
    c = CASES[name]
    M = 1.0; R = 2 * M / c["nu"]; t_dyn = np.sqrt(R**3 / M)
    rng = np.random.default_rng(seed)
    g = Grid(r_out=6.0 * R, K=K)
    P = shell_family(M, R, n_part, rng, flatness=c["flatness"], L_rms=c["L_rms"])
    for _ in range(4):
        P.N *= M / solve_metric(g, P).M
    cfg = RunConfig(Gamma=c["Gamma_tdyn"] / t_dyn, t_end=5.0 * t_dyn, cfl=0.4, diag_every=25, seed=seed,
                    emit_fraction=0.25, store_profiles=True, store_particles=3000)
    run = Run(g, P, cfg)
    t0 = time.time(); outcome = run.run()
    out = dict(name=name, params=c, M=M, R=R, t_dyn=t_dyn, r_edge=g.r_edge, r_c=g.r_c, dr=g.dr,
               outcome=outcome, diag=run.diag, profiles=run.profiles, wall=time.time() - t0)
    Path("runs/viz").mkdir(parents=True, exist_ok=True)
    with open(f"runs/viz/{name}.pkl", "wb") as f:
        pickle.dump(out, f)
    return name, outcome, out["wall"], len(run.profiles)

if __name__ == "__main__":
    names = sys.argv[1:] or list(CASES)
    with Pool(2) as pool:
        for r in pool.imap_unordered(run_case, names):
            print(r, flush=True)