"""The timing benchmarks: a few sets, each started from a model rugnux has to place by molecular
replacement, timed per step and checked for a right answer.

timing.json lists the sets of each benchmark (README, 'The timing benchmark'): 'standard', three
open-arm EIGER 16M sets and 4nwv, needs nothing but the open data; 'inhouse' is two in-house standard
proteins and 4nwv. A set is started from its AlphaFold DB prediction (models/, committed with its
predicted aligned error) or from its deposited model moved by a fixed rigid-body transform, made at
run time from the entry committed in models/pdb/ (`start_model`). Each runs the battery's own command
for its arm with --model <start model>. A set passes (`check`) when every one of these holds:

  space_group   the battery's own verdict (score.judge: lattice and space group against the
                reference, the hand excepted) is a pass
  placed        MODEL_PLACEMENT is MR: molecular replacement placed the model, and the R-factors,
                maps and model written are of that placement
  mr_group      MR_SPACE_GROUP is the expected group, hand included - the model decides it
  r_free        R_FREE <= rfree_max (a random structure: 0.586 acentric, Wilson 1950)
  confirmed     CC_MODEL_CONFIRMED_TO_D_MIN is not NONE: some shell's CC(model, data) clears
                rugnux's empirical null
  placement_cc  the placed model (p_model.pdb) has the reference model's intensities: their
                correlation, shell by shell normalised, 10-4 A, the best over the reindexing
                operators the two lattices allow, is >= cc_min. Independent of rugnux's own scores:
                |F| is the same for every origin the space group allows and every symmetry copy, so
                a right placement correlates highly and a wrong one (or the wrong hand) near zero.
"""
import gzip
import json
import math
import os
import re

import gemmi
import numpy as np

import model_check
import score

HERE = os.path.dirname(os.path.abspath(__file__))
MANIFEST = os.path.join(HERE, "timing.json")

CC_D_MAX, CC_D_MIN = 10.0, 4.0     # the placement check's resolution range (A)
CC_SHELLS = 8                      # equal-count shells for the normalisation
STEP_LINE = re.compile(r"TIMING step (.+?)(?:, \d+ images)? ([\d.]+)$")


def names():
    return list(json.load(open(MANIFEST))["benchmarks"])


def load(benchmark):
    return json.load(open(MANIFEST))["benchmarks"][benchmark]


def read_model(path):
    """A model with strict NCS expanded (as rugnux's validation does) and no hydrogens or waters."""
    st = gemmi.read_structure(path)
    st.setup_entities()
    if any(not op.given for op in st.ncs):
        st.expand_ncs(gemmi.HowToNameCopiedChain.Short)
    st.remove_hydrogens()
    st.remove_waters()
    return st


def rotation(axis, deg):
    x, y, z = np.asarray(axis, float) / np.linalg.norm(axis)
    c, s = math.cos(math.radians(deg)), math.sin(math.radians(deg))
    t = 1 - c
    return [[c + x * x * t, x * y * t - z * s, x * z * t + y * s],
            [y * x * t + z * s, c + y * y * t, y * z * t - x * s],
            [z * x * t - y * s, z * y * t + x * s, c + z * z * t]]


def start_model(start, wd, pdb_cache):
    """The file given to --model: the committed prediction, or the deposited model (models/pdb/,
    else the pdb_cache) moved by the manifest's transform, written into the set's work directory."""
    if start.get("alphafold"):
        return os.path.join(HERE, start["alphafold"])
    mv = start["moved"]
    meta = model_check.deposition(mv["pdb"], pdb_cache)
    if meta is None:
        return None
    st = read_model(meta["xyz"])
    pos = np.array([a.pos.tolist() for ch in st[0] for r in ch for a in r])
    centre = pos.mean(axis=0)
    rot = np.array(rotation(mv["axis"], mv["deg"]))
    tr = gemmi.Transform(gemmi.Mat33(rot.tolist()),
                         gemmi.Vec3(*(centre - rot @ centre + np.asarray(mv["shift"], float))))
    st[0].transform_pos_and_adp(tr)
    out = os.path.join(wd, f"{mv['pdb']}_moved.cif.gz")
    with gzip.open(out, "wt") as fh:
        fh.write(st.make_mmcif_document().as_string())
    return out


def intensities(st):
    hkl, f = model_check.fcalc(st, CC_D_MIN)
    return hkl, f.astype(float) ** 2


def normalised(i, d):
    """Each intensity over its shell's mean (equal-count shells in d)."""
    out = np.empty_like(i)
    for sel in np.array_split(np.argsort(d), CC_SHELLS):
        out[sel] = i[sel] / max(i[sel].mean(), 1e-30)
    return out


def placement_cc(placed_path, reference_path):
    """Correlation of the placed model's intensities with the reference model's, each in its own cell
    and space group, the placed model's indices carried into the reference's frame by every change of
    basis the two lattices allow; the best one."""
    placed, ref = read_model(placed_path), read_model(reference_path)
    sg_p, sg_r = placed.find_spacegroup(), ref.find_spacegroup()
    hp, ip = intensities(placed)
    hr, ir = intensities(ref)
    d_r = ref.cell.calculate_d_array(hr)
    keep = d_r <= CC_D_MAX
    hr, ir = hr[keep], ir[keep]
    best = None
    for t in model_check.basis_changes(placed.cell, sg_p, ref.cell, sg_r):
        hm = hp @ t
        whole = np.all(np.abs(hm - np.round(hm)) < 1e-3, axis=1)
        n = int(whole.sum())
        h, i, _, _ = model_check.to_group(np.round(hm[whole]).astype(int), ip[whole], np.ones(n),
                                          np.zeros(n, int), ref.cell, sg_r, np.arange(n))
        _, a, b = np.intersect1d(model_check.hkl_key(h), model_check.hkl_key(hr), return_indices=True)
        if len(a) < 100:
            continue
        d = ref.cell.calculate_d_array(hr[b])
        cc = float(np.corrcoef(normalised(i[a], d), normalised(ir[b], d))[0, 1])
        if best is None or cc > best[0]:
            best = (cc, len(a), model_check.op_string(t))
    return best


