#!/usr/bin/env python3
"""model_check.py -- R-factors of the deposited model against Rugnux's merged data.

An external check for the open arm, run AFTER Rugnux has written its MTZ, so it cannot bias the
processing (Rugnux is never given the model). For one set it:

  * fetches the entry's mmCIF (and, for the baseline, its structure factors) from RCSB into a
    cache, and reads the published R-work / R-free / d_min from the entry's `_refine`;
  * matches the data to the model's setting: the two lattices are compared through their Niggli
    cells, every change of basis that carries one onto the other is tried (this covers an
    alternative setting, a permuted or centred-vs-primitive description and the merohedral
    indexing ambiguity), and the one whose reindexed data agree best with the model's Fcalc on a
    mid-resolution shell is taken;
  * re-expresses the data in the DEPOSITED space group: expanded to P1 with Rugnux's own group,
    reindexed, reduced to the deposited group's asymmetric unit (an over-call just splits
    equivalents apart, an under-call averages them) and given the DEPOSITED cell;
  * has REFMAC compute R-work / R-free of the model AS DEPOSITED, with its bulk solvent and
    overall anisotropic scaling, at the lower of the two high-resolution limits. This is REFMAC's
    rigid-body mode read at its first cycle, before any shift is applied: rigid body reads no
    restraint dictionaries, so no ligand is stripped, no atom is moved and no cell is imposed on
    the model. Refining is deliberately not done - the data were moved into the model's frame, so
    there is nothing for a rigid body to correct, and at low resolution REFMAC's rigid body can
    walk a correct model away (R 0.30 -> 0.49 over ten cycles on one 5 A set);
  * scores the SAME model with the SAME protocol against the depositor's own structure factors
    and free set, as the baseline for "same method, other data", and - when that free set can be
    carried over - also scores Rugnux's data on the depositor's free reflections. The depositor's
    data go through the same setting choice as Rugnux's (they can be indexed in the other branch
    of a merohedral ambiguity); where only intensities were deposited, amplitudes come from
    ctruncate's French-Wilson;
  * scores an entry that declares twinning (_pdbx_reflns_twin) with REFMAC's twin refinement,
    both data sets alike, and keeps the untwinned numbers beside it;
  * and, as a second protocol (refined_*), refines the same deposited model against each data set
    at its OWN resolution (restrained, same protocol for both) with R-free on a SHARED free set:
    the depositor's free reflections present in both data sets (see refined_shared).

The first-cycle numbers are taken at d_min_used = max(Rugnux's d_min, the deposited d_min).

On free sets: the depositor's model was refined against every reflection except THEIR free set,
so most of Rugnux's free reflections were work reflections to it. `rfree` (Rugnux's own flags)
therefore reads like a work R. `rfree_depflags` (depositor's flags on Rugnux's data) is the
number that is comparable with `rfree_depdata` and, loosely, with `rfree_deposited`.

  model_check.py rugnux.mtz 1abc [--workdir DIR] [--cache DIR] [--no-baseline]
"""
import argparse
import itertools
import json
import os
import re
import subprocess
import time

import gemmi
import numpy as np

DEFAULT_CACHE = "/data/battery/pdb_cache"
CCP4_SETUP = os.environ.get("CCP4_SETUP", "/opt/xtal/ccp4-9/bin/ccp4.setup-sh")
RCSB = "https://files.rcsb.org/download/"
PDB_ID = re.compile(r"^[0-9][a-z0-9]{3}$")
MAX_ATOMS = 200000          # REFMAC time grows with atoms x reflections
CELL_TOL = 0.05             # Niggli cells compared loosely on purpose; the R score decides
ANGLE_TOL = 3.0
REFINE_CYCLES = 10          # restrained refinement of the shared-free-set protocol


# ------------------------------------------------------------------ deposition

def fetch(url, dest):
    if not os.path.exists(dest) or os.path.getsize(dest) == 0:
        tmp = dest + ".part"
        r = subprocess.run(["curl", "-sfL", "--max-time", "300", "-o", tmp, url])
        if r.returncode != 0 or not os.path.exists(tmp):
            if os.path.exists(tmp):
                os.remove(tmp)
            return None
        os.replace(tmp, dest)
    return dest


