#!/usr/bin/env python3
"""R0ice Resource Team -- physics worker (pseudo-experiments for public ATLAS Open Data results).

What a job is: a frozen binned model (the published H->gamma-gamma fit to 36 fb^-1 of 2015+2016 ATLAS
Open Data) plus a seed range. The worker draws that many pseudo-experiments ("toys") from the model,
refits each one exactly like the real analysis (background = exp(cubic), signal = Gaussian, width
1.7 GeV, scanned 118-132 GeV) and returns a few numbers per toy. Thousands of toys measure things the
single real fit can't: whether the quoted significance formula holds, how much scanning the mass
inflates a bump (look-elsewhere effect), and how much the background shape choice can fake a signal.

Only public numbers go in or out: bin means from a published fit and the toy statistics.
Needs only numpy. Every toy is reproducible from (seed, toy index), so the server re-runs a sample of
each worker's toys and rejects results that don't match.

  pip install numpy
  python physics-worker.py --join "Your Name"         # first run: join, then work (name goes on the leaderboard)
  python physics-worker.py [--procs 2]                # later runs reuse the saved token
  python physics-worker.py --token-file "<Reader app folder>/config.json"   # Reader volunteers: same identity
  python physics-worker.py --replay 40                # re-run a published job here and compare every number
  python physics-worker.py --selftest OBSERVED_JSON   # fitter check on the real data, no network
  python physics-worker.py --kinds hhbbyy --once      # one claim of Higgs-pair (HH->bb gamma gamma) limit toys only

Since physics-1.1 the worker also takes "hhbbyy-cls-toys" jobs when no H->gamma-gamma job is waiting: CLs
pseudo-experiments for a planned upper limit on Higgs-pair production. Those jobs name a statistics engine by its
sha256; the worker downloads that exact file from r0ice.com/resource-team/physics/engines/, refuses it if the
hash differs, and runs it. The job carries only model numbers (expected yields) and seeds, no collision data.
"""
import json, math, os, sys, time, urllib.request, urllib.error

import numpy as np

WORKER_VERSION = "physics-1.1"
API_BASE = os.environ.get("RT_API_BASE", "https://r0ice.com/resource-team/api")
if not (API_BASE.startswith("https://") or API_BASE.startswith("http://127.0.0.1:")):
    raise SystemExit("RT_API_BASE must be https:// (or http://127.0.0.1:<port> for local testing)")
MAX_TOYS_PER_JOB = 2000

# ---------------------------------------------------------------- deterministic toy generation
def _poisson_tables(means):
    """Per-bin inverse-CDF tables. Built once per job; exact Poisson out to +-12 sigma."""
    tables = []
    for lam in means:
        lam = float(lam)
        sd = math.sqrt(max(lam, 1.0))
        lo = max(0, int(math.floor(lam - 12 * sd - 10)))
        hi = int(math.ceil(lam + 12 * sd + 10))
        k = np.arange(lo, hi + 1, dtype=np.float64)
        logp = k * math.log(max(lam, 1e-300)) - lam - np.array([math.lgamma(v + 1.0) for v in k])
        cdf = np.cumsum(np.exp(logp))
        tables.append((lo, cdf))
    return tables

def draw_toy(tables, seed, index):
    """Counts for toy `index` of job `seed`. PCG64 raw output is stable across numpy versions."""
    bg = np.random.PCG64(np.random.SeedSequence([int(seed), int(index)]))
    u = (bg.random_raw(len(tables)) >> np.uint64(11)).astype(np.float64) * (1.0 / 9007199254740992.0)
    out = np.empty(len(tables), dtype=np.float64)
    for i, (lo, cdf) in enumerate(tables):
        j = int(np.searchsorted(cdf, u[i], side="right"))
        out[i] = lo + min(j, len(cdf) - 1)
    return out

