#!/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.

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


# ------------------------------------------------------------------ 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 write_mtz(path, hkl, f, sigf, flags, cell, sg, dep_flags=None):
    mtz = gemmi.Mtz(with_base=True)
    mtz.spacegroup, mtz.cell = sg, cell
    mtz.add_dataset("data")
    cols = [hkl.astype(float), f, sigf, flags]
    for lab, typ in (("F", "F"), ("SIGF", "Q"), ("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))
    mtz.write_to_file(path)


def refmac(workdir, tag, mtz, xyz, d_min, free_label="FreeR_flag"):
    """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)."""
    log = os.path.join(workdir, tag + ".log")
    keywords = (f"labin FP=F SIGFP=SIGF FREE={free_label}\nrefi type rigid\n"
                f"rigid ncycle 1\nreso 999 {d_min:.3f}\nend\n")
    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]
        raise RuntimeError("REFMAC gave no R factors: " + (err[-1] if err else "see " + log))
    return rw[0], rf[0]


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

def depositor_data(pdb_id, cache_dir, cell_m, sg_m):
    """The first reflection block of the entry's -sf.cif, reduced to the model group's ASU.
    Free flags from _refln.status ('f' free) or a numeric _refln.pdbx_r_free_flag (minority
    value = free). 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()
    if "F_meas_au" in labels and "F_meas_sigma_au" in labels:
        f, s = b.make_float_array("F_meas_au"), b.make_float_array("F_meas_sigma_au")
    else:     # intensities only: no amplitude conversion here that could bias the baseline
        return None, "deposited structure factors carry no amplitudes"
    if "status" in labels:
        st = np.array(list(b.block.find_values("_refln.status")) if b.block.find_values("_refln.status")
                      else [], dtype=object)
        st = np.array([x.strip("'\"") for x in st])
        if len(st) != len(hkl) or not np.any(st == "f"):
            return None, "deposited structure factors carry no free set"
        use = (st == "f") | (st == "o")
        flags = np.where(st == "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 = 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 = hkl[use], f[use], s[use], flags[use]
    return to_group(hkl, f, s, flags, cell_m, sg_m, np.arange(len(hkl))), None


# ------------------------------------------------------------------ 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,
           "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:
        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 = []
        for t in basis_changes(mtz.cell, mtz.spacegroup, cell_m, sg_m):
            (hkl, f, s, fl), frac = reindexed(mtz, t, cell_m, sg_m)
            r = shell_r(hkl, f, fc_hkl, fc, cell_m, d_lo, d_hi)
            if r is not None:
                scored.append((r, frac, t))
        if not scored:
            return done("failed", "no change of basis carries the data lattice onto the model's")
        scored.sort(key=lambda x: x[0])
        r_best, frac, t = 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)

        (hkl, f, s, fl), _ = reindexed(mtz, t, cell_m, sg_m)
        dep = None
        if baseline:
            dep, why = depositor_data(pdb_id, cache_dir, cell_m, sg_m)
            res["baseline_note"] = why
        dep_flags = None
        if dep is not None:
            _, i, j = np.intersect1d(hkl_key(hkl), hkl_key(dep[0]), return_indices=True)
            dep_flags = np.full(len(hkl), -1)        # -1: absent from the deposited set -> work
            dep_flags[i] = dep[3][j]
            dep_flags[dep_flags < 0] = 1
        data_mtz = os.path.join(workdir, "data.mtz")
        write_mtz(data_mtz, hkl, f, s, fl, cell_m, sg_m, dep_flags)
        xyz = os.path.join(workdir, "model.cif")
        st.setup_entities()
        st.make_mmcif_document().write_file(xyz)

        res["rwork"], res["rfree"] = refmac(workdir, "rugnux", data_mtz, xyz, d_used)
        if dep is not None:
            res["rfree_depflags"] = refmac(workdir, "rugnux_depflags", data_mtz, xyz, d_used,
                                           "FreeR_dep")[1]
            dep_mtz = os.path.join(workdir, "deposited.mtz")
            write_mtz(dep_mtz, *dep, cell_m, sg_m)
            res["rwork_depdata"], res["rfree_depdata"] = refmac(workdir, "deposited", dep_mtz, xyz,
                                                                d_used)
        for k in ("rwork", "rfree", "rfree_depflags",
                  "rwork_depdata", "rfree_depdata"):
            if res[k] is not None:
                res[k] = round(res[k], 4)
        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, wd, a.cache, baseline=not a.no_baseline), indent=1))


if __name__ == "__main__":
    main()