#!/usr/bin/env python3
"""Rugnux still battery: run serial (XFEL stills) datasets, score the deposited model against each
merge, report, compare.

The rotation battery (battery.py) scores a lattice and a space group. For stills both are given
(-C/-S, the way a serial user runs), so this battery reports what the merge is worth and what it
costs instead: indexing rate, R_split, the R-free of the deposited model against our merge, not
refined (model_check.py: REFMAC's first cycle, the deposition's own free set) beside the R-free of
the same protocol on the deposited data (the target the verdict is taken against; the published
R-free is shown for information), and the wall time of the rugnux command with the input's I/O floor
next to it.

Usage (README.md in this directory, section "The still battery", is the guide):
  still_battery.py run --rugnux BIN [--label L] [--only a,b | --tier smoke] [--threads N]
                       [--cold] [--io-gbps G]
  still_battery.py compare RUN_A RUN_B
  still_battery.py report RUN [--out DIR]          re-render a run's report.md (a run scored by an
                                                   older protocol is rescored: give --out)
  still_battery.py list                            still runs under the runs root
  still_battery.py keep RUN --as workstation|server  copy report.md, results.json, manifest.json
                                                   into results/<role>/<run>/ (committed); refuses a
                                                   private run

The site config is the rotation battery's (--site, $JFJOCH_BATTERY_SITE, tools/battery/site.json),
with a "still" section: {"manifest": "still_manifest.json", "data_root": DIR} and optionally
"private_manifest" for local-only sets. A run writes <runs_root>/<YYYYMMDD-HHMM>_<git>_still-<label>/
(manifest.json, results.json, report.md, work/<set>/) and is made read-only when done; a run directory
is never reused.
"""
import argparse
import datetime
import glob
import json
import os
import re
import shutil
import socket
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import battery  # noqa: E402  binary_info, hardware_info, run_logged, write_json, make_read_only, ...
import model_check  # noqa: E402  check, fetch, RCSB

MANIFEST = os.path.join(HERE, "still_manifest.json")
KEPT_FILES = ("report.md", "results.json", "manifest.json")
IO_GBPS = 5.0          # sequential read speed the I/O floor is quoted at (a PCIe 4 NVMe SSD)
DRAW = 0.01            # R-free difference to the target still called a draw

COLUMNS = [("facility", "facility"), ("beamline", "beamline"), ("detector", "detector"),
           ("dataset", "set"), ("PDB", "pdb"), ("images", "images"), ("images (deposition)", "images_deposited"),
           ("indexed %", "indexed_pct"), ("d_min", "d_min"), ("R_split", "r_split"),
           ("R_free rugnux", "r_free"), ("R_free deposition", "r_free_depdata"),
           ("R_free published", "r_free_deposited"), ("difference", "r_free_diff"), ("quality", "quality"),
           ("time s", "wall_s"), ("ms/image", "ms_per_image"), ("I/O floor s", "io_floor_s"),
           ("page cache", "cache")]


# ------------------------------------------------------------------ configuration

def load_site(path):
    path = path or os.environ.get("JFJOCH_BATTERY_SITE") or battery.LOCAL_SITE
    if not os.path.exists(path):
        sys.exit(f"no site config at {path}: see site.example.json and README.md")
    site = json.load(open(path))
    if "still" not in site:
        sys.exit(f"{path}: no \"still\" section (see site.example.json)")
    base = os.path.dirname(os.path.abspath(path))
    st = site["still"]
    st["data_root"] = os.path.join(base, st["data_root"])
    st["manifest"] = os.path.join(HERE, st.get("manifest", "still_manifest.json"))
    if st.get("private_manifest"):
        st["private_manifest"] = os.path.join(base, st["private_manifest"])
    st["runs_root"] = os.path.join(base, st.get("runs_root") or site["runs_root"])
    st["pdb_cache"] = os.path.join(base, site.get("pdb_cache") or model_check.DEFAULT_CACHE)
    return st