# ---------------------------------------------------------------- the fit (mirrors higgs-open analyze.py)
class Fitter:
    def __init__(self, model):
        lo, hi, n = model["bins"]["lo"], model["bins"]["hi"], model["bins"]["n"]
        w = (hi - lo) / n
        self.centers = lo + w * (np.arange(n) + 0.5)
        x = (self.centers - model["x_offset"]) / model["x_scale"]
        self.X = np.vstack([x ** 3, x ** 2, x, np.ones_like(x)]).T          # same order as np.polyval
        self.sigma = float(model["signal_sigma"])
        self.width = w
        self.grid = np.arange(model["scan"]["lo"], model["scan"]["hi"] + 1e-9, model["scan"]["step"])
        self.m_fixed = float(model["m_fixed"])
        self.shapes = {round(float(m), 4): self.shape(m) for m in list(self.grid) + [self.m_fixed]}

    def shape(self, mu):
        return self.width * np.exp(-0.5 * ((self.centers - mu) / self.sigma) ** 2) / (self.sigma * math.sqrt(2 * math.pi))

    @staticmethod
    def nll(pred, n):
        pred = np.maximum(pred, 1e-9)
        return float(np.sum(pred - n * np.log(pred)))

    def fit(self, n, g=None, c0=None, s0=0.0, iters=60):
        """Newton-Raphson on the Poisson NLL. g=None -> background only. Returns (nll, c, s, ok)."""
        X = self.X
        c = np.array(c0 if c0 is not None else np.polyfit(X[:, 2], np.log(np.maximum(n, 0.5)), 3), dtype=np.float64)
        s = float(s0) if g is not None else 0.0
        def pred_of(c, s):
            b = np.exp(np.clip(X @ c, -700, 700))
            return b, (b + s * g if g is not None else b)
        b, p = pred_of(c, s)
        cur = self.nll(p, n)
        for _ in range(iters):
            p = np.maximum(p, 1e-9)
            r = 1.0 - n / p
            q = n / (p * p)
            gc = X.T @ (r * b)
            Hcc = X.T @ (((r * b) + q * b * b)[:, None] * X)
            if g is None:
                grad, H = gc, Hcc
            else:
                gs = float(np.sum(r * g))
                Hcs = X.T @ (q * b * g)
                Hss = float(np.sum(q * g * g))
                grad = np.append(gc, gs)
                H = np.zeros((5, 5)); H[:4, :4] = Hcc; H[:4, 4] = Hcs; H[4, :4] = Hcs; H[4, 4] = Hss
            try:
                step = np.linalg.solve(H + 1e-9 * np.eye(len(grad)), grad)
            except np.linalg.LinAlgError:
                return cur, c, s, False
            t = 1.0
            while t > 1e-6:
                c_new = c - t * step[:4]
                s_new = s - t * step[4] if g is not None else 0.0
                b_new, p_new = pred_of(c_new, s_new)
                if np.all(p_new > 0):
                    v = self.nll(p_new, n)
                    if v <= cur + 1e-12:
                        break
                t *= 0.5
            if t <= 1e-6:
                return cur, c, s, True          # no descent possible: at the minimum to numerical precision
            dec = cur - v
            c, s, b, p, cur = c_new, s_new, b_new, p_new, v
            if dec < 1e-10 and np.max(np.abs(t * step)) < 1e-8:
                return cur, c, s, True
        return cur, c, s, True

    def analyse(self, n):
        """Everything one toy reports. q = 2*(NLL_b - NLL_sb), one-sided (0 if the fitted signal < 0)."""
        nb, cb, _, okb = self.fit(n)
        q_at = {}
        s_fixed = None
        ok = okb
        c_prev, s_prev = cb, 0.0
        qs = []
        for m in self.grid:
            v, c, s, okm = self.fit(n, self.shapes[round(float(m), 4)], c_prev, s_prev)
            ok = ok and okm
            q = max(0.0, 2.0 * (nb - v)) if s > 0 else 0.0
            qs.append(q)
            c_prev, s_prev = c, s
        v, c, s_fixed, okf = self.fit(n, self.shapes[round(self.m_fixed, 4)], cb, 0.0)
        ok = ok and okf
        q_fixed = max(0.0, 2.0 * (nb - v)) if s_fixed > 0 else 0.0
        qs = np.array(qs)
        imax = int(np.argmax(qs))
        # upcrossings of q(m) through a few low levels (Gross & Vitells 2010 look-elsewhere method)
        ups = {}
        for c0 in (0.5, 1.0, 2.0):
            above = qs > c0
            ups[str(c0)] = int(np.sum(above[1:] & ~above[:-1]))
        return {"q_fixed": q_fixed, "s_fixed": float(s_fixed), "q_max": float(qs[imax]),
                "m_max": float(self.grid[imax]), "up": ups, "ok": bool(ok)}

