Jungfraujoch/rugnux_anomalous.py

#!/usr/bin/env python3
"""
rugnux_anomalous.py -- score one or more rugnux merge "arms" by anomalous peak height
at known anomalous scatterers, using SHELXC + ANODE and a model that is placed once
and then held fixed.

WHY THIS EXISTS.  Most of the statistics we reach for when judging a processing change
are compromised for anything that touches partiality:

  * ISa, R_meas and the error-model b are one measurement, dominated by the low
    resolution shells, and none of them is invariant to the uniform intensity rescale
    that every partiality change produces;
  * last-shell R_meas moves with the denominator, i.e. in the wrong direction by
    construction;
  * `rugnux --model` R-free tracks its own zero-information floor, which moves ~22x
    more than R-free itself over the same sweep;
  * per-shell agreement with XDS is biased, because XDS never divides by partiality --
    "divide less" mechanically moves us toward it.

Anomalous peak height has none of those problems.  It is read in units of the map's own
sigma, so a uniform intensity rescale cancels exactly, and it references the STRUCTURE
rather than another program's partiality model.

WHAT IT REPORTS.  For every dataset and every arm: the mean anomalous density at the
known sites, the off-site noise floor, and the baseline sanity gate.  Between arms: the
PAIRED per-site change (the same sites see the same data at every arm, so the honest
test is paired, not a difference of two means).  With numeric arm labels it also prints
the normalised curve and each dataset's optimum, which is how a sigma_M-style scan is
read.

FOUR TRAPS THIS TOOL ENCAPSULATES -- all of them have already cost real time:

  1. The phasing space group comes from the CONFIG, never from the merged file.  I23 and
     I2_1 3 have identical absences, so no data can separate them, and phaser's automatic
     space-group test only tries the ENANTIOMORPH -- which for I23 is itself.  Phasing
     an I-centred cubic crystal in the I23 that rugnux and XDS both report makes four
     perfectly good datasets look signal-free (TFZ 7-11 instead of 30-50).
  2. The model is placed ONCE, from a reference arm, and reused unchanged everywhere.
     Re-phasing per arm lets the model move and contaminates the comparison.
  3. A dataset whose baseline has no signal is EXCLUDED, never reported as a zero.  The
     difference between two noise measurements is not a measurement.
  4. The off-site floor skips peaks on SPECIAL POSITIONS.  A peak at (0,0,0) is a Fourier
     ripple of the phases, not background; leaving it in inflated one dataset's floor to
     6.8 sigma and made a passing dataset look like a failing one.

USAGE

    rugnux_anomalous.py --config <datasets.json> [label=]<arm> [[label=]<arm> ...]

An <arm> is either a directory holding one subdirectory per dataset (the usual rugnux
output layout, `<arm>/<id>/<id>.hkl` + `.mtz`), or a path template containing `{name}`
for anything else, e.g. a scan point `'.../sscan/{name}/s1p2.hkl'`.  The first arm is
the reference unless `--reference` says otherwise; with numeric labels the `1.0` arm is.

The dataset table (ids, phasing space groups, models, crystal identities) lives in the
JSON config OUTSIDE this repository -- see `--write-config-template`.

Self-contained: on first run it builds a private venv, pip-installs gemmi and re-execs.
"""

import os
import sys
import subprocess
import pathlib
import venv

_SELF = os.path.abspath(__file__)


# --------------------------------------------------------------------------- #
#  venv bootstrap (gemmi is the only third-party dependency)
# --------------------------------------------------------------------------- #
def _venv_dir():
    if os.environ.get("RUGNUX_CMP_VENV"):
        return pathlib.Path(os.environ["RUGNUX_CMP_VENV"])
    cache = os.environ.get("XDG_CACHE_HOME") or os.path.join(os.path.expanduser("~"), ".cache")
    return pathlib.Path(cache) / "rugnux_vs_xds" / "venv"


