"""The wrong-model null: a wrong --model must change nothing and must be recognised as wrong.

nullmodel.json lists sets of the open arm, each with five wrong models - deposited coordinates of
other open-arm entries (README, 'The wrong-model null', says how they were chosen). The runner
(battery.py) runs each set once without --model and once with each wrong model; `score` compares a
wrong-model run with the no-model run of the same set:

  (a) data      p.hkl byte for byte; p.mtz reflection by reflection (group, cell, columns, every
                value - the file itself differs by the command line in its history); p.cif with its
                command line left out
  (b) decisions every KEY= line of p_report.txt outside the MODEL VALIDATION section (space group,
                SOHNCKE group and SPACE_GROUP_ENANTIOMORPH, cell, resolution cut, twinning,
                warnings, ...), provenance lines aside
  (c) null      the model is recognised as wrong: MODEL_FIT is not ACCEPTED, MODEL_DECISIONS_TAKEN is
                NONE, no enantiomorph adopted, CC_MODEL_CONFIRMED_TO_D_MIN is NONE (no shell's
                CC(model, data) clears its empirical null), and R_FREE
                is where a random structure's is: 0.586 for acentric reflections, 0.828 for centric
                (Wilson 1950); the scale, anisotropic B, bulk solvent and rigid body fitted to the
                working set can only pull it a few points below that, so R_FREE >= RANDOM_R_FREE_MIN.
"""
import hashlib
import json
import os
import re

import gemmi
import numpy

import score

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

# Wilson's random-structure R is 0.586 (acentric); a fit of scale, anisotropic B, bulk solvent and a
# rigid body on the working set takes a few points off it, never fifteen
RANDOM_R_FREE_MIN = 0.50
# lines of p_report.txt that say how and when the run was made, not what it decided
PROVENANCE = {"DATE", "COMMAND_LINE", "WALL_TIME"}
MODEL_SECTION = "5. MODEL VALIDATION"


def f(x, fmt):
    return "-" if x is None else fmt.format(x)


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


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def decision_keys(path):
    """KEY= lines of a report outside the model section, provenance left out."""
    out, in_model = {}, False
    for line in open(path, errors="replace"):
        if line.startswith("#  ") and line[3:].strip()[:1].isdigit():
            in_model = line[3:].strip() == MODEL_SECTION
        m = score.REPORT_LINE.match(line.rstrip("\n"))
        if m and not in_model and m.group(1) not in PROVENANCE:
            out[m.group(1)] = m.group(2).strip()
    return out


def same_mtz(a, b):
    ma, mb = gemmi.read_mtz_file(a), gemmi.read_mtz_file(b)
    if (ma.spacegroup.number != mb.spacegroup.number or ma.cell.parameters != mb.cell.parameters
            or ma.column_labels() != mb.column_labels()):
        return False
    return numpy.array_equal(numpy.asarray(ma), numpy.asarray(mb), equal_nan=True)


def same_cif(a, b):
    def text(path):
        return re.sub(r"_jfjoch_provenance\.command_line\s*\n;.*?\n;", "", open(path).read(), flags=re.S)
    return text(a) == text(b)