def run_job(model, job):
    """job = {seed, toy_start, n_toys, mode}. mode picks the generation means from the model."""
    n_toys = int(job["n_toys"])
    if not (1 <= n_toys <= MAX_TOYS_PER_JOB):
        raise ValueError("n_toys out of range")
    means = model["generate"][job["mode"]]
    tables = _poisson_tables(means)
    fitter = Fitter(model)
    rows = {"q_fixed": [], "s_fixed": [], "q_max": [], "m_max": [], "up05": [], "up1": [], "up2": []}
    bad = 0
    t0 = time.time()
    for i in range(int(job["toy_start"]), int(job["toy_start"]) + n_toys):
        r = fitter.analyse(draw_toy(tables, job["seed"], i))
        if not r["ok"]:
            bad += 1
        rows["q_fixed"].append(round(r["q_fixed"], 5)); rows["s_fixed"].append(round(r["s_fixed"], 3))
        rows["q_max"].append(round(r["q_max"], 5)); rows["m_max"].append(r["m_max"])
        rows["up05"].append(r["up"]["0.5"]); rows["up1"].append(r["up"]["1.0"]); rows["up2"].append(r["up"]["2.0"])
    return {"v": 1, "worker_version": WORKER_VERSION, "model_id": model["id"], "mode": job["mode"],
            "seed": int(job["seed"]), "toy_start": int(job["toy_start"]), "n_toys": n_toys,
            "fit_failures": bad, "seconds": round(time.time() - t0, 2), "toys": rows}

# ---------------------------------------------------------------- self-test on the real data
def selftest(path):
    obs = json.load(open(path))
    model = obs["model"]
    f = Fitter(model)
    n = np.array(obs["counts"], dtype=np.float64)
    t0 = time.time()
    r = f.analyse(n)
    z = math.sqrt(r["q_fixed"])
    print(json.dumps({"z_at_125": round(z, 3), "signal_events": round(r["s_fixed"], 1), "best_mass": r["m_max"],
                      "reference_z": obs.get("reference_z"), "reference_s": obs.get("reference_s"),
                      "seconds_per_fit": round(time.time() - t0, 3), "ok": r["ok"]}))
    tables = _poisson_tables(model["generate"]["bkg"])
    a, b = draw_toy(tables, 12345, 7), draw_toy(tables, 12345, 7)
    assert np.array_equal(a, b), "toy draw is not deterministic"
    t0 = time.time()
    res = run_job(model, {"seed": 1, "toy_start": 0, "n_toys": 20, "mode": "bkg"})
    print(json.dumps({"toys": 20, "seconds": round(time.time() - t0, 2), "fit_failures": res["fit_failures"],
                      "mean_q_fixed": round(float(np.mean(res["toys"]["q_fixed"])), 3)}))


# ---------------------------------------------------------------- hhbbyy-cls-toys (physics-1.1)
# A job names the statistics engine by sha256. The file is fetched once, hash-checked, cached under
# ~/.r0ice-physics-engines/<sha256>.py and imported from there; a file with any other hash is refused,
# so a worker can only ever run the exact engine the job was frozen with.
ENGINE_BASE = os.environ.get("RT_ENGINE_BASE", "https://r0ice.com/resource-team/physics/engines")
ENGINE_DIR = os.path.join(os.path.expanduser("~"), ".r0ice-physics-engines")
_ENGINES = {}