def _bootstrap():
    vdir = _venv_dir()
    py = vdir / "bin" / "python"
    if not py.exists():
        venv.create(vdir, with_pip=True)
    have_gemmi = subprocess.run([str(py), "-c", "import gemmi"],
                                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0
    if not have_gemmi:
        subprocess.run([str(py), "-m", "pip", "install", "-q", "--disable-pip-version-check", "gemmi"],
                       check=True)
    os.environ["RUGNUX_ANO_BOOTSTRAPPED"] = "1"
    os.execv(str(py), [str(py), _SELF, *sys.argv[1:]])


if __name__ == "__main__" and not os.environ.get("RUGNUX_ANO_BOOTSTRAPPED"):
    _bootstrap()

# ---- past this point we are running inside the venv -----------------------
import argparse
import json
import math
import re
import shutil
import statistics
import datetime
import gemmi


CCP4_SETUP_DEFAULT = "/opt/xtal/ccp4-9/bin/ccp4.setup-sh"

# Elements that carry usable anomalous signal in an MX experiment.  Sites are looked up
# by ELEMENT from the PDB, never by atom name -- SG and SD are sulfur, and any
# leading-letters rule gets them wrong.
ANOM = {"S", "CL", "P", "CA", "K", "FE", "ZN", "SE", "BR", "CU", "MN", "NI", "CO", "I"}

# ANODE names its ranked peaks after the SFAC element it was given, so the labels read
# S1, S2, ... on a sulfur case but FE1, MN1, ZN1, ... as soon as the model carries a
# heavier scatterer.  Matching only "S<n>" silently drops the whole peak list on exactly
# those datasets, which then fail the gate as "no peaks" however strong they are.
PEAK_LABEL = re.compile(r"^[A-Za-z]{1,2}\d+$")

# ANODE is run with -s4.0, so its ranked peak list stops at ~4 sigma.  An empty off-site
# list therefore means "below 4", not "zero".
FLOOR_MIN = 4.0

# The baseline sanity gate.  Both conditions must hold on the reference arm.
GATE_PEAK_RATIO = 1.5
GATE_MIN_SITES_5SIG = 3


# --------------------------------------------------------------------------- #
#  config
# --------------------------------------------------------------------------- #
CONFIG_TEMPLATE = {
    "workdir": "/path/to/anomalous_work",
    "ccp4_setup": CCP4_SETUP_DEFAULT,
    "datasets": [
        {
            "id": "<directory name of the dataset under each arm>",
            "crystal": "<identity of the physical crystal; two energies of the same "
                       "crystal share it, so the report can count crystals and datasets "
                       "separately>",
            "spacegroup": "<space group to PHASE in, e.g. I213 -- NOT necessarily the "
                          "one the merged file reports>",
            "model": "<placed (and optionally refined) PDB, reused for every arm>",
            "energy_kev": 0.0,
            "search_model": "<optional: unplaced PDB, only used by --place>",
            "mw": 0,
            "ncopy": 1
        }
    ]
}


def load_config(path):
    cfg = json.loads(pathlib.Path(path).read_text())
    seen = set()
    for d in cfg.get("datasets", []):
        for key in ("id", "spacegroup"):
            if not d.get(key):
                sys.exit(f"config: dataset entry missing '{key}': {d}")
        if d["id"] in seen:
            sys.exit(f"config: duplicate dataset id {d['id']}")
        seen.add(d["id"])
        d.setdefault("crystal", d["id"])
    return cfg


# --------------------------------------------------------------------------- #
#  space groups: the indistinguishable classes
# --------------------------------------------------------------------------- #
def same_absences(sg_a, sg_b, hmax=6):
    """Do two space groups have identical systematic absences?

    13 classes of chiral space group cannot be separated by any absence test.  Most are
    enantiomorph pairs, which phaser covers; I23/I2_1 3 and I222/I2_12_12_1 are NOT --
    they are genuinely different groups, and phaser will never try the other member.
    """
    def absent(sg):
        ops = gemmi.SpaceGroup(sg).operations()
        return {(h, k, l)
                for h in range(-hmax, hmax + 1)
                for k in range(-hmax, hmax + 1)
                for l in range(-hmax, hmax + 1)
                if ops.is_systematically_absent([h, k, l])}
    return absent(sg_a) == absent(sg_b)


def indistinguishable(sg_hm):
    """Every chiral space group that this data could equally well be in."""
    sg = gemmi.SpaceGroup(sg_hm)
    out, seen = [], set()
    for cand in gemmi.spacegroup_table():
        if cand.number in seen:
            continue                      # one setting per group is enough
        if cand.point_group_hm() != sg.point_group_hm() or cand.hm[0] != sg.hm[0]:
            continue
        if cand.number == sg.number or same_absences(sg.hm, cand.hm):
            seen.add(cand.number)
            out.append(cand)
    return out


def shelx_spag(hm):
    """SHELX/ANODE spell space groups without spaces: 'I 21 3' -> 'I213'."""
    return gemmi.SpaceGroup(hm).hm.replace(" ", "")


# --------------------------------------------------------------------------- #
#  ANODE
# --------------------------------------------------------------------------- #
def ccp4_run(cmd, cwd, setup, **kw):
    """Run a CCP4 program, sourcing the CCP4 environment unless it is already there."""
    if shutil.which("shelxc"):
        return subprocess.run(cmd, cwd=str(cwd), **kw)
    line = " ".join(f"'{c}'" for c in cmd)
    return subprocess.run(["bash", "-c", f". '{setup}' >/dev/null 2>&1; exec {line}"],
                          cwd=str(cwd), **kw)


def run_anode(hkl, model, outdir, cell, spag, dmin, setup, find=10):
    """SHELXC (to F_A) then ANODE (anomalous density at the model's atoms).

    Returns the path of the .lsa, or raises RuntimeError.  Everything is written into
    `outdir`, so a run can be inspected afterwards.
    """
    outdir = pathlib.Path(outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(hkl, outdir / "x.hkl")
    # ANISOU records confuse the ANODE PDB reader; the B(iso) it prints is all we need.
    (outdir / "x.pdb").write_text(
        "".join(l for l in pathlib.Path(model).read_text(errors="replace").splitlines(True)
                if not l.startswith("ANISOU")))
    (outdir / "x.inp").write_text(
        f"CELL {cell}\nSPAG {spag}\nSAD x.hkl\nSHEL 999 {dmin:.2f}\nFIND {find}\nNTRY 20\n")

    with open(outdir / "x.inp") as inp, open(outdir / "shelxc.log", "w") as log:
        ccp4_run(["shelxc", "x"], outdir, setup, stdin=inp, stdout=log, stderr=subprocess.STDOUT)
    fa = outdir / "x_fa.hkl"
    if not fa.exists() or not fa.stat().st_size:
        raise RuntimeError(f"shelxc found no usable anomalous differences in this arm "
                           f"(0 reflections written to x_fa.hkl; see {outdir}/shelxc.log)")

    with open(outdir / "anode.log", "w") as log:
        ccp4_run(["anode", "-a", "x"], outdir, setup, stdout=log, stderr=subprocess.STDOUT)
    lsa = outdir / "x.lsa"
    if not lsa.exists() or not lsa.stat().st_size:
        raise RuntimeError(f"anode produced no x.lsa (see {outdir}/anode.log)")
    return lsa


def element_map(pdb):
    """ANODE labels an atom NAME_CHAIN:RESNAMESEQ.  Map that label to its element."""
    st = gemmi.read_structure(str(pdb))
    st.setup_entities()
    return {f"{a.name}_{ch.name}:{r.name}{r.seqid.num}": a.element.name.upper()
            for ch in st[0] for r in ch for a in r}


def on_special_position(xyz, sg):
    """Is this fractional position fixed by more than the identity?

    A peak on a special position -- the cell origin above all -- is a ripple of the
    calculated phases, not a sample of the background, so it must not set the noise
    floor.
    """
    ops = gemmi.SpaceGroup(sg).operations()
    n = 0
    for op in ops:
        d = [a - b for a, b in zip(op.apply_to_xyz(list(xyz)), xyz)]
        if all(abs(v - round(v)) < 1e-3 for v in d):
            n += 1
    return n > 1


def parse_lsa(lsa, pdb, spag):
    """Read an ANODE .lsa: density at every anomalous scatterer, plus the noise floor."""
    elem = element_map(pdb)
    atoms, peaks, mode = [], [], None
    for line in pathlib.Path(lsa).read_text(errors="replace").splitlines():
        if "Density  Occupancy" in line:
            mode = "atoms"
            continue
        if "Strongest unique anomalous peaks" in line:
            mode = "peaks"
            continue
        f = line.split()
        if mode == "atoms" and len(f) == 4:
            try:
                atoms.append((f[3], float(f[0])))
            except ValueError:
                pass
        elif mode == "peaks" and len(f) == 8 and PEAK_LABEL.match(f[0]):
            try:
                peaks.append({"xyz": tuple(float(v) for v in f[1:4]), "h": float(f[4]),
                              "dist": float(f[6]), "near": f[7]})
            except ValueError:
                pass

    # Alternate conformations share one ANODE label (SG ACYS 159 and SG BCYS 159 both print
    # as SG_C:CYS159), so disambiguate rather than silently dropping a site.
    sites = {}
    for label, dens in atoms:
        if elem.get(label, "?") not in ANOM:
            continue
        key, n = label, 1
        while key in sites:
            n += 1
            key = f"{label}#{n}"
        sites[key] = dens

    off, special = [], []
    for p in peaks:
        if elem.get(p["near"], "?") in ANOM and p["dist"] < 1.5:
            continue                     # the site itself
        (special if on_special_position(p["xyz"], spag) else off).append(p["h"])

    h = list(sites.values())
    return {
        "sites": sites,
        "n_sites": len(h),
        "mean": statistics.fmean(h) if h else None,
        "sem": statistics.stdev(h) / math.sqrt(len(h)) if len(h) > 1 else None,
        "top_site": max(h) if h else None,
        "n4": sum(1 for v in h if v > 4.0),
        "n5": sum(1 for v in h if v > 5.0),
        "floor": max(off) if off else None,
        "special": max(special) if special else None,
        "top_peak": peaks[0]["h"] if peaks else None,
    }


def gate(m):
    """The baseline sanity gate, evaluated on the reference arm.

    A dataset that fails it carries no information about any arm, so it is EXCLUDED --
    reporting it as "no change" would present the difference of two noise measurements
    as a measurement.
    """
    if not m or m.get("top_peak") is None:
        return False, "no peaks"
    floor = m.get("floor") or FLOOR_MIN
    if m["top_peak"] <= GATE_PEAK_RATIO * floor:
        return False, f"top peak {m['top_peak']:.2f} <= {GATE_PEAK_RATIO}x floor {floor:.2f}"
    if m["n5"] < GATE_MIN_SITES_5SIG:
        return False, f"only {m['n5']} site(s) over 5 sigma"
    return True, ""


# --------------------------------------------------------------------------- #
#  arms
# --------------------------------------------------------------------------- #
def parse_arm(spec):
    """'label=path' or 'path'.  A path may contain {name} for the dataset id."""
    if "=" in spec and not os.path.exists(spec.split("=", 1)[0]):
        label, path = spec.split("=", 1)
    else:
        label, path = pathlib.Path(spec.rstrip("/")).name, spec
    return label, path


def resolve_hkl(path, name):
    """Locate the merged reflection file for one dataset under one arm."""
    if "{name}" in path:
        p = pathlib.Path(path.format(name=name))
        return p if p.exists() else None
    d = pathlib.Path(path)
    # rugnux writes both <prefix>.hkl and <prefix>_01.hkl; <prefix>.hkl is the stable one.
    for cand in (d / name / f"{name}.hkl", d / f"{name}.hkl"):
        if cand.exists():
            return cand
    hits = [h for h in sorted((d / name).glob("*.hkl")) if not h.name.endswith("_fa.hkl")]
    return hits[0] if hits else None


def cell_and_dmin(hkl):
    """Cell and resolution from the .mtz rugnux wrote next to the .hkl."""
    mtz = hkl.with_suffix(".mtz")
    if not mtz.exists():
        raise RuntimeError(f"no {mtz.name} beside {hkl.name} (needed for the cell)")
    m = gemmi.read_mtz_file(str(mtz))
    c = m.cell
    return (f"{c.a:.4f} {c.b:.4f} {c.c:.4f} {c.alpha:.3f} {c.beta:.3f} {c.gamma:.3f}",
            m.resolution_high(), m.spacegroup.hm)


# --------------------------------------------------------------------------- #
#  placing the model (run once, then the result goes in the config)
# --------------------------------------------------------------------------- #
def place(ds, arm_path, workdir, setup):
    """Place `search_model` with phaser, trying EVERY space group the data cannot
    distinguish, and keep the best.  Print what to put in the config."""
    name = ds["id"]
    hkl = resolve_hkl(arm_path, name)
    if hkl is None:
        return f"{name}: no reflection file in the reference arm"
    mtz = hkl.with_suffix(".mtz")
    if not ds.get("search_model"):
        return f"{name}: config has no 'search_model'"

    best = None
    for cand in indistinguishable(ds.get("spacegroup") or gemmi.read_mtz_file(str(mtz)).spacegroup.hm):
        out = pathlib.Path(workdir) / name / f"mr_{shelx_spag(cand.hm)}"
        out.mkdir(parents=True, exist_ok=True)
        m = gemmi.read_mtz_file(str(mtz))
        m.spacegroup = cand
        m.write_to_file(str(out / "in.mtz"))
        script = (f"TITLE place\nMODE MR_AUTO\nHKLIN in.mtz\nLABIN F=F SIGF=SIGF\n"
                  f"ENSEMBLE ens PDBFILE {ds['search_model']} IDENTITY 0.95\n"
                  f"COMPOSITION PROTEIN MW {ds.get('mw', 0)} NUM {ds.get('ncopy', 1)}\n"
                  f"SEARCH ENSEMBLE ens NUM {ds.get('ncopy', 1)}\nRESOLUTION 3.0\nROOT placed\n")
        (out / "phaser.inp").write_text(script)
        with open(out / "phaser.inp") as inp, open(out / "phaser.log", "w") as log:
            ccp4_run(["phaser"], out, setup, stdin=inp, stdout=log, stderr=subprocess.STDOUT)
        sol = out / "placed.sol"
        llg = tfz = None
        if sol.exists():
            for line in sol.read_text(errors="replace").splitlines():
                for tok in line.split():
                    if tok.startswith("LLG="):
                        llg = float(tok[4:])
                    elif tok.startswith("TFZ=="):
                        tfz = float(tok[5:])
        pdb = out / "placed.1.pdb"
        print(f"    {name:22s} {shelx_spag(cand.hm):10s} LLG={llg if llg else 0:8.0f} "
              f"TFZ={tfz if tfz else 0:6.1f} {'ok' if pdb.exists() else 'NO SOLUTION'}")
        if pdb.exists() and (best is None or (llg or 0) > best[0]):
            best = ((llg or 0), cand.hm, pdb)
    if best is None:
        return f"{name}: no solution in any space group"
    print(f"    -> config: \"spacegroup\": \"{shelx_spag(best[1])}\", \"model\": \"{best[2]}\"")
    return None


# --------------------------------------------------------------------------- #
#  statistics
# --------------------------------------------------------------------------- #
def _betacf(a, b, x, itmax=200, eps=3e-14):
    qab, qap, qam = a + b, a + 1.0, a - 1.0
    c, d = 1.0, 1.0 - qab * x / qap
    d = 1.0 / (d if abs(d) > 1e-30 else 1e-30)
    h = d
    for m in range(1, itmax + 1):
        m2 = 2 * m
        for aa in (m * (b - m) * x / ((qam + m2) * (a + m2)),
                   -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2))):
            d = 1.0 + aa * d
            d = 1.0 / (d if abs(d) > 1e-30 else 1e-30)
            c = 1.0 + aa / (c if abs(c) > 1e-30 else 1e-30)
            h *= d * c
        if abs(d * c - 1.0) < eps:
            break
    return h


