#!/usr/bin/env python3
"""Does a passage's score depend on which passages happened to share its batch?

`score.py` tokenizes with `padding=True`, so the pad width of a batch is set by the longest
passage in it, which depends on the shuffle. A correctly-masked roberta is invariant to trailing
pad tokens — but "should be" is not "is", and if it were false then every number in this study
would be conditional on an arbitrary grouping rather than on the text.

The test: re-score a sample **one passage at a time**. A batch of one has nothing to pad to, so
its score is the canonical value. Compare against what the batched run recorded.

The sample is not random-only. Float-level noise can only matter where a passage sits near a
decision threshold, so every passage within 0.01 of any published threshold is included by
construction, plus a seeded random sample for the general case. Taking only a random sample would
test the question exactly where the answer cannot matter.

Runs over the control corpus too. The first version of this script read `passages.jsonl` and
nothing else, which would have discharged the padding question for the headline table and left it
open for the control column of the same sweep — half a check, presented as a whole one.

  python3 padding_check.py <model-key> [--control] [--dry]
"""
import json, os, random, sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from score import MODELS, SEED

HERE = os.path.dirname(os.path.abspath(__file__))
THRESHOLDS = (0.5, 0.7, 0.9, 0.95, 0.99)
NEAR = 0.01
N_RANDOM = 150


def load_corpus(control):
    """Same key scheme score.py uses, for the same reason: the two corpora must differ by corpus
    and by nothing else, and a second way of naming a row is where that guarantee goes to die."""
    rows = {}
    if control:
        for line in open(os.path.join(HERE, "control.jsonl")):
            d = json.loads(line)
            rows["%s#%d" % (d["id"], d["seq"])] = d
    else:
        for line in open(os.path.join(HERE, "passages.jsonl")):
            d = json.loads(line)
            rows[d["n"]] = d
    return rows


def select_sample(scored):
    """Which passages get re-scored singly, as a function of the recorded scores alone.

    Pulled out of `main` because the page states this rule as a claim about a *set* — "every
    passage within 0.01 of any published threshold is included by construction" — and a claim
    about a set is only true if something recomputes it. `build_page.py` imports this function and
    demands that the output file hold exactly the ids it returns. Writing the rule out a second
    time there would put two hand-written copies of one rule in the tree, and they drift toward
    whichever branch nothing tests.

    Returns (sample, near) — both lists, `near` being the near-threshold part, for the log line.
    """
    near = [n for n, p in scored.items() if any(abs(p - t) < NEAR for t in THRESHOLDS)]
    rest = sorted(set(scored) - set(near), key=str)
    random.Random(SEED).shuffle(rest)
    return sorted(set(near) | set(rest[:N_RANDOM]), key=str), near


def main():
    key = sys.argv[1] if len(sys.argv) > 1 else "hello"
    control = "--control" in sys.argv[2:]
    dry = "--dry" in sys.argv[2:]
    name, ai_label = MODELS[key]
    tag = ("control_" if control else "") + key

    rows = load_corpus(control)
    spath = os.path.join(HERE, "scores_%s.jsonl" % tag)
    if not os.path.exists(spath):
        print("no scores at %s — nothing to check" % os.path.basename(spath), file=sys.stderr)
        return
    scored = {}
    for l in open(spath):
        d = json.loads(l)
        scored[d["n"]] = d["p_ai"]

    missing = [n for n in scored if n not in rows]
    if missing:
        raise SystemExit("%d scored ids are absent from the corpus file, e.g. %r"
                         % (len(missing), missing[:3]))

    sample, near = select_sample(scored)
    print("re-scoring %d passages singly (%d near a threshold, %d random) of %d scored in %s"
          % (len(sample), len(near), len(sample) - len(near), len(scored),
             "control.jsonl" if control else "passages.jsonl"), file=sys.stderr, flush=True)

    # --dry exercises everything up to the model load: which score file, which key scheme, that
    # every sampled id resolves to text. This whole script runs as the last stage of a six-hour
    # chain, which is the worst possible place to discover a typo.
    if dry:
        print("  DRY: near-threshold counts " + "  ".join(
            "%.2f:%d" % (t, sum(1 for p in scored.values() if abs(p - t) < NEAR))
            for t in THRESHOLDS), file=sys.stderr)
        print("  DRY: text present on all sampled rows: %s"
              % all(isinstance(rows[n].get("text"), str) and rows[n]["text"] for n in sample),
              file=sys.stderr)
        print("  DRY: first %r  last %r" % (sample[0], sample[-1]) if sample else "  DRY: empty",
              file=sys.stderr)
        return

    import torch
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    torch.set_num_threads(1)
    tok = AutoTokenizer.from_pretrained(name)
    model = AutoModelForSequenceClassification.from_pretrained(name).eval()
    id2label = {int(k): v for k, v in model.config.id2label.items()}
    ai_idx = [i for i, v in id2label.items() if v == ai_label][0]

    out = open(os.path.join(HERE, "padding_%s.jsonl" % tag), "w")
    worst = 0.0
    flips = []
    for i, n in enumerate(sample):
        enc = tok([rows[n]["text"]], return_tensors="pt", truncation=True, max_length=512)
        with torch.no_grad():
            p = torch.softmax(model(**enc).logits, dim=-1)[0, ai_idx].item()
        was = scored[n]
        d = abs(p - was)
        worst = max(worst, d)
        for t in THRESHOLDS:
            if (was > t) != (p > t):
                flips.append((n, t, was, p))
        out.write(json.dumps({"n": n, "batched": was, "single": round(p, 6),
                              "delta": round(p - was, 9)}) + "\n")
        if i % 25 == 0:
            print("  %d/%d  worst delta so far %.3g" % (i, len(sample), worst),
                  file=sys.stderr, flush=True)
    out.close()

    print("\nlargest absolute difference: %.3g" % worst, file=sys.stderr)
    print("verdict flips at any published threshold: %d" % len(flips), file=sys.stderr)
    for n, t, a, b in flips:
        print("  n=%s  threshold %.2f  batched %.6f -> single %.6f" % (n, t, a, b), file=sys.stderr)
    print("done", file=sys.stderr, flush=True)


if __name__ == "__main__":
    main()