def deposition(pdb_id, cache_dir):
    """Coordinates path and the published refinement statistics, cached as <id>.json."""
    os.makedirs(cache_dir, exist_ok=True)
    meta_path = os.path.join(cache_dir, pdb_id + ".json")
    xyz = fetch(RCSB + pdb_id + ".cif.gz", os.path.join(cache_dir, pdb_id + ".cif.gz"))
    if xyz is None:
        return None
    if os.path.exists(meta_path):
        return json.load(open(meta_path))
    block = gemmi.cif.read(xyz).sole_block()
    methods = [m.strip("'\"") for m in block.find_values("_exptl.method")]
    rows = block.find("_refine.", ["pdbx_refine_id", "?ls_R_factor_R_free", "?ls_R_factor_R_work",
                                   "?ls_d_res_high"])
    ref = {}
    for row in rows:       # a joint X-ray/neutron entry has one row per method; take the X-ray one
        if "X-RAY" in row.str(0).upper():
            ref = {k: (float(row[i]) if row.has(i) and gemmi.cif.is_null(row[i]) is False else None)
                   for k, i in (("rfree", 1), ("rwork", 2), ("dmin", 3))}
    meta = {"id": pdb_id, "methods": methods, "rfree": ref.get("rfree"), "rwork": ref.get("rwork"),
            "dmin": ref.get("dmin"), "xyz": xyz}
    json.dump(meta, open(meta_path, "w"), indent=1)
    return meta


# ------------------------------------------------------------------ lattice / setting matching

def metric(cell):
    a, b, c, al, be, ga = cell.parameters
    ca, cb, cg = np.cos(np.radians([al, be, ga]))
    return np.array([[a * a, a * b * cg, a * c * cb],
                     [a * b * cg, b * b, b * c * ca],
                     [a * c * cb, b * c * ca, c * c]])


def niggli(cell, sg):
    """Matrix M with (reduced primitive basis) = (conventional basis) @ M, and the reduced cell."""
    gv = gemmi.GruberVector(cell, sg.centring_type(), True)
    gv.niggli_reduce()
    op = gv.change_of_basis
    return np.array(op.rot, float) / op.DEN, gv.get_cell()


UNIMODULAR = np.array([u for u in itertools.product((-1, 0, 1), repeat=9)
                       if round(np.linalg.det(np.reshape(u, (3, 3)))) == 1], float).reshape(-1, 3, 3)


def basis_changes(cell_d, sg_d, cell_m, sg_m):
    """Every T (hkl_model = hkl_data @ T) whose lattice matches the model's, one per distinct
    reindexing modulo the model's point group (the member nearest the identity names it)."""
    md, red_d = niggli(cell_d, sg_d)
    mm, red_m = niggli(cell_m, sg_m)
    gd, gm = metric(red_d), metric(red_m)
    g = np.einsum("nji,jk,nkl->nil", UNIMODULAR, gd, UNIMODULAR)
    lengths = np.sqrt(np.einsum("nii->ni", g))
    lm = np.sqrt(np.diag(gm))
    ok = np.all(np.abs(lengths / lm - 1) < CELL_TOL, axis=1)
    cos = lambda g_, i, j, l_: g_[..., i, j] / (l_[..., i] * l_[..., j])
    for i, j in ((1, 2), (0, 2), (0, 1)):
        ang = np.degrees(np.arccos(np.clip(cos(g, i, j, lengths), -1, 1)))
        ok &= np.abs(ang - np.degrees(np.arccos(cos(gm, i, j, lm)))) < ANGLE_TOL
    asu, gops = gemmi.ReciprocalAsu(sg_m), sg_m.operations()
    probe = np.array([[1, 2, 3], [3, 1, 7], [5, 11, 2]]) * 12
    found = {}
    for u in UNIMODULAR[ok]:
        t = md @ u @ np.linalg.inv(mm)
        key = tuple(tuple(asu.to_asu([int(round(x)) for x in h], gops)[0]) for h in probe @ t)
        key = tuple(sorted(key))
        if key not in found or np.abs(t - np.eye(3)).sum() < np.abs(found[key] - np.eye(3)).sum():
            found[key] = t
    return list(found.values())


def op_string(t):
    """hkl_model = hkl_data @ T written as the usual h,k,l triplet."""
    from fractions import Fraction
    out = []
    for j in range(3):
        terms = ""
        for i, name in enumerate("hkl"):
            f = Fraction(t[i, j]).limit_denominator(12)
            if f == 0:
                continue
            sign = "-" if f < 0 else ("+" if terms else "")
            mag = abs(f)
            terms += sign + (name if mag == 1 else f"{mag}{name}" if mag.denominator == 1
                             else f"{mag.numerator}{name}/{mag.denominator}")
        out.append(terms)
    return ",".join(out)


# ------------------------------------------------------------------ reflection data

def hkl_key(hkl):
    h = hkl.astype(np.int64) + 1024
    return (h[:, 0] * 2048 + h[:, 1]) * 2048 + h[:, 2]


