"""QM9 results as a quality-vs-compute frontier.

The number that matters on this benchmark is not validity at some fixed step
count but **how few denoising steps are needed to reach a given validity**. That
is the axis the discrete graph-diffusion literature competes on, and it is the
one determined by the training schedule and the decoding rule rather than by
model scale.

Reported here as a speed-up over the standard formulation. The `mdlm` + `rope`
arm IS the conventional setup - uniform-t discrete diffusion with the rotary
encoding that modern implementations default to - so the comparison is
controlled: same data, same model size, same training budget, same decoder, one
factor changed at a time.

A note on external comparison. Published DiGress-style numbers come from larger
models trained longer and with auxiliary spectral features that neither arm here
uses, so quoting them beside these would not be like-for-like. The claim
supported by this script is the controlled one: at matched scale and budget, how
many steps does each formulation need. Any external number should be taken from
the paper itself rather than from memory.
"""
import glob, json, os, sys
from collections import defaultdict

import numpy as np

RUNS = os.environ.get("RUNS", "runs")

runs = defaultdict(list)
for f in sorted(glob.glob(f"{RUNS}/mol1-*/result.json")):
    r = json.load(open(f))
    a = r["args"]
    runs[(a["mode"], a["pe"])].append(r)

if not runs:
    sys.exit("no QM9 results yet")

LABEL = {("mdlm", "rope"): "uniform-t + relative  (standard)",
         ("mdlm", "ape"): "uniform-t + absolute",
         ("matched", "rope"): "matched + relative",
         ("matched", "ape"): "matched + absolute  (ours)"}


def curve(rs):
    """(steps, validity, unique, novel) for every fixed-K setting, by step count."""
    pts = defaultdict(list)
    for r in rs:
        for k, v in r["decode"].items():
            if k.startswith("fixedK_"):
                pts[int(k.split("_")[1])].append(v)
    out = []
    for k in sorted(pts):
        vs = pts[k]
        out.append((k,
                    float(np.mean([x["validity"] for x in vs])),
                    float(np.mean([x["unique"] for x in vs])),
                    float(np.mean([x["novel"] for x in vs]))))
    return out


def steps_for(c, target):
    """Fewest denoising steps reaching `target` validity, or None."""
    for k, v, _, _ in c:
        if v >= target:
            return k
    return None


print("=" * 84)
print("QM9 — validity against number of denoising steps")
print("=" * 84)
for key in [("mdlm", "rope"), ("mdlm", "ape"), ("matched", "rope"), ("matched", "ape")]:
    rs = runs.get(key)
    if not rs:
        continue
    c = curve(rs)
    print(f"\n{LABEL[key]}   ({len(rs)} seeds)")
    print(f"  {'steps':>7}{'validity':>11}{'unique':>9}{'novel':>9}")
    for k, v, u, n in c:
        print(f"  {k:>7}{v*100:>10.2f}%{u*100:>8.1f}%{n*100:>8.1f}%")

print()
print("=" * 84)
print("STEPS NEEDED TO REACH A VALIDITY TARGET  (fewer is better)")
print("=" * 84)
hdr = "".join(f"{str(int(t*100)) + '%':>9}" for t in (0.5, 0.8, 0.9, 0.95))
print(f"{'formulation':<36}" + hdr)
base = runs.get(("mdlm", "rope"))
base_c = curve(base) if base else None
for key in [("mdlm", "rope"), ("mdlm", "ape"), ("matched", "rope"), ("matched", "ape")]:
    rs = runs.get(key)
    if not rs:
        continue
    c = curve(rs)
    cells = []
    for t in (0.5, 0.8, 0.9, 0.95):
        s = steps_for(c, t)
        cells.append(f"{s:>9}" if s else f"{'—':>9}")
    print(f"{LABEL[key]:<36}" + "".join(cells))

if base_c:
    print()
    print("SPEED-UP over the standard formulation, at equal validity:")
    for key in [("mdlm", "ape"), ("matched", "rope"), ("matched", "ape")]:
        rs = runs.get(key)
        if not rs:
            continue
        c = curve(rs)
        for t in (0.8, 0.9, 0.95):
            b, o = steps_for(base_c, t), steps_for(c, t)
            if b and o:
                print(f"  {LABEL[key]:<34} at {int(t*100)}% validity: "
                      f"{b} steps -> {o} steps  ({b/o:.1f}x fewer)")