#!/usr/bin/env python3
"""paper_plots.py -- three battery figures of the open arm: R-free and high-resolution limit
(deposition vs Rugnux) and the processing time per dataset.

    tools/battery/paper_plots.py RUN [RUN ...] --out DIR [--time-runs RUN [RUN ...]] [--time-arms ARM [ARM ...]]
                                 [--uncontended RUN]

RUN is a battery run directory, or its name under the site's runs_root or among the kept runs
(tools/battery/results/<role>/<run>/, battery.py keep); later runs fill in sets
an earlier one lacks or that failed to run there (a `-redo` of CUDA casualties). --time-runs draws
the time histogram from other runs than the two quality figures - a run whose sets had the GPU to
themselves, where the runs for the quality figures had to share it. The quality figures read only the
open arm; --time-arms widens the time figures (e.g. `open inhouse`). Everything the figures are drawn from is written
beside them, so a figure can be audited and redrawn: rfree.dat, dmin.dat (value pairs with the set
and its beamline), time.dat (set, arm, seconds, images, other GPU processes), excluded.txt (every
set left out, with the reason), summary.txt (the runs, their binary, the counts and the timing
statistics), figures.gp (the gnuplot script) and the rendered figures: a single-column vector PDF
with embedded fonts and the same figure as a 600 d.p.i. PNG. gnuplot and pdftoppm (poppler) must be
on PATH; gemmi reads the facility and beamline from each entry's cached mmCIF (_diffrn_source).

Left figure: R-free of the deposited model after the battery's REFMAC check on the depositor's data
(x) and on Rugnux's data (y): the model as deposited, not refined, REFMAC's first cycle on the
depositor's free set and both data sets' common reflections (model_check.py).
A run without REFMAC values (the check is opt-in, --model-check) gets no R-free figure; the rest are drawn.
Right figure: the deposited high-resolution limit (x) against Rugnux's own cut
(y), log axes so the low-resolution entries show. Small molecules (no deposited model) and sets
whose verdict is `fail` (a wrong lattice: nothing to compare) are left out.

Third figure: a histogram of the processing time of every open-arm set that ran to a report, whatever its verdict, with the median marked. The time is `rugnux_wall_s`, the
WALL_TIME rugnux itself reports over the whole invocation (reading the images through to the
written merge) - not the runner's `elapsed_s`, which also counts the wait for a GPU slot under
--gpulock (`wall_s` is the same value with `elapsed_s` as the fallback). Linear 25 s bins: the
distribution is long-tailed but spans little more than a decade, and a linear axis keeps "how long
does a dataset take" readable. Two caveats the reader must know, both written to summary.txt:
each set runs once, so a time includes reading the images from disk unless they were still in the
page cache (a cold read of a 10 GB dataset from the workstation's HDD costs of the order of a
minute), and a run under --gpulock can share the GPU with other processes - `gpu_others` counts
them as seen just before the set started. Sets that ran with a shared GPU are drawn as the lighter
top segment of their bar and are never dropped; with --uncontended RUN (a run of the same sets
with the GPU to itself) summary.txt also reports how much slower they were than there. Sets that
crashed have no WALL_TIME and go to excluded.txt. A fourth figure, time_zoom, is the same histogram
over the first 60 s in 5 s bins, where most sets fall.
"""
import argparse
import collections
import json
import math
import os
import re
import shutil
import subprocess
import sys

import gemmi

RUNS_ROOT = "/data/battery/runs"
KEPT_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")
RFREE_BAND = 0.05     # |R-free Rugnux - R-free depositor| counted as "within"
TIME_BIN = 25         # s, histogram bin width
ZOOM_MAX = 60         # s, range of the zoomed time histogram
ZOOM_BIN = 5          # s, its bin width