def to_group(hkl, f, sigf, flags, cell, sg, src):
    """Reduce reflections to the asymmetric unit of sg; independent measurements of one
    reflection are averaged (1/sigma^2), copies of the same one (same `src`) are not. A reflection
    is free if any contributor is."""
    mtz = gemmi.Mtz()
    mtz.spacegroup, mtz.cell = sg, cell
    mtz.add_dataset("d")
    for lab, typ in (("H", "H"), ("K", "H"), ("L", "H"), ("S", "R")):
        mtz.add_column(lab, typ)
    mtz.set_data(np.column_stack([hkl, np.arange(len(hkl))]).astype(np.float32))
    mtz.ensure_asu()
    a = np.array(mtz, copy=False)
    order = a[:, 3].astype(np.int64)
    hkl, f, sigf, flags, src = a[:, :3].astype(int), f[order], sigf[order], flags[order], src[order]
    keep = ~sg.operations().systematic_absences(hkl)
    hkl, f, sigf, flags, src = hkl[keep], f[keep], sigf[keep], flags[keep], src[keep]
    key = hkl_key(hkl)
    _, first = np.unique(np.column_stack([key, src]), axis=0, return_index=True)
    hkl, f, sigf, flags, key = hkl[first], f[first], sigf[first], flags[first], key[first]
    _, idx, inv = np.unique(key, return_index=True, return_inverse=True)
    w = 1.0 / np.maximum(sigf, 1e-6) ** 2
    sw = np.bincount(inv, w)
    fm = np.bincount(inv, w * f) / sw
    flag = np.minimum.reduceat(flags[np.argsort(inv, kind="stable")],
                               np.r_[0, np.cumsum(np.bincount(inv))[:-1]])
    return hkl[idx], fm, 1.0 / np.sqrt(sw), flag


def read_rugnux(mtz_path):
    mtz = gemmi.read_mtz_file(mtz_path)
    labels = mtz.column_labels()
    need = ("F", "SIGF", "FreeR_flag")
    if not all(l in labels for l in need):
        raise RuntimeError(f"MTZ lacks one of {need}")
    return mtz


def expand_p1(mtz):
    """Rugnux's merge expanded to P1 (one hemisphere), with each row's source index."""
    m = gemmi.Mtz()
    m.spacegroup, m.cell = mtz.spacegroup, mtz.cell
    m.add_dataset("d")
    for lab, typ in (("H", "H"), ("K", "H"), ("L", "H"), ("S", "R")):
        m.add_column(lab, typ)
    m.set_data(np.column_stack([mtz.make_miller_array(), np.arange(mtz.nreflections)]).astype(np.float32))
    m.expand_to_p1()
    a = np.array(m, copy=False)
    return a[:, :3].astype(int), a[:, 3].astype(int)


def reindexed(mtz, t, cell_m, sg_m):
    f = np.array(mtz.column_with_label("F"), copy=False).astype(float)
    s = np.array(mtz.column_with_label("SIGF"), copy=False).astype(float)
    fl = np.array(mtz.column_with_label("FreeR_flag"), copy=False).astype(int)
    good = np.isfinite(f) & np.isfinite(s) & (s > 0)
    hkl, src = expand_p1(mtz)
    hm = hkl @ t
    integral = np.all(np.abs(hm - np.round(hm)) < 1e-3, axis=1) & good[src]
    hm, src = np.round(hm[integral]).astype(int), src[integral]
    return to_group(hm, f[src], s[src], fl[src], cell_m, sg_m, src), integral.mean()


def fcalc(st, d_min):
    dc = gemmi.DensityCalculatorX()
    dc.d_min = d_min
    dc.grid.setup_from(st)
    dc.set_refmac_compatible_blur(st[0])
    dc.put_model_density_on_grid(st[0])
    asu = gemmi.transform_map_to_f_phi(dc.grid).prepare_asu_data(dmin=d_min, unblur=dc.blur)
    return asu.miller_array, np.abs(asu.value_array)


def shell_r(hkl, fo, fc_hkl, fc, cell, d_lo, d_hi):
    """R after an overall scale and B, on the shell used only to choose the setting (low
    resolution excluded: unmodelled bulk solvent dominates there)."""
    _, i, j = np.intersect1d(hkl_key(hkl), hkl_key(fc_hkl), return_indices=True)
    fo, fc = fo[i], fc[j]
    d = cell.calculate_d_array(hkl[i])
    sel = (d <= d_lo) & (d >= d_hi) & (fo > 0) & (fc > 0)
    if sel.sum() < 50:
        return None
    fo, fc, s2 = fo[sel], fc[sel], 1.0 / d[sel] ** 2
    slope, lnk = np.polyfit(s2, np.log(fo / fc), 1)
    fm = np.exp(lnk + slope * s2) * fc
    return float(np.sum(np.abs(fo - fm)) / np.sum(fo))


# ------------------------------------------------------------------ REFMAC

_ENV = None


def ccp4_env():
    global _ENV
    if _ENV is None:
        out = subprocess.run(["bash", "-c", f"source {CCP4_SETUP} >/dev/null 2>&1; env -0"],
                             capture_output=True).stdout
        _ENV = dict(kv.split("=", 1) for kv in out.decode().split("\0") if "=" in kv)
    return _ENV


