Jungfraujoch/tools/battery/battery.py

#!/usr/bin/env python3
"""Rugnux validation battery: run, score, report, compare.

Arms and where their truth comes from:
  open     public PDB depositions (and a few published small-molecule sets); scored against the
           deposited space group, cell and resolution. Manifest: open.json (committed).
  inhouse  standard test crystals measured at PSI (lysozyme, thaumatin, insulin, cytochrome C,
           myoglobin) plus no-crystal controls; scored against XDS (point group only - XDS never
           tests a screw axis). Manifest: inhouse.json (committed).
  private  user data, scored like inhouse. Its manifest lives OUTSIDE the repository (path in the
           site config) and a private run gets its own run directory and report.

Usage (README.md in this directory is the full guide):
  battery.py run --rugnux BIN [--label L] [--arm open,inhouse|private] [--only a,b | --tier smoke|ci]
                 [--build-dir DIR] [--gpulock] [--threads N]
                 [--model-check] [--nullmodel | --no-nullmodel] [--baseline RUN|none]
                                                  --tier ci also runs the wrong-model null (nullmodel.py)
  battery.py nullmodel --rugnux BIN [--label L] [--build-dir DIR] [--gpulock] [--threads N]
                                                  the wrong-model null on its own, in a run directory of
                                                  its own: exit 1 if a pair failed (CI runs it as its own
                                                  step, after run --tier ci --no-nullmodel)
  battery.py timing --rugnux BIN [--benchmark standard|inhouse] [--label L] [--build-dir DIR] [--gpulock]
                    [--threads N] [--baseline RUN|none]
                                                  the timing benchmark (timing.json). standard (the default,
                                                  open data only): 9zlo, 9s02 and 9gjx from their AlphaFold DB
                                                  models, 4nwv from its deposited model moved away; inhouse:
                                                  myob_x10sa and cytc_x10sa from AlphaFold DB models, and 4nwv.
                                                  All placed by molecular replacement; checks each answer and
                                                  prints wall and step times; exit 1 if one is wrong
  battery.py keep RUN --as workstation|server     copy a finished open/inhouse run's report.md, results.json
                                                  and manifest.json into results/<role>/<run>/ (committed);
                                                  refuses a private run
  battery.py stress --rugnux BIN (--only a,b | --tier T) [--repeat N] [--concurrent K]
                 [--sanitizer memcheck|racecheck|synccheck|initcheck] [--launch-blocking]
                                                  crash hunt: exit 1 if any run crashed
  battery.py list                                 runs under the runs root, with their state
  battery.py compare RUN_A RUN_B [--all] [--rerun-changed]
  battery.py report RUN --out DIR [--baseline RUN] re-render a run's report into DIR/<run>.{md,html}
  battery.py table RUN                            the one-line-per-set table a run ends with
  battery.py abort RUN [--reason TEXT]            mark an unfinished run as aborted
  battery.py recheck RUN [--only a,b] [--jobs N]  rerun the model checks of a finished run's open
                                                  arm on its p.mtz files; recorded in the manifest
  battery.py discover [--refresh]                 open arm: propose inputs for new directories
  battery.py refs --arm inhouse|private [--write] re-read the XDS references from CORRECT.LP
  battery.py remap --arm A MAP.json [--write]     follow a rename of the data directories

The site config (data roots, runs root, lock, gpulock, private manifest, baselines) is --site, else
$JFJOCH_BATTERY_SITE, else site.json beside this script (not committed; site.example.json is).

Every set runs once, with one command per arm (COMMAND_DOC): the open arm with the deposited model
(--model), the XDS arms with XDS's range as a report-only second
table (--report-resolution) and, where a set names the PDB entry of its crystal form, that model
(--model). Neither option changes the processing, so the one run is both what a user gets and the
like-for-like comparison.

A run writes <runs_root>/<YYYYMMDD-HHMM>_<rugnux git>_<label>/ holding a copy of the binary,
manifest.json (what was run, on what, with which references), results.json (one row per set),
report.md + report.html and work/<arm>/<set>/ (log, report, cif). It is made read-only when done.
A full-arm run (no --only / --tier) takes the site lock and refuses to start if another holds it.
A run exits with status 1 when any set crashed (score.crash_note: killed by a signal, "Processing
failed", a device fault even where rugnux went on) or a set of the ci tier it ran failed (lattice,
symmetry, no input). `stress` runs sets repeatedly, side by side or under compute-sanitizer and
only asks whether any of them crashed.
A run that did not finish is never compared or reported on unless asked (--allow-incomplete).
"""
import argparse
import concurrent.futures
import datetime
import hashlib
import json
import os
import re
import shutil
import signal
import socket
import stat
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import inputs  # noqa: E402
import report  # noqa: E402
import score   # noqa: E402

import model_check  # noqa: E402  the deposited model (--model); REFMAC check (--model-check)
import nullmodel  # noqa: E402  a wrong --model changes nothing and is recognised as wrong
import shelx_check  # noqa: E402  SHELXL refinement of a small molecule's COD structure
import timing  # noqa: E402  the timing benchmarks: sets placed by molecular replacement
import depdata_check  # noqa: E402  our merge vs the depositor's data, shell by shell

PDB_ID = re.compile(r"^[0-9][A-Za-z0-9]{3}$")

# The one command each set runs, by its arm's reference. --model and --report-resolution only add
# to the report (the merge is the same with and without them), so this one run is what a user
# gets and carries the like-for-like numbers too.
COMMAND_DOC = {
    "deposition": "`rugnux --model <deposited coordinates> <input>` (without --model where there is no "
                  "deposited model: small molecules, unpublished sets)",
    "xds": "`rugnux --report-resolution <XDS d_min>,<XDS d_max> [--model <reference model>] <input>`: "
           "the second statistics table at XDS's range is report-only; where XDS kept Friedel mates "
           "apart the scorer reads the run's own per-hand table, which needs no flag. --model where "
           "the set names a PDB entry of its crystal form (the standard proteins)",
}

LOCAL_SITE = os.path.join(HERE, "site.json")
# what a site config may leave out for the two committed arms: then it needs only their data roots
PUBLIC_ARM_DEFAULTS = {"open": {"manifest": "open.json", "reference": "deposition"},
                       "inhouse": {"manifest": "inhouse.json", "reference": "xds"}}


# ------------------------------------------------------------------ configuration and manifests

def load_site(path):
    path = path or os.environ.get("JFJOCH_BATTERY_SITE") or LOCAL_SITE
    if not os.path.exists(path):
        sys.exit(f"no site config at {path}: copy {os.path.join(HERE, 'site.example.json')} to "
                 f"{LOCAL_SITE} (or anywhere, and point JFJOCH_BATTERY_SITE or --site at it) and "
                 "fill in this machine's paths")
    site = json.load(open(path))
    # a relative path is relative to the site config itself, so one shipped beside its data works
    # wherever it is unpacked and whatever directory rugnux runs in; an absolute path stays as it is
    base = os.path.dirname(os.path.abspath(path))
    for k in ("runs_root", "lock_file", "pdb_cache", "cod_cache"):
        if site.get(k):
            site[k] = os.path.join(base, site[k])
    for name, arm in site["arms"].items():
        for k, v in PUBLIC_ARM_DEFAULTS.get(name, {}).items():
            arm.setdefault(k, v)
        arm["manifest"] = os.path.join(HERE, arm["manifest"])   # relative to tools/battery/
        if arm.get("data_root"):
            arm["data_root"] = os.path.join(base, arm["data_root"])
    return site


def load_sets(site, arm):
    cfg = site["arms"][arm]
    man = json.load(open(cfg["manifest"]))
    for e in man["sets"]:
        e["arm"] = arm
    return man["sets"]


def select(site, arms, only, tier):
    sets = [e for arm in arms for e in load_sets(site, arm)]
    if only:
        want = [s.strip() for s in only.split(",") if s.strip()]
        unknown = [w for w in want if w not in {e["id"] for e in sets}]
        if unknown:
            sys.exit(f"not in the {'/'.join(arms)} manifest(s): {', '.join(unknown)}")
        sets = [e for e in sets if e["id"] in want]
    if tier:
        sets = [e for e in sets if tier in e.get("tiers", {})]
    return sets


# ------------------------------------------------------------------ the binary

