"""Did the baseline fail to PARALLELISE, or fail to LEARN?

At 20 digits the baseline scores 0.0% at one pass and 0.0% at sixteen. It is
tempting to read that as a very large parallel-decoding gap, but the standard
pass grid stops at sixteen while a 20-digit answer occupies 21 slots - so that
number was never a fully sequential decode, and the two explanations are not
distinguished by it:

  (a) the baseline learned the task and merely cannot commit it in parallel
      -> it should succeed once decoding is fully sequential;
  (b) the baseline never learned the task at all
      -> it fails at every budget, and the claim is about sample efficiency
         rather than about parallelism.

These are different papers. This script decodes every run at a budget equal to
the number of answer positions, which is genuinely one token per pass, and
reports accuracy at that budget next to the one-pass number.
"""
import argparse, glob, json, os, sys
import numpy as np
import torch

sys.path.insert(0, os.path.dirname(__file__))
from data import build_addition
from model import Transformer
import diffusion as dfn
from train_add import exact_match

p = argparse.ArgumentParser()
p.add_argument("--runs", default="runs")
p.add_argument("--pattern", default="g13-*")
p.add_argument("--n_eval", type=int, default=200)
p.add_argument("--eval_bs", type=int, default=50)
p.add_argument("--out", default="figures/seq_check.json")
a = p.parse_args()

dev = "cuda" if torch.cuda.is_available() else "cpu"
rows = []

for run_dir in sorted(glob.glob(os.path.join(a.runs, a.pattern))):
    rf = os.path.join(run_dir, "result.json")
    mf = os.path.join(run_dir, "model.pt")
    if not (os.path.exists(rf) and os.path.exists(mf)):
        continue
    r = json.load(open(rf))
    cfg = r["args"]
    d, SL = cfg["digits"], cfg["seq_len"]

    X, A, tok = build_addition(a.n_eval, d, SL, seed=91_000, exact=True)
    X, A = torch.from_numpy(X), torch.from_numpy(A)
    n_slots = int(A[0].sum())

    m = Transformer(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"],
                    causal=False, max_len=SL + 8).to(dev)
    m.load_state_dict(torch.load(mf, map_location=dev))
    m.eval()

    accs = {}
    for steps in (1, n_slots):
        vals = []
        for i in range(0, len(X), a.eval_bs):
            xb, ab = X[i:i + a.eval_bs], A[i:i + a.eval_bs]
            out = dfn.cond_decode(m, xb, ab, tok, steps, dev)
            vals.append(exact_match(out.cpu(), xb, ab, tok))
        accs[steps] = float(np.mean(vals))

    rows.append({"run": os.path.basename(run_dir), "digits": d,
                 "mode": cfg["mode"], "seed": cfg["seed"], "slots": n_slots,
                 "acc_1pass": accs[1], "acc_sequential": accs[n_slots]})
    print(f"{rows[-1]['run']:<22} {d:>3}d {cfg['mode']:<6} "
          f"1-pass {accs[1]*100:6.2f}   fully-sequential ({n_slots} passes) "
          f"{accs[n_slots]*100:6.2f}", flush=True)

print("\n=== does the baseline learn the task at ANY budget? ===")
agg = {}
for r in rows:
    agg.setdefault((r["digits"], r["mode"]), []).append(r)
print(f"{'digits':>7} {'mode':<7}{'1 pass':>9}{'sequential':>12}  verdict")
for k in sorted(agg):
    v = agg[k]
    one = np.mean([x["acc_1pass"] for x in v]) * 100
    seq = np.mean([x["acc_sequential"] for x in v]) * 100
    if seq < 5:
        verdict = "never learned the task (sample efficiency, not parallelism)"
    elif one < seq - 5:
        verdict = "learned it, cannot parallelise (a true parallelism gap)"
    else:
        verdict = "learned it and parallelises"
    print(f"{k[0]:>7} {k[1]:<7}{one:8.1f} {seq:11.1f}  {verdict}")

os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
json.dump(rows, open(a.out, "w"), indent=2)
print(f"\nwrote {a.out}")