def make_mtz(hkl, f, sigf, flags, cell, sg, dep_flags=None, labels=("F", "SIGF")):
    mtz = gemmi.Mtz(with_base=True)
    mtz.spacegroup, mtz.cell = sg, cell
    mtz.add_dataset("data")
    cols = [hkl.astype(float), f, sigf, flags]
    types = ("F", "Q") if labels[0] == "F" else ("J", "Q")
    for lab, typ in ((labels[0], types[0]), (labels[1], types[1]), ("FreeR_flag", "I")):
        mtz.add_column(lab, typ)
    if dep_flags is not None:
        mtz.add_column("FreeR_dep", "I")
        cols.append(dep_flags)
    mtz.set_data(np.column_stack(cols).astype(np.float32))
    return mtz


def run_refmac(workdir, tag, mtz, xyz, keywords):
    """REFMAC's R-work and R-free as it prints them, one per cycle."""
    log = os.path.join(workdir, tag + ".log")
    with open(log, "w") as fh:
        subprocess.run(["refmac5", "hklin", mtz, "xyzin", xyz,
                        "hklout", os.path.join(workdir, tag + "_out.mtz"),
                        "xyzout", os.path.join(workdir, tag + "_out.pdb")],
                       input=keywords.encode(), stdout=fh, stderr=subprocess.STDOUT,
                       cwd=workdir, env=ccp4_env())
    text = open(log, errors="replace").read()
    rw = [float(x) for x in re.findall(r"^Overall R factor\s+=\s+([\d.]+)", text, re.M)]
    rf = [float(x) for x in re.findall(r"^Free R factor\s+=\s+([\d.]+)", text, re.M)]
    if not rw or not rf:
        err = [l.strip() for l in text.splitlines() if "ERROR" in l.upper() or "Stop" in l
               or "Problem" in l or "does not seem" in l]
        raise RuntimeError("REFMAC gave no R factors: " + (err[-1] if err else "see " + log))
    return rw, rf


def refmac(workdir, tag, mtz, xyz, d_min, free_label="FreeR_flag", twin=False):
    """R-work, R-free of the model as given: the first cycle of REFMAC's rigid body, whose
    R-factors come before any shift is applied (bulk solvent + anisotropic overall scale are
    REFMAC's defaults). With `twin`, REFMAC finds the twin operators and fractions itself."""
    rw, rf = run_refmac(workdir, tag, mtz, xyz,
                        f"labin FP=F SIGFP=SIGF FREE={free_label}\nrefi type rigid\nrigid ncycle 1\n"
                        + ("twin\n" if twin else "") + f"reso 999 {d_min:.3f}\nend\n")
    return rw[0], rf[0]


def refine(workdir, tag, mtz, xyz, twin=False):
    """R-work, R-free after REFMAC's restrained refinement of the model, over the data's whole
    resolution range: REFINE_CYCLES cycles, automatic weight, isotropic B for every atom (a
    deposited ANISOU is often TLS-derived, and refining it atom by atom diverged: -LL rising every
    cycle at 1.4 A on a model deposited at 1.7 A); a ligand without a dictionary entry gets REFMAC's
    own restraints instead of stopping the run."""
    rw, rf = run_refmac(workdir, tag, mtz, xyz,
                        f"labin FP=F SIGFP=SIGF FREE=FreeR_flag\nncyc {REFINE_CYCLES}\nweight auto\n"
                        "refi bref iso\nmake newligand continue\nmake check none\n"
                        + ("twin\n" if twin else "") + "end\n")
    return rw[-1], rf[-1]


def rename_misfit_ligands(workdir, tag, xyz):
    """A model whose ligand carries a CCD code that REFMAC's monomer library describes with other
    atoms ("atom ... is absent in the library": a code reused for another compound, or a model
    built with its own) stops a restrained refinement. Such a residue is given a code the library
    does not have, so REFMAC makes its restraints from the model's own coordinates, as for any new
    ligand. Returns the renamed model and a note, or None when that is not the failure."""
    log = open(os.path.join(workdir, tag + ".log"), errors="replace").read()
    names = sorted(set(re.findall(r"ERROR : atom :\S+\s+(\S+)\s.*is absent in the library", log)))
    if not names:
        return None
    lib = ccp4_env().get("CLIBD_MON", "")
    free = (c for c in ("".join(t) for t in itertools.product("DQX", "0123456789", "0123456789"))
            if not os.path.exists(os.path.join(lib, c[0].lower(), c + ".cif")))
    new = {n: next(free) for n in names}
    st = gemmi.read_structure(xyz)
    for ch in st[0]:
        for res in ch:
            if res.name in new:
                res.name = new[res.name]
    out = os.path.join(workdir, "model_ligands_renamed.pdb")
    st.write_pdb(out)
    return out, ", ".join(f"{k}->{v}" for k, v in new.items())


