"""Head-to-head against inference-time decoding methods.

This is the comparison that decides whether the paper's claim holds. Roughly a
dozen 2026 papers accelerate diffusion decoding by choosing *which* positions to
reveal together - confidence ranking, entropy gating, dependency-aware ordering.
All of them take the trained marginals as given and only reorder commits.

The argument of this project is that no such method can close the gap, because
the marginals themselves were never trained for the one-pass regime. That is
testable directly: take a BASELINE model, decode it with every strategy at every
budget, and check whether any of them reaches the training-time method.

For a fair fight each inference-time method is given its best configuration and
is allowed MORE passes than our method uses.
"""
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("--digits", type=int, default=12)
p.add_argument("--n_eval", type=int, default=400)
p.add_argument("--eval_bs", type=int, default=100)
p.add_argument("--out", default="figures/beat_baselines.json")
a = p.parse_args()

dev = "cuda" if torch.cuda.is_available() else "cpu"
SL = 3 * a.digits + 4
X, A, tok = build_addition(a.n_eval, a.digits, SL, seed=90_000, exact=True)
X, A = torch.from_numpy(X), torch.from_numpy(A)


def load(run_dir):
    cfg = json.load(open(os.path.join(run_dir, "result.json")))["args"]
    m = Transformer(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"],
                    causal=False, max_len=cfg["seq_len"] + 8).to(dev)
    m.load_state_dict(torch.load(os.path.join(run_dir, "model.pt"), map_location=dev))
    m.eval()
    return m, cfg


@torch.no_grad()
def run(model, steps, strategy, tau=None):
    accs = []
    for i in range(0, len(X), a.eval_bs):
        xb, ab = X[i:i + a.eval_bs].to(dev), A[i:i + a.eval_bs].to(dev)
        x = torch.where(ab, torch.full_like(xb, tok.mask), xb)
        n_ans = int(ab[0].sum())
        pred = None
        for s in range(steps, 0, -1):
            masked = (x == tok.mask) & ab
            if not masked.any():
                break
            logits = model(x)
            probs = dfn.real_probs(logits, tok)
            conf, pred = probs.max(-1)
            if strategy == "random":
                score = torch.rand_like(conf)
            elif strategy == "entropy":
                ent = -(probs * probs.clamp_min(1e-9).log()).sum(-1)
                score = -ent
            elif strategy == "margin":
                top2 = probs.topk(2, dim=-1).values
                score = top2[..., 0] - top2[..., 1]
            else:
                score = conf
            score = score.masked_fill(~masked, -1e9)
            left = int(n_ans * (s - 1) / steps)
            for b in range(x.shape[0]):
                k = max(0, int(masked[b].sum()) - left)
                if strategy == "entropy" and tau is not None and k > 0:
                    ent_b = -(probs[b] * probs[b].clamp_min(1e-9).log()).sum(-1)
                    allowed = int(((ent_b < tau) & masked[b]).sum())
                    k = max(1, min(k, allowed)) if allowed else 1
                if k:
                    idx = score[b].topk(k).indices
                    x[b, idx] = pred[b, idx]
        rem = (x == tok.mask) & ab
        if rem.any() and pred is not None:
            x[rem] = pred[rem]
        accs.extend(per_example(x.cpu(), xb.cpu(), ab.cpu()))
    return np.array(accs, dtype=float)


def per_example(pred, gold, amask):
    """Correct/incorrect per evaluation item, not a batch mean.

    A paired test needs to know which items each system got right, on the SAME
    items. Batch means throw that away.
    """
    out = []
    for p, g, a in zip(pred.tolist(), gold.tolist(), amask.tolist()):
        pa = [p[i] for i in range(len(a)) if a[i]]
        ga = [g[i] for i in range(len(a)) if a[i]]
        out.append(float(pa == ga))
    return out