def select(st, only, tier, private):
    path = st.get("private_manifest") if private else st["manifest"]
    if not path:
        sys.exit("--private: the site config names no still private_manifest")
    sets = json.load(open(path))["sets"]
    if only:
        want = [s.strip() for s in only.split(",") if s.strip()]
        unknown = [w for w in want if w not in {e["id"] for e in sets}]
        if unknown:
            sys.exit(f"not in the still manifest: {', '.join(unknown)}")
        sets = [e for e in sets if e["id"] in want]
    if tier:
        sets = [e for e in sets if tier in e.get("tiers", {})]
    return sets


# ------------------------------------------------------------------ one set

def inputs(e, data_root):
    """The input files: each glob of the row in order, each glob's matches sorted by name."""
    files = []
    for pattern in e["inputs"]:
        files += sorted(glob.glob(os.path.join(data_root, pattern)))
    return files


def page_cache(files):
    """Fraction of the input bytes already in the page cache (fincore), or None without fincore."""
    if not shutil.which("fincore") or not files:
        return None
    r = subprocess.run(["fincore", "--bytes", "--noheadings", "--output", "RES,SIZE"] + files,
                       capture_output=True, text=True)
    res = size = 0
    for line in r.stdout.splitlines():
        parts = line.split()
        if len(parts) == 2:
            res, size = res + int(parts[0]), size + int(parts[1])
    return res / size if size else None


def evict(files):
    """Drop the input's pages from the page cache, so the run reads it from the disk."""
    for f in files:
        fd = os.open(f, os.O_RDONLY)
        os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
        os.close(fd)


def cache_state(fraction):
    if fraction is None:
        return "unknown"
    return "cold" if fraction < 0.05 else "warm" if fraction > 0.95 else f"partial {100 * fraction:.0f}%"


def frame_limit(e, opts):
    if opts["tier"]:
        return (e.get("tiers") or {}).get(opts["tier"], {}).get("images")
    return e.get("max_images")   # a set kept to a subset for now (see the manifest)


def command(binary, e, data_root, opts, model):
    """The command run, and the same written for a reader: rugnux, the data root as <data>, the input
    globs instead of the file list, the model as <PDB>.pdb."""
    head = ["-o", "p", "--no-export-unmerged"] + (["-N", str(opts["threads"])] if opts["threads"] else [])
    if frame_limit(e, opts):
        head += ["-e", str(frame_limit(e, opts))]
    tail = ["--model", model] if model else []
    cmd = [binary] + head + [a.replace("{data}", data_root) for a in e["args"]] + tail + inputs(e, data_root)
    short = (["rugnux"] + head + [a.replace("{data}", "<data>") for a in e["args"]]
             + (["--model", e["pdb"].lower() + ".pdb"] if model else []) + ["<data>/" + p for p in e["inputs"]])
    return cmd, " ".join(short)


def report_keys(path):
    if not os.path.exists(path):
        return {}
    return dict(re.findall(r"^([A-Z_]+)= *(.*?)\s*$", open(path, errors="replace").read(), re.M))


def num(s):
    try:
        return float(str(s).split()[0].rstrip("%"))
    except (ValueError, IndexError):
        return None


def deposited_model(pdb, cache):
    """The deposited coordinates: PDB format where the RCSB serves it, else mmCIF (entries too large or
    too new for the PDB format exist only as mmCIF; rugnux reads both)."""
    return (model_check.fetch(model_check.RCSB + pdb + ".pdb", os.path.join(cache, pdb + ".pdb"))
            or model_check.fetch(model_check.RCSB + pdb + ".cif", os.path.join(cache, pdb + ".cif")))


def score(wd, mtz, e, cache):
    """R-free of the deposited model, not refined, against our merge and against the deposited data
    through the same protocol (model_check.py: REFMAC's first cycle, the deposition's free set, both
    data sets' common reflections at the lower of the two resolutions)."""
    res = model_check.check(mtz, e["pdb"], os.path.join(wd, "model_check"), cache)
    return {"r_work": res["rwork"], "r_free": res["rfree_depflags"], "r_free_depdata": res["rfree_depdata"],
            "d_min_used": res["d_min_used"],
            "score_note": res["reason"] or res["baseline_note"] or res["model_modification"]}


def quality(r_free, goal):
    """win / draw / loss of our R-free against the deposited data's."""
    if r_free is None or goal is None:
        return None
    d = r_free - goal
    return "win" if d < -DRAW else "draw" if d <= DRAW else "loss"