def _betai(a, b, x):
    """Regularized incomplete beta I_x(a,b) (Numerical Recipes)."""
    if x <= 0.0:
        return 0.0
    if x >= 1.0:
        return 1.0
    front = math.exp(math.lgamma(a + b) - math.lgamma(a) - math.lgamma(b)
                     + a * math.log(x) + b * math.log(1.0 - x))
    if x < (a + 1.0) / (a + b + 2.0):
        return front * _betacf(a, b, x) / a
    return 1.0 - front * _betacf(b, a, 1.0 - x) / b


def t_pvalue(t, df):
    """Two-sided p for Student's t."""
    if df <= 0:
        return None
    return _betai(df / 2.0, 0.5, df / (df + t * t))


def paired(a_sites, b_sites):
    """Paired per-site change b - a.  The same sites see the same data in both arms, so
    a difference of means throws away most of the information (and most of the power)."""
    keys = sorted(set(a_sites) & set(b_sites))
    d = [b_sites[k] - a_sites[k] for k in keys]
    if len(d) < 2:
        return None
    mean = statistics.fmean(d)
    sem = statistics.stdev(d) / math.sqrt(len(d))
    t = mean / sem if sem > 0 else 0.0
    return {"n": len(d), "mean": mean, "sem": sem, "t": t, "p": t_pvalue(t, len(d) - 1)}