# Sized and lettered to the IUCr artwork guide (journals.iucr.org/services/help/artwork/guide.html):
# one column is 8.85 cm, lettering upright in a standard font (Helvetica, embedded by cairo) and in
# the 1.5-3 mm range of the notes for authors (10 pt: capitals ~2.5 mm),
# line weights 0.35-1.5 pt at final size (pdfcairo lw 1 = 0.5 pt), no grid, a colour that stays
# legible in greyscale. Ticks are short, outward and not mirrored.
GNUPLOT = r"""
set terminal pdfcairo size 8.8cm,8.8cm font "Helvetica,10"
set encoding utf8
set size ratio 1
set border lw 1.5
set tics out scale 0.4 nomirror
set key off
set style line 1 lc rgb "black" lw 1.5 dt 2
set style line 2 pt 7 ps 0.35 lc rgb "#0072b2"
# counts: names left-aligned, numbers right-aligned in a column beside them
set label 11 "Structures\nFacilities\nBeamlines" at graph 0.05, graph 0.93 left front
set label 12 "%(n)d\n%(nf)d\n%(nb)d" at graph 0.31, graph 0.93 right front
"""

# Drawn only where the runs carry REFMAC R-free values; the battery's REFMAC check is opt-in.
GNUPLOT_RFREE = r"""
# --- R-free
set output "%(out)s/rfree.pdf"
set xlabel "{/:Italic R}_{free}, depositor data"
set ylabel "{/:Italic R}_{free}, Rugnux data"
set xrange [%(rlo).2f:%(rhi).2f]; set yrange [%(rlo).2f:%(rhi).2f]
set xtics 0.1; set ytics 0.1
set format x "%%.1f"; set format y "%%.1f"
set label 2 "Within ±%(band).2f: %(within)d/%(n)d" at graph 0.95, graph 0.07 right front
plot x ls 1, "%(out)s/rfree.dat" using 1:2 ls 2
unset label 2
"""

GNUPLOT_REST = r"""
# --- high-resolution limit
set output "%(out)s/dmin.pdf"
set logscale xy
set xlabel "{/:Italic d}_{min}, deposition (Å)"
set ylabel "{/:Italic d}_{min}, Rugnux (Å)"
set xrange [%(dlo).2f:%(dhi).2f]; set yrange [%(dlo).2f:%(dhi).2f]
set xtics (0.5, 1, 2, 4, 8); set ytics (0.5, 1, 2, 4, 8)
set format x "%%.1f"; set format y "%%.1f"
set label 12 "%(dn)d\n%(dnf)d\n%(dnb)d" at graph 0.31, graph 0.93 right front
set label 3 "Rugnux deeper: %(deeper)d" at graph 0.95, graph 0.07 right front
set label 4 "Deposition deeper: %(shallower)d" at graph 0.05, graph 0.68 left front
plot x ls 1, "%(out)s/dmin.dat" using 1:2 ls 2
unset label 3; unset label 4

# --- processing time: all sets in the lighter fill, those with the GPU to themselves over them in
# full colour, so the lighter top of a bar is the sets that shared the GPU
set output "%(out)s/time.pdf"
unset logscale
set xlabel "Processing time (s)"
set ylabel "Datasets"
set xrange [0:%(thi)d]; set yrange [0:%(tyhi)d]
set xtics 100; set ytics 10
set format x "%%.0f"; set format y "%%.0f"
set boxwidth %(bin)d absolute
set style fill solid noborder
set label 11 "Datasets" at graph 0.64, graph 0.93 left front
set label 12 "%(tn)d" at graph 0.95, graph 0.93 right front
%(tkey)s
set arrow 1 from %(tmed).1f, graph 0 to %(tmed).1f, graph 1 nohead ls 1 front
set label 5 "median %(tmed).0f s" at %(tmed).1f + 0.02 * %(thi)d, graph 0.93 left front
bin(x) = (floor(x / %(bin)d) + 0.5) * %(bin)d
plot "%(out)s/time.dat" using (bin($3)):(1) smooth frequency with boxes lc rgb "#8fc0e0" title "GPU shared", \
     "" using (bin($3)):($5 > 0 ? 0 : 1) smooth frequency with boxes lc rgb "#0072b2" title "GPU exclusive"

# --- the same over the first %(zhi)d s; sets beyond it are counted in the label, not drawn
set output "%(out)s/time_zoom.pdf"
set xrange [0:%(zhi)d]; set yrange [0:%(zyhi)d]
set xtics 10; set ytics 5
set boxwidth %(zbin)d absolute
set label 11 "Datasets\nbeyond %(zhi)d s" at graph 0.64, graph 0.93 left front
set label 12 "%(tn)d\n%(zbeyond)d" at graph 0.95, graph 0.93 right front
set label 5 "median %(tmed).0f s" at %(tmed).1f + 0.02 * %(zhi)d, graph 0.80 left front
zbin(x) = (floor(x / %(zbin)d) + 0.5) * %(zbin)d
plot "%(out)s/time.dat" using (zbin($3)):($3 < %(zhi)d ? 1 : 0) smooth frequency with boxes lc rgb "#8fc0e0" title "GPU shared", \
     "" using (zbin($3)):($5 > 0 || $3 >= %(zhi)d ? 0 : 1) smooth frequency with boxes lc rgb "#0072b2" title "GPU exclusive"
"""