def run_one(run_dir, binary, e, st, opts):
    wd = os.path.join(run_dir, "work", e["id"])
    os.makedirs(wd)
    files = inputs(e, st["data_root"])
    row = {k: e.get(k) for k in ("facility", "beamline", "detector", "pdb", "r_free_deposited",
                                 "images_deposited", "format")}
    row.update(set=e["id"], dataset=e.get("cxidb"))
    if not files:
        return dict(row, verdict="not_run", reason="no input files under the data root")
    model = None
    if e.get("pdb"):
        pdb = e["pdb"].lower()
        model = deposited_model(pdb, st["pdb_cache"])
    cmd, short = command(binary, e, st["data_root"], opts, model)
    io_bytes = sum(os.path.getsize(f) for f in files)
    if opts["cold"]:
        evict(files)
    cached = page_cache(files)
    t0 = time.time()
    rc = battery.run_logged(cmd, wd, opts["timeout"])
    wall = time.time() - t0
    keys = report_keys(os.path.join(wd, "p_report.txt"))
    images = num(keys.get("IMAGES_PROCESSED"))
    rng = keys.get("INCLUDE_RESOLUTION_RANGE", "").split()
    row.update(command=short, command_run=" ".join(cmd), exit_code=rc, wall_s=round(wall, 1), io_bytes=io_bytes,
               # with a frame limit (a smoke tier) only part of each file is read: no floor then
               io_floor_s=None if frame_limit(e, opts) else round(io_bytes / (opts["io_gbps"] * 1e9), 1),
               cache=cache_state(cached),
               input_files=len(files), images=int(images) if images else None,
               ms_per_image=round(1000 * wall / images, 2) if images else None,
               indexed_pct=round(100 * num(keys["INDEXING_RATE"]), 1) if keys.get("INDEXING_RATE") else None,
               d_min=num(rng[-1]) if rng else None, r_split=num(keys.get("R_SPLIT")),
               cc_half_split=num(keys.get("CC_HALF_SPLIT")), r_free_model_check=num(keys.get("R_FREE")))
    if rc != 0 or not keys:
        return dict(row, verdict="crash" if rc else "fail", reason=battery.failure_note(os.path.join(wd, "run.log")))
    if e.get("pdb"):
        row.update(score(wd, os.path.join(wd, "p.mtz"), e, st["pdb_cache"]))
    return dict(row, verdict="done", reason="")


# ------------------------------------------------------------------ report

def fmt(v, key):
    if v is None:
        return "-"
    if key == "r_split":
        return f"{100 * v:.1f}%"
    if key in ("r_free", "r_free_deposited", "r_free_depdata"):
        return f"{v:.3f}"
    if key == "r_free_diff":
        return f"{v:+.3f}"
    if key == "images_deposited" and isinstance(v, (int, float)):
        return f"{int(v):,}"
    if key == "images":
        return f"{v:,}"
    return str(v)