def score_pair(none_wd, wd):
    """The (a)(b)(c) verdict of one wrong-model run in wd against the no-model run in none_wd."""
    rep = score.read_report(os.path.join(wd, "p_report.txt"))
    files = {}
    for name, same in (("p.hkl", lambda a, b: sha256(a) == sha256(b)), ("p.mtz", same_mtz),
                       ("p.cif", same_cif)):
        a, b = os.path.join(none_wd, name), os.path.join(wd, name)
        files[name] = os.path.exists(a) and os.path.exists(b) and same(a, b)
    ka, kb = (decision_keys(os.path.join(d, "p_report.txt")) for d in (none_wd, wd))
    moved = sorted(k for k in set(ka) | set(kb) if ka.get(k) != kb.get(k))
    r_free = score.fnum(rep.get("R_FREE"))
    row = {
        "data_same": all(files.values()), "files_differ": [n for n, s in files.items() if not s],
        "decisions_moved": moved,
        "sg": kb.get("SPACE_GROUP_NAME"), "enantiomorph": kb.get("SPACE_GROUP_ENANTIOMORPH"),
        "model_validation": rep.get("MODEL_VALIDATION"), "model_fit": rep.get("MODEL_FIT"),
        "model_fit_sigma": score.fnum(rep.get("MODEL_FIT_SIGMA")),
        "model_decisions": rep.get("MODEL_DECISIONS_TAKEN"),
        "enantiomorph_adopted": rep.get("MODEL_ENANTIOMORPH_ADOPTED"),
        "indexing_margin_sigma": score.fnum(rep.get("MODEL_INDEXING_MARGIN_SIGMA")),
        "change_of_basis": rep.get("MODEL_CHANGE_OF_BASIS"),
        "r_free": r_free, "r_work": score.fnum(rep.get("R_WORK")),
        "cc_model": score.fnum(rep.get("CC_MODEL_OVERALL")),
        "cc_confirmed": rep.get("CC_MODEL_CONFIRMED_TO_D_MIN"),
    }
    why = []
    if not row["data_same"]:
        why.append("data differ: " + ", ".join(row["files_differ"]))
    if moved:
        why.append("decisions moved: " + ", ".join(moved[:6]) + (" ..." if len(moved) > 6 else ""))
    if row["model_validation"] != "PERFORMED":
        why.append(f"model validation {row['model_validation'] or 'missing'}")
    else:
        if row["model_fit"] == "ACCEPTED":
            why.append(f"MODEL_FIT ACCEPTED ({row['model_fit_sigma']:+.1f} sigma)")
        if row["model_decisions"] != "NONE" or row["enantiomorph_adopted"] != "FALSE":
            why.append(f"model decided {row['model_decisions']}")
        if r_free is None or r_free < RANDOM_R_FREE_MIN:
            why.append(f"R_free {r_free} below a random structure's")
        if row["cc_confirmed"] not in (None, "NONE"):
            why.append(f"CC_MODEL_CONFIRMED_TO_D_MIN {row['cc_confirmed']} (a significant CC with a wrong model)")
    row["verdict"] = "fail" if why else "pass"
    row["reason"] = "; ".join(why)
    return row


def section(doc, rows):
    """The report's own pass/fail section for the wrong-model null (nullmodel_results.json in a run directory)."""
    bad = [r for r in rows if r["verdict"] != "pass"]
    doc.h(2, f"Wrong-model null: {len(rows) - len(bad)} of {len(rows)} pass")
    doc.p("Each set run without --model and with five wrong models (nullmodel.json). A pair passes when "
          "(a) p.hkl, the p.mtz reflections and p.cif are identical to the no-model run's, (b) no "
          "KEY= line outside MODEL VALIDATION moved, and (c) the model is recognised as wrong: "
          "MODEL_FIT not ACCEPTED, no decision taken, CC_MODEL_CONFIRMED_TO_D_MIN NONE and R_free "
          f">= {RANDOM_R_FREE_MIN} (a random structure: 0.586 acentric, 0.828 centric). "
          "NOT_TESTED means the model claimed no hand, indexing or setting, so no null was needed.")
    head = ["set", "wrong model", "(a) data", "(b) decisions", "space group", "hand", "MODEL_FIT",
            "sigma", "decided", "R_free", "CC model", "CC confirmed to", "verdict", "why"]
    doc.table(head, [[r["set"], r["model"], "same" if r["data_same"] else "DIFFER",
                      "same" if not r["decisions_moved"] else f"{len(r['decisions_moved'])} moved",
                      r["sg"] or "-", r["enantiomorph"] or "-", r["model_fit"] or "-",
                      f(r["model_fit_sigma"], "{:+.1f}"), r["model_decisions"] or "-",
                      f(r["r_free"], "{:.3f}"), f(r["cc_model"], "{:.3f}"),
                      r["cc_confirmed"] or "-", r["verdict"], r["reason"]] for r in rows],
              num=(7, 9, 10),
              row_class=lambda row: "fail" if row[12] != "pass" else "",
              cell_class=lambda i, x: {"pass": "pass", "fail": "fail", "crash": "fail"}.get(x) if i == 12 else None)