def paired_test(ours, theirs, n_perm=20000, seed=0):
    """Paired permutation test on the same evaluation items.

    The two systems are scored on identical inputs, so the items are paired and
    a paired test is the right one: under the null that the systems are
    interchangeable, the sign of each item's difference is exchangeable.
    Reports the observed difference, a bootstrap 95% interval on it, and the
    two-sided p-value.
    """
    d = ours - theirs
    obs = float(d.mean())
    rng = np.random.default_rng(seed)
    flips = rng.choice([-1.0, 1.0], size=(n_perm, len(d)))
    null = (flips * d).mean(axis=1)
    p = float((np.abs(null) >= abs(obs) - 1e-12).mean())
    idx = rng.integers(0, len(d), size=(10000, len(d)))
    bs = d[idx].mean(axis=1)
    return {"diff": obs, "ci95": [float(np.percentile(bs, 2.5)),
                                  float(np.percentile(bs, 97.5))],
            "p_value": p, "n": int(len(d))}


base_dirs = sorted(glob.glob(f"{a.runs}/g11-d{a.digits}-mdlm-s*"))
ours_dirs = sorted(glob.glob(f"{a.runs}/g11-d{a.digits}-full-s*"))
if not base_dirs or not ours_dirs:
    sys.exit(f"need both baseline and method runs for {a.digits} digits")

results = {"digits": a.digits, "inference_time_on_baseline": {}, "ours": {}}

print(f"\n=== {a.digits}-digit addition: can any inference-time method rescue "
      f"the baseline? ===")
print(f"{'strategy':<22}{'1 pass':>9}{'2':>8}{'4':>8}{'8':>8}{'16':>8}")
ITEMS = {}
for strat, tau in [("confidence", None), ("entropy", None), ("entropy-gated", 0.5),
                   ("margin", None), ("random", None)]:
    row, row_items = [], []
    for steps in (1, 2, 4, 8, 16):
        vals = [run(load(d)[0], steps,
                    "entropy" if strat == "entropy-gated" else strat, tau)
                for d in base_dirs]
        per = np.mean(np.stack(vals), axis=0)   # per-item, averaged over seeds
        row.append(float(per.mean())); row_items.append(per)
    results["inference_time_on_baseline"][strat] = row
    ITEMS[strat] = row_items
    print(f"{strat:<22}" + "".join(f"{v*100:8.1f}" for v in row))

row, ours_items = [], []
for steps in (1, 2, 4, 8, 16):
    vals = [run(load(d)[0], steps, "confidence") for d in ours_dirs]
    per = np.mean(np.stack(vals), axis=0)
    row.append(float(per.mean())); ours_items.append(per)
results["ours"]["confidence"] = row
print(f"{'OURS (training-time)':<22}" + "".join(f"{v*100:8.1f}" for v in row))

print(f"\n=== is the difference significant? (paired permutation test) ===")
print("Our method at ONE pass against each inference-time strategy at its BEST")
print("budget, scored on identical evaluation items.\n")
print(f"{'vs strategy':<22}{'their best':>12}{'ours @1':>10}{'diff':>9}"
      f"{'95% CI':>18}{'p':>10}")
results["significance"] = {}
for strat, rows in ITEMS.items():
    best_i = int(np.argmax([r.mean() for r in rows]))
    st = paired_test(ours_items[0], rows[best_i])
    results["significance"][strat] = dict(st, their_best_budget=[1, 2, 4, 8, 16][best_i])
    print(f"{strat:<22}{rows[best_i].mean()*100:11.1f}%{ours_items[0].mean()*100:9.1f}%"
          f"{st['diff']*100:+8.1f}"
          f"   [{st['ci95'][0]*100:+.1f},{st['ci95'][1]*100:+.1f}]"
          f"{st['p_value']:10.2g}")
print(f"\n(n = {len(ours_items[0])} paired evaluation items; "
      f"p is two-sided, 20000 permutations)")

best_inf = max(v[0] for v in results["inference_time_on_baseline"].values())
best_inf_any = max(max(v) for v in results["inference_time_on_baseline"].values())
print(f"\nbest inference-time method at 1 pass : {best_inf*100:.1f}%")
print(f"best inference-time at ANY budget     : {best_inf_any*100:.1f}%")
print(f"ours at 1 pass                        : {row[0]*100:.1f}%")
results["summary"] = {"best_inference_1pass": best_inf,
                      "best_inference_any_budget": best_inf_any,
                      "ours_1pass": row[0]}
os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
json.dump(results, open(a.out, "w"), indent=2)
print(f"\nwrote {a.out}")