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]
                 [--build-dir DIR] [--gpulock] [--threads N]
                 [--model-check] [--baseline RUN|none]
  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 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). 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 that did not finish is never compared or reported on unless asked (--allow-incomplete).
"""
import argparse
import concurrent.futures
import datetime
import fcntl
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 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> <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",
}

LOCAL_SITE = os.path.join(HERE, "site.json")


# ------------------------------------------------------------------ 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))
    for arm in site["arms"].values():
        arm["manifest"] = os.path.join(HERE, arm["manifest"])   # an absolute path stays as it is
    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()])


# ------------------------------------------------------------------ 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"]:
        cmd = [opts["gpulock"]] + cmd
    return cmd


def failure_note(log):
    """rugnux's own last error message, without the spdlog timestamp and severity prefix."""
    try:
        lines = [l.strip() for l in open(log, errors="replace") if l.strip()]
    except OSError:
        return ""
    for l in reversed(lines):
        if "[error]" in l or "failed" in l.lower():
            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."""
    pdb = 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_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})
    if not os.path.exists(e["input_path"]):
        row.update(score.judge(e, {}, "no input"))
        return row
    meta = None
    if e["reference"] == "deposition":
        # 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"]
    cmd = command(binary, e, opts, row["model"])
    row["cmd"] = " ".join(cmd)
    t0 = time.time()
    note = ""
    with open(os.path.join(wd, "run.log"), "w") as log:
        try:
            rc = subprocess.run(cmd, cwd=wd, stdout=log, stderr=subprocess.STDOUT,
                                timeout=opts["timeout"]).returncode
        except subprocess.TimeoutExpired:
            rc, note = None, "timeout"
    row["elapsed_s"] = round(time.time() - t0, 1)
    row["exit_code"] = rc
    if 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 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))
    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 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:
        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
    binary = os.path.join(run_dir, "bin", "rugnux")
    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(), "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"]), "sets": sets}
    write_json(os.path.join(run_dir, "manifest.json"), man)
    print(f"run directory {run_dir}", flush=True)

    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)
    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)
    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,
            "gpulock": site["gpulock"] if a.gpulock else None, "extra_args": a.extra.split(),
            "runs_root": a.runs_root or site["runs_root"]}
    if a.gpulock and not site.get("gpulock"):
        sys.exit("--gpulock: the site config names no gpulock")
    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"])


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)
    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)
    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_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),
             "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, e.g. smoke")
    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("--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")

    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")

    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.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, "list": cmd_list, "compare": cmd_compare, "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()