def shared_set_r(out_mtz, shared_keys):
    """R = sum|Fo - Fc| / sum Fo over the shared free reflections, from REFMAC's output MTZ (F
    against FC_ALL_LS, which carries REFMAC's own scaling and bulk solvent: over REFMAC's free set
    it reproduces REFMAC's R-free). Not for a twinned refinement, whose twinned Fc is not written."""
    m = gemmi.read_mtz_file(out_mtz)
    fo = np.array(m.column_with_label("F"), copy=False).astype(float)
    fc = np.array(m.column_with_label("FC_ALL_LS"), copy=False).astype(float)
    sel = np.isin(asu_key(m.make_miller_array(), m.spacegroup), shared_keys) & np.isfinite(fo) & np.isfinite(fc)
    return float(np.sum(np.abs(fo[sel] - fc[sel])) / np.sum(fo[sel])), int(sel.sum())


def asu_key(hkl, sg):
    """hkl_key of the reflections mapped into gemmi's asymmetric unit (REFMAC writes its output
    in its own convention, which is not always the same one)."""
    asu, ops = gemmi.ReciprocalAsu(sg), sg.operations()
    return hkl_key(np.array([asu.to_asu([int(x) for x in h], ops)[0] for h in hkl]).reshape(-1, 3))


def refined_shared(workdir, data, dep, cell_m, sg_m, xyz, twin):
    """Both data sets refined from the deposited model by the same protocol, each over its own
    resolution range, and R-free on a SHARED free set: the depositor's free reflections present in
    both data sets. Depositor free reflections outside that intersection are removed from both, so
    neither refinement uses a free reflection as a work reflection; every other reflection of a
    data set is a work reflection of that data set."""
    hkl, f, s = data
    hkl_d, f_d, s_d, flags_d = dep
    k, k_d = asu_key(hkl, sg_m), asu_key(hkl_d, sg_m)
    shared = np.intersect1d(k_d[flags_d == 0], k)
    keep_d = (flags_d != 0) | np.isin(k_d, shared)
    free = np.isin(k, shared)
    out = {"shared_free_n": int(len(shared))}
    tags = (("", hkl, f, s, free, "data_refine.mtz"),
            ("_depdata", hkl_d[keep_d], f_d[keep_d], s_d[keep_d], flags_d[keep_d] == 0,
             "deposited_refine.mtz"))
    for suffix, h, fo, so, fr, name in tags:
        mtz = os.path.join(workdir, name)
        make_mtz(h, fo, so, np.where(fr, 0, 1), cell_m, sg_m).write_to_file(mtz)
        tag = "refine" + (suffix or "_rugnux")
        try:
            r = refine(workdir, tag, mtz, xyz, twin)
        except RuntimeError:
            renamed = rename_misfit_ligands(workdir, tag, xyz)
            if renamed is None:
                raise
            xyz, out["ligands_renamed"] = renamed
            r = refine(workdir, tag, mtz, xyz, twin)
        out["refined_rwork" + suffix], out["refined_rfree_refmac" + suffix] = r
        if not twin:
            out["refined_rfree_shared" + suffix], out["shared_free_n_used" + suffix] = shared_set_r(
                os.path.join(workdir, tag + "_out.mtz"), shared)
        else:     # REFMAC's own R-free is over exactly the shared set (it is these data's free set)
            out["refined_rfree_shared" + suffix] = out["refined_rfree_refmac" + suffix]
    return {k_: (round(v, 4) if isinstance(v, float) else v) for k_, v in out.items()}


def model_for_refmac(st, workdir):
    """The deposited model written in a form REFMAC reads, unchanged in everything it scores;
    returns the file and what had to be changed ("none" normally)."""
    st = st.clone()
    st.setup_entities()
    changes = []
    # An atom of unknown element (UNX, element X) stops REFMAC outright ("Atom does not seem to be
    # an atom"); it has no scattering factor, so leaving it out loses nothing REFMAC could use.
    unknown = 0
    for ch in st[0]:
        for i in reversed(range(len(ch))):
            for k in reversed(range(len(ch[i]))):
                if ch[i][k].element == gemmi.Element("X"):
                    del ch[i][k]
                    unknown += 1
            if len(ch[i]) == 0:
                del ch[i]
    if unknown:
        st.remove_empty_chains()
        changes.append(f"{unknown} atoms of unknown element dropped")
    # PDB format wherever it can hold the model: REFMAC's mmCIF reader rejects many models whose
    # atoms carry ANISOU only in part ("rdaniso_cif: Atom symbol mismatch"), while in the PDB format
    # each ANISOU follows its own ATOM line. Chain names longer than two characters (e.g. 'AAA')
    # are what usually rules the PDB format out; renaming chains changes nothing REFMAC scores.
    if any(len(ch.name) > 2 for ch in st[0]):
        st.shorten_chain_names()
    if st[0].count_atom_sites() < 100000 and all(len(ch.name) <= 2 for ch in st[0]):
        xyz = os.path.join(workdir, "model.pdb")
        st.write_pdb(xyz)
    else:
        # too large for the PDB format: mmCIF with isotropic B only, which REFMAC reads reliably
        # (B_iso holds B_eq for an anisotropic atom)
        for ch in st[0]:
            for res in ch:
                for atom in res:
                    atom.aniso = gemmi.SMat33f(0, 0, 0, 0, 0, 0)
        xyz = os.path.join(workdir, "model.cif")
        st.make_mmcif_document().write_file(xyz)
        changes.append("ANISOU dropped (mmCIF)")
    return xyz, "; ".join(changes) or "none"