# --------------------------------------------------------------------------- #
#  measuring
# --------------------------------------------------------------------------- #
def measure(ds, label, path, workdir, setup, dmin, cache, find, refresh):
    """One (dataset, arm) point, cached on disk -- ANODE is fast but not free."""
    name = ds["id"]
    hkl = resolve_hkl(path, name)
    if hkl is None:
        return None, "no reflection file"
    cell, own_dmin, merged_sg = cell_and_dmin(hkl)
    spag = shelx_spag(ds["spacegroup"])
    key = f"{name}/{label}"
    stamp = [str(hkl), int(hkl.stat().st_mtime), ds["model"], spag,
             round(dmin or own_dmin, 3), cell, find]
    if not refresh and key in cache and cache[key].get("stamp") == stamp:
        return cache[key]["m"], None

    out = pathlib.Path(workdir) / name / f"anode_{label}"
    try:
        lsa = run_anode(hkl, ds["model"], out, cell, spag, dmin or own_dmin, setup, find)
        m = parse_lsa(lsa, out / "x.pdb", ds["spacegroup"])
    except Exception as e:
        return None, str(e)
    m["merged_sg"] = merged_sg
    m["own_dmin"] = own_dmin
    cache[key] = {"stamp": stamp, "m": m}
    return m, None


