#!/usr/bin/env python3
"""Merge the pangram corpora into one deduplicated scoring set.

Dedup is on exact passage text, not on (item, seq): the corpora were built by overlapping
queries and the same book was fetched by more than one of them. Counting a passage twice would
inflate every rate downstream by however much the queries overlapped, which is not a constant.

Emits corpus.jsonl with one row per distinct passage, carrying the source corpus file so the
per-genre breakdown promised in PREREG.md is computable later.
"""
import json, glob, os, re, sys

SRC = "/home/agent/work/pangram"
OUT = "/home/agent/work/aidetect-fpr/passages.jsonl"
FUTURE_YEAR = re.compile(r"\b(20[0-3]\d)\b")
# corpora 7-10 were built after the book-level own-text date filter existed
STRICT = {"corpus7.jsonl", "corpus8.jsonl", "corpus9.jsonl", "corpus10.jsonl"}

def yr(v):
    try:
        return int(str(v)[:4])
    except Exception:
        return None

# Everything below is the merge itself, behind a main guard so that FUTURE_YEAR and yr() above can
# be imported by another script without re-running the merge and overwriting a corpus that is
# pinned by digest. drift.py needs exactly those two rules and nothing else; a second copy of the
# year regex is how the corpus and a measurement about the corpus start disagreeing.
def main():
  seen = {}
  per_file = {}
  skipped_future = 0
  for path in sorted(glob.glob(os.path.join(SRC, "corpus*.jsonl"))):
      fname = os.path.basename(path)
      n = 0
      for line in open(path):
          try:
              d = json.loads(line)
          except Exception:
              continue
          text = (d.get("text") or "").strip()
          if not text:
              continue
          # B3: no post-1999 year in the passage's own text
          if FUTURE_YEAR.search(text):
              skipped_future += 1
              continue
          if text in seen:
              r = seen[text]
              r["files"].add(fname)
              # the same passage reached us through more than one query, and the two rows can
              # disagree about the catalogue year. Keep the earliest: a passage is pre-2001 if
              # ANY source row says so. Deduping before the year test silently dropped 7.
              if (yr(d.get("year")) or 9999) < (yr(r["year"]) or 9999):
                  r["year"], r["item"], r["title"], r["seq"] = (
                      d.get("year"), d.get("id"), d.get("title"), d.get("seq"))
              continue
          seen[text] = {"item": d.get("id"), "title": d.get("title"), "year": d.get("year"),
                        "seq": d.get("seq"), "text": text, "files": {fname}}
          n += 1
      per_file[fname] = n

  rows = list(seen.values())
  pre2001 = [r for r in rows if (yr(r["year"]) or 9999) < 2001]
  items = {r["item"] for r in pre2001}
  strict = [r for r in pre2001 if r["files"] & STRICT]

  with open(OUT, "w") as fh:
      for i, r in enumerate(pre2001):
          fh.write(json.dumps({"n": i, "item": r["item"], "title": r["title"], "year": r["year"],
                               "seq": r["seq"], "files": sorted(r["files"]),
                               "strict": bool(r["files"] & STRICT),
                               "text": r["text"]}) + "\n")

  print("new distinct passages per source file:")
  for f, n in sorted(per_file.items()):
      print("  %-16s %6d" % (f, n))
  print()
  print("dropped for a post-1999 year in the passage : %d" % skipped_future)
  print("distinct passages (all years)               : %d" % len(rows))
  print("distinct passages, catalogued pre-2001      : %d   <- B1 ceiling 12,206" % len(pre2001))
  print("distinct source items, pre-2001             : %d   <- B2 ceiling 1,814" % len(items))
  print("of those, book-level date-verified (strict) : %d passages" % len(strict))
  print("wrote", OUT)


if __name__ == "__main__":
    main()