# ------------------------------------------------------------------ depositor's structure factors

def refln_value(b, labels, value, sigma, plus, minus):
    """A value and its sigma per reflection, and where they came from: the columns
    `value`/`sigma`, else the mean of the Friedel pair `plus`/`minus` (either member alone where
    the other is missing); None if the block has neither."""
    if value in labels and sigma in labels:
        return b.make_float_array(value), b.make_float_array(sigma), value
    if not all(x in labels for x in (plus, plus + "_sigma", minus, minus + "_sigma")):
        return None
    vp, sp = b.make_float_array(plus), b.make_float_array(plus + "_sigma")
    vm, sm = b.make_float_array(minus), b.make_float_array(minus + "_sigma")
    both, p_only = np.isfinite(vp) & np.isfinite(vm), np.isfinite(vp) & ~np.isfinite(vm)
    v = np.where(both, (vp + vm) / 2, np.where(p_only, vp, vm))
    s = np.where(both, np.sqrt(sp ** 2 + sm ** 2) / 2, np.where(p_only, sp, sm))
    return v, s, f"mean of {plus}/{minus}"


def french_wilson(workdir, hkl, i, sigi, flags, cell, sg):
    """Amplitudes from merged intensities by ctruncate (French & Wilson), as the depositor's own
    CCP4 route would have made them."""
    hklin = os.path.join(workdir, "deposited_i.mtz")
    hklout = os.path.join(workdir, "deposited_fw.mtz")
    make_mtz(hkl, i, sigi, flags, cell, sg, labels=("I", "SIGI")).write_to_file(hklin)
    with open(os.path.join(workdir, "ctruncate.log"), "w") as fh:
        subprocess.run(["ctruncate", "-hklin", hklin, "-hklout", hklout, "-colin", "/*/*/[I,SIGI]",
                        "-freein", "/*/*/[FreeR_flag]"],
                       stdout=fh, stderr=subprocess.STDOUT, cwd=workdir, env=ccp4_env())
    if not os.path.exists(hklout):
        raise RuntimeError("ctruncate failed: see " + os.path.join(workdir, "ctruncate.log"))
    out = gemmi.read_mtz_file(hklout)
    col = lambda lab: np.array(out.column_with_label(lab), copy=False).astype(float)
    return out.make_miller_array(), col("F"), col("SIGF"), col("FreeR_flag").astype(int)


def depositor_data(pdb_id, cache_dir, workdir, cell_m, sg_m):
    """The first reflection block of the entry's -sf.cif as an MTZ (F, SIGF, FreeR_flag) in the
    model's cell and group, and what the amplitudes were made from. Amplitudes: F_meas, else the
    mean of F(+)/F(-); where only intensities were deposited (I, else the mean of I(+)/I(-)) they
    are converted by French-Wilson. Free flags from _refln.status ('f' free), else from a numeric
    _refln.pdbx_r_free_flag (0 when it takes more than two values - the CCP4 convention - else the
    minority value). Returns None, reason when there is nothing usable."""
    path = fetch(RCSB + pdb_id + "-sf.cif.gz", os.path.join(cache_dir, pdb_id + "-sf.cif.gz"))
    if path is None:
        return None, "no deposited structure factors"
    rb = gemmi.as_refln_blocks(gemmi.cif.read(path))
    if not rb:
        return None, "no reflection block"
    b = rb[0]
    labels = b.column_labels()
    hkl = b.make_miller_array()
    val = refln_value(b, labels, "F_meas_au", "F_meas_sigma_au", "pdbx_F_plus", "pdbx_F_minus")
    amplitudes = val is not None
    if not amplitudes:
        val = refln_value(b, labels, "intensity_meas", "intensity_sigma", "pdbx_I_plus", "pdbx_I_minus")
    if val is None:
        return None, "deposited structure factors carry neither amplitudes nor intensities"
    f, s, kind = val
    status = np.array([x.strip("'\"") for x in b.block.find_values("_refln.status")]) \
        if "status" in labels else np.array([])
    if len(status) == len(hkl) and np.any(status == "f"):
        use = (status == "f") | (status == "o")
        flags = np.where(status == "f", 0, 1)
    elif "pdbx_r_free_flag" in labels:
        v = b.make_int_array("pdbx_r_free_flag", -1)
        vals, counts = np.unique(v[v >= 0], return_counts=True)
        if len(vals) < 2:
            return None, "deposited structure factors carry no free set"
        free = 0 if len(vals) > 2 and 0 in vals else vals[np.argmin(counts)]
        use = v >= 0
        flags = np.where(v == free, 0, 1)
    else:
        return None, "deposited structure factors carry no free set"
    use &= np.isfinite(f) & np.isfinite(s) & (s > 0)
    if not np.allclose(b.cell.parameters, cell_m.parameters, rtol=0.02, atol=0.5):
        return None, "deposited structure factors in another cell"
    hkl, f, s, flags = to_group(hkl[use], f[use], s[use], flags[use], cell_m, sg_m,
                                np.arange(use.sum()))
    if not amplitudes:
        hkl, f, s, flags = french_wilson(workdir, hkl, f, s, flags, cell_m, sg_m)
        kind += ", French-Wilson (ctruncate)"
    return (make_mtz(hkl, f, s, flags, cell_m, sg_m), kind), None