def render(run_dir, cache, src=None):
    """report.md of a run. A row scored by an older protocol (no r_free_depdata) is rescored from the
    merge in src's work/<set>/ (src: the run itself unless the report is written elsewhere), into
    run_dir's work/<set>/model_check/, and results.json is rewritten."""
    man = json.load(open(os.path.join(run_dir, "manifest.json")))
    res = json.load(open(os.path.join(run_dir, "results.json")))
    sets = {e["id"]: e for e in man["sets"]}
    stale = [r for r in res if r.get("verdict") == "done" and r.get("pdb") and "r_free_depdata" not in r]
    for r in stale:
        r.update(score(os.path.join(run_dir, "work", r["set"]),
                       os.path.join(src or run_dir, "work", r["set"], "p.mtz"), sets[r["set"]], cache))
    if stale:
        battery.write_json(os.path.join(run_dir, "results.json"), res)
    for r in res:
        if r.get("r_free") is not None and r.get("r_free_depdata") is not None:
            r["r_free_diff"] = r["r_free"] - r["r_free_depdata"]
        r["quality"] = quality(r.get("r_free"), r.get("r_free_depdata"))
    b = man["binary"]
    hw = man.get("hardware") or {}
    lines = [f"# Rugnux still battery - {os.path.basename(os.path.normpath(run_dir))}", "",
             f"- binary: rugnux {b.get('version')} (git {b.get('git')}), {b.get('flags', '?')}",
             f"- host: {man['host']}; {hw.get('cpu_model', '?')}, {hw.get('threads', '?')} threads, "
             f"{(hw.get('memory_gib') or 0):.0f} GiB; GPUs: "
             f"{', '.join(g['name'] for g in hw.get('gpus', [])) or 'none'}",
             f"- start {man['start']}, end {man.get('end') or 'unfinished'}; subset: {man.get('subset') or 'all sets'}"
             + ("; ABORTED" if man.get("aborted") else ""),
             f"- R_free rugnux: the deposited model against our merge, not refined (REFMAC's first cycle, bulk "
             f"solvent and anisotropic scale, the deposition's free set, the common reflections at the lower of "
             f"the two resolutions; model_check.py); R_free deposition: the same protocol on the deposited data, "
             f"the target; R_free published: the entry's own value, for information; difference: ours minus "
             f"the deposition's; quality: win where ours is lower by more than {DRAW}, "
             f"draw within {DRAW}, loss otherwise. Time: wall clock of the rugnux "
             f"command (to model validation); I/O floor: input bytes / {man['options']['io_gbps']} GB/s; page "
             f"cache: share of the input already in memory when the run started.", "",
             "| " + " | ".join(c for c, _ in COLUMNS) + " |", "|" + "---|" * len(COLUMNS)]
    for r in res:
        lines.append("| " + " | ".join(fmt(r.get(k), k) for _, k in COLUMNS) + " |")
    bad = [r for r in res if r.get("verdict") != "done"]
    if bad:
        lines += ["", "## Not processed", ""] + [f"- {r['set']}: {r.get('verdict')} - {r.get('reason')}" for r in bad]
    lines += ["", "## Commands", ""] + [f"- `{r['set']}`: `{r.get('command', '-')}`" for r in res]
    name = "report_PRIVATE.md" if man.get("private") else "report.md"
    open(os.path.join(run_dir, name), "w").write("\n".join(lines) + "\n")
    return os.path.join(run_dir, name)


# ------------------------------------------------------------------ commands

def cmd_run(a, st):
    sets = select(st, a.only, a.tier, a.private)
    if not sets:
        sys.exit("nothing selected")
    info = battery.binary_info(a.rugnux, a.build_dir)
    label = "still-" + (a.label or info.get("version") or "run")
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
    run_dir = os.path.join(st["runs_root"], f"{stamp}_{info['git']}_{label}" + ("_private" if a.private else ""))
    os.makedirs(run_dir)                      # never reuse or overwrite a run directory
    binary = os.path.join(run_dir, "bin", "rugnux")
    os.makedirs(os.path.dirname(binary))
    shutil.copy2(a.rugnux, binary)
    info.update(copied_from=os.path.realpath(a.rugnux), sha256=battery.sha256(binary))
    opts = {"threads": a.threads, "timeout": a.timeout, "tier": a.tier,
            "io_gbps": a.io_gbps, "cold": a.cold}
    man = {"kind": "still", "label": label, "private": a.private, "subset": a.only or (f"tier {a.tier}" if a.tier else None),
           "binary": info, "runner_git": battery.git(HERE, "rev-parse", "--short", "HEAD") or "unknown",
           "host": socket.gethostname(), "hardware": battery.hardware_info(),
           "start": datetime.datetime.now().isoformat(timespec="seconds"), "end": None, "complete": False,
           "options": opts, "sets": sets}
    battery.write_json(os.path.join(run_dir, "manifest.json"), man)
    print(f"run directory {run_dir}", flush=True)
    results = []
    try:
        for e in sets:
            row = run_one(run_dir, binary, e, st, opts)
            results.append(row)
            battery.write_json(os.path.join(run_dir, "results.json"), results)
            print(f"[{len(results)}/{len(sets)}] {e['id']:16s} {row['verdict']:8s} img={row.get('images')} "
                  f"idx={row.get('indexed_pct')}% Rsplit={row.get('r_split')} Rfree={row.get('r_free')} "
                  f"t={row.get('wall_s')}s io={row.get('io_floor_s')}s {row.get('cache')} {row.get('reason') or ''}",
                  flush=True)
    except KeyboardInterrupt:
        battery.mark_aborted(run_dir, f"interrupted after {len(results)} of {len(sets)} sets")
        sys.exit(f"aborted; {run_dir} is marked as such")
    man.update(end=datetime.datetime.now().isoformat(timespec="seconds"), complete=True)
    battery.write_json(os.path.join(run_dir, "manifest.json"), man)
    print(f"report {render(run_dir, st['pdb_cache'])}", flush=True)
    battery.make_read_only(run_dir)
    if any(r["verdict"] in ("crash", "fail") for r in results):
        sys.exit("some sets failed: " + ", ".join(r["set"] for r in results if r["verdict"] in ("crash", "fail")))