def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for block in iter(lambda: fh.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def binary_info(path, build_dir):
    """Version and git hash as the binary itself prints them, plus the build flags from the
    CMakeCache of its build tree (given, or found above the binary)."""
    out = subprocess.run([path], capture_output=True, text=True, timeout=120).stdout
    m = re.search(r"Version (\S+) \(git (\w+)", out)
    info = {"version": m.group(1) if m else None, "git": m.group(2) if m else "unknown"}
    cache = None
    d = build_dir or os.path.dirname(os.path.realpath(path))
    for _ in range(4):
        if os.path.exists(os.path.join(d, "CMakeCache.txt")):
            cache = os.path.join(d, "CMakeCache.txt")
            break
        d = os.path.dirname(d)
    if cache:
        vals = {}
        for line in open(cache, errors="replace"):
            m = re.match(r"^(CMAKE_BUILD_TYPE|CMAKE_CXX_FLAGS|CMAKE_HOME_DIRECTORY|JFJOCH_USE_CUDA):\w+=(.*)$",
                         line.strip())
            if m:
                vals[m.group(1)] = m.group(2)
        info["cmake_cache"] = cache
        info["flags"] = (f"{vals.get('CMAKE_BUILD_TYPE', '?')} CXX_FLAGS='{vals.get('CMAKE_CXX_FLAGS', '')}' "
                         f"CUDA={vals.get('JFJOCH_USE_CUDA', '?')}")
        src = vals.get("CMAKE_HOME_DIRECTORY")
        if src:
            info["source"] = src
            info["source_head"] = git(src, "rev-parse", "HEAD")
            info["source_dirty"] = bool(git(src, "status", "--porcelain", "--untracked-files=no"))
    return info


def git(d, *args):
    r = subprocess.run(["git", "-C", d] + list(args), capture_output=True, text=True)
    return r.stdout.strip() if r.returncode == 0 else None


def gpu_others():
    """Number of processes other than ours on the GPU right now (None if nvidia-smi is absent)."""
    try:
        out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid", "--format=csv,noheader"],
                             capture_output=True, text=True, timeout=30).stdout
    except (OSError, subprocess.TimeoutExpired):
        return None
    return len([x for x in out.split() if x.strip().isdigit()])


def gpu_hardware():
    """[{"name", "memory"}] from nvidia-smi, [] if there are none or it is absent."""
    try:
        out = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
                             capture_output=True, text=True, timeout=30).stdout
    except (OSError, subprocess.TimeoutExpired):
        return []
    gpus = []
    for line in out.splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) == 2:
            gpus.append({"name": parts[0], "memory": parts[1]})
    return gpus


def linux_hardware():
    blocks = []
    try:
        blocks = [b for b in open("/proc/cpuinfo").read().split("\n\n") if b.strip()]
    except OSError:
        pass
    model = None
    sockets, cores = set(), set()
    for b in blocks:
        m = re.search(r"model name\s*:\s*(.+)", b)
        if m and model is None:
            model = m.group(1).strip()
        p, c = re.search(r"physical id\s*:\s*(\d+)", b), re.search(r"core id\s*:\s*(\d+)", b)
        if p and c:
            sockets.add(p.group(1))
            cores.add((p.group(1), c.group(1)))

    # Nominal (base) clock, never the boost clock: the model-name string, else cpufreq's own
    # base_frequency (intel_pstate, kHz), else ACPI CPPC's nominal_freq (MHz; what AMD's
    # acpi-cpufreq exposes instead - cpuinfo_max_freq on those boxes is the boost clock, not this).
    freq_ghz, freq_source = None, "unknown"
    m = re.search(r"@\s*([\d.]+)\s*GHz", model or "")
    if m:
        freq_ghz, freq_source = float(m.group(1)), "model name"
    else:
        try:
            freq_ghz = int(open("/sys/devices/system/cpu/cpu0/cpufreq/base_frequency").read()) / 1e6
            freq_source = "base_frequency"
        except (OSError, ValueError):
            try:
                freq_ghz = int(open("/sys/devices/system/cpu/cpu0/acpi_cppc/nominal_freq").read()) / 1e3
                freq_source = "acpi_cppc"
            except (OSError, ValueError):
                pass

    # Reported separately, and never as "nominal": cpuinfo_max_freq is the highest boost state.
    max_boost_ghz = None
    try:
        max_boost_ghz = int(open("/sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq").read()) / 1e6
    except (OSError, ValueError):
        pass

    mem_gib = None
    try:
        for line in open("/proc/meminfo"):
            if line.startswith("MemTotal:"):
                mem_gib = int(line.split()[1]) / 1024 / 1024   # kiB -> GiB
                break
    except OSError:
        pass

    return {"cpu_model": model or "unknown", "cpu_base_ghz": freq_ghz, "cpu_base_source": freq_source,
            "cpu_max_boost_ghz": max_boost_ghz, "sockets": len(sockets) or None,
            "physical_cores": len(cores) or None, "threads": os.cpu_count(), "memory_gib": mem_gib,
            "gpus": gpu_hardware()}


def windows_hardware():
    """From WMI through PowerShell (wmic is gone from current Windows). Win32_Processor has one
    entry per socket; its MaxClockSpeed is the nominal clock, not the boost."""
    def cim(query):
        r = subprocess.run(["powershell", "-NoProfile", "-Command", query + " | ConvertTo-Json"],
                           capture_output=True, text=True, errors="replace", timeout=60)
        out = json.loads(r.stdout) if r.returncode == 0 and r.stdout.strip() else []
        return out if isinstance(out, list) else [out]

    cpus = cim("Get-CimInstance Win32_Processor | "
               "Select-Object Name, NumberOfCores, NumberOfLogicalProcessors, MaxClockSpeed")
    mem = cim("(Get-CimInstance Win32_ComputerSystem).TotalPhysicalMemory")
    model = cpus[0]["Name"].strip() if cpus else None
    freq_ghz, freq_source = None, "unknown"
    m = re.search(r"@\s*([\d.]+)\s*GHz", model or "")
    if m:
        freq_ghz, freq_source = float(m.group(1)), "model name"
    elif cpus and cpus[0].get("MaxClockSpeed"):
        freq_ghz, freq_source = cpus[0]["MaxClockSpeed"] / 1000, "Win32_Processor.MaxClockSpeed"
    return {"cpu_model": model or "unknown", "cpu_base_ghz": freq_ghz, "cpu_base_source": freq_source,
            "cpu_max_boost_ghz": None, "sockets": len(cpus) or None,
            "physical_cores": sum(c.get("NumberOfCores") or 0 for c in cpus) or None,
            "threads": sum(c.get("NumberOfLogicalProcessors") or 0 for c in cpus) or os.cpu_count(),
            "memory_gib": int(mem[0]) / 1024 ** 3 if mem else None, "gpus": gpu_hardware()}


def macos_hardware():
    def sysctl(key):
        r = subprocess.run(["sysctl", "-n", key], capture_output=True, text=True)
        return r.stdout.strip() if r.returncode == 0 else None

    model = sysctl("machdep.cpu.brand_string")
    freq_ghz, freq_source = None, "unknown"
    m = re.search(r"@\s*([\d.]+)\s*GHz", model or "")
    if m:
        freq_ghz, freq_source = float(m.group(1)), "model name"
    cores = sysctl("hw.physicalcpu")
    # Apple silicon has two or three core tiers (an M6: Super, Performance, Efficiency), each a perflevel
    # with its own name, fastest first: written as e.g. 2S+4P+6E.
    tiers = []
    for level in range(int(sysctl("hw.nperflevels") or 0)):
        n, name = sysctl(f"hw.perflevel{level}.physicalcpu"), sysctl(f"hw.perflevel{level}.name")
        if n and name:
            tiers.append(f"{n}{name[0]}")
    if len(tiers) > 1:
        cores = "+".join(tiers)
    mem = sysctl("hw.memsize")
    return {"cpu_model": model or "unknown", "cpu_base_ghz": freq_ghz, "cpu_base_source": freq_source,
            "cpu_max_boost_ghz": None, "sockets": int(sysctl("hw.packages") or 1), "physical_cores": cores,
            "threads": int(sysctl("hw.logicalcpu") or os.cpu_count()),
            "memory_gib": int(mem) / 1024 ** 3 if mem else None, "gpus": []}


def hardware_info():
    """The run host's hardware, for the manifest: CPU model, nominal/base frequency (and where it
    came from) plus max boost frequency separately, physical cores, threads, sockets, memory, GPUs.
    Linux is detailed; macOS and Windows are a short best-effort; anywhere else this is just empty,
    and the report says so."""
    try:
        if sys.platform == "darwin":
            return macos_hardware()
        if sys.platform.startswith("linux"):
            return linux_hardware()
        if sys.platform == "win32":
            return windows_hardware()
    except (OSError, ValueError, subprocess.TimeoutExpired):
        pass
    return {}


# ------------------------------------------------------------------ running

