#!/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 both data sets on their common reflections only: the hkl present in both, at
    d_min_used = max(Rugnux's d_min, the deposited d_min), so a cut either data set made (a
    low-resolution limit, an anisotropic cut) applies to both;
  * 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.

The model is never refined: every number is REFMAC's first cycle, before any shift.

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`.
Where the deposited structure factors mark no free set (SHELXL entries often do not), both data
sets are scored on a synthetic one instead (synthetic_free; `baseline_note` says so). Their model
was refined against those reflections too, so both numbers read like work Rs, but alike.

  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/"
COMMITTED = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models", "pdb")   # the timing benchmark's entries
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
SYNTHETIC_FREE = 0.05       # free fraction given to deposited structure factors that carry no free set


# ------------------------------------------------------------------ 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: from the copy committed in
    models/pdb/ where there is one (read from the file itself, no cache and no network), else
    fetched from the RCSB into cache_dir and cached as <id>.json."""
    committed = os.path.join(COMMITTED, pdb_id + ".cif.gz")
    if os.path.exists(committed):
        return refinement_record(pdb_id, committed)
    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):
        # the coordinates where this cache is now: a cache copied elsewhere keeps the old path in its json
        meta = dict(json.load(open(meta_path)), xyz=xyz)
        if "program" not in meta:          # a cache written before the program was recorded
            meta["program"] = refinement_program(gemmi.cif.read(xyz).sole_block())
        return meta
    meta = refinement_record(pdb_id, xyz)
    json.dump(meta, open(meta_path, "w"), indent=1)
    return meta


def refinement_record(name, xyz):
    """The refinement statistics an mmCIF coordinate file states about itself (_refine)."""
    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))}
    return {"id": name, "methods": methods, "rfree": ref.get("rfree"), "rwork": ref.get("rwork"),
            "dmin": ref.get("dmin"), "program": refinement_program(block), "xyz": xyz}


def refinement_program(block):
    """The refinement program(s) the entry names (_software, classification 'refinement'), or None."""
    names = [row.str(1) for row in block.find("_software.", ["classification", "name"])
             if row.str(0).lower() == "refinement"]
    return " ".join(names) or None


# ------------------------------------------------------------------ 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 synthetic_free(hkl):
    """Free flags (0 free, 1 work) for deposited structure factors that carry no free set: a fixed
    hash of the asymmetric-unit hkl (splitmix64), so the same reflection is free in every run and in
    both data sets, independent of reflection order."""
    with np.errstate(over="ignore"):
        x = hkl_key(hkl).astype(np.uint64) + np.uint64(0x9E3779B97F4A7C15)
        x = (x ^ (x >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9)
        x = (x ^ (x >> np.uint64(27))) * np.uint64(0x94D049BB133111EB)
        x ^= x >> np.uint64(31)
    return np.where(x < np.uint64(SYNTHETIC_FREE * 2.0 ** 64), 0, 1)


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 refmac_version():
    """REFMAC's version as its banner states it ("REFMAC 5.8.0431, CCP4 9.0.014"), for the run's record;
    None where refmac5 does not run."""
    try:
        out = subprocess.run(["refmac5"], input=b"END\n", capture_output=True, timeout=60, env=ccp4_env()).stdout
    except (OSError, subprocess.TimeoutExpired):
        return None
    m = re.search(r"CCP4\s+(\S+):\s+Refmac\s+version\s+(\S+)", out.decode(errors="replace"))
    return f"REFMAC {m.group(2)}, CCP4 {m.group(1)}" if m else None


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 read_model(xyz):
    """The deposited model with its strict-NCS copies generated: an entry can hold part of the
    asymmetric unit and the operators that make the rest (_struct_ncs_oper 'generate', MTRIX without
    the given flag) - an icosahedral virus deposits a few chains of dozens."""
    st = gemmi.read_structure(xyz)
    if any(not op.given for op in st.ncs):
        st.expand_ncs(gemmi.HowToNameCopiedChain.Short)
    return st


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). Where neither marks a free set (e.g. a SHELXL entry whose free set was not
    deposited), a synthetic one is made by synthetic_free and noted. Returns None, reason when there
    is nothing usable, else the data and a note (None, or what was synthesised)."""
    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([])
    note = None
    v = b.make_int_array("pdbx_r_free_flag", -1) if "pdbx_r_free_flag" in labels else None
    if len(status) == len(hkl) and np.any(status == "f"):
        use = (status == "f") | (status == "o")
        flags = np.where(status == "f", 0, 1)
    elif v is not None and len(np.unique(v[v >= 0])) >= 2:
        vals, counts = np.unique(v[v >= 0], return_counts=True)
        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:
        use = np.ones(len(hkl), bool) if len(status) != len(hkl) else (status == "o")
        flags = np.ones(len(hkl), int)
        note = f"deposited structure factors carry no free set: a synthetic {SYNTHETIC_FREE:.0%} one is used"
    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 note:
        flags = synthetic_free(hkl)
    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), note


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):
    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,
           "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 = read_model(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]
                # both data sets scored on their common reflections only
                common = np.isin(hkl_key(hkl), hkl_key(dep[0]))
                hkl, f, s, fl = hkl[common], f[common], s[common], fl[common]
                dep = tuple(x[np.isin(hkl_key(dep[0]), hkl_key(hkl))] for x in dep)
                _, i, j = np.intersect1d(hkl_key(hkl), hkl_key(dep[0]), return_indices=True)
                dep_flags = np.full(len(hkl), 1)
                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")
        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")
    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),
                     indent=1))


if __name__ == "__main__":
    main()