def load_pinned_engine(spec):
    import hashlib, importlib.util, re
    sha = str(spec.get("sha256", ""))
    if not re.fullmatch(r"[0-9a-f]{64}", sha):
        raise ValueError("job has no valid engine sha256")
    if sha in _ENGINES:
        return _ENGINES[sha]
    path = os.path.join(ENGINE_DIR, sha + ".py")
    raw = open(path, "rb").read() if os.path.exists(path) else b""
    if hashlib.sha256(raw).hexdigest() != sha:
        req = urllib.request.Request(f"{ENGINE_BASE}/{sha}.py", headers={"User-Agent": "r0ice-physics-worker/" + WORKER_VERSION})
        with urllib.request.urlopen(req, timeout=60) as r:
            raw = r.read()
        if hashlib.sha256(raw).hexdigest() != sha:
            raise ValueError(f"downloaded engine does not hash to {sha[:8]}; refusing to run it")
        os.makedirs(ENGINE_DIR, exist_ok=True)
        with open(path + ".tmp", "wb") as fh:
            fh.write(raw)
        os.replace(path + ".tmp", path)
    mspec = importlib.util.spec_from_file_location(f"hhbbyy_engine_{sha[:12]}", path)
    mod = importlib.util.module_from_spec(mspec); mspec.loader.exec_module(mod)
    if spec.get("version") and getattr(mod, "ENGINE_VERSION", None) != spec["version"]:
        raise ValueError("engine version does not match the job")
    _ENGINES[sha] = mod
    return mod

def hhbbyy_diagnostics(eng, M, mu, toys):
    """Per-job fit diagnostics, computed from the toy inputs (the verifier recomputes the same numbers).
    The pinned engines expose no fit status, so this wrapper counts a toy as a FIT FAILURE when q is not finite,
    a conditional fit (at mu and at 0) returns a non-finite NLL or b, or (engine 0.2) a golden-section search ends on
    its +-10 bracket edge. zero_q counts toys with q~_mu exactly 0 (mu_hat >= mu or clipped), reported separately."""
    q = np.asarray(toys["q"], dtype=np.float64)
    fail = ~np.isfinite(q)
    method = "non-finite q"
    ss = bool(getattr(M, "ss", False))                      # engine 0.3 spurious-signal path (sigma_ss > 0)
    cf = getattr(M, "conditional_ss" if ss else "conditional_full", None)
    if cf is not None and all(k in toys for k in ("n", "m", "a")):
        n, m, a = (np.asarray(toys[k], dtype=np.float64) for k in ("n", "m", "a"))
        c = np.asarray(toys["c"], dtype=np.float64) if "c" in toys else np.zeros_like(n)
        d = np.asarray(toys["d"], dtype=np.float64) if "d" in toys else np.zeros_like(n)
        # +-10 search ranges are those of engines 0.2 and 0.3 (golden section over alpha/beta; 0.3's theta Newton is bracketed
        # to [d-10, d+10] and returns the range end when the root lies outside). Other engines: finiteness checks only.
        bounds = getattr(eng, "ENGINE_VERSION", "") in ("hhbbyy-stats-0.2", "hhbbyy-stats-0.3")
        with np.errstate(all="ignore"):
            for mu_fit in (float(mu), 0.0):
                out = cf(mu_fit, n, m, a, c, d) if ss else cf(mu_fit, n, m, a, c)
                val, b, al, be = (np.broadcast_to(np.asarray(x, dtype=np.float64), n.shape) for x in out[:4])
                fail |= ~np.isfinite(val) | ~np.isfinite(b) | ~np.isfinite(al) | ~np.isfinite(be)
                if ss:
                    th = np.broadcast_to(np.asarray(out[4], dtype=np.float64), n.shape)
                    fail |= ~np.isfinite(th)
                    if bounds:
                        fail |= np.abs(np.abs(th - d) - 10.0) < 1e-3
                if bounds and getattr(M, "h_nuis", False):
                    fail |= np.abs(np.abs(be - c) - 10.0) < 1e-3
                if bounds and getattr(M, "sigma_s", 0.0) != 0.0 and mu_fit != 0.0:
                    a_lo = np.maximum(a - 10.0, -1.0 / M.sigma_s + 1e-6)
                    fail |= (np.abs(al - a_lo) < 1e-3) | (np.abs(al - (a + 10.0)) < 1e-3)
        method = ("non-finite q/NLL/b/nuisances at mu and 0" + (" + search-range edge hits (+-10 around a, c" + (", d" if ss else "") + ")" if bounds else "")
                  + "; Newton non-convergence is not exposed by the engine (v0.4 item)")
    idx = [int(i) for i in np.flatnonzero(fail)]
    return {"fit_failures": len(idx), "fit_failure_index": idx[:100], "zero_q": int(np.sum(np.isfinite(q) & (q == 0.0))),
            "nonfinite_q": int(np.sum(~np.isfinite(q))), "method": method}