def run_path(run):
    """A run directory as given, else by name under RUNS_ROOT, else among the kept runs."""
    if os.path.isdir(run):
        return run
    if os.path.isdir(os.path.join(RUNS_ROOT, run)):
        return os.path.join(RUNS_ROOT, run)
    kept = [os.path.join(KEPT_ROOT, h, run) for h in sorted(os.listdir(KEPT_ROOT))
            if os.path.isdir(os.path.join(KEPT_ROOT, h, run))] if os.path.isdir(KEPT_ROOT) else []
    if not kept:
        sys.exit(f"{run}: not a directory, not under {RUNS_ROOT} and not a kept run")
    return kept[0]


def load_rows(runs, arms=("open",)):
    """Rows of the given arms, a later run's row replacing one that did not run."""
    rows, manifests = {}, []
    for run in runs:
        path = run_path(run)
        manifests.append((path, json.load(open(os.path.join(path, "manifest.json")))))
        for r in json.load(open(os.path.join(path, "results.json"))):
            if r["arm"] not in arms:
                continue
            if r["set"] not in rows or rows[r["set"]].get("exit_code"):
                rows[r["set"]] = r
    return rows, manifests


def percentile(values, p):
    v = sorted(values)
    return v[round(p * (len(v) - 1))]


def source(model_path):
    """(facility, beamline) from the entry's _diffrn_source, facility without its operator suffix
    ("PETRA III, EMBL c/o DESY" -> "PETRA III"); None where the entry does not record one."""
    block = gemmi.cif.read(model_path).sole_block()

    def value(tag):
        v = block.find_value(tag)
        return None if v in (None, "?", ".") else gemmi.cif.as_string(v).strip()

    site = value("_diffrn_source.pdbx_synchrotron_site")
    if not site:
        return None
    beamline = value("_diffrn_source.pdbx_synchrotron_beamline") or "?"
    facility = site.split(",")[0].strip().upper()
    return facility, facility + " " + re.sub(r"\s+", "", beamline.upper())