def declared_twin(xyz):
    """The twin operators the entry declares (_pdbx_reflns_twin), [] when it declares none."""
    block = gemmi.cif.read(xyz).sole_block()
    ops = [o.strip("'\"").replace(" ", "").upper()
           for o in block.find_values("_pdbx_reflns_twin.operator")]
    return ops if len(ops) > 1 else []


def choose_setting(mtz, fc_hkl, fc, cell_m, sg_m, d_lo, d_hi):
    """Every change of basis carrying the data's lattice onto the model's, as (shell R, fraction
    indexed, T, reflections in the model's setting), best agreement with the model's Fcalc first."""
    scored = []
    for t in basis_changes(mtz.cell, mtz.spacegroup, cell_m, sg_m):
        data, frac = reindexed(mtz, t, cell_m, sg_m)
        r = shell_r(data[0], data[1], fc_hkl, fc, cell_m, d_lo, d_hi)
        if r is not None:
            scored.append((r, frac, t, data))
    return sorted(scored, key=lambda x: x[0])


# ------------------------------------------------------------------ the check

def check(mtz_path, pdb_id, workdir, cache_dir=DEFAULT_CACHE, baseline=True, refined=True):
    t0 = time.time()
    res = {"status": "failed", "reason": "", "rwork": None, "rfree": None,
           "rfree_deposited": None, "rwork_deposited": None, "d_min_used": None, "sg_used": None,
           "reindex_op": None, "cycles": 0, "seconds": None,
           "rfree_depflags": None,
           "rwork_depdata": None, "rfree_depdata": None, "baseline_note": None,
           "depdata_kind": None, "reindex_op_depdata": None, "twin_declared": None,
           "refined_rwork": None, "refined_rfree_shared": None, "refined_rfree_refmac": None,
           "refined_rwork_depdata": None, "refined_rfree_shared_depdata": None,
           "refined_rfree_refmac_depdata": None, "shared_free_n": None, "refined_reason": None,
           "sg_data": None, "d_min_data": None, "d_min_deposited": None, "cell_diff_pct": None,
           "setting_r": None, "setting_r_next": None, "fraction_indexed": None,
           "model_modification": "none", "protocol": "REFMAC R-factors of the model as deposited, bulk solvent, no refinement"}

    def done(status, reason=""):
        res["status"], res["reason"] = status, reason
        res["seconds"] = round(time.time() - t0, 1)
        return res

    pdb_id = pdb_id.lower()
    if not PDB_ID.match(pdb_id):
        return done("n/a", "not a PDB entry (small molecule or unpublished set)")
    try:
        workdir = os.path.abspath(workdir)          # REFMAC runs inside it
        os.makedirs(workdir, exist_ok=True)
        meta = deposition(pdb_id, cache_dir)
        if meta is None:
            return done("failed", "entry could not be downloaded")
        res["rfree_deposited"], res["rwork_deposited"] = meta["rfree"], meta["rwork"]
        res["d_min_deposited"] = meta["dmin"]
        if not any("X-RAY" in m.upper() for m in meta["methods"]):
            return done("n/a", "not an X-ray structure: " + ", ".join(meta["methods"]))
        if meta["rfree"] is None:
            return done("n/a", "no deposited R-free")
        st = gemmi.read_structure(meta["xyz"])
        if len(st) > 1:
            return done("skipped", f"{len(st)} models in the entry")
        natoms = st[0].count_atom_sites()
        if natoms > MAX_ATOMS:
            return done("skipped", f"{natoms} atoms, above {MAX_ATOMS}")
        sg_m, cell_m = st.find_spacegroup(), st.cell
        if sg_m is None:
            return done("failed", f"model space group '{st.spacegroup_hm}' not recognised")

        mtz = read_rugnux(mtz_path)
        res["sg_data"] = mtz.spacegroup.hm
        res["d_min_data"] = round(mtz.resolution_high(), 3)
        d_used = max(mtz.resolution_high(), meta["dmin"] or 0)
        res["d_min_used"] = round(d_used, 3)
        res["sg_used"] = sg_m.hm

        # choose the setting on a mid-resolution shell against the model's Fcalc
        d_hi = max(3.0, d_used)
        d_lo = max(8.0, 1.6 * d_hi)
        fc_hkl, fc = fcalc(st, d_hi)
        scored = choose_setting(mtz, fc_hkl, fc, cell_m, sg_m, d_lo, d_hi)
        if not scored:
            return done("failed", "no change of basis carries the data lattice onto the model's")
        r_best, frac, t, (hkl, f, s, fl) = scored[0]
        res["setting_r"] = round(r_best, 4)
        res["setting_r_next"] = round(scored[1][0], 4) if len(scored) > 1 else None
        res["fraction_indexed"] = round(float(frac), 4)
        res["reindex_op"] = op_string(t)
        gd = t.T @ metric(mtz.cell) @ t
        res["cell_diff_pct"] = round(100 * float(np.max(np.abs(
            np.sqrt(np.diag(gd)) / np.array(cell_m.parameters[:3]) - 1))), 2)

        # the depositor's data go through the same setting choice: their indexing can be the
        # other branch of a merohedral ambiguity (a twinned entry refined in its minor domain's)
        dep = None
        if baseline:
            dep, why = depositor_data(pdb_id, cache_dir, workdir, cell_m, sg_m)
            res["baseline_note"] = why
        dep_flags = None
        if dep is not None:
            dep_mtz, res["depdata_kind"] = dep
            dep_scored = choose_setting(dep_mtz, fc_hkl, fc, cell_m, sg_m, d_lo, d_hi)
            if dep_scored:
                res["reindex_op_depdata"] = op_string(dep_scored[0][2])
                dep = dep_scored[0][3]
                _, i, j = np.intersect1d(hkl_key(hkl), hkl_key(dep[0]), return_indices=True)
                dep_flags = np.full(len(hkl), 1)     # absent from the deposited set -> work
                dep_flags[i] = dep[3][j]
            else:
                dep, res["baseline_note"] = None, "deposited structure factors too few to place"
        data_mtz = os.path.join(workdir, "data.mtz")
        make_mtz(hkl, f, s, fl, cell_m, sg_m, dep_flags).write_to_file(data_mtz)
        xyz, res["model_modification"] = model_for_refmac(st, workdir)
        if dep is not None:
            dep_mtz = os.path.join(workdir, "deposited.mtz")
            make_mtz(*dep, cell_m, sg_m).write_to_file(dep_mtz)

        def score(twin, suffix):
            out = {}
            out["rwork"], out["rfree"] = refmac(workdir, "rugnux" + suffix, data_mtz, xyz, d_used,
                                                twin=twin)
            if dep is not None:
                out["rfree_depflags"] = refmac(workdir, "rugnux_depflags" + suffix, data_mtz, xyz,
                                               d_used, "FreeR_dep", twin=twin)[1]
                out["rwork_depdata"], out["rfree_depdata"] = refmac(workdir, "deposited" + suffix,
                                                                    dep_mtz, xyz, d_used, twin=twin)
            return {k: round(v, 4) for k, v in out.items()}

        # An entry refined as a twin is scored as one, both data sets alike: its published R-free
        # is a twinned one, and an untwinned score of twinned data says more about the twin
        # fraction than about the data. The untwinned numbers are kept beside it.
        res["twin_declared"] = declared_twin(meta["xyz"]) or None
        res.update(score(bool(res["twin_declared"]), ""))
        if res["twin_declared"]:
            res["protocol"] += ", REFMAC twin refinement (the entry declares twinning)"
            res["untwinned"] = score(False, "_untwinned")
        if dep is not None and refined:
            try:
                res.update(refined_shared(workdir, (hkl, f, s), dep, cell_m, sg_m, xyz,
                                          bool(res["twin_declared"])))
            except Exception as e:           # the first-cycle numbers stand without it
                if os.environ.get("MODEL_CHECK_RAISE"):
                    raise
                res["refined_reason"] = f"{type(e).__name__}: {e}"
        return done("ok")
    except Exception as e:                           # one bad entry must not stop a battery
        if os.environ.get("MODEL_CHECK_RAISE"):
            raise
        return done("failed", f"{type(e).__name__}: {e}")


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("mtz", help="Rugnux merged MTZ (needs F, SIGF, FreeR_flag)")
    ap.add_argument("pdb_id")
    ap.add_argument("--workdir", default=None, help="default: model_check_<id> beside the MTZ")
    ap.add_argument("--cache", default=DEFAULT_CACHE)
    ap.add_argument("--no-baseline", action="store_true",
                    help="skip the depositor's structure factors")
    ap.add_argument("--no-refine", action="store_true",
                    help="first-cycle R factors only (skip the refined shared-free-set protocol)")
    a = ap.parse_args()
    wd = a.workdir or os.path.join(os.path.dirname(os.path.abspath(a.mtz)), "model_check_" + a.pdb_id)
    print(json.dumps(check(a.mtz, a.pdb_id, os.path.abspath(wd), a.cache, baseline=not a.no_baseline,
                           refined=not a.no_refine), indent=1))


if __name__ == "__main__":
    main()