def hhbbyy_summary(q):
    q = np.asarray(q, dtype=np.float64)
    return {"q_mean": float(q.mean()), "q_zero_fraction": float(np.mean(q == 0.0)),
            "q_quantiles": {str(p): float(np.quantile(q, p)) for p in (0.025, 0.16, 0.5, 0.84, 0.975)}}

def run_hhbbyy_job(item):
    """item = {id, job:{mode, mu_index, mu, gen, seed, toy_start, n_toys}, model:{id, model}, engine:{sha256, version}}"""
    job, eng_spec = item["job"], item["engine"]
    n_toys = int(job["n_toys"])
    if not (1 <= n_toys <= MAX_TOYS_PER_JOB):
        raise ValueError("n_toys out of range")
    eng = load_pinned_engine(eng_spec)
    model = eng.OnOffModel.from_json(item["model"]["model"])
    t0 = time.time()
    out = eng.run_batch(model, float(job["mu"]), dict(job["gen"]), int(job["seed"]), int(job["toy_start"]), n_toys)
    toys = {}
    for k, v in out.items():
        arr = np.asarray(v, dtype=np.float64)
        toys[k] = [int(x) for x in arr] if k in ("n", "m") else [float(x) for x in arr]
    seconds = round(time.time() - t0, 3)
    diag = hhbbyy_diagnostics(eng, model, float(job["mu"]), toys)
    return {"v": 1, "kind": "hhbbyy-cls-toys", "worker_version": WORKER_VERSION, "model_id": item["model"]["id"],
            "model_sha256": item["model"].get("sha256"), "diag": diag,
            "engine_sha256": eng_spec["sha256"], "engine_version": eng.ENGINE_VERSION, "mode": job["mode"],
            "mu_index": int(job["mu_index"]), "mu": float(job["mu"]), "seed": str(job["seed"]),
            "toy_start": int(job["toy_start"]), "n_toys": n_toys, "seconds": seconds,
            "summary": hhbbyy_summary(toys["q"]), "toys": toys}

def hhbbyy_once(token):
    """One claim/run/submit cycle of hhbbyy-cls-toys jobs. Returns seconds to wait, or None if nothing was claimed."""
    got = _http("GET", "/jobs/hhbbyy-claim?n=2", token)
    items = got.get("items") or []
    if not items:
        return None
    out = []
    for it in items:
        try:
            res = run_hhbbyy_job(it)
            out.append({"id": it["id"], "result": res})
            print(f"[physics] hhbbyy job {it['id']} {res['mode']} mu={res['mu']} {res['n_toys']} toys in {res['seconds']}s", flush=True)
        except Exception as e:
            out.append({"id": it["id"], "error": str(e)[:200]})
            print(f"[physics] hhbbyy job {it['id']} failed: {e}", flush=True)
    for attempt in range(4):
        try:
            r = _http("POST", "/jobs/hhbbyy-result", token, out)
            print(f"[physics] hhbbyy submitted {r.get('submitted')} rejected {r.get('rejected')}", flush=True)
            break
        except (urllib.error.URLError, OSError) as e:
            print("[physics] hhbbyy submit failed, retrying:", e, flush=True)
            time.sleep(15 * (attempt + 1))
    return 1