# --------------------------------------------------------------------------- #
#  rendering
# --------------------------------------------------------------------------- #
def fmt(v, w, prec=2, dash="-"):
    return dash.rjust(w) if v is None else f"{v:{w}.{prec}f}"


def print_baseline(cfg, results, ref, order):
    print(f"  BASELINE ({ref}) -- the gate: top peak > {GATE_PEAK_RATIO}x the off-site floor "
          f"AND >= {GATE_MIN_SITES_5SIG} sites over 5 sigma")
    print()
    print(f"  {'dataset':24s} {'keV':>5s} {'dmin':>5s} {'SG(merge)':>10s} {'SG(phase)':>10s} "
          f"{'n':>4s} {'mean':>6s} {'s.e.':>5s} {'top':>6s} {'>4s':>4s} {'>5s':>4s} "
          f"{'floor':>6s} {'spec':>6s} {'peak':>6s} {'ratio':>6s}  gate")
    for name in order:
        ds, m = cfg[name], results.get((name, ref))
        if not m:
            print(f"  {name:24s}  (missing in the reference arm)")
            continue
        ok, why = gate(m)
        floor = m.get("floor")
        print(f"  {name:24s} {fmt(ds.get('energy_kev'), 5, 2)} {fmt(m['own_dmin'], 5, 2)} "
              f"{shelx_spag(m['merged_sg']):>10s} {shelx_spag(ds['spacegroup']):>10s} "
              f"{m['n_sites']:4d} {fmt(m['mean'], 6)} {fmt(m['sem'], 5)} {fmt(m['top_site'], 6)} "
              f"{m['n4']:4d} {m['n5']:4d} "
              f"{(f'{floor:6.2f}' if floor else '  <4.0')} {fmt(m.get('special'), 6)} "
              f"{fmt(m['top_peak'], 6)} {fmt((m['top_peak'] or 0) / (floor or FLOOR_MIN), 6)}  "
              f"{'PASS' if ok else 'EXCLUDED: ' + why}")
    print()


