#!/usr/bin/env python3
"""Check the determinism-cost record against its own readings, and every figure the text
quotes from it.

    python3 -I wai_determinism_cost_verify.py <record.json> [<text.md> ...]

Written from the procedure in research/determinism-cost.md, sharing no code with
tools/wai_determinism_cost.py, which wrote the record. From the record's raw readings
alone it recomputes every value of `summary`: the BD-rates (its own monotone cubic
interpolant, integrated by two-point Gauss-Legendre on each knot interval, exact for a
cubic), the per-quality figures, the cross-platform tallies
and the timing statistics with their verdicts. Then each `<!--det:PATH-->` marker in the
text files must follow a number equal to the record's value at PATH, rounded as printed.
Standard library only. Exit 0 when everything holds, 1 with one line per finding.
"""
import json
import math
import re
import sys

FINDINGS = []


def bad(msg):
    FINDINGS.append(msg)


def median(xs):
    xs = sorted(xs)
    n = len(xs)
    if n == 0:
        return None
    return xs[(n - 1) // 2] if n % 2 == 1 else 0.5 * (xs[n // 2 - 1] + xs[n // 2])


def describe(xs):
    xs = list(xs)
    if not xs:
        return None
    return {"n": len(xs), "median": median(xs), "mean": math.fsum(xs) / len(xs), "min": min(xs), "max": max(xs)}


def psnr(sse, samples):
    return 10 * math.log10(65025.0 * samples / sse)


def msssim_db(v):
    return -10 * math.log10(1 - v)


def hermite_slopes(xs, ys):
    """Shape-preserving slopes (Fritsch-Carlson with Brodlie weights; one-sided
    three-point ends, clipped to keep the shape)."""
    n = len(xs)
    dx = [xs[i + 1] - xs[i] for i in range(n - 1)]
    sec = [(ys[i + 1] - ys[i]) / dx[i] for i in range(n - 1)]
    sl = []
    for i in range(n):
        if i == 0 or i == n - 1:
            j = 0 if i == 0 else n - 2
            k = 1 if i == 0 else n - 3
            a, b = dx[j], dx[k]
            s = ((2 * a + b) * sec[j] - a * sec[k]) / (a + b)
            if math.copysign(1, s) != math.copysign(1, sec[j]) or s == 0 or sec[j] == 0:
                s = 0.0
            elif math.copysign(1, sec[j]) != math.copysign(1, sec[k]) and abs(s) > 3 * abs(sec[j]):
                s = 3 * sec[j]
            sl.append(s)
        else:
            a, b = sec[i - 1], sec[i]
            if a == 0 or b == 0 or (a > 0) != (b > 0):
                sl.append(0.0)
            else:
                wa, wb = 2 * dx[i] + dx[i - 1], dx[i] + 2 * dx[i - 1]
                sl.append((wa + wb) / (wa / a + wb / b))
    return sl


def hermite_eval(xs, ys, sl, x):
    i = 0
    while i < len(xs) - 2 and x > xs[i + 1]:
        i += 1
    h = xs[i + 1] - xs[i]
    t = (x - xs[i]) / h
    return ((2 * t ** 3 - 3 * t ** 2 + 1) * ys[i] + (t ** 3 - 2 * t ** 2 + t) * h * sl[i]
            + (-2 * t ** 3 + 3 * t ** 2) * ys[i + 1] + (t ** 3 - t ** 2) * h * sl[i + 1])


def integral(xs, ys, sl, a, b):
    """The interpolant's integral over [a, b]: on each knot interval's share of it, two-point
    Gauss-Legendre, which is exact for the cubic that interval carries."""
    total = 0.0
    g = 1 / math.sqrt(3)
    for i in range(len(xs) - 1):
        lo, hi = max(a, xs[i]), min(b, xs[i + 1])
        if hi <= lo:
            continue
        mid, half = (lo + hi) / 2, (hi - lo) / 2
        total += half * (hermite_eval(xs, ys, sl, mid - half * g) + hermite_eval(xs, ys, sl, mid + half * g))
    return total


def bd(r_a, q_a, r_b, q_b):
    """Percentage more rate curve b needs than curve a at equal quality (VCEG-M33),
    log10(rate) a monotone cubic of quality, over the common quality range."""
    curves = []
    for r, q in ((r_a, q_a), (r_b, q_b)):
        pairs = sorted(zip(q, r))
        qs = [p[0] for p in pairs]
        if any(qs[i + 1] <= qs[i] for i in range(len(qs) - 1)):
            return None
        ls = [math.log10(p[1]) for p in pairs]
        curves.append((qs, ls, hermite_slopes(qs, ls)))
    lo = max(c[0][0] for c in curves)
    hi = min(c[0][-1] for c in curves)
    if hi <= lo:
        return None
    ia, ib = (integral(c[0], c[1], c[2], lo, hi) for c in curves)
    return 100 * (10 ** ((ib - ia) / (hi - lo)) - 1)


def close(a, b, rel=1e-9, ab=1e-9):
    if a is None or b is None:
        return a is b
    if isinstance(a, (list, str, bool)) or isinstance(b, (list, str, bool)):
        return a == b
    return abs(a - b) <= max(ab, rel * max(abs(a), abs(b)))


def compare(path, mine, theirs, tol=None):
    """Recursively compare a recomputed value with the record's."""
    if isinstance(mine, dict):
        if not isinstance(theirs, dict):
            bad(f"{path}: the record has {type(theirs).__name__}, recomputed a mapping")
            return
        for k in set(mine) | set(theirs):
            if k not in theirs:
                bad(f"{path}.{k}: missing from the record")
            elif k not in mine:
                bad(f"{path}.{k}: in the record, not recomputed")
            else:
                compare(f"{path}.{k}", mine[k], theirs[k], tol)
        return
    ok = close(mine, theirs, *(tol or ())) if not isinstance(mine, int) or isinstance(mine, bool) else mine == theirs
    if not ok:
        bad(f"{path}: the record says {theirs!r}, the readings give {mine!r}")


def recompute(rec):
    s = {}
    items = rec["pictures"]["items"]
    qnums = sorted(int(q) for q in rec["models"])
    P = rec["points"]
    curves = {"psnr": {}, "ms_ssim": {}}
    for it in items:
        px = it["width"] * it["height"]
        keys = [it["id"] + "/q" + str(q) for q in qnums]
        rate_f = [P[k]["float_stream_bytes"] * 8.0 / px for k in keys]
        rate_i = [P[k]["int_stream_bytes"] * 8.0 / px for k in keys]
        curves["psnr"][it["id"]] = bd(rate_f, [psnr(P[k]["float"]["sse"], P[k]["float"]["samples"]) for k in keys],
                                      rate_i, [psnr(P[k]["int"]["sse"], P[k]["int"]["samples"]) for k in keys])
        curves["ms_ssim"][it["id"]] = bd(rate_f, [msssim_db(P[k]["float"]["ms_ssim"]) for k in keys],
                                         rate_i, [msssim_db(P[k]["int"]["ms_ssim"]) for k in keys])
    s["bd_rate_percent"] = {}
    for m, per in curves.items():
        d = {"per_picture": per, "undefined": sorted(k for k, v in per.items() if v is None)}
        d.update(describe([v for v in per.values() if v is not None]) or {})
        s["bd_rate_percent"][m] = d
    pq = {}
    for q in qnums:
        keys = [it["id"] + "/q" + str(q) for it in items]
        g = lambda f: describe([f(P[k]) for k in keys])
        pq["q%d" % q] = {
            "float_bpp": g(lambda p: p["float_stream_bytes"] * 24.0 / p["float"]["samples"]),
            "rate_ratio": g(lambda p: p["int_stream_bytes"] / p["float_stream_bytes"]),
            "float_psnr_db": g(lambda p: psnr(p["float"]["sse"], p["float"]["samples"])),
            "delta_psnr_db": g(lambda p: psnr(p["int"]["sse"], p["int"]["samples"]) - psnr(p["float"]["sse"], p["float"]["samples"])),
            "delta_ms_ssim_db": g(lambda p: msssim_db(p["int"]["ms_ssim"]) - msssim_db(p["float"]["ms_ssim"])),
        }
    s["per_quality"] = pq
    xi = {}
    for tag, d in rec["integer_decodes"].items():
        res = d["results"]
        same = refused = 0
        for k, h in res.items():
            if h.startswith("refused:"):
                refused += 1
            elif k in P and h == P[k]["int"]["rgb_sha256"]:
                same += 1
        xi[tag] = {"decodes": len(res), "identical": same, "refused": refused, "differ": len(res) - same - refused,
                   "missing": sorted(k for k in P if k not in res)}
    s["integer_cross_platform"] = xi
    tot = {"platforms": len(xi)}
    for f in ("decodes", "identical", "differ", "refused"):
        tot[f] = 0
        for v in xi.values():
            tot[f] += v[f]
    s["integer_cross_platform_total"] = tot
    xf = {}
    for tag, d in rec["float_decodes"].items():
        res = d["results"]
        t = {"decodes": len(res), "identical": 0, "rgb_differ": 0, "indexes_differ": 0, "symbols_differ": 0}
        deltas, desync = [], []
        for k, v in res.items():
            ref = P[k]
            change = psnr(v["sse"], ref["float"]["samples"]) - psnr(ref["float"]["sse"], ref["float"]["samples"])
            if v["rgb_sha256"] == ref["float"]["rgb_sha256"]:
                t["identical"] += 1
            else:
                t["rgb_differ"] += 1
                deltas.append(change)
            if v["indexes_sha256"] != ref["float_cpu_indexes_sha256"]:
                t["indexes_differ"] += 1
            if v["symbols_sha256"] != ref["float_cpu_symbols_sha256"]:
                t["symbols_differ"] += 1
                desync.append(change)
        t["delta_psnr_db_where_differ"] = describe(deltas)
        t["delta_psnr_db_where_symbols_differ"] = describe(desync)
        t["missing"] = sorted(k for k in P if k not in res)
        t["tables_differ"] = sorted(q for q, h in d.get("tables_sha256", {}).items() if h != rec["float_tables_sha256"].get(q))
        xf[tag] = t
    s["float_cross_platform"] = xf
    tm = {}
    for run in rec["timing"]:
        R = run["rounds"]
        t = {"rounds": len(R), "decodes_per_round": sorted(set(r["decodes"] for r in R)),
             "ms_per_decode": describe([r["child_wall_ns"] / 1e6 / r["decodes"] for r in R]),
             "load_avg_1min": describe([r["load_avg_before"][0] for r in R] + [r["load_avg_after"][0] for r in R])}
        if R and all("task" in r for r in R):
            t["task_j_per_decode"] = describe([r["task"]["energy_nj"] / 1e9 / r["decodes"] for r in R])
            t["task_cpu_ms_per_decode"] = describe([r["task"]["cpu_ns"] / 1e6 / r["decodes"] for r in R])
            t["billed_j_per_decode"] = describe([r["task"]["billed_nj"] / 1e9 / r["decodes"] for r in R])
        for cn in ("rapl", "ioreport-gpu"):
            if not R or not all(cn in r for r in R):
                continue
            verdicts = []
            usable = []
            lives = 0
            for r in R:
                e = r[cn]
                ratio = e.get("ratio")
                v = None if ratio is None else ("clean" if ratio >= 10 else "relative-only" if ratio >= 3 else "noisy" if ratio >= 1 else "unusable")
                if v != e.get("verdict"):
                    bad(f"timing {run['label']}: a round's {cn} verdict is {e.get('verdict')!r}, its ratio {ratio} gives {v!r}")
                verdicts.append(v)
                share_ok = cn != "ioreport-gpu" or (r.get("gpu_time", {}).get("share") or 0) >= 0.99
                alive = is_live(e.get("busy")) and is_live(e.get("idle_before")) and is_live(e.get("idle_after"))
                lives += alive
                if v in ("clean", "relative-only") and share_ok and alive:
                    usable.append(r)
            t[cn] = {"verdicts": verdicts, "live_rounds": lives,
                     "ratio": describe([r[cn]["ratio"] for r in R if r[cn].get("ratio") is not None]),
                     "usable_rounds": len(usable),
                     "marginal_j_per_decode": describe([r[cn]["marginal_nj"] / 1e9 / r["decodes"] for r in usable]),
                     "publishable": len(usable) >= 7 and all(r[cn]["verdict"] == "clean" for r in usable)}
            if cn == "ioreport-gpu":
                t[cn]["gpu_share"] = describe([r["gpu_time"]["share"] for r in R if r.get("gpu_time", {}).get("share") is not None])
        tm[run["label"]] = t
    s["timing"] = tm
    return s


def is_live(b):
    """A bracket is live (energy-measurement §3.1) when the counter advanced in two
    intervals in a row, in five or more, in at least 80% of them, over 300 ms or more."""
    if not b:
        return False
    return (b["longest_run"] >= 2 and b["advanced"] >= 5 and b["advanced"] >= 0.8 * b["intervals"]
            and b["span_ns"] >= 300e6)


def structure(rec):
    items = rec["pictures"]["items"]
    ids = [it["id"] for it in items]
    if len(set(ids)) != len(ids):
        bad("two pictures share an id")
    for it in items:
        if it["id"] != it["file_sha256"][:16]:
            bad(f"picture {it['id']}: its id is not the first 16 hex digits of its file's SHA-256")
    for q in rec["models"]:
        for it in items:
            k = f"{it['id']}/q{q}"
            p = rec["points"].get(k)
            if p is None:
                bad(f"{k}: no reading")
                continue
            n = it["width"] * it["height"] * 3
            for side in ("float", "int"):
                if p[side]["samples"] != n:
                    bad(f"{k} {side}: {p[side]['samples']} samples, the picture has {n}")
                if not 0 < p[side]["ms_ssim"] < 1:
                    bad(f"{k} {side}: MS-SSIM {p[side]['ms_ssim']} outside (0, 1)")
                if p[side]["sse"] <= 0:
                    bad(f"{k} {side}: SSE {p[side]['sse']}")
            if p["int_payload_bytes"] != p["int_stream_bytes"] + 40:
                bad(f"{k}: a WIH1 payload is its two streams plus 40 bytes, not {p['int_payload_bytes']} for {p['int_stream_bytes']}")
    if len(rec["points"]) != len(items) * len(rec["models"]):
        bad(f"{len(rec['points'])} readings for {len(items)} pictures at {len(rec['models'])} qualities")


def lookup(rec, path):
    v = rec
    for part in path.split("."):
        if isinstance(v, dict) and part in v:
            v = v[part]
        else:
            return KeyError
    return v


def check_text(rec, files):
    marker = re.compile(r"([+\-−]?\d[\d,]*(?:\.\d+)?)\s*(%|dB|ms|mJ|J)?\**\s*<!--det:([\w.\-/]+)-->")
    n = 0
    for f in files:
        text = open(f, encoding="utf-8").read()
        for m in marker.finditer(text):
            n += 1
            printed, path = m.group(1).replace("−", "-").replace(",", ""), m.group(3)
            v = lookup(rec, path)
            if v is KeyError or not isinstance(v, (int, float)) or isinstance(v, bool):
                bad(f"{f}: <!--det:{path}--> names no number in the record")
                continue
            dec = len(printed.split(".")[1]) if "." in printed else 0
            want = round(v, dec)
            if abs(float(printed) - want) > 0.5 * 10 ** -dec + 1e-12 or (dec == 0 and float(printed) != round(v)):
                bad(f"{f}: prints {m.group(1)} for {path}, the record has {v} ({want} at {dec} decimals)")
        found = len(marker.findall(text))
        every = len(re.findall(r"<!--det:", text))
        if every != found:
            bad(f"{f}: {every - found} det marker(s) do not directly follow a number")
    return n


def main(argv):
    if len(argv) < 2:
        sys.exit(__doc__)
    rec = json.load(open(argv[1]))
    if rec.get("format") != "wai-determinism-cost/1":
        sys.exit(f"not a determinism-cost record: {rec.get('format')!r}")
    structure(rec)
    compare("summary", recompute(rec), rec.get("summary", {}), tol=(1e-9, 1e-12))
    quoted = check_text(rec, argv[2:])
    for f in FINDINGS:
        print(f)
    if FINDINGS:
        print(f"{len(FINDINGS)} finding(s)")
        return 1
    s = rec["summary"]
    print(f"ok: {len(rec['points'])} readings, {len(rec['integer_decodes'])} integer and "
          f"{len(rec['float_decodes'])} float platform decodes, {len(rec['timing'])} timed runs; "
          f"{quoted} quoted figure(s) match; BD-rate (PSNR) median {s['bd_rate_percent']['psnr']['median']:.2f} %")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))