KINDS = ("hyy", "hhbbyy")

# ---------------------------------------------------------------- work loop
def _http(method, path, token, body=None, timeout=60):
    req = urllib.request.Request(API_BASE + path, method=method,
                                 data=None if body is None else json.dumps(body).encode(),
                                 headers={"Authorization": "Bearer " + token, "Content-Type": "application/json",
                                          "User-Agent": "r0ice-physics-worker/" + WORKER_VERSION})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode())

def work_once(token, kinds=KINDS):
    """Claim one job, run it, submit it. Returns seconds to wait before the next claim.
    H->gamma-gamma jobs first; when none is waiting, hhbbyy-cls-toys jobs (physics-1.1)."""
    items, got = [], {}
    if "hyy" in kinds:
        try:
            got = _http("GET", "/jobs/physics-claim?n=1", token)
        except (urllib.error.URLError, OSError, ValueError) as e:
            print("[physics] claim failed:", e, flush=True)
            return 120
        items = got.get("items") or []
    if not items:
        if "hhbbyy" in kinds:
            try:
                w = hhbbyy_once(token)
                if w is not None:
                    return w
            except (urllib.error.URLError, OSError, ValueError, KeyError) as e:
                print("[physics] hhbbyy claim failed:", e, flush=True)
        return int(got.get("retry_after", 300))
    out = []
    for it in items:
        try:
            res = run_job(it["model"], it["job"])
            out.append({"id": it["id"], "result": res})
            print(f"[physics] job {it['id']} {it['job']['mode']} {res['n_toys']} toys in {res['seconds']}s", flush=True)
        except Exception as e:  # report, never crash the host worker
            out.append({"id": it["id"], "error": str(e)[:200]})
    for attempt in range(4):
        try:
            _http("POST", "/jobs/physics-result", token, out)
            break
        except (urllib.error.URLError, OSError) as e:
            print("[physics] submit failed, retrying:", e, flush=True)
            time.sleep(15 * (attempt + 1))
    return 1

def work_loop(token, minutes=None, stop=None, kinds=KINDS):
    end = time.time() + minutes * 60 if minutes else None
    while (end is None or time.time() < end) and not (stop and stop.is_set()):
        wait = work_once(token, kinds)
        if stop:
            stop.wait(wait)
        else:
            time.sleep(wait)

PUBLIC_BASE = "https://r0ice.com/resource-team/physics"

def replay(job_id, n_check):
    """Re-run a published job on THIS computer and compare every number with what the fleet submitted."""
    def get(path):
        req = urllib.request.Request(PUBLIC_BASE + path, headers={"User-Agent": "r0ice-physics-worker/" + WORKER_VERSION})
        with urllib.request.urlopen(req, timeout=60) as r:
            return json.loads(r.read().decode())
    res = get(f"/results/{int(job_id)}.json")
    model = get(f"/models/{res['model_id']}.json")
    n = res["n_toys"] if n_check <= 0 else min(n_check, res["n_toys"])
    print(f"job {job_id}: {res['mode']} toys {res['toy_start']}..{res['toy_start'] + res['n_toys'] - 1}, seed {res['seed']}, "
          f"run by {res.get('machine', '?')} in {res.get('seconds')}s; re-running {n} toy(s) here")
    tables = _poisson_tables(model["generate"][res["mode"]])
    fitter = Fitter(model)
    bad = 0
    for k in range(n):
        r = fitter.analyse(draw_toy(tables, res["seed"], res["toy_start"] + k))
        t = res["toys"]
        same = (abs(r["q_fixed"] - t["q_fixed"][k]) < 2e-3 and abs(r["q_max"] - t["q_max"][k]) < 2e-3
                and abs(r["s_fixed"] - t["s_fixed"][k]) < 2e-2 and r["m_max"] == t["m_max"][k])
        bad += 0 if same else 1
        print(f"  toy {res['toy_start'] + k:>7}: q125 {r['q_fixed']:9.5f} vs {t['q_fixed'][k]:9.5f}   "
              f"qmax {r['q_max']:9.5f} vs {t['q_max'][k]:9.5f} @ {r['m_max']:.2f} GeV   {'MATCH' if same else 'DIFFERENT'}")
    print("ALL MATCH" if bad == 0 else f"{bad} toy(s) differ")
    return bad