def print_arms(cfg, results, ref, arms, usable):
    """Mean site height per arm, and the paired per-site change against the reference."""
    labels = [l for l, _ in arms]
    print(f"  MEAN SITE HEIGHT (sigma) and PAIRED per-site change vs {ref}")
    print()
    head = f"  {'dataset':24s} " + " ".join(f"{l:>16s}" for l in labels)
    print(head)
    for name in usable:
        cells = []
        for label in labels:
            m = results.get((name, label))
            if not m:
                cells.append(f"{'-':>16s}")
                continue
            if label == ref:
                cells.append(f"{m['mean']:8.2f}          ")
            else:
                d = paired(results[(name, ref)]["sites"], m["sites"])
                cells.append(f"{m['mean']:8.2f} {d['mean']:+6.2f} " if d else f"{m['mean']:8.2f}      ")
        print(f"  {name:24s} " + " ".join(cells))
    print()
    print(f"  paired change, mean +- s.e. over the sites of one dataset (t, two-sided p):")
    for label in labels:
        if label == ref:
            continue
        print(f"    vs {label}:")
        per_dataset = {}
        for name in usable:
            m = results.get((name, label))
            d = paired(results[(name, ref)]["sites"], m["sites"]) if m else None
            if not d:
                print(f"      {name:24s}  -")
                continue
            per_dataset[name] = d["mean"]
            print(f"      {name:24s} n={d['n']:3d}  {d['mean']:+6.3f} +- {d['sem']:5.3f} "
                  f"sigma   t={d['t']:+6.2f}  p={d['p']:.3f}"
                  f"{'' if d['p'] >= 0.05 else '   *'}")
        if per_dataset:
            # Two datasets of the SAME crystal are not two independent votes, so pool the
            # crystal first and count crystals, not datasets.
            by_crystal = {}
            for name, v in per_dataset.items():
                by_crystal.setdefault(cfg[name]["crystal"], []).append(v)
            per_crystal = [statistics.fmean(v) for v in by_crystal.values()]
            up = sum(1 for v in per_crystal if v > 0)
            mean = statistics.fmean(per_crystal)
            sem = (statistics.stdev(per_crystal) / math.sqrt(len(per_crystal))
                   if len(per_crystal) > 1 else 0.0)
            print(f"      {'-> over crystals':24s} n={len(per_crystal):3d}  {mean:+6.3f} "
                  f"+- {sem:5.3f} sigma   {up} better / {len(per_crystal) - up} worse")
    print()