def command(binary, e, opts, model=None):
    cmd = [binary, "-o", "p", "--no-export-unmerged"]
    if opts["threads"]:
        cmd += ["-N", str(opts["threads"])]
    ref = e.get("ref") or {}
    if e["reference"] == "xds":
        # No -A, whatever the reference did. A rotation run decides with Friedel's mates merged
        # either way, and reports the merge counted BOTH ways from that one run; -A only says which
        # of the two tables was asked for. Where XDS merged with FRIEDEL'S_LAW=FALSE the scorer
        # reads rugnux's per-hand block (ANOM_REFRES_*) instead, which is the like-for-like table.
        # a second statistics table over XDS's own range, from the same merge (report only). XDS's
        # own limit, not the derived reference d_min: the manifest's XDS totals are over XDS's range,
        # so only that makes the table like for like. Shells past rugnux's own cut print as not
        # merged.
        dmin = ref.get("dmin_xds") or ref.get("dmin")
        if dmin:
            cmd += ["--report-resolution",
                    f"{dmin:.3f}" + (f",{ref['dmax']:.3f}" if ref.get("dmax") else "")]
    if model:
        cmd += ["--model", model]
    cmd += opts["extra_args"]
    cmd.append(e["input_path"])
    if opts["gpulock"] and "--no-gpu" not in cmd:     # a CPU run has no card to share
        cmd = [opts["gpulock"]] + cmd
    return cmd


def run_logged(cmd, wd, timeout, env=None):
    """Run cmd in wd with its output in wd/run.log; its exit code, or None on a timeout. It runs in a
    process group of its own, and a timeout or an interrupt kills the whole group: killing only the
    direct child - the GPU lock wrapper, or compute-sanitizer - leaves rugnux running on the card."""
    with open(os.path.join(wd, "run.log"), "w") as log:
        proc = subprocess.Popen(cmd, cwd=wd, stdout=log, stderr=subprocess.STDOUT, env=env,
                                start_new_session=True)
        try:
            return proc.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            return None
        finally:
            if proc.poll() is None:
                if hasattr(os, "killpg"):
                    os.killpg(proc.pid, signal.SIGKILL)
                else:                         # Windows: no process groups
                    proc.kill()
                proc.wait()


def failure_note(log):
    """rugnux's own last error message, without the spdlog timestamp and severity prefix. Read from
    the log file rugnux writes beside its outputs (p_log.txt), whose lines are whole; the terminal
    output in run.log wraps them."""
    full = os.path.join(os.path.dirname(log), "p_log.txt")
    if os.path.exists(full):
        log = full
    try:
        lines = [l.strip() for l in open(log, errors="replace") if l.strip()]
    except OSError:
        return ""
    # the error itself before a "Step N FAILED" line that follows it
    for test in (lambda l: "[error]" in l or "[critical]" in l, lambda l: "failed" in l.lower()):
        for l in reversed(lines):
            if test(l):
                return l.split("] ")[-1][:300]
    return lines[-1].split("] ")[-1][:300] if lines else ""


def deposited_model(e, opts):
    """(deposition record from model_check's RCSB cache, None) for the entry a set belongs to, or
    (None, why not). A set id may carry a suffix (6h2p_native); the entry is the part before it.
    An in-house set names the entry of its crystal form in its `model` field instead, or an mmCIF
    file beside its data: a reference model of our own, refined against that set."""
    local = os.path.join(os.path.dirname(e["input_path"]), e.get("model") or "")
    if e.get("model") and os.path.isfile(local):
        return model_check.refinement_record(e["model"], local), None
    pdb = (e.get("model") or e["id"].split("_")[0]).lower()
    if not PDB_ID.match(pdb):
        return None, "no model: not a PDB entry (small molecule or unpublished set)"
    meta = model_check.deposition(pdb, opts["pdb_cache"])
    if meta is None:
        return None, f"no model: coordinates of {pdb} could not be downloaded"
    return meta, None


def run_one(run_dir, binary, e, opts):
    wd = os.path.join(run_dir, "work", e["arm"], e["id"])
    os.makedirs(wd)
    row = {"set": e["id"], "arm": e["arm"], "tags": e.get("tags", []), "input": e["input_path"],
           "gpu_others": gpu_others(), "model": None, "model_note": None, "rfree_deposited": None,
           "rfree_deposited_program": None,
           "rfree_ratio": None, "cmd": None, "exit_code": None, "elapsed_s": None, "wall_s": None}
    row.update({k: None for k in REFMAC_KEYS + depdata_check.KEYS + shelx_check.KEYS})
    if not os.path.exists(e["input_path"]):
        row.update(score.judge(e, {}, "no input"))
        return row
    meta = None
    if e["reference"] == "deposition" or e.get("model"):
        # no deposited model (small molecule, unpublished): the same command without --model
        meta, row["model_note"] = deposited_model(e, opts)
        if meta:
            row["model"], row["rfree_deposited"] = meta["xyz"], meta["rfree"]
            row["rfree_deposited_program"] = meta.get("program")
    reference = row["model"]
    if e.get("timing"):
        # the timing benchmark: rugnux starts from a model it has to place, and the set's own model
        # (above) is the reference the placement is checked against
        row["model"] = timing.start_model(e["timing"]["start"], wd, opts["pdb_cache"])
    cmd = command(binary, e, opts, row["model"])
    row["cmd"] = " ".join(cmd)
    t0 = time.time()
    note = ""
    rc = run_logged(cmd, wd, opts["timeout"])
    if rc is None:
        note = "timeout"
    row["elapsed_s"] = round(time.time() - t0, 1)
    row["exit_code"] = rc
    crash = score.crash_note(rc, wd)
    if crash:
        note = "crash: " + crash
    elif rc:
        note = f"exit {rc}: {failure_note(os.path.join(wd, 'run.log'))}"
    row.update(score.judge(e, score.read_report(os.path.join(wd, "p_report.txt")), note))
    # rugnux's own WALL_TIME where it wrote one: the runner's clock also counts the wait for a
    # GPU slot under --gpulock
    row["wall_s"] = row["rugnux_wall_s"] if row["rugnux_wall_s"] is not None else row["elapsed_s"]
    if row["rfree"] and row["rfree_deposited"]:
        row["rfree_ratio"] = round(row["rfree"] / row["rfree_deposited"], 4)
    if e.get("timing"):
        row.update(timing.check(e["timing"], row, wd, reference))
    elif e["reference"] == "deposition" and row["model"] and os.path.exists(os.path.join(wd, "p.mtz")):
        row.update(depdata_check.check_set(wd, e["id"], opts["pdb_cache"]))
    if opts["model_check"] and e["arm"] == "open":
        row.update(check_model(e, wd))
    if e.get("cod"):
        # small molecule: its published structure refined against our p.hkl (report only)
        row.update(shelx_check.check(wd, e["cod"], opts.get("cod_cache") or shelx_check.DEFAULT_CACHE,
                                     opts.get("shelxl")))
    return row


# what the REFMAC model check (model_check.py, --model-check) adds to an open-arm row
REFMAC_REFINED = ("refined_rfree_shared", "refined_rfree_shared_depdata", "refined_rwork",
                  "refined_rwork_depdata", "refined_rfree_refmac", "refined_rfree_refmac_depdata",
                  "shared_free_n", "refined_reason")
REFMAC_KEYS = ("refmac_rfree", "refmac_rwork", "refmac_rfree_depflags", "refmac_rfree_depdata",
               "refmac_rfree_ratio", "refmac_twin", "refmac_untwinned", "refmac_status",
               "refmac_reason", "refmac_refined_ratio") + tuple("refmac_" + k for k in REFMAC_REFINED)


def check_model(e, wd):
    """R-free of the deposited model against this run's merge, by REFMAC (model_check.py)."""
    out = {k: None for k in REFMAC_KEYS}
    pdb = e["id"].split("_")[0]
    mtz = os.path.join(wd, "p.mtz")
    if not PDB_ID.match(pdb):
        return dict(out, refmac_status="skipped", refmac_reason="not a PDB entry")
    if not os.path.exists(mtz):
        return dict(out, refmac_status="skipped", refmac_reason="no merged MTZ")
    try:
        res = model_check.check(mtz, pdb, os.path.join(wd, "model_check"))
    except Exception as ex:          # a failing check must not lose the set's processing result
        return dict(out, refmac_status="error", refmac_reason=f"{type(ex).__name__}: {ex}"[:300])
    out.update({"refmac_" + k: res.get(k) for k in ("rfree", "rwork", "rfree_depflags", "rfree_depdata",
                                                    "untwinned") + REFMAC_REFINED},
               refmac_twin=" ".join(res.get("twin_declared") or []) or None,
               refmac_status=res.get("status"),
               refmac_reason=res.get("reason") or res.get("baseline_note"))
    # The fair ratio: our merge and the depositor's structure factors, both scored on the
    # depositor's free set by the same protocol. Our own free set was mostly work reflections
    # for the depositor, so R_free on it reads low.
    if out["refmac_rfree_depflags"] and out["refmac_rfree_depdata"]:
        out["refmac_rfree_ratio"] = round(out["refmac_rfree_depflags"] / out["refmac_rfree_depdata"], 4)
    # the same after refinement of each data set at its own resolution, on the shared free set
    if out["refmac_refined_rfree_shared"] and out["refmac_refined_rfree_shared_depdata"]:
        out["refmac_refined_ratio"] = round(out["refmac_refined_rfree_shared"]
                                            / out["refmac_refined_rfree_shared_depdata"], 4)
    return out