def main(argv):
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--selftest", metavar="OBSERVED_JSON")
    ap.add_argument("--token", default=os.environ.get("RT_TOKEN"))
    ap.add_argument("--token-file", help="read the token from this file (keeps it out of the process list)")
    ap.add_argument("--join", metavar="NAME", help="first run only: join the Resource Team under this name (shown on the leaderboard)")
    ap.add_argument("--procs", type=int, default=1)
    ap.add_argument("--minutes", type=float, default=None)
    ap.add_argument("--replay", metavar="JOB_ID", help="re-run a published job on this computer and compare with the fleet's numbers")
    ap.add_argument("--check", type=int, default=5, help="with --replay: how many toys to re-run (0 = all)")
    ap.add_argument("--kinds", default=",".join(KINDS), help="job kinds to take: hyy (H->gamma-gamma), hhbbyy (Higgs-pair limit toys)")
    ap.add_argument("--once", action="store_true", help="one claim/run/submit cycle, then exit")
    a = ap.parse_args(argv)
    a.kinds = tuple(k.strip() for k in a.kinds.split(",") if k.strip() in KINDS) or KINDS
    if a.replay:
        sys.exit(1 if replay(a.replay, a.check) else 0)
    if a.selftest:
        return selftest(a.selftest)
    if a.token_file:
        with open(a.token_file) as fh:
            raw = fh.read().strip()
        # Also accepts the Resource Team Reader app's config.json, so a Reader volunteer earns physics
        # credit under the identity they already have instead of joining again under a new name.
        a.token = json.loads(raw).get("token", "") if raw.startswith("{") else raw
    saved = os.path.join(os.path.expanduser("~"), ".r0ice-physics-token")
    if not a.token and a.join:
        req = urllib.request.Request(API_BASE + "/jobs/join", method="POST", data=json.dumps({"handle": a.join[:40]}).encode(),
                                     headers={"Content-Type": "application/json", "User-Agent": "r0ice-physics-worker/" + WORKER_VERSION})
        with urllib.request.urlopen(req, timeout=60) as r:
            got = json.loads(r.read().decode())
        a.token = got["token"]
        if got.get("handle") and got["handle"] != a.join[:40]:
            print(f"that name was taken, so you joined as '{got['handle']}'", flush=True)
            a.join = got["handle"]
        with open(saved, "w") as fh:
            fh.write(a.token)
        try:
            os.chmod(saved, 0o600)
        except OSError:
            pass
        print(f"joined the R0ice Resource Team as '{a.join[:40]}' -- token saved to {saved}; next time just run with no arguments", flush=True)
    if not a.token and os.path.exists(saved):
        with open(saved) as fh:
            a.token = fh.read().strip()
    if not a.token:
        ap.error('first run: --join "Your Name" (or --token / --token-file / RT_TOKEN)')
    if a.once:
        return work_once(a.token, a.kinds)
    if a.procs <= 1:
        return work_loop(a.token, a.minutes, kinds=a.kinds)
    import multiprocessing as mp
    ps = [mp.Process(target=work_loop, args=(a.token, a.minutes, None, a.kinds)) for _ in range(a.procs)]
    for p in ps: p.start()
    for p in ps: p.join()

if __name__ == "__main__":
    main(sys.argv[1:])