def print_curve(cfg, results, ref, arms, usable):
    """With numeric arm labels the arms are a scan: print the normalised curve, each
    dataset's own optimum and the pooled optimum.  A per-dataset optimum at the edge of
    the grid is not an optimum -- extend the grid instead of reporting it."""
    facs = [float(l) for l, _ in arms]
    labels = [l for l, _ in arms]
    rel = {}
    for name in usable:
        base = results[(name, ref)]["mean"]
        row = [results[(name, l)]["mean"] / base if results.get((name, l)) else None
               for l in labels]
        if all(v is not None for v in row):
            rel[name] = row
    if not rel:
        return
    print("  SCAN -- mean site height normalised to the reference arm")
    print()
    print(f"  {'dataset':24s} " + " ".join(f"{f:6.2f}" for f in facs) + "     opt")
    for name, row in rel.items():
        i = row.index(max(row))
        opt = facs[i]
        # A per-dataset optimum is only worth quoting if that arm actually beats the
        # reference on the paired per-site test; otherwise it is where the noise happened
        # to peak, and it will move when anything else about the analysis moves.
        d = paired(results[(name, ref)]["sites"], results[(name, labels[i])]["sites"])
        flags = []
        if opt in (facs[0], facs[-1]):
            flags.append("grid edge")
        if opt != float(ref) and (d is None or abs(d["t"]) < 2.0):
            flags.append("not significant vs ref")
        note = ("  (" + "; ".join(flags) + ")") if flags else ""
        print(f"  {name:24s} " + " ".join(f"{v:6.3f}" for v in row) + f"  {opt:6.2f}{note}")
    n = len(rel)
    mean = [statistics.fmean(r[i] for r in rel.values()) for i in range(len(facs))]
    sem = [statistics.stdev([r[i] for r in rel.values()]) / math.sqrt(n) if n > 1 else 0.0
           for i in range(len(facs))]
    print(f"  {'pooled mean (n=' + str(n) + ')':24s} " + " ".join(f"{v:6.3f}" for v in mean))
    print(f"  {'s.e.':24s} " + " ".join(f"{v:6.3f}" for v in sem))
    best = mean.index(max(mean))
    at_ref = mean[facs.index(float(ref))] if float(ref) in facs else None
    print()
    print(f"  pooled optimum {facs[best]:.2f}"
          + (f";  the reference arm ({ref}) is {100 * (1 - at_ref / mean[best]):.2f}% below it"
             if at_ref else ""))
    opts = sorted(facs[r.index(max(r))] for r in rel.values())
    print(f"  per-dataset optima {opts}  median {statistics.median(opts):.2f}")
    print()


def print_counts(cfg, usable, excluded):
    crystals = sorted({cfg[n]["crystal"] for n in usable})
    sites = sum(cfg[n].get("_n_sites", 0) for n in usable)
    print(f"  n = {len(usable)} dataset(s) from {len(crystals)} distinct crystal(s), "
          f"{sites} sites in total")
    for c in crystals:
        members = [n for n in usable if cfg[n]["crystal"] == c]
        if len(members) > 1:
            print(f"      the same crystal appears as {len(members)} datasets: "
                  f"{', '.join(members)} -- they are NOT independent")
    if excluded:
        print(f"  excluded (baseline gate): {', '.join(excluded)}")
    print()