def run_nullmodel(run_dir, binary, site, opts):
    """The wrong-model null (nullmodel.py): each of its sets once without --model and once with each
    of its wrong models, every wrong-model run scored against the no-model run. Its rows go to
    nullmodel_results.json in the run directory, apart from results.json, so the sets' own rows and scores
    are those of a run without it."""
    sets = {e["id"]: e for e in load_sets(site, "open")}
    rows = []
    out = os.path.join(run_dir, "nullmodel_results.json")
    for n in nullmodel.load():
        e = dict(sets[n["set"]], reference="deposition",
                 input_path=os.path.join(site["arms"]["open"]["data_root"], sets[n["set"]]["input"]))
        base = os.path.join(run_dir, "work", "nullmodel", e["id"])
        none_wd = os.path.join(base, "none")
        os.makedirs(none_wd)
        rc_none = run_logged(command(binary, e, opts), none_wd, opts["timeout"])
        for m in n["models"]:
            wd = os.path.join(base, m["pdb"])
            os.makedirs(wd)
            meta = model_check.deposition(m["pdb"], opts["pdb_cache"])
            row = {"set": e["id"], "model": m["pdb"], "verdict": "fail"}
            if meta is None:
                row["reason"] = f"coordinates of {m['pdb']} could not be read or downloaded"
            else:
                rc = run_logged(command(binary, e, opts, meta["xyz"]), wd, opts["timeout"])
                crash = score.crash_note(rc_none, none_wd) or score.crash_note(rc, wd)
                if crash or rc_none != 0 or rc != 0:
                    row.update(verdict="crash" if crash else "fail",
                               reason=crash or f"exit {rc_none} without --model, {rc} with it")
                else:
                    row.update(nullmodel.score_pair(none_wd, wd))
                row["wall_s"] = score.fnum(score.read_report(os.path.join(wd, "p_report.txt")).get("WALL_TIME"))
            rows.append(row)
            write_json(out, rows)
            print(f"[null {len(rows)}] {e['id']:8s} --model {m['pdb']}  {row['verdict']:5s} "
                  f"{row.get('model_fit') or '-':10s} R_free={report.f(row.get('r_free'), '{:.3f}')}  "
                  f"{row.get('reason') or ''}", flush=True)
    return rows


def make_read_only(path):
    for dp, dirs, files in os.walk(path):
        for fn in files:
            p = os.path.join(dp, fn)
            if not os.path.islink(p):
                os.chmod(p, os.stat(p).st_mode & ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH))
    for dp, dirs, files in os.walk(path, topdown=False):
        os.chmod(dp, os.stat(dp).st_mode & ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH))


def write_json(path, obj):
    with open(path + ".tmp", "w") as fh:
        json.dump(obj, fh, indent=1)
    os.replace(path + ".tmp", path)


def write_manifest(path, man):
    """A manifest with one set per line, so a diff of it reads set by set."""
    lines = [json.dumps(e) for e in man["sets"]]
    head = json.dumps({k: v for k, v in man.items() if k != "sets"})[:-1]
    with open(path, "w") as fh:
        fh.write(head + ', "sets": [\n ' + ",\n ".join(lines) + "\n]}\n")


def write_report(run_dir, baseline=None, out_dir=None, allow_incomplete=False):
    """report.md and report.html (report_PRIVATE.* for a private run) in the run directory, or
    <run name>.md/.html (with _PRIVATE) in out_dir."""
    man = json.load(open(os.path.join(run_dir, "manifest.json")))
    name = os.path.basename(os.path.normpath(run_dir)) if out_dir else "report"
    base = os.path.join(out_dir or run_dir, name + ("_PRIVATE" if man.get("private") else ""))
    doc = report.build(run_dir, baseline, allow_incomplete=allow_incomplete or not out_dir)
    open(base + ".md", "w").write(doc.markdown())
    open(base + ".html", "w").write(doc.html())
    return base + ".html"