def run_path(st, run):
    return run if os.path.isdir(run) else os.path.join(st["runs_root"], run)


def cmd_compare(a, st):
    ra = {r["set"]: r for r in json.load(open(os.path.join(run_path(st, a.run_a), "results.json")))}
    rb = {r["set"]: r for r in json.load(open(os.path.join(run_path(st, a.run_b), "results.json")))}
    keys = [("indexed %", "indexed_pct"), ("d_min", "d_min"), ("R_split", "r_split"), ("R_free", "r_free"),
            ("time s", "wall_s"), ("cache", "cache")]
    print("| set | " + " | ".join(f"{c} A -> B" for c, _ in keys) + " |")
    print("|---|" + "---|" * len(keys))
    for s in [s for s in ra if s in rb] + [s for s in rb if s not in ra]:
        x, y = ra.get(s, {}), rb.get(s, {})
        print(f"| {s} | " + " | ".join(f"{x.get(k, '-')} -> {y.get(k, '-')}" for _, k in keys) + " |")


def cmd_list(a, st):
    for d in sorted(glob.glob(os.path.join(st["runs_root"], "*_still-*"))):
        man = json.load(open(os.path.join(d, "manifest.json")))
        state = "aborted" if man.get("aborted") else "complete" if man.get("complete") else "unfinished"
        print(f"{os.path.basename(d):60s} {state}")


def cmd_report(a, st):
    src = run_path(st, a.run)
    dest = src
    if a.out:
        dest = os.path.join(a.out, os.path.basename(os.path.normpath(src)))
        os.makedirs(dest, exist_ok=True)
        for f in ("manifest.json", "results.json"):
            shutil.copyfile(os.path.join(src, f), os.path.join(dest, f))
    print(render(dest, st["pdb_cache"], src))


def cmd_keep(a, st):
    src = run_path(st, a.run)
    man = json.load(open(os.path.join(src, "manifest.json")))
    if man.get("private") or man.get("kind") != "still":
        sys.exit(f"{src}: refused - not a public still run")
    dest = os.path.join(battery.KEPT_ROOT, a.role, os.path.basename(os.path.normpath(src)))
    os.makedirs(dest, exist_ok=True)
    for name in KEPT_FILES:
        shutil.copyfile(os.path.join(src, name), os.path.join(dest, name))
    print(dest)


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--site")
    sub = ap.add_subparsers(dest="cmd", required=True)
    r = sub.add_parser("run")
    r.add_argument("--rugnux", required=True)
    r.add_argument("--label")
    r.add_argument("--only")
    r.add_argument("--tier")
    r.add_argument("--private", action="store_true", help="the site's local-only still manifest")
    r.add_argument("--threads", type=int)
    r.add_argument("--cold", action="store_true", help="evict each set's input from the page cache before its run")
    r.add_argument("--timeout", type=int, default=6 * 3600)
    r.add_argument("--build-dir")
    r.add_argument("--io-gbps", type=float, default=IO_GBPS)
    c = sub.add_parser("compare")
    c.add_argument("run_a")
    c.add_argument("run_b")
    p = sub.add_parser("report")
    p.add_argument("run")
    p.add_argument("--out")
    sub.add_parser("list")
    k = sub.add_parser("keep")
    k.add_argument("run")
    k.add_argument("--as", dest="role", required=True, choices=["workstation", "server"])
    a = ap.parse_args()
    st = load_site(a.site)
    {"run": cmd_run, "compare": cmd_compare, "list": cmd_list, "report": cmd_report, "keep": cmd_keep}[a.cmd](a, st)


if __name__ == "__main__":
    main()