def steps(wd):
    """{step: seconds} from the TIMING step lines of rugnux's log, in run order."""
    out = {}
    log = os.path.join(wd, "p_log.txt")
    if os.path.exists(log):
        for line in open(log, errors="replace"):
            m = STEP_LINE.search(line.rstrip())
            if m:
                out[m.group(1)] = out.get(m.group(1), 0.0) + float(m.group(2))
    return out


def check(t, row, wd, reference_path):
    """The benchmark's verdict on one finished set: bench_checks {name: [ok, value, limit]},
    bench_verdict pass/fail, bench_reason (the checks that failed) and steps."""
    rep = score.read_report(os.path.join(wd, "p_report.txt"))
    exp = t["expect"]
    mr_sg = gemmi.find_spacegroup_by_name(rep.get("MR_SPACE_GROUP") or "")
    checks = {
        "space_group": [row["verdict"] == "pass", row.get("sg"), row.get("sg_ref")],
        "placed": [rep.get("MODEL_PLACEMENT") == "MR", rep.get("MODEL_PLACEMENT"), "MR"],
        "mr_group": [bool(mr_sg) and mr_sg.number == exp["mr_sgno"], rep.get("MR_SPACE_GROUP"),
                     gemmi.find_spacegroup_by_number(exp["mr_sgno"]).hm],
        "r_free": [row.get("rfree") is not None and row["rfree"] <= exp["rfree_max"], row.get("rfree"),
                   exp["rfree_max"]],
        "confirmed": [rep.get("CC_MODEL_CONFIRMED_TO_D_MIN") not in (None, "NONE"),
                      rep.get("CC_MODEL_CONFIRMED_TO_D_MIN"), "not NONE"],
    }
    placed = os.path.join(wd, "p_model.pdb")
    cc = None
    if os.path.exists(placed) and reference_path:
        cc = placement_cc(placed, reference_path)
    checks["placement_cc"] = [cc is not None and cc[0] >= exp["cc_min"], cc and round(cc[0], 3), exp["cc_min"]]
    bad = [k for k, (ok, _, _) in checks.items() if not ok]
    return {"bench_checks": checks, "bench_verdict": "fail" if bad else "pass",
            "bench_reason": ", ".join(f"{k} {checks[k][1]} (want {checks[k][2]})" for k in bad) or None,
            "bench_cc_op": cc and cc[2], "bench_cc_n": cc and cc[1],
            "mr_seconds": score.fnum(rep.get("MR_SECONDS")), "mr_decision": rep.get("MR_DECISION"),
            "steps": steps(wd)}


def console(rows):
    """One block per set: verdict, wall time, the checks, the step times."""
    out = []
    for r in rows:
        out.append(f"{r['set']:12s} {r.get('bench_verdict') or '-':5s} wall {r.get('wall_s') or 0:7.1f} s"
                   f"  MR {r.get('mr_seconds') or 0:.1f} s  {r.get('bench_reason') or ''}")
        for k, (ok, v, lim) in (r.get("bench_checks") or {}).items():
            out.append(f"    {'ok ' if ok else 'BAD'} {k:13s} {v} (want {lim})")
        out.append("    " + ", ".join(f"{k} {v:.1f}" for k, v in (r.get("steps") or {}).items()))
    return "\n".join(out)


def section(doc, rows):
    """The report's section for a timing run (rows of results.json carrying bench_* keys)."""
    rows = [r for r in rows if r.get("bench_checks")]
    if not rows:
        return
    bad = [r for r in rows if r["bench_verdict"] != "pass"]
    doc.h(2, f"Timing benchmark: {len(rows) - len(bad)} of {len(rows)} pass")
    doc.p("Each set started from a model it has to place by molecular replacement (timing.json). A set "
          "passes when the battery's verdict passes, MODEL_PLACEMENT is MR, MR_SPACE_GROUP is the expected "
          "group (hand included), R_FREE is at most its limit, CC_MODEL_CONFIRMED_TO_D_MIN is not NONE, "
          f"and the placed model's intensities correlate with the reference model's ({CC_D_MAX:.0f}-{CC_D_MIN:.0f} A, "
          "shell-normalised) at least as much as its limit.")
    names = list(rows[0]["bench_checks"])
    doc.table(["set", "verdict", "wall s", "MR s"] + names + ["why"],
              [[r["set"], r["bench_verdict"], f"{r['wall_s']:.1f}" if r.get("wall_s") else "-",
                f"{r['mr_seconds']:.1f}" if r.get("mr_seconds") else "-"]
               + [("" if r["bench_checks"][k][0] else "BAD ") + str(r["bench_checks"][k][1]) for k in names]
               + [r.get("bench_reason") or ""] for r in rows],
              row_class=lambda row: "fail" if row[1] != "pass" else "")
    step_names = []
    for r in rows:
        step_names += [k for k in r.get("steps", {}) if k not in step_names]
    doc.table(["step"] + [r["set"] for r in rows],
              [[k] + [f"{r['steps'][k]:.1f}" if k in r.get("steps", {}) else "-" for r in rows]
               for k in step_names])