# --------------------------------------------------------------------------- #
#  main
# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser(
        description="Score rugnux merge arms by anomalous peak height (SHELXC + ANODE).")
    ap.add_argument("arms", nargs="*", metavar="[label=]arm",
                    help="rugnux output directory, or a path template containing {name}")
    ap.add_argument("--config", help="dataset table (JSON), kept OUTSIDE this repository")
    ap.add_argument("--write-config-template", metavar="PATH",
                    help="write a commented config skeleton and exit")
    ap.add_argument("--reference", help="arm to gate on and compare against "
                                        "(default: the 1.0 arm if labels are numeric, else the first)")
    ap.add_argument("--workdir", help="where ANODE runs go (default: config 'workdir')")
    ap.add_argument("--only", help="comma-separated substrings; only matching datasets are used")
    ap.add_argument("--find", type=int, default=10, help="ANODE peaks to list (default: 10)")
    ap.add_argument("--per-arm-dmin", action="store_true",
                    help="use each arm's own resolution instead of the reference arm's "
                         "(the map resolution then differs between arms)")
    ap.add_argument("--refresh", action="store_true", help="re-run ANODE even if cached")
    ap.add_argument("--place", action="store_true",
                    help="place each dataset's 'search_model' from the reference arm, trying "
                         "every space group the data cannot distinguish, and exit")
    ap.add_argument("--ccp4-setup", default=os.environ.get("CCP4_SETUP", CCP4_SETUP_DEFAULT))
    args = ap.parse_args()

    if args.write_config_template:
        pathlib.Path(args.write_config_template).write_text(
            json.dumps(CONFIG_TEMPLATE, indent=2) + "\n")
        print(f"wrote {args.write_config_template}")
        return
    if not args.config:
        sys.exit("--config is required (the dataset table lives outside this repository)")
    if not args.arms:
        sys.exit("give at least one arm")

    cfg_file = load_config(args.config)
    setup = cfg_file.get("ccp4_setup", args.ccp4_setup)
    workdir = pathlib.Path(args.workdir or cfg_file.get("workdir") or "rugnux_anomalous")
    workdir.mkdir(parents=True, exist_ok=True)

    filters = [s for s in (args.only or "").split(",") if s]
    datasets = [d for d in cfg_file["datasets"]
                if not filters or any(f in d["id"] for f in filters)]
    if not datasets:
        sys.exit("no datasets selected")
    cfg = {d["id"]: d for d in datasets}
    order = [d["id"] for d in datasets]

    arms = [parse_arm(a) for a in args.arms]
    numeric = all(_isfloat(l) for l, _ in arms)
    if args.reference:
        ref = args.reference
    elif numeric and any(float(l) == 1.0 for l, _ in arms):
        ref = next(l for l, _ in arms if float(l) == 1.0)
    else:
        ref = arms[0][0]
    if ref not in [l for l, _ in arms]:
        sys.exit(f"reference arm '{ref}' is not among the arms given")

    if args.place:
        print("\n  placing models from the reference arm "
              f"({dict(arms)[ref]}) -- one placement per dataset, reused everywhere\n")
        for d in datasets:
            err = place(d, dict(arms)[ref], workdir, setup)
            if err:
                print(f"    {err}")
        print()
        return

    for d in datasets:
        if not d.get("model") or not os.path.exists(d["model"]):
            sys.exit(f"{d['id']}: config 'model' missing or not found ({d.get('model')}) -- "
                     f"run with --place first")

    cache_file = workdir / "anode_cache.json"
    cache = json.loads(cache_file.read_text()) if cache_file.exists() else {}

    # The reference arm sets the resolution for every arm, so the maps are computed on
    # the same data range and only the intensities differ.
    dmin = {}
    for name in order:
        hkl = resolve_hkl(dict(arms)[ref], name)
        dmin[name] = None if args.per_arm_dmin or hkl is None else cell_and_dmin(hkl)[1]

    results, failures = {}, []
    for name in order:
        for label, path in arms:
            m, err = measure(cfg[name], label, path, workdir, setup, dmin[name],
                             cache, args.find, args.refresh)
            if m:
                results[(name, label)] = m
            elif err != "no reflection file":
                failures.append(f"{name}/{label}: {err}")
    cache_file.write_text(json.dumps(cache, indent=1))

    for name in order:
        m = results.get((name, ref))
        cfg[name]["_n_sites"] = m["n_sites"] if m else 0

    usable, excluded = [], []
    for name in order:
        m = results.get((name, ref))
        ok, _ = gate(m) if m else (False, "")
        (usable if ok else excluded).append(name)

    today = datetime.date.today().isoformat()
    print()
    print(f"  anomalous peak height at known sites   .   {today}")
    print(f"  config {args.config}   .   ANODE work in {workdir}")
    print(f"  arms: " + ",  ".join(f"{l} = {p}" for l, p in arms))
    print(f"  Peak height is in map sigma, so a uniform intensity rescale cancels exactly -- "
          f"which is why")
    print(f"  this and not XDS agreement or R-free is the arbiter for a partiality change.")
    print(f"  The model is placed once and held fixed; the space group is the config's, "
          f"not the merged file's.")
    print()
    print_baseline(cfg, results, ref, order)
    print_counts(cfg, usable, excluded)
    if usable:
        print_arms(cfg, results, ref, arms, usable)
        if numeric:
            print_curve(cfg, results, ref, arms, usable)
    if failures:
        print("  failures:")
        for f in failures:
            print(f"    {f}")
        print()


def _isfloat(s):
    try:
        float(s)
        return True
    except ValueError:
        return False


if __name__ == "__main__":
    main()