#!/usr/bin/env python3
"""Score every passage with one open AI-text detector. Append-only, resumable.

  python3 score.py <model-key>

Design notes, all of which are scar tissue:

* **Append-JSONL + skip-set.** A 200k-item scan of mine once died silently two-thirds of the
  way through and I did not notice, because a crashed writer and a finished writer leave the
  same file. Every result is flushed as its own line, keyed by passage id, and a rerun skips
  what is already there.
* **Deterministic shuffle before scoring.** The corpus is ordered by source file, i.e. by
  genre. If this is killed halfway, an ordered prefix would be a biased sample of genres while
  looking like a sample of the corpus. Shuffled with a fixed seed, any prefix is a uniform
  random sample and the partial run is still publishable.
* **One model in memory at a time.** This box is 1 core / ~2 GB; two roberta-base checkpoints
  in fp32 is most of the RAM, and starvation here looks exactly like a hang.
"""
import json, os, random, sys, time

MODELS = {
    "hello": ("Hello-SimpleAI/chatgpt-detector-roberta", "ChatGPT"),
    "openai": ("openai-community/roberta-base-openai-detector", "Fake"),
}
HERE = os.path.dirname(os.path.abspath(__file__))
SEED = 20260828          # fixed before the run; never re-drawn


def main():
    key = sys.argv[1]
    control = "--control" in sys.argv[2:]
    name, ai_label = MODELS[key]

    # The control corpus goes through THIS function, not a copy of it. The two rates have to
    # differ by corpus and by nothing else, and a second implementation is where that guarantee
    # goes to die — the drift always lands on the branch the check doesn't reach.
    if control:
        out_path = os.path.join(HERE, f"scores_control_{key}.jsonl")
        rows = []
        for line in open(os.path.join(HERE, "control.jsonl")):
            d = json.loads(line)
            d["n"] = "%s#%d" % (d["id"], d["seq"])      # control rows carry no corpus-wide id
            rows.append(d)
    else:
        out_path = os.path.join(HERE, f"scores_{key}.jsonl")
        rows = [json.loads(l) for l in open(os.path.join(HERE, "passages.jsonl"))]
    random.Random(SEED).shuffle(rows)

    done = set()
    if os.path.exists(out_path):
        for line in open(out_path):
            try:
                done.add(json.loads(line)["n"])
            except Exception:
                pass
    todo = [r for r in rows if r["n"] not in done]
    print(f"{name}: {len(rows)} passages, {len(done)} already scored, {len(todo)} to go",
          file=sys.stderr, flush=True)
    if not todo:
        return

    # --dry exercises everything up to the model load: which file, which key scheme, how many
    # rows resume skips. Worth having because the control path only runs hours from now, at the
    # end of a chain, and a typo there costs the whole paired comparison rather than announcing
    # itself. It also runs in a few MB, which matters when the box has one core and no swap.
    if "--dry" in sys.argv[2:]:
        print("  DRY: would score %d of %d rows from %s -> %s" % (len(todo), len(rows),
              "control.jsonl" if control else "passages.jsonl", os.path.basename(out_path)),
              file=sys.stderr)
        print("  DRY: first key %r  last key %r" % (todo[0]["n"], todo[-1]["n"]), file=sys.stderr)
        print("  DRY: text present on all rows: %s"
              % all(isinstance(r.get("text"), str) and r["text"] for r in todo), 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()

    # which logit index means "machine-written" — read from the checkpoint's own config rather
    # than assumed. Guessing this inverts every number in the study.
    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]
    if len(ai_idx) != 1:
        raise SystemExit(f"cannot locate '{ai_label}' in id2label={id2label}")
    ai_idx = ai_idx[0]
    print(f"  id2label={id2label}  ai_idx={ai_idx}", file=sys.stderr, flush=True)

    BATCH = 8
    t0 = time.time()
    with open(out_path, "a") as fh:
        for i in range(0, len(todo), BATCH):
            chunk = todo[i:i + BATCH]
            enc = tok([r["text"] for r in chunk], return_tensors="pt",
                      truncation=True, max_length=512, padding=True)
            with torch.no_grad():
                probs = torch.softmax(model(**enc).logits, dim=-1)[:, ai_idx].tolist()
            for r, p in zip(chunk, probs):
                fh.write(json.dumps({"n": r["n"], "p_ai": round(p, 6)}) + "\n")
            fh.flush()
            if (i // BATCH) % 25 == 0:
                el = time.time() - t0
                rate = (i + len(chunk)) / max(el, 1e-9)
                eta = (len(todo) - i - len(chunk)) / max(rate, 1e-9)
                print(f"  {i+len(chunk)}/{len(todo)}  {rate:.2f}/s  eta {eta/60:.0f}m",
                      file=sys.stderr, flush=True)
    print("done", file=sys.stderr, flush=True)


if __name__ == "__main__":
    main()