def run_sets(site, sets, arms, binary_src, build_dir, label, opts, subset, baseline=None, lock=None):
    """Run the given manifest rows into a new run directory. Returns the directory."""
    private = any(site["arms"][a].get("private") for a in arms)
    if private and len(arms) > 1:
        sys.exit("the private arm runs on its own: its results must not share a run or a report")
    if baseline:
        report.check_baseline(baseline, private)
    lock_fh = None
    if lock:
        import fcntl                          # POSIX only; a subset run (and Windows) takes no lock
        os.makedirs(os.path.dirname(lock), exist_ok=True)
        lock_fh = open(lock, "w")
        try:
            fcntl.flock(lock_fh, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            sys.exit(f"another full battery run holds {lock}; refusing to start a second one")

    info = binary_info(binary_src, build_dir)
    label = label or info.get("version") or "run"
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
    run_dir = os.path.join(opts["runs_root"],
                           f"{stamp}_{info['git']}_{label}" + ("_private" if private else ""))
    os.makedirs(run_dir)                      # never reuse or overwrite a run directory
    # rugnux.exe keeps its extension, which Windows needs to run it
    binary = os.path.join(run_dir, "bin", "rugnux" + os.path.splitext(binary_src)[1])
    os.makedirs(os.path.dirname(binary))
    shutil.copy2(binary_src, binary)
    info.update(copied_from=os.path.realpath(binary_src), sha256=sha256(binary))

    for e in sets:
        e["reference"] = site["arms"][e["arm"]]["reference"]
        e.setdefault("input_path", os.path.join(site["arms"][e["arm"]]["data_root"], e["input"]))
    man = {"label": label, "arms": arms, "private": private, "subset": subset, "baseline": baseline,
           "command_doc": {arm: COMMAND_DOC[site["arms"][arm]["reference"]] for arm in arms},
           "binary": info,
           "runner_git": git(HERE, "rev-parse", "--short", "HEAD") or "unknown",
           "runner_dirty": bool(git(HERE, "status", "--porcelain", "--", ".")),
           "host": socket.gethostname(), "hardware": hardware_info(), "pid": os.getpid(),
           "start": datetime.datetime.now().isoformat(timespec="seconds"),
           "end": None, "complete": False, "results_schema": report.RESULTS_SCHEMA, "options": {k: v for k, v in opts.items() if k != "runs_root"},
           "gpu_shared": bool(opts["gpulock"]),
           "refmac": model_check.refmac_version() if opts["model_check"] else None, "sets": sets}
    write_json(os.path.join(run_dir, "manifest.json"), man)
    print(f"run directory {run_dir}", flush=True)

    results = []
    write_json(os.path.join(run_dir, "results.json"), results)
    # a kill by SIGTERM is recorded like Ctrl-C: the run directory then says it was aborted
    signal.signal(signal.SIGTERM, interrupted)
    try:
        for e in sets:
            row = run_one(run_dir, binary, e, opts)
            results.append(row)
            if row.get("gpu_others"):
                man["gpu_shared"] = True
            write_json(os.path.join(run_dir, "results.json"), results)
            print(f"[{len(results)}/{len(sets)}] {e['arm']:8s} {e['id']:24s} {row['verdict']:8s} "
                  f"{row.get('sg') or '-':12s} d={report.f(row.get('d_min'))} "
                  f"t={report.f(row.get('wall_s'), '{:.0f}')}s  {row.get('reason') or ''}", flush=True)
        if opts.get("nullmodel"):
            run_nullmodel(run_dir, binary, site, opts)
    except KeyboardInterrupt:
        mark_aborted(run_dir, f"interrupted after {len(results)} of {len(sets)} sets")
        sys.exit(f"aborted; {run_dir} is marked as such")

    man.update(end=datetime.datetime.now().isoformat(timespec="seconds"), complete=True)
    write_json(os.path.join(run_dir, "manifest.json"), man)
    print(f"report {write_report(run_dir, baseline)}", flush=True)
    print("\n" + report.console_table(results), flush=True)
    make_read_only(run_dir)
    if lock_fh:
        lock_fh.close()
    return run_dir


def interrupted(*_):
    raise KeyboardInterrupt


def mark_aborted(run_dir, reason):
    man = json.load(open(os.path.join(run_dir, "manifest.json")))
    man["aborted"] = {"at": datetime.datetime.now().isoformat(timespec="seconds"), "reason": reason}
    write_json(os.path.join(run_dir, "manifest.json"), man)
    make_read_only(run_dir)


def default_baseline(site, arms, given):
    """--baseline, else the site's persisted baseline for this kind of run ('none' = no delta)."""
    if given:
        return None if given == "none" else given
    private = any(site["arms"][a].get("private") for a in arms)
    return site.get("baseline_private" if private else "baseline") or None


# ------------------------------------------------------------------ subcommands

def cmd_run(a, site):
    arms = [x.strip() for x in a.arm.split(",")]
    for arm in arms:
        if arm not in site["arms"]:
            sys.exit(f"unknown arm {arm}; the site config has {', '.join(site['arms'])}")
    sets = select(site, arms, a.only, a.tier)
    if not sets:
        sys.exit("nothing selected")
    subset = a.only or (f"tier {a.tier}" if a.tier else None)
    opts = {"threads": a.threads, "timeout": a.timeout, "model_check": a.model_check and "open" in arms,
            "pdb_cache": site.get("pdb_cache") or model_check.DEFAULT_CACHE,
            "cod_cache": site.get("cod_cache") or shelx_check.DEFAULT_CACHE,
            "shelxl": shelx_check.find_shelxl(site.get("shelxl")),
            "gpulock": site["gpulock"] if a.gpulock else None, "extra_args": a.extra.split(),
            "runs_root": a.runs_root or site["runs_root"],
            # the wrong-model null rides along with every ci-tier run (and any run that asks for it)
            "nullmodel": (a.tier == "ci" or a.nullmodel) and not a.no_nullmodel and "open" in arms}
    if a.gpulock and not site.get("gpulock"):
        sys.exit("--gpulock: the site config names no gpulock")
    run_dir = run_sets(site, sets, arms, a.rugnux, a.build_dir, a.label, opts, subset,
                       default_baseline(site, arms, a.baseline), lock=None if subset else site["lock_file"])
    # The ci tier holds sets every recent release passed: one of them failing (wrong lattice or
    # symmetry, a missing input) fails the run, whatever else the run included. A crash fails the
    # run on any set, in or out of the tier.
    ci = {(e["arm"], e["id"]) for e in sets if "ci" in e.get("tiers", {})}
    res = json.load(open(os.path.join(run_dir, "results.json")))
    crashed = [r for r in res if r["verdict"] == "crash"]
    bad = [r for r in res if (r["arm"], r["set"]) in ci and r["verdict"] in ("fail", "not_run")]
    null_path = os.path.join(run_dir, "nullmodel_results.json")
    null_bad = [r for r in json.load(open(null_path)) if r["verdict"] != "pass"] if os.path.exists(null_path) else []
    if crashed or bad or null_bad:
        sys.exit("; ".join(([f"{len(crashed)} set(s) CRASHED: " +
                             ", ".join(f"{r['set']} ({r['reason']})" for r in crashed)] if crashed else [])
                           + ([f"{len(bad)} of {len(ci)} ci-tier sets failed: " +
                               ", ".join(f"{r['set']} ({r['reason']})" for r in bad)] if bad else [])
                           + ([f"wrong-model null: {len(null_bad)} pair(s) failed: " +
                               ", ".join(f"{r['set']} --model {r['model']} ({r['reason']})" for r in null_bad)]
                              if null_bad else [])))


def cmd_nullmodel(a, site):
    """The wrong-model null on its own: a run directory with no sets of its own and only
    nullmodel_results.json, its report holding just that section. Exits 1 when a pair failed."""
    opts = {"threads": a.threads, "timeout": a.timeout, "model_check": False,
            "pdb_cache": site.get("pdb_cache") or model_check.DEFAULT_CACHE,
            "cod_cache": site.get("cod_cache") or shelx_check.DEFAULT_CACHE, "shelxl": None,
            "gpulock": site["gpulock"] if a.gpulock else None, "extra_args": [],
            "runs_root": a.runs_root or site["runs_root"], "nullmodel": True}
    if a.gpulock and not site.get("gpulock"):
        sys.exit("--gpulock: the site config names no gpulock")
    run_dir = run_sets(site, [], ["open"], a.rugnux, a.build_dir, a.label, opts, "wrong-model null", None)
    rows = json.load(open(os.path.join(run_dir, "nullmodel_results.json")))
    bad = [r for r in rows if r["verdict"] != "pass"]
    print(f"\nwrong-model null: {len(rows) - len(bad)} of {len(rows)} pairs pass", flush=True)
    if bad:
        sys.exit(f"wrong-model null: {len(bad)} pair(s) failed: " +
                 ", ".join(f"{r['set']} --model {r['model']} ({r['reason']})" for r in bad))


def cmd_timing(a, site):
    """A timing benchmark (timing.json): its sets, each with --model <a model it has to place>, in a
    run directory of its own; prints the checks and the step times, exits 1 when a set failed."""
    sets = []
    for t in timing.load(a.benchmark):
        e = select(site, [t["arm"]], t["set"], None)[0]
        e["timing"] = t
        sets.append(e)
    opts = {"threads": a.threads, "timeout": a.timeout, "model_check": False,
            "pdb_cache": site.get("pdb_cache") or model_check.DEFAULT_CACHE,
            "cod_cache": site.get("cod_cache") or shelx_check.DEFAULT_CACHE, "shelxl": None,
            "gpulock": site["gpulock"] if a.gpulock else None, "extra_args": a.extra.split(),
            "runs_root": a.runs_root or site["runs_root"], "nullmodel": False}
    if a.gpulock and not site.get("gpulock"):
        sys.exit("--gpulock: the site config names no gpulock")
    baseline = None if a.baseline == "none" else a.baseline or (site.get("baseline_timing") or {}).get(a.benchmark)
    run_dir = run_sets(site, sets, sorted({e["arm"] for e in sets}), a.rugnux, a.build_dir,
                       a.label or f"timing-{a.benchmark}", opts, f"timing benchmark {a.benchmark}", baseline)
    res = json.load(open(os.path.join(run_dir, "results.json")))
    print("\n" + timing.console(res), flush=True)
    bad = [r for r in res if r.get("bench_verdict") != "pass" or r["verdict"] == "crash"]
    if bad:
        sys.exit(f"timing benchmark: {len(bad)} of {len(res)} sets failed: " +
                 "; ".join(f"{r['set']} ({r.get('bench_reason') or r.get('reason')})" for r in bad))


KEPT_ROOT = os.path.join(HERE, "results")
KEPT_FILES = ("report.md", "results.json", "manifest.json")
PUBLIC_ARMS = {"open", "inhouse"}


def cmd_keep(a, site):
    """Copy a run's report.md, results.json and manifest.json, as written, into
    results/<role>/<run name>/, where they are committed for the paper tooling. The role names the
    machine (workstation, server); the files themselves record its host and hardware. Open and in-house
    runs only: a run its manifest marks private, or with any row or arm outside those two, is refused.
    A run of which only report.md exists (another host's, its other files still to come) is kept with
    what there is, checked on the report's own arms line."""
    src = a.run if os.path.isdir(a.run) else os.path.join(site["runs_root"], a.run)
    report_md = os.path.join(src, "report.md")
    if not os.path.exists(report_md):
        sys.exit(f"{src}: no report.md (a private run's is report_PRIVATE.md, and is never kept)")
    text = open(report_md, errors="replace").read()
    title = re.match(r"# Rugnux battery - (\S+)", text)
    arms_line = re.search(r"^- arms: ([^;]+);", text, re.M)
    if not (title and arms_line):
        sys.exit(f"{report_md}: not a battery report (no title or arms line)")
    arms = {x.strip() for x in arms_line.group(1).split(",")}
    if "PRIVATE ARM" in text:
        arms.add("private")
    man_path, res_path = os.path.join(src, "manifest.json"), os.path.join(src, "results.json")
    if os.path.exists(man_path):
        man = json.load(open(man_path))
        if man.get("private"):
            arms.add("private")
        arms |= set(man.get("arms", [])) | {e.get("arm") for e in man.get("sets", [])}
    if os.path.exists(res_path):
        arms |= {r.get("arm") for r in json.load(open(res_path))}
    if not arms <= PUBLIC_ARMS:
        sys.exit(f"{src}: refused - it holds arm(s) {', '.join(sorted(map(str, arms - PUBLIC_ARMS)))}; "
                 "only open and inhouse runs are kept")
    dest = os.path.join(KEPT_ROOT, a.role, title.group(1))
    os.makedirs(dest, exist_ok=True)
    for name in KEPT_FILES:
        if os.path.exists(os.path.join(src, name)):
            shutil.copyfile(os.path.join(src, name), os.path.join(dest, name))
        else:
            print(f"note: {src} has no {name}; kept without it")
    print(dest)


def stress_one(stress_dir, binary, e, i, opts):
    """One stress job: set e, repetition i. Returns its line of stress.json."""
    wd = os.path.join(stress_dir, "work", e["arm"], e["id"], str(i))
    os.makedirs(wd)
    model = None
    if e["reference"] == "deposition" or e.get("model"):
        meta, _ = deposited_model(e, opts)
        model = meta and meta["xyz"]
    cmd = command(binary, e, dict(opts, gpulock=None), model)
    if opts["sanitizer"]:
        # API errors rugnux handles itself (a probe for a driver entry point, an allocation it retries)
        # are not the defect looked for here; a kernel's bad access, race or bad sync is.
        cmd = ["compute-sanitizer", "--tool", opts["sanitizer"], "--report-api-errors", "no"] + cmd
    env = dict(os.environ, CUDA_LAUNCH_BLOCKING="1") if opts["launch_blocking"] else None
    t0 = time.time()
    rc = run_logged(cmd, wd, opts["timeout"], env)
    hkl = os.path.join(wd, "p.hkl")
    return {"set": e["id"], "arm": e["arm"], "rep": i, "cmd": " ".join(cmd), "exit_code": rc,
            "elapsed_s": round(time.time() - t0, 1), "timeout": rc is None,
            "crash": score.crash_note(rc, wd) or None,
            "hkl_sha256": sha256(hkl) if os.path.exists(hkl) else None}


def cmd_stress(a, site):
    """Run the selected sets --repeat times, --concurrent at once, optionally under compute-sanitizer,
    and ask one question: did any run crash (score.crash_note)? A race between GPU users of one
    process shows up rarely, so a set that passed once proves little; running it many times, beside
    another process that changes its timing, or under racecheck, is what finds one before a release."""
    arms = [x.strip() for x in a.arm.split(",")]
    sets = select(site, arms, a.only, a.tier)
    if not sets or not (a.only or a.tier):
        sys.exit("stress: name the sets (--only or --tier)")
    for e in sets:
        e["reference"] = site["arms"][e["arm"]]["reference"]
        e.setdefault("input_path", os.path.join(site["arms"][e["arm"]]["data_root"], e["input"]))
    info = binary_info(a.rugnux, None)
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    stress_dir = os.path.join(a.runs_root or site["runs_root"], "stress",
                              f"{stamp}_{info['git']}_{a.label or info.get('version') or 'stress'}")
    binary = os.path.join(stress_dir, "bin", "rugnux")
    os.makedirs(os.path.dirname(binary))
    shutil.copy2(a.rugnux, binary)
    opts = {"threads": a.threads, "timeout": a.timeout, "extra_args": a.extra.split(),
            "sanitizer": a.sanitizer, "launch_blocking": a.launch_blocking,
            "pdb_cache": site.get("pdb_cache") or model_check.DEFAULT_CACHE}
    jobs = [(e, i) for i in range(a.repeat) for e in sets]   # sets side by side, not one set K times
    print(f"stress directory {stress_dir}: {len(jobs)} runs, {a.concurrent} at a time"
          + (f", under compute-sanitizer --tool {a.sanitizer}" if a.sanitizer else ""), flush=True)
    out = []
    with concurrent.futures.ThreadPoolExecutor(a.concurrent) as pool:
        futures = [pool.submit(stress_one, stress_dir, binary, e, i, opts) for e, i in jobs]
        for fut in concurrent.futures.as_completed(futures):
            r = fut.result()
            out.append(r)
            write_json(os.path.join(stress_dir, "stress.json"), out)
            state = "CRASH" if r["crash"] else "timeout" if r["timeout"] else "ok"
            print(f"[{len(out)}/{len(jobs)}] {r['arm']:8s} {r['set']:24s} #{r['rep']} {state:7s} "
                  f"t={r['elapsed_s']:.0f}s  {r['crash'] or ''}", flush=True)
    print()
    for e in sets:
        rs = [r for r in out if (r["arm"], r["set"]) == (e["arm"], e["id"])]
        hkls = {r["hkl_sha256"] for r in rs if r["hkl_sha256"]}
        print(f"{e['id']:24s} {len(rs)} runs, {sum(bool(r['crash']) for r in rs)} crashed, "
              f"{sum(r['timeout'] for r in rs)} timed out, {len(hkls)} distinct p.hkl")
    crashed = [r for r in out if r["crash"] or r["timeout"]]
    if crashed:
        sys.exit(f"{len(crashed)} of {len(out)} runs CRASHED or timed out: " +
                 "; ".join(f"{r['set']} #{r['rep']} ({r['crash'] or 'timeout'})" for r in crashed))


def cmd_list(a, site):
    root = a.runs_root or site["runs_root"]
    for d in sorted(os.listdir(root)):
        run = os.path.join(root, d)
        if not os.path.exists(os.path.join(run, "manifest.json")):
            continue
        man = json.load(open(os.path.join(run, "manifest.json")))
        state = report.run_state(man)
        n = len(json.load(open(os.path.join(run, "results.json")))) \
            if os.path.exists(os.path.join(run, "results.json")) else 0
        mark = " <- baseline" if run in (site.get("baseline"), site.get("baseline_private")) else ""
        print(f"{d:48s} {state:10s} {n:3d}/{report.expected_rows(man):3d} rows  "
              f"{man.get('subset') or 'full arm':20.20s} {','.join(man['arms'])}{mark}")


def cmd_abort(a, site):
    man = json.load(open(os.path.join(a.run, "manifest.json")))
    state = report.run_state(man)
    if state not in ("unfinished", "running") or (state == "running" and not a.force):
        sys.exit(f"{a.run} is {state}; only an unfinished run can be marked aborted "
                 "(a live one needs --force)")
    mark_aborted(a.run, a.reason)
    print(f"{a.run}: marked aborted ({a.reason})")


def cmd_compare(a, site):
    man_a, res_a = report.load_run(a.run_a, a.allow_incomplete)
    man_b, res_b = report.load_run(a.run_b, a.allow_incomplete)
    report.fill_shelx(a.run_a, man_a, res_a)
    report.fill_shelx(a.run_b, man_b, res_b)
    rows = report.compare_rows(res_a, res_b)
    doc = report.Doc(f"{os.path.basename(os.path.normpath(a.run_a))}  vs  "
                     f"{os.path.basename(os.path.normpath(a.run_b))}")
    doc.p(f"A: rugnux {man_a['binary'].get('version')} ({man_a['binary']['sha256'][:12]}); "
          f"B: rugnux {man_b['binary'].get('version')} ({man_b['binary']['sha256'][:12]}). "
          f"{len(rows)} sets, {sum(1 for r in rows if r[3])} changed beyond noise.")
    report.compare_table(doc, rows, only_changed=not a.all)
    print(doc.markdown())
    if not a.rerun_changed:
        return
    # Rerun the changed sets with A's own binary and options: a set that moves again under the
    # same binary was noise; one that reproduces A was changed by B for real.
    changed = {k for k, ra, rb, ch in rows if ra and rb and [c for c in ch if c != "time"]}
    if not changed:
        print("nothing changed beyond noise; no rerun")
        return
    if man_a.get("results_schema", 1) < report.RESULTS_SCHEMA:
        sys.exit("--rerun-changed: A ran the older variant commands; rerun it with this runner instead")
    sets = [e for e in man_a["sets"] if (e["arm"], e["id"]) in changed]
    opts = dict(man_a["options"], runs_root=a.runs_root or site["runs_root"])
    opts["gpulock"] = opts["gpulock"] and site.get("gpulock")
    opts.setdefault("pdb_cache", site.get("pdb_cache") or model_check.DEFAULT_CACHE)
    opts.setdefault("cod_cache", site.get("cod_cache") or shelx_check.DEFAULT_CACHE)
    opts.setdefault("shelxl", shelx_check.find_shelxl(site.get("shelxl")))
    label = "rerunA-" + man_a["label"]
    rerun = run_sets(site, sets, man_a["arms"], os.path.join(a.run_a, "bin", "rugnux"), None, label,
                     opts, ",".join(e["id"] for e in sets))
    _, res_a2 = report.load_run(rerun)
    again = {k: ch for k, ra, ra2, ch in report.compare_rows(res_a, res_a2) if ra and ra2}
    doc = report.Doc("Changed sets, rerun with A's binary")
    rows_out = []
    for k, ra, rb, ch in rows:
        if k not in changed:
            continue
        noise = [c for c in again.get(k, []) if c != "time"]
        rows_out.append([k[1], k[0], ", ".join(c for c in ch if c != "time"),
                         ", ".join(noise) or "-",
                         "NOISE - A itself moved" if noise else "REAL - A reproduces"])
    doc.table(["set", "arm", "A -> B", "A -> A again", "reading"], rows_out)
    doc.p(f"rerun directory: {rerun}")
    print(doc.markdown())


def cmd_table(a, site):
    _, res = report.load_run(a.run, a.allow_incomplete)
    print(report.console_table(res))


def cmd_report(a, site):
    # a finished run directory is read-only, so a re-render goes elsewhere, named after the run
    man, _ = report.load_run(a.run, a.allow_incomplete)
    repo = git(HERE, "rev-parse", "--show-toplevel")
    out = os.path.realpath(a.out)
    if man.get("private") and repo and (out + os.sep).startswith(os.path.realpath(repo) + os.sep):
        sys.exit("a private run's report must not be written into the repository")
    os.makedirs(a.out, exist_ok=True)
    print(write_report(a.run, a.baseline, a.out, a.allow_incomplete))


def recheck_one(e, wd, pdb_cache):
    """The post-processing checks of one open-arm row, on the run's own p.mtz (a worker process)."""
    out = {}
    if e["model"] and os.path.exists(os.path.join(wd, "p.mtz")):
        out.update(depdata_check.check_set(wd, e["id"], pdb_cache))
    out.update(check_model(e, wd))
    return out


def set_writable(paths):
    for p in paths:
        os.chmod(p, os.stat(p).st_mode | stat.S_IWUSR)


def cmd_recheck(a, site):
    """Rerun model_check / depdata_check on a finished run's open-arm merges and update its
    results.json in place. The previous results.json is kept beside it and the recheck is recorded
    in manifest.json ("rechecks": which runner version produced the refmac_* / dep_* numbers)."""
    run = os.path.abspath(a.run)
    report.load_run(run, False)       # refuses a run that did not finish
    # the raw files: load_run rescores rows, and results.json keeps the verdicts as they were scored
    man = json.load(open(os.path.join(run, "manifest.json")))
    results = json.load(open(os.path.join(run, "results.json")))
    only = set(a.only.split(",")) if a.only else None
    rows = [r for r in results if r["arm"] == "open" and (only is None or r["set"] in only)]
    if only and len(rows) != len(only):
        sys.exit("not open-arm rows of this run: " + ", ".join(sorted(only - {r["set"] for r in rows})))
    if not rows:
        sys.exit("nothing to recheck")
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    entry = {"date": datetime.datetime.now().isoformat(timespec="seconds"),
             "runner_git": git(HERE, "rev-parse", "--short", "HEAD") or "unknown",
             "runner_dirty": bool(git(HERE, "status", "--porcelain", "--", ".")),
             "sets": [r["set"] for r in rows],
             "fields": list(REFMAC_KEYS + depdata_check.KEYS),
             "refmac": model_check.refmac_version(),
             "previous_results": "results.json.pre-recheck-" + stamp,
             "command": " ".join(sys.argv)}
    pdb_cache = site.get("pdb_cache") or model_check.DEFAULT_CACHE
    # a finished run directory is read-only: open up only what the recheck writes
    wds = [os.path.join(run, "work", "open", r["set"]) for r in rows]
    set_writable([run, os.path.join(run, "results.json"), os.path.join(run, "manifest.json")] + wds)
    for wd in wds:                    # the previous check's files stay, renamed
        if os.path.isdir(os.path.join(wd, "model_check")):
            set_writable([os.path.join(wd, "model_check")])
            os.rename(os.path.join(wd, "model_check"), os.path.join(wd, "model_check.pre-recheck-" + stamp))
    shutil.copy2(os.path.join(run, "results.json"), os.path.join(run, "results.json.pre-recheck-" + stamp))
    with concurrent.futures.ProcessPoolExecutor(a.jobs) as pool:
        futures = {pool.submit(recheck_one, {"id": r["set"], "model": r.get("model")}, wd, pdb_cache): r
                   for r, wd in zip(rows, wds)}
        for n, fut in enumerate(concurrent.futures.as_completed(futures), 1):
            r = futures[fut]
            r.update(fut.result())
            print(f"[{n}/{len(rows)}] {r['set']:24s} {r['refmac_status'] or '-':8s} "
                  f"R_free depflags {report.f(r.get('refmac_rfree_depflags'), '{:.4f}')} "
                  f"depdata {report.f(r.get('refmac_rfree_depdata'), '{:.4f}')}  "
                  f"{r.get('refmac_reason') or ''}",
                  flush=True)
    write_json(os.path.join(run, "results.json"), results)
    man.setdefault("rechecks", []).append(dict(entry, end=datetime.datetime.now().isoformat(timespec="seconds")))
    write_json(os.path.join(run, "manifest.json"), man)
    for name in ("report", "report_PRIVATE"):
        for ext in (".md", ".html"):
            if os.path.exists(os.path.join(run, name + ext)):
                set_writable([os.path.join(run, name + ext)])
                os.rename(os.path.join(run, name + ext),
                          os.path.join(run, f"{name}{ext}.pre-recheck-{stamp}"))
    print(f"report {write_report(run, man.get('baseline'))}", flush=True)
    make_read_only(run)


def cmd_discover(a, site):
    """Propose an input for every open-arm directory the manifest does not have yet; with
    --refresh, also show where today's discovery disagrees with the manifest's (unpinned) input."""
    cfg = site["arms"]["open"]
    root = cfg["data_root"]
    by_dir = {}                       # a directory can hold several sets (6h2p_native, 6h2p_1p89A)
    for e in load_sets(site, "open"):
        by_dir.setdefault(e["input"].split("/")[0], []).append(e)
    for d in sorted(os.listdir(root)):
        if d.startswith(".") or not os.path.isdir(os.path.join(root, d)):
            continue
        sets = by_dir.get(d, [])
        if (sets and not a.refresh) or len(sets) > 1 or (sets and sets[0].get("pinned")):
            continue
        path, fmt, note = inputs.find_input(os.path.join(root, d))
        rel = os.path.relpath(path, root) if path else None
        if not sets:
            print(f"NEW   {d:10s} {rel}   [{note}]")
        elif rel != sets[0]["input"]:
            print(f"DIFF  {d:10s} manifest {sets[0]['input']}\n      {'':10s} discover {rel}   [{note}]")
    for d in sorted(set(by_dir) - set(os.listdir(root))):
        print(f"GONE  {d:10s} in the manifest, not on disk")


def cmd_refs(a, site):
    """Re-read every set's XDS reference from the CORRECT.LP beside its input."""
    cfg = site["arms"][a.arm]
    if cfg["reference"] != "xds":
        sys.exit(f"the {a.arm} arm is not referenced to XDS")
    man = json.load(open(cfg["manifest"]))
    for e in man["sets"]:
        lp = os.path.join(cfg["data_root"], os.path.dirname(e["input"]), "CORRECT.LP")
        if not os.path.exists(lp):
            print(f"{e['id']:24s} no CORRECT.LP")
            continue
        new = inputs.parse_correct_lp(lp)
        old = e.get("ref") or {}
        diff = [f"{k} {old.get(k)} -> {new.get(k)}" for k in sorted(set(old) | set(new))
                if old.get(k) != new.get(k)]
        if diff:
            print(f"{e['id']:24s} " + "; ".join(diff))
            e["ref"] = new
    if a.write:
        write_manifest(cfg["manifest"], man)


def cmd_remap(a, site):
    """Follow a rename of data directories. MAP is a JSON object {old: new}: a string renames the
    set of that id and, wherever it is the first directory of an input, that directory; an object
    {"id": ..., "input": ...} gives one set's new id and input outright. A renamed set's old id is
    kept under "aliases", so runs made before the rename still compare set by set."""
    cfg = site["arms"][a.arm]
    man = json.load(open(cfg["manifest"]))
    mp = json.load(open(a.map))
    aliases = man.setdefault("aliases", {})
    for e in man["sets"]:
        new = mp.get(e["id"])
        head, _, rest = e["input"].partition("/")
        if isinstance(new, dict):
            new_id, new_input = new.get("id", e["id"]), new.get("input", e["input"])
        else:
            new_id = new or e["id"]
            new_input = (mp[head] if isinstance(mp.get(head), str) else head) + "/" + rest
        if (new_id, new_input) == (e["id"], e["input"]):
            continue
        there = os.path.exists(os.path.join(cfg["data_root"], new_input))
        print(f"{e['id']:24s} -> {new_id:24s} {new_input}" + ("" if there else "   NOT FOUND"))
        if new_id != e["id"]:
            for old in [k for k, v in aliases.items() if v == e["id"]]:
                aliases[old] = new_id
            aliases[e["id"]] = new_id
        e["id"], e["input"] = new_id, new_input
    if a.write:
        write_manifest(cfg["manifest"], man)


def load_aliases(site):
    """{arm: {old set id: current id}} from the manifests, applied when runs are compared."""
    out = {}
    for arm, cfg in site["arms"].items():
        if os.path.exists(cfg["manifest"]):
            out[arm] = json.load(open(cfg["manifest"])).get("aliases", {})
    return out


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--site", help="site config (default: $JFJOCH_BATTERY_SITE, else site.json "
                                   "beside this script)")
    sub = ap.add_subparsers(dest="cmd", required=True)

    r = sub.add_parser("run")
    r.add_argument("--rugnux", required=True)
    r.add_argument("--label", help="short name, part of the run directory's name "
                                   "(default: the rugnux version)")
    r.add_argument("--arm", default="open,inhouse", help="comma-separated (default: open,inhouse)")
    r.add_argument("--only", help="comma-separated set ids")
    r.add_argument("--tier", help="named subset from the manifests: smoke or ci")
    r.add_argument("--build-dir", help="build tree whose CMakeCache.txt records the flags")
    r.add_argument("--gpulock", action="store_true",
                   help="queue each rugnux behind the site's GPU lock (timing is then not a reference)")
    r.add_argument("--threads", type=int, default=0, help="-N for rugnux (default: its own choice)")
    r.add_argument("--timeout", type=int, default=10800, help="per rugnux run, in s")
    r.add_argument("--model-check", action="store_true",
                   help="open arm: also score each merge against the deposited model with "
                        "REFMAC (model_check.py; needs CCP4)")
    r.add_argument("--extra", default="", help="extra rugnux arguments, applied to every set")
    r.add_argument("--nullmodel", action="store_true",
                   help="also run the wrong-model null (nullmodel.json; on by default with --tier ci)")
    r.add_argument("--no-nullmodel", action="store_true",
                   help="not even with --tier ci (CI runs it as its own step: battery.py nullmodel)")
    r.add_argument("--baseline", help="earlier run directory to report the delta against (default: "
                                      "the site's persisted baseline; 'none' for no delta)")
    r.add_argument("--runs-root")

    nm = sub.add_parser("nullmodel", help="the wrong-model null on its own; fail on any failed pair")
    nm.add_argument("--rugnux", required=True)
    nm.add_argument("--label", help="short name, part of the run directory's name")
    nm.add_argument("--build-dir", help="build tree whose CMakeCache.txt records the flags")
    nm.add_argument("--gpulock", action="store_true", help="queue each rugnux behind the site's GPU lock")
    nm.add_argument("--threads", type=int, default=0, help="-N for rugnux (default: its own choice)")
    nm.add_argument("--timeout", type=int, default=10800, help="per rugnux run, in s")
    nm.add_argument("--runs-root")

    tb = sub.add_parser("timing", help="the timing benchmark: sets placed by molecular replacement, "
                                       "checked and timed per step (timing.json); fail on any wrong answer")
    tb.add_argument("--rugnux", required=True)
    tb.add_argument("--benchmark", default="standard", choices=timing.names(),
                    help="which set list of timing.json: standard (open data only, the default) or inhouse")
    tb.add_argument("--label", help="short name, part of the run directory's name (default: timing-<benchmark>)")
    tb.add_argument("--build-dir", help="build tree whose CMakeCache.txt records the flags")
    tb.add_argument("--gpulock", action="store_true",
                    help="queue each rugnux behind the site's GPU lock (timing is then not a reference)")
    tb.add_argument("--threads", type=int, default=0, help="-N for rugnux (default: its own choice)")
    tb.add_argument("--timeout", type=int, default=10800, help="per rugnux run, in s")
    tb.add_argument("--extra", default="", help="extra rugnux arguments, applied to every set")
    tb.add_argument("--baseline", help="earlier timing run to report the delta against (default: the "
                                       "site's baseline_timing for this benchmark; 'none' for no delta)")
    tb.add_argument("--runs-root")

    kp = sub.add_parser("keep", help="copy a run's report.md, results.json and manifest.json into "
                                     "results/<role>/<run>/ (open and inhouse runs only)")
    kp.add_argument("run")
    kp.add_argument("--as", dest="role", required=True, choices=("workstation", "server"),
                    help="the machine the run is from: workstation (1 socket, RTX 5080) or server "
                         "(4 sockets, 4x T4)")

    st = sub.add_parser("stress", help="run sets repeatedly / side by side / under compute-sanitizer; "
                                       "fail on any crash")
    st.add_argument("--rugnux", required=True)
    st.add_argument("--label")
    st.add_argument("--arm", default="open,inhouse")
    st.add_argument("--only", help="comma-separated set ids")
    st.add_argument("--tier", help="named subset from the manifests: smoke or ci")
    st.add_argument("--repeat", type=int, default=1, help="runs of each set (default 1)")
    st.add_argument("--concurrent", type=int, default=1, help="rugnux processes at once (default 1)")
    st.add_argument("--sanitizer", choices=("memcheck", "racecheck", "synccheck", "initcheck"),
                    help="run every rugnux under compute-sanitizer with this tool (10-100x slower)")
    st.add_argument("--launch-blocking", action="store_true",
                    help="CUDA_LAUNCH_BLOCKING=1 for rugnux: a kernel's fault is reported by its own launch")
    st.add_argument("--threads", type=int, default=0, help="-N for rugnux (default: its own choice)")
    st.add_argument("--timeout", type=int, default=10800, help="per rugnux run, in s")
    st.add_argument("--extra", default="", help="extra rugnux arguments, applied to every run")
    st.add_argument("--runs-root")

    c = sub.add_parser("compare")
    c.add_argument("run_a")
    c.add_argument("run_b")
    c.add_argument("--all", action="store_true", help="list unchanged sets too")
    c.add_argument("--rerun-changed", action="store_true",
                   help="rerun the changed sets with A's saved binary to separate real change from noise")
    c.add_argument("--runs-root")
    c.add_argument("--allow-incomplete", action="store_true",
                   help="compare a run that did not finish (aborted, killed or still running)")

    p = sub.add_parser("report")
    p.add_argument("run")
    p.add_argument("--baseline")
    p.add_argument("--out", required=True, help="directory for the re-rendered report")
    p.add_argument("--allow-incomplete", action="store_true",
                   help="render a run that did not finish; the report says so")

    t = sub.add_parser("table", help="print a run's one-line-per-set table again")
    t.add_argument("run")
    t.add_argument("--allow-incomplete", action="store_true")

    ls = sub.add_parser("list")
    ls.add_argument("--runs-root")

    ab = sub.add_parser("abort")
    ab.add_argument("run")
    ab.add_argument("--reason", default="aborted by hand")
    ab.add_argument("--force", action="store_true", help="even if its runner still seems alive")

    rc = sub.add_parser("recheck")
    rc.add_argument("run")
    rc.add_argument("--only", help="comma-separated set ids (default: every open-arm row)")
    rc.add_argument("--jobs", type=int, default=8, help="sets checked in parallel (CPU; REFMAC)")

    d = sub.add_parser("discover")
    d.add_argument("--refresh", action="store_true")

    m = sub.add_parser("remap")
    m.add_argument("--arm", required=True)
    m.add_argument("map", help="JSON file {old: new}")
    m.add_argument("--write", action="store_true")

    f = sub.add_parser("refs")
    f.add_argument("--arm", required=True)
    f.add_argument("--write", action="store_true")

    a = ap.parse_args()
    site = load_site(a.site)
    report.ALIASES = load_aliases(site)
    report.PDB_CACHE = site.get("pdb_cache") or model_check.DEFAULT_CACHE
    report.COD_CACHE = site.get("cod_cache") or shelx_check.DEFAULT_CACHE
    report.SHELXL = shelx_check.find_shelxl(site.get("shelxl"))
    report.MANIFESTS = {arm: {e["id"]: e for e in load_sets(site, arm)}
                        for arm, cfg in site["arms"].items() if os.path.exists(cfg["manifest"])}
    {"run": cmd_run, "nullmodel": cmd_nullmodel, "timing": cmd_timing, "keep": cmd_keep, "stress": cmd_stress, "list": cmd_list, "compare": cmd_compare, "table": cmd_table, "report": cmd_report,
     "abort": cmd_abort, "recheck": cmd_recheck, "discover": cmd_discover, "refs": cmd_refs, "remap": cmd_remap}[a.cmd](a, site)


if __name__ == "__main__":
    main()