def counts(sets, where):
    known = [where[s] for s in sets if where[s]]
    return len(sets), len({f for f, _ in known}), len({b for _, b in known})


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("runs", nargs="+", help="battery run directories or names (later ones fill in)")
    ap.add_argument("--out", required=True, help="output directory")
    ap.add_argument("--time-runs", nargs="+", help="runs to draw the time histogram from "
                                                  "(default: the runs above)")
    ap.add_argument("--time-arms", nargs="+", default=["open"],
                    help="arms the time figures count (default: open)")
    ap.add_argument("--uncontended", help="a run of the same sets with the GPU to itself, to size the "
                                          "slowdown of the sets that shared it")
    a = ap.parse_args()
    for tool in ("gnuplot", "pdftoppm"):
        if not shutil.which(tool):
            sys.exit(tool + " not found on PATH")
    os.makedirs(a.out, exist_ok=True)
    out = os.path.abspath(a.out)

    rows, manifests = load_rows(a.runs)
    time_rows, time_manifests = load_rows(a.time_runs or a.runs, a.time_arms)
    if not a.time_runs:
        time_manifests = []
    excluded, rfree, dmin, where, times = [], [], [], {}, []
    for s, r in sorted(time_rows.items()):
        if r.get("rugnux_wall_s") is None:
            excluded.append((s, "time", "no WALL_TIME: " + str(r.get("reason"))))
        else:
            times.append((s, r["arm"], r["rugnux_wall_s"], r.get("images"), r.get("gpu_others")))
    for s, r in sorted(rows.items()):
        if not r.get("model"):
            excluded.append((s, "both", "no deposited model (small molecule or unpublished)"))
            continue
        if r.get("verdict") == "fail":
            excluded.append((s, "both", "verdict fail: " + str(r.get("reason"))))
            continue
        where[s] = source(r["model"])
        dep, rug = r.get("refmac_rfree_depdata"), r.get("refmac_rfree_depflags")
        if dep and rug:
            rfree.append((dep, rug, s))
        else:
            excluded.append((s, "rfree", "no REFMAC pair: " + str(r.get("refmac_reason")
                                                                   or r.get("dep_reason") or "depositor data not scored")))
        if r.get("d_min_ref") and r.get("d_min"):
            dmin.append((r["d_min_ref"], r["d_min"], s))
        else:
            excluded.append((s, "dmin", "no deposited or Rugnux d_min"))

    bl = lambda s: (where[s][1] if where[s] else "not-recorded").replace(" ", "_")
    with open(os.path.join(out, "rfree.dat"), "w") as fh:
        fh.write("# rfree_depositor_data rfree_rugnux_data set beamline\n")
        fh.writelines(f"{x:.4f} {y:.4f} {s} {bl(s)}\n" for x, y, s in rfree)
    with open(os.path.join(out, "dmin.dat"), "w") as fh:
        fh.write("# dmin_deposition dmin_rugnux set beamline\n")
        fh.writelines(f"{x:.3f} {y:.3f} {s} {bl(s)}\n" for x, y, s in dmin)
    with open(os.path.join(out, "time.dat"), "w") as fh:
        fh.write("# set arm seconds images gpu_others\n")
        fh.writelines(f"{s} {arm} {t:.2f} {'nan' if im is None else round(im)} {'nan' if g is None else g}\n"
                      for s, arm, t, im, g in times)
    with open(os.path.join(out, "excluded.txt"), "w") as fh:
        fh.writelines(f"{s}\t{which}\t{why}\n" for s, which, why in excluded)

    n, nf, nb = counts([s for _, _, s in rfree], where)
    dn, dnf, dnb = counts([s for _, _, s in dmin], where)
    within = sum(abs(y - x) <= RFREE_BAND for x, y, _ in rfree)
    deeper = sum(y < x for x, y, _ in dmin)
    values = [v for x, y, _ in rfree for v in (x, y)]
    secs = [t for _, _, t, _, _ in times]
    tmed = percentile(secs, 0.5)
    shared = [s for s, _, _, _, g in times if g]
    tallest = max(collections.Counter(math.floor(t / TIME_BIN) for t in secs).values())
    # the shared/exclusive key only means something when some set did share the GPU
    tkey = ("set key at graph 0.95, graph 0.76 right Left reverse samplen 1.5 spacing 1.1" if shared
            else "set key off")
    params = dict(out=out, tkey=tkey, n=n, nf=nf, nb=nb, dn=dn, dnf=dnf, dnb=dnb, within=within, band=RFREE_BAND,
                  deeper=deeper, shallower=sum(y > x for x, y, _ in dmin),
                  rlo=0.05 * math.floor((min(values, default=0.2) - 0.01) / 0.05),
                  rhi=0.05 * math.ceil((max(values, default=0.3) + 0.01) / 0.05),
                  dlo=0.5, dhi=8.0,
                  bin=TIME_BIN, thi=TIME_BIN * math.ceil(max(secs) / TIME_BIN), tyhi=10 * math.ceil((tallest + 1) / 10),
                  tmed=tmed, tn=len(times),
                  zhi=ZOOM_MAX, zbin=ZOOM_BIN, zbeyond=sum(t >= ZOOM_MAX for t in secs),
                  zyhi=5 * math.ceil((max(collections.Counter(math.floor(t / ZOOM_BIN) for t in secs
                                                              if t < ZOOM_MAX).values()) + 1) / 5))
    gp = os.path.join(out, "figures.gp")
    open(gp, "w").write((GNUPLOT + (GNUPLOT_RFREE if rfree else "") + GNUPLOT_REST) % params)
    subprocess.run(["gnuplot", gp], check=True)
    for fig in (["rfree"] if rfree else []) + ["dmin", "time", "time_zoom"]:    # the PNG is the PDF rasterised at 600 d.p.i.
        subprocess.run(["pdftoppm", "-r", "600", "-png", "-singlefile", os.path.join(out, fig + ".pdf"),
                        os.path.join(out, fig)], check=True)

    with open(os.path.join(out, "summary.txt"), "w") as fh:
        for path, m in manifests + [(p, dict(m, label=f"{m.get('label')} (time figure only)"))
                                    for p, m in time_manifests]:
            b = m.get("binary", {})
            fh.write(f"run {path}\n  label {m.get('label')}  binary git {b.get('source_head')} "
                     f"dirty {b.get('source_dirty')}  sha256 {b.get('sha256')}\n  flags {b.get('flags')}  "
                     f"runner git {m.get('runner_git')} dirty {m.get('runner_dirty')}\n")
        fh.write((f"R-free: {n} structures, {nf} facilities, {nb} beamlines; within +-{RFREE_BAND}: {within}\n"
                  if rfree else "R-free: no REFMAC values in these runs (run the battery with --model-check); "
                                "no rfree figure\n")
                 + f"d_min: {dn} structures, {dnf} facilities, {dnb} beamlines; Rugnux reaches further on "
                 f"{deeper}\n")
        queue_wait = max(time_rows[s]["elapsed_s"] - t for s, _, t, _, _ in times)
        fh.write(f"time: {len(times)} sets of arm(s) {', '.join(a.time_arms)}; field rugnux_wall_s "
                 f"(rugnux's WALL_TIME over the whole invocation, not the runner's elapsed_s, which adds the wait "
                 f"for a GPU slot: here at most {queue_wait:.1f} s)\n"
                 f"  median {tmed:.1f} s, mean {sum(secs) / len(secs):.1f} s, p10 {percentile(secs, 0.1):.1f}, "
                 f"p25 {percentile(secs, 0.25):.1f}, p75 {percentile(secs, 0.75):.1f}, p90 {percentile(secs, 0.9):.1f}, "
                 f"min {min(secs):.1f}, max {max(secs):.1f}\n"
                 f"  each set ran once: a time includes reading the images from disk unless they were still in the "
                 f"page cache\n"
                 f"  GPU shared with other processes (gpu_others > 0) on {len(shared)} sets, kept and drawn as the "
                 f"lighter segment; median without them "
                 f"{percentile([t for s, _, t, _, _ in times if s not in shared], 0.5):.1f} s\n")
        if a.uncontended:
            other, _ = load_rows([a.uncontended])
            ratio = lambda names: percentile([t / other[s]["rugnux_wall_s"] for s, _, t, _, _ in times
                                              if s in names and other.get(s, {}).get("rugnux_wall_s")], 0.5)
            exclusive = [s for s, _, _, _, g in times if not g]
            fh.write(f"  against {a.uncontended}: median time ratio {ratio(shared):.2f} on the shared-GPU sets, "
                     f"{ratio(exclusive):.2f} on the rest\n")
        fh.write("excluded (set, figure, reason): excluded.txt\n")
    print(open(os.path.join(out, "summary.txt")).read(), end="")


if __